82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
from pydantic import BaseModel, Field
|
|
from typing import List, Optional
|
|
from pydantic_ai import Agent, RunContext
|
|
import os
|
|
|
|
# --- Data Models ---
|
|
|
|
class Character(BaseModel):
|
|
name: str = Field(description="The name of the character.")
|
|
description: str = Field(description="A concise visual description of the character (e.g., 'Blonde hair, blue dress, young').")
|
|
|
|
class CharacterAnalysisResult(BaseModel):
|
|
characters: List[Character] = Field(description="List of main characters identified in the text.")
|
|
|
|
class MangaSimplePrompt(BaseModel):
|
|
prompt: str = Field(description="The generated English manga image prompt.")
|
|
|
|
# --- Wrapper Functions ---
|
|
|
|
async def analyze_characters_with_agent(text: str, api_key: str, base_url: Optional[str] = None, model: str = "gpt-4o") -> str:
|
|
"""
|
|
Uses PydanticAI Agent to analyze characters and returns a formatted string context.
|
|
"""
|
|
from pydantic_ai.models.openai import OpenAIModel
|
|
|
|
# Allow model override
|
|
model_name = model if model else "gpt-4o"
|
|
|
|
# Create the model instance with the specific API Key and Base URL
|
|
openai_model = OpenAIModel(
|
|
model_name,
|
|
api_key=api_key,
|
|
base_url=base_url
|
|
)
|
|
|
|
# Create a temporary agent for this run
|
|
agent = Agent(
|
|
openai_model,
|
|
result_type=CharacterAnalysisResult,
|
|
system_prompt="You are a professional manga editor. Analyze the provided novel text and extract the visual descriptions of the main characters to ensure consistency in manga adaptation."
|
|
)
|
|
|
|
try:
|
|
result = await agent.run(text)
|
|
|
|
# Format the result into a string context
|
|
context_str = ""
|
|
for char in result.data.characters:
|
|
context_str += f"- {char.name}: {char.description}\n"
|
|
return context_str
|
|
except Exception as e:
|
|
print(f"Agent Error (Characters): {e}")
|
|
return ""
|
|
|
|
async def generate_single_prompt_with_agent(paragraph: str, character_context: str, api_key: str, base_url: Optional[str] = None, model: str = "gpt-4o") -> str:
|
|
from pydantic_ai.models.openai import OpenAIModel
|
|
|
|
model_name = model if model else "gpt-4o"
|
|
openai_model = OpenAIModel(
|
|
model_name,
|
|
api_key=api_key,
|
|
base_url=base_url
|
|
)
|
|
|
|
agent = Agent(
|
|
openai_model,
|
|
result_type=MangaSimplePrompt,
|
|
deps_type=str,
|
|
system_prompt=(
|
|
"You are a professional manga artist assistant. Convert the novel text into a detailed manga image prompt. "
|
|
"The prompt MUST be in English. Focus on visual details, character appearance, setting, and style (monochrome, manga style, high quality). "
|
|
f"Maintain character consistency:\n{character_context}"
|
|
)
|
|
)
|
|
|
|
try:
|
|
result = await agent.run(paragraph, deps=character_context)
|
|
return result.data.prompt
|
|
except Exception as e:
|
|
print(f"Agent Error (Prompt): {e}")
|
|
return f"Error generation prompt: {e}"
|