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."""

main.py

@@ -19,12 +19,12 @@ def main():
             iterations=10,  # Reduced for testing
             patience=8,
         )
-        print(f"Initial text: {results['initial_text']}")
+        print(f"Initial text: {results.initial_text}")
-        print(f"Final tweet: {results['final_tweet']}")
+        print(f"Final tweet: {results.final_tweet}")
-        print(f"Best score: {results['best_score']:.2f}")
+        print(f"Best score: {results.best_score:.2f}")
-        print(f"Iterations run: {results['iterations_run']}")
+        print(f"Iterations run: {results.iterations_run}")
-        print(f"Improvements found: {results['improvement_count']}")
+        print(f"Improvements found: {results.improvement_count}")
-        print(f"Early stopped: {results['early_stopped']}")
+        print(f"Early stopped: {results.early_stopped}")
     except Exception as e:
         print(f"Error in optimization: {e}")