Refactor: Replace dataclasses with Pydantic models

agents/impl_generator.py

@@ -3,21 +3,10 @@
 from dotenv import load_dotenv
 import textwrap
 from . import utils
-import dataclasses
-
-
-@dataclasses.dataclass(frozen=True)
-class GenerationAttempt:
-    code: str
-    errors: str
-
-
-@dataclasses.dataclass(frozen=True)
-class ImplGenerationRequest:
-    interface_str: str
-    test_str: str
-    prior_attempts: list[GenerationAttempt] = dataclasses.field(default_factory=list)
+from .models import GenerationAttempt, ImplGenerationRequest # Import Pydantic models
 
+# Removed dataclass definitions for GenerationAttempt and ImplGenerationRequest
+# Removed import dataclasses
 
 class ImplGenerator:
 

agents/models.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel, Field
+from typing import List
+
+class GenerationAttempt(BaseModel):
+    code: str
+    errors: str
+
+class ImplGenerationRequest(BaseModel):
+    interface_str: str
+    test_str: str
+    prior_attempts: List[GenerationAttempt] = Field(default_factory=list)
+
+class TestGenerationRequest(BaseModel):
+    interface_str: str
+    prior_attempts: List[GenerationAttempt] = Field(default_factory=list)

agents/test_generator.py

@@ -3,20 +3,10 @@
 from dotenv import load_dotenv
 import textwrap
 from . import utils
-from dataclasses import dataclass, field
-
-
-@dataclass(frozen=True)
-class GenerationAttempt:
-    code: str
-    errors: str
-
-
-@dataclass(frozen=True)
-class TestGenerationRequest:
-    interface_str: str
-    prior_attempts: list[GenerationAttempt] = field(default_factory=list)
+from .models import GenerationAttempt, TestGenerationRequest # Import Pydantic models
 
+# Removed dataclass definitions for GenerationAttempt and TestGenerationRequest
+# Removed import for dataclasses
 
 class TestGenerator:
 

demo.py

@@ -1,6 +1,6 @@
 from agents import impl_generator
 from agents import test_generator
-from dataclasses import dataclass
+from pydantic import BaseModel # Import BaseModel for Pydantic
 from dotenv import load_dotenv
 import os
 import py_compile
@@ -10,12 +10,13 @@
 import logfire
 
 
-@dataclass
-class Example:
+# Define the Pydantic model Example
+class Example(BaseModel):
     code: str
     project_name: str
     filename: str
 
+# Removed old dataclass Example definition
 
 math_utils = Example(
     project_name='math-utils',

requirements.txt

@@ -5,3 +5,4 @@ pytest
 modelcontextprotocol
 uvicorn[standard]
 fastapi
+pydantic