diff --git a/imagecreator/backend/app/services/llm.py b/imagecreator/backend/app/services/llm.py index 7ec364a..e7ed806 100644 --- a/imagecreator/backend/app/services/llm.py +++ b/imagecreator/backend/app/services/llm.py @@ -1 +1,14 @@ -"""LLM service placeholder.""" +from typing import Optional +from google.genai import Client +from app.config import settings + +class LLMService: + def __init__(self): + self.client = Client(api_key=settings.GEMINI_API_KEY) + + def generate_text(self, prompt: str, system_instruction: Optional[str] = None) -> str: + response = self.client.models.generate_content( + model="gemini-2.5-flash", + contents=prompt + ) + return response.text \ No newline at end of file diff --git a/imagecreator/backend/app/services/prompt.py b/imagecreator/backend/app/services/prompt.py new file mode 100644 index 0000000..5928923 --- /dev/null +++ b/imagecreator/backend/app/services/prompt.py @@ -0,0 +1,28 @@ +from app.services.llm import LLMService + +TEXT2IMAGE_SYSTEM_PROMPT = """You are an expert at optimizing text-to-image prompts. +Your task is to take user's simple or vague prompt and enhance it with detailed, vivid descriptions. +Include details about: subject, setting, lighting, mood, style, composition, and quality modifiers. +Output ONLY the optimized prompt in English, no explanations.""" + +IMAGE2IMAGE_SYSTEM_PROMPT = """You are an expert at optimizing image-to-image prompts. +The user wants to modify an existing image. Your task is to clearly describe what should be changed. + +Important rules: +1. Identify if user wants to ADD, REMOVE, REPLACE, or ENHANCE something +2. Be specific about what to keep vs what to change +3. Describe the desired style, mood, and quality +4. Output ONLY the optimized prompt in English, no explanations.""" + +class PromptOptimizationService: + def __init__(self): + self.llm = LLMService() + + def optimize_text2image(self, user_prompt: str) -> str: + full_prompt = f"{TEXT2IMAGE_SYSTEM_PROMPT}\n\nUser's original prompt: {user_prompt}" + return self.llm.generate_text(full_prompt) + + def optimize_image2image(self, user_prompt: str, original_description: str = "") -> str: + context = f"Original image description: {original_description}\n\n" if original_description else "" + full_prompt = f"{IMAGE2IMAGE_SYSTEM_PROMPT}\n\n{context}User's modification request: {user_prompt}" + return self.llm.generate_text(full_prompt) \ No newline at end of file diff --git a/imagecreator/backend/tests/test_prompt.py b/imagecreator/backend/tests/test_prompt.py new file mode 100644 index 0000000..e4569a9 --- /dev/null +++ b/imagecreator/backend/tests/test_prompt.py @@ -0,0 +1,17 @@ +import pytest +from app.services.prompt import PromptOptimizationService + +def test_optimize_text2image_prompt(): + service = PromptOptimizationService() + result = service.optimize_text2image("a cat") + + assert result is not None + assert len(result) > len("a cat") + assert "cat" in result.lower() + +def test_optimize_image2image_prompt(): + service = PromptOptimizationService() + result = service.optimize_image2image("make it colorful", "a cat image") + + assert result is not None + assert len(result) > 0 \ No newline at end of file