diff --git a/opus_orchestrator/__init__.py b/opus_orchestrator/__init__.py index 221de18..9b2c0d1 100644 --- a/opus_orchestrator/__init__.py +++ b/opus_orchestrator/__init__.py @@ -41,6 +41,21 @@ from opus_orchestrator.crews import ( create_fiction_crew, create_nonfiction_crew, ) +from opus_orchestrator.pydanticai_agent import ( + OpusPydanticAgent, + StorySeed, + CharacterProfile, + ChapterOutline, + ChapterDraft, + CritiqueResult, + StyleGuide, + create_story_seed_agent, + create_character_agent, + create_chapter_outline_agent, + create_chapter_draft_agent, + create_critique_agent, + create_style_guide_agent, +) __all__ = [ # Config @@ -71,12 +86,26 @@ __all__ = [ "OpusGraphState", "run_opus", "StoryFramework", - # Crews (NEW!) + # Crews "OpusCrew", "FictionCrew", "NonfictionCrew", "create_fiction_crew", "create_nonfiction_crew", + # PydanticAI (NEW!) + "OpusPydanticAgent", + "StorySeed", + "CharacterProfile", + "ChapterOutline", + "ChapterDraft", + "CritiqueResult", + "StyleGuide", + "create_story_seed_agent", + "create_character_agent", + "create_chapter_outline_agent", + "create_chapter_draft_agent", + "create_critique_agent", + "create_style_guide_agent", # Main "OpusOrchestrator", "CritiqueCrew", diff --git a/opus_orchestrator/pydanticai_agent.py b/opus_orchestrator/pydanticai_agent.py new file mode 100644 index 0000000..97384f0 --- /dev/null +++ b/opus_orchestrator/pydanticai_agent.py @@ -0,0 +1,251 @@ +"""PydanticAI Agent for Opus Orchestrator. + +Provides structured output validation and agent expertise using PydanticAI. +""" + +import os +from typing import Any, Optional, Type + +from pydantic import BaseModel +from pydantic_ai import Agent +from dotenv import load_dotenv + +load_dotenv("/home/solaria/.openclaw/workspace/opus-orchestrator-ai/.env") + +from opus_orchestrator.config import get_config + + +def get_pydanticai_model() -> str: + """Get the model name for PydanticAI. + + Returns: + Model string for PydanticAI + """ + config = get_config() + + # Map our config to PydanticAI model names + if config.agent.provider == "openai": + return "openai:gpt-4o" + elif config.agent.provider == "anthropic": + return "anthropic:claude-3-5-sonnet-20241022" + else: + return "openai:gpt-4o" + + +class OpusPydanticAgent: + """PydanticAI-powered agent with structured output validation. + + This agent ensures all outputs conform to defined schemas, + providing type-safe, validated responses. + """ + + def __init__( + self, + result_type: Optional[Type[BaseModel]] = None, + model: Optional[str] = None, + system_prompt: Optional[str] = None, + ): + """Initialize the PydanticAI agent. + + Args: + result_type: Pydantic model for structured output + model: Override model name + system_prompt: System prompt for the agent + """ + self.config = get_config() + self.model = model or get_pydanticai_model() + self.result_type = result_type + self.system_prompt = system_prompt + + # Create the PydanticAI agent + self._agent: Optional[Agent] = None + self._setup_agent() + + def _setup_agent(self) -> None: + """Set up the PydanticAI agent.""" + model = self.model + + # Build system prompt + system_prompt = self.system_prompt or """You are an expert writer and editor for Opus Orchestrator. +You produce high-quality, structured output that conforms to the given schema. +Always follow best practices for the content type you're creating.""" + + if self.result_type: + self._agent = Agent( + model=model, + output_type=self.result_type, + system_prompt=system_prompt, + ) + else: + self._agent = Agent( + model=model, + system_prompt=system_prompt, + ) + + async def run(self, prompt: str) -> BaseModel | str: + """Run the agent with a prompt. + + Args: + prompt: User prompt + + Returns: + Structured result or string + """ + if not self._agent: + self._setup_agent() + + if self.result_type: + result = await self._agent.run(prompt) + return result.output + else: + result = await self._agent.run(prompt) + return result.output + + def run_sync(self, prompt: str) -> BaseModel | str: + """Run the agent synchronously. + + Args: + prompt: User prompt + + Returns: + Structured result or string + """ + import asyncio + + if not self._agent: + self._setup_agent() + + if self.result_type: + result = asyncio.run(self._agent.run(prompt)) + return result.output + else: + result = asyncio.run(self._agent.run(prompt)) + return result.output + + +# ============================================================================= +# SCHEMAS FOR STRUCTURED OUTPUT +# ============================================================================= + +class StorySeed(BaseModel): + """Structured story seed output.""" + one_sentence: str + one_paragraph: str + genre: str + target_audience: str + tone: str + themes: list[str] + + +class CharacterProfile(BaseModel): + """Character profile schema.""" + name: str + role: str # protagonist, antagonist, supporting + description: str + motivations: list[str] + flaws: list[str] + backstory: str + arc: str # character transformation + + +class ChapterOutline(BaseModel): + """Chapter outline schema.""" + chapter_number: int + title: str + summary: str + pov_character: str + key_events: list[str] + setting: str + chapter_goal: str + chapter_resolution: str + + +class ChapterDraft(BaseModel): + """Chapter draft schema.""" + chapter_number: int + title: str + content: str + word_count: int + pov: str + tone: str + key_moments: list[str] + dialogue_snippets: list[str] + + +class CritiqueResult(BaseModel): + """Critique result schema.""" + score: float # 0.0 - 1.0 + strengths: list[str] + weaknesses: list[str] + suggestions: list[str] + verdict: str # APPROVED, MINOR_REVISIONS, MAJOR_REVISIONS + + +class StyleGuide(BaseModel): + """Style guide schema.""" + tone: str + voice: str + vocabulary_level: str + sentence_structure: str + pacing: str + dialogue_style: str + description_style: str + prohibited_elements: list[str] + + +# ============================================================================= +# FACTORY FUNCTIONS +# ============================================================================= + +def create_story_seed_agent() -> OpusPydanticAgent: + """Create agent for generating story seeds.""" + return OpusPydanticAgent( + result_type=StorySeed, + system_prompt="""You are a story architect expert. Create detailed story seeds +that capture the essence of a story concept. Include genre, audience, tone, and themes.""", + ) + + +def create_character_agent() -> OpusPydanticAgent: + """Create agent for generating character profiles.""" + return OpusPydanticAgent( + result_type=CharacterProfile, + system_prompt="""You are a character development expert. Create rich, detailed +character profiles with clear motivations, flaws, and arcs. Make characters feel real.""", + ) + + +def create_chapter_outline_agent() -> OpusPydanticAgent: + """Create agent for generating chapter outlines.""" + return OpusPydanticAgent( + result_type=ChapterOutline, + system_prompt="""You are a plot structure expert. Create detailed chapter outlines +with clear goals, events, and resolutions. Ensure pacing works for the genre.""", + ) + + +def create_chapter_draft_agent() -> OpusPydanticAgent: + """Create agent for writing chapter drafts.""" + return OpusPydanticAgent( + result_type=ChapterDraft, + system_prompt="""You are an expert fiction writer. Write compelling, well-paced +chapters that follow the outline while adding rich detail, dialogue, and description.""", + ) + + +def create_critique_agent() -> OpusPydanticAgent: + """Create agent for critiquing content.""" + return OpusPydanticAgent( + result_type=CritiqueResult, + system_prompt="""You are an expert literary critic. Evaluate content objectively, +providing constructive feedback. Score honestly - don't inflate scores. Be specific.""", + ) + + +def create_style_guide_agent() -> OpusPydanticAgent: + """Create agent for generating style guides.""" + return OpusPydanticAgent( + result_type=StyleGuide, + system_prompt="""You are an expert editor. Create comprehensive style guides +that capture the voice and tone of the intended work. Be specific and actionable.""", + )