from __future__ import annotations

from app.core.agent.processes.v1.workflow.flow_main.context import V1FlowContext
from app.core.agent.processes.v1.workflow.flow_main.steps.finalize_answer_step import FinalizeAnswerStep
from app.core.agent.processes.v1.workflow.flow_main.steps.generate_answer_step import GenerateAnswerStep
from app.core.agent.processes.v1.workflow.flow_main.steps.prepare_user_message_step import PrepareUserMessageStep
from app.core.agent.utils.llm import AgentLlmService
from app.core.agent.utils.workflow import WorkflowGraph


class V1FlowMainGraph:
    def __init__(self, llm: AgentLlmService) -> None:
        self._graph = WorkflowGraph(
            workflow_id="v1.flow_main",
            source="workflow.v1",
            steps=(
                PrepareUserMessageStep(),
                GenerateAnswerStep(llm),
                FinalizeAnswerStep(),
            ),
        )

    async def run(self, context: V1FlowContext) -> V1FlowContext:
        return await self._graph.run(context)
