Background
Currently, the project's state classes are defined as TypedDicts. While TypedDict provides some structural typing, it allows too much flexibility at runtime, leading to unpredictable object shapes being passed to API.py and downstream code — which in turn makes LLM prompts and outputs less consistent.
Goal
Replace all current TypedDict-based state classes with Pydantic BaseModel classes to enforce strict validation and serialization. This will ensure API.py receives well-structured, validated objects and will make LLM interactions more reliable.
Scope
- Identify all TypedDict state classes across the repository (likely in modules named state.py, types.py, models.py, or similar).
- Replace each TypedDict with an equivalent Pydantic model (inherit from pydantic.BaseModel).
- Preserve field names and types where possible; when types are ambiguous, choose the safest stricter type and add comments TODOs for ambiguous cases.
- Add validation logic for fields where necessary (e.g., constraining strings, lists, optional fields).
- Update imports and any type checks/usages that depended on TypedDicts.
- Update tests and add new tests asserting Pydantic validation behavior (invalid inputs raise ValidationError, serialization produces dicts matching expected schema).
- Update API.py to accept/construct Pydantic models if it currently expects plain dicts.
Acceptance Criteria
- All previous TypedDict definitions are removed or converted to Pydantic models.
- API.py no longer accepts unvalidated dicts for state objects; instead, it receives Pydantic model instances or validated dicts from model.dict().
- Existing unit tests are updated and passing (or new tests added) verifying validation and serialization behavior.
- A migration note is added to the repository README or a docs file describing the change and how to construct the new models.
Implementation Notes
- Use pydantic==1.x or pydantic v2 depending on the project's existing dependencies. If the repo already pins pydantic in requirements/pyproject, follow that.
- Consider adding a centralized module (e.g., project_ai.models.state_models) to keep models organized.
- For backward compatibility, consider providing helper functions (from_dict/to_dict) that convert between legacy dicts and Pydantic models during transition.
Risks
- This is a breaking change for any external consumers relying on dict-shaped state objects.
- Tests and other code paths will need thorough updates.
Background
Currently, the project's state classes are defined as TypedDicts. While TypedDict provides some structural typing, it allows too much flexibility at runtime, leading to unpredictable object shapes being passed to API.py and downstream code — which in turn makes LLM prompts and outputs less consistent.
Goal
Replace all current TypedDict-based state classes with Pydantic BaseModel classes to enforce strict validation and serialization. This will ensure API.py receives well-structured, validated objects and will make LLM interactions more reliable.
Scope
Acceptance Criteria
Implementation Notes
Risks