feat: implement LLM and prompt optimization service
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -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
|
||||||
Reference in New Issue
Block a user