Complete Migration

agent/index.py

@@ -1,6 +1,6 @@
 from modaic import PrecompiledAgent, PrecompiledConfig
 from .modules import TweetGeneratorModule, TweetEvaluatorModule
-from .models import EvaluationResult
+from .models import EvaluationResult, FinalResult
 from .hill_climbing import HillClimbingOptimizer
 from typing import Optional, List
 from .utils import get_dspy_lm
@@ -50,7 +50,7 @@ class TweetOptimizerAgent(PrecompiledAgent):
         input_text: str,
         iterations: Optional[int] = None,
         patience: Optional[int] = None,
-    ) -> str:
+    ) -> FinalResult:
         """Run full optimization process."""
         max_iterations = iterations or self.max_iterations
         patience_limit = patience or self.patience
@@ -97,6 +97,7 @@ class TweetOptimizerAgent(PrecompiledAgent):
         results.update({"final_tweet": best_tweet, "best_score": best_score})
 
         self.reset()
+        results = FinalResult(**results)
 
         return results
 

agent/models.py

@@ -23,6 +23,16 @@ class CategoryEvaluation(BaseModel):
             )
         return score
 
+class FinalResult(BaseModel):
+    """Pydantic model for final tweet optimization result."""
+    initial_text: str = Field(description="The initial tweet text")
+    final_tweet: str = Field(description="The final optimized tweet")
+    best_score: float = Field(description="The best score for the final tweet")
+    iterations_run: int = Field(description="The number of iterations run")
+    early_stopped: bool = Field(description="Whether the optimization early stopped")
+    scores_history: List[List[int]] = Field(description="The history of scores")
+    improvement_count: int = Field(description="The number of improvements found")
+
 
 class EvaluationResult(BaseModel):
     """Pydantic model for tweet evaluation results."""