from __future__ import annotations

import asyncio

from app.core.agent.processes.v1.workflow.flow_main.context import V1FlowContext
from app.core.agent.utils.llm import AgentLlmService
from app.core.agent.utils.workflow import WorkflowStep


class GenerateAnswerStep(WorkflowStep[V1FlowContext]):
    step_id = "generate_answer"
    title = "Вызов LLM"

    def __init__(self, llm: AgentLlmService) -> None:
        self._llm = llm

    async def run(self, context: V1FlowContext) -> V1FlowContext:
        request_id = context.runtime.request.request_id
        context.answer = await asyncio.to_thread(
            self._llm.generate,
            context.prompt_name,
            context.prepared_message,
            log_context=f"agent:{request_id}",
            trace=context.runtime.trace.module("workflow.v1.llm"),
        )
        return context

    def trace_input(self, context: V1FlowContext) -> dict[str, object]:
        return {"prompt_name": context.prompt_name, "prepared_message_length": len(context.prepared_message)}

    def trace_output(self, context: V1FlowContext) -> dict[str, object]:
        return {"answer_length": len(context.answer)}
