diff --git a/agent/__pycache__/constants.cpython-310.pyc b/agent/__pycache__/constants.cpython-310.pyc
index 46f0c3e..8b7a0a1 100644
Binary files a/agent/__pycache__/constants.cpython-310.pyc and b/agent/__pycache__/constants.cpython-310.pyc differ
diff --git a/agent/__pycache__/helpers.cpython-310.pyc b/agent/__pycache__/helpers.cpython-310.pyc
index 212e7bd..8824895 100644
Binary files a/agent/__pycache__/helpers.cpython-310.pyc and b/agent/__pycache__/helpers.cpython-310.pyc differ
diff --git a/agent/__pycache__/hill_climbing.cpython-310.pyc b/agent/__pycache__/hill_climbing.cpython-310.pyc
index 9346839..3329097 100644
Binary files a/agent/__pycache__/hill_climbing.cpython-310.pyc and b/agent/__pycache__/hill_climbing.cpython-310.pyc differ
diff --git a/agent/__pycache__/index.cpython-310.pyc b/agent/__pycache__/index.cpython-310.pyc
new file mode 100644
index 0000000..f2316f1
Binary files /dev/null and b/agent/__pycache__/index.cpython-310.pyc differ
diff --git a/agent/__pycache__/models.cpython-310.pyc b/agent/__pycache__/models.cpython-310.pyc
index d6f8f20..372dac3 100644
Binary files a/agent/__pycache__/models.cpython-310.pyc and b/agent/__pycache__/models.cpython-310.pyc differ
diff --git a/agent/__pycache__/modules.cpython-310.pyc b/agent/__pycache__/modules.cpython-310.pyc
index 30b5443..47e0184 100644
Binary files a/agent/__pycache__/modules.cpython-310.pyc and b/agent/__pycache__/modules.cpython-310.pyc differ
diff --git a/agent/__pycache__/utils.cpython-310.pyc b/agent/__pycache__/utils.cpython-310.pyc
index d5d51f0..b003ec8 100644
Binary files a/agent/__pycache__/utils.cpython-310.pyc and b/agent/__pycache__/utils.cpython-310.pyc differ
diff --git a/agent/index.py b/agent/index.py
index 5fe2607..8e4e9d8 100644
--- a/agent/index.py
+++ b/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
 
diff --git a/agent/models.py b/agent/models.py
index 81b5308..fea59e8 100644
--- a/agent/models.py
+++ b/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."""
diff --git a/main.py b/main.py
index 90b10eb..5fa4dde 100644
--- a/main.py
+++ b/main.py
@@ -19,12 +19,12 @@ def main():
             iterations=10,  # Reduced for testing
             patience=8,
         )
-        print(f"Initial text: {results['initial_text']}")
-        print(f"Final tweet: {results['final_tweet']}")
-        print(f"Best score: {results['best_score']:.2f}")
-        print(f"Iterations run: {results['iterations_run']}")
-        print(f"Improvements found: {results['improvement_count']}")
-        print(f"Early stopped: {results['early_stopped']}")
+        print(f"Initial text: {results.initial_text}")
+        print(f"Final tweet: {results.final_tweet}")
+        print(f"Best score: {results.best_score:.2f}")
+        print(f"Iterations run: {results.iterations_run}")
+        print(f"Improvements found: {results.improvement_count}")
+        print(f"Early stopped: {results.early_stopped}")
     except Exception as e:
         print(f"Error in optimization: {e}")