Artifex/project_brain/sol.py

from __future__ import annotations

from dataclasses import dataclass


from model_router.router import ModelRequestContract, ModelRouter
from project_brain.planning import ProjectPlanContract, parse_project_plan_response


@dataclass(frozen=True)
class PlanningPrompt:
    idea: str
    constraints: dict[str, object]


class SolProjectBrain:
    def __init__(self, router: ModelRouter) -> None:
        self.router = router

    def plan_project(self, prompt: PlanningPrompt) -> str:
        response = self.router.complete(
            ModelRequestContract(purpose="PLANNING", prompt=prompt.idea)
        )
        return response.content

    def create_project_plan_contract(self, prompt: PlanningPrompt) -> ProjectPlanContract:
        response = self.router.complete(ModelRequestContract(purpose="PLANNING", prompt=prompt.idea))
        return parse_project_plan_response(response.content)