src/routers/analyze_generate_graph.py

from fastapi import APIRouter, Response
from src.models.analysis_models import InputModel
from src.workflows.analysis_workflow import MLAnalysisWorkflow, MLImplementationPlanner
from datetime import datetime
from phi.storage.workflow.sqlite import SqlWorkflowStorage
from phi.utils.pprint import pprint_run_response
from phi.utils.log import logger
from typing import Iterator
from llama_index.core.settings import Settings
from src.models.workflow_graph import GraphInputSchema, GraphOutputSchema
from src.workflows.graph_workflow import DesignGraphWorkflow

router = APIRouter()

async def analyze_problem(problem_statement: str):

    analysis_workflow = MLAnalysisWorkflow(
        session_id=f"ml-analysis-{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        storage=SqlWorkflowStorage(
            table_name="ml_analysis_workflows",
            db_file="storage/workflows.db"
        )
    )

    analysis_response: Iterator[RunResponse] = analysis_workflow.run(problem_statement)
    
    pprint_run_response(analysis_response, markdown=True)

    requirements_result = analysis_workflow.requirements_analyst.run_response.content if analysis_workflow.requirements_analyst.run_response else None
    research_result = analysis_workflow.technical_researcher.run_response.content if analysis_workflow.technical_researcher.run_response else None

    if requirements_result:
        logger.info("===Planning Phase===")
        planning_workflow = MLImplementationPlanner(
            session_id=f"ml-planning-{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            storage=SqlWorkflowStorage(
                table_name="ml_planning_workflows",
                db_file="storage/workflows.db"
            )
        )

        planning_response_stream: Iterator[RunResponse] = planning_workflow.run(requirements_result, research_result)
        
        pprint_run_response(planning_response_stream, markdown=True)

        return planning_workflow.writer.run_response.content

    else:
        return "Requirements analysis did not complete successfully."

@router.post("/", response_model=GraphOutputSchema)
async def analyzer_generate_graph(data: InputModel):

    task_description = await analyze_problem(data.problem_statement)

    try:
        graph_workflow = DesignGraphWorkflow(timeout=60, verbose=True)
        graph_result = await graph_workflow.run(_project_description=task_description, llm=Settings._llm)

        return GraphOutputSchema(graph=graph_result)

    except Exception as e:
        return {"detail": f"Error processing {e}"}