from smolagents import CodeAgent
from smolagents.memory import ActionStep, FinalAnswerStep, PlanningStep


class StreamingAgentRunner:
    def __init__(self, agent: CodeAgent):
        self.agent = agent

    def run(self, query: str):
        """Генератор прогресса агента и финального ответа."""
        parts = ["Агент думает...\n"]
        yield parts[0], 0

        try:
            for step in self.agent.run(query, stream=True, reset=True):
                if isinstance(step, PlanningStep):
                    parts.append(f"**План:** {step.plan.strip()}")
                elif isinstance(step, ActionStep):
                    if step.tool_calls:
                        names = ", ".join(tc.name for tc in step.tool_calls)
                        parts.append(f"**Шаг {step.step_number}:** вызов `{names}`")
                    if step.error:
                        parts.append(f"**Ошибка:** {step.error}")
                elif isinstance(step, FinalAnswerStep):
                    parts.append(f"## Ответ\n\n{step.output}")

                progress = min(len(parts) / (self.agent.max_steps + 1), 0.95)
                yield "\n\n".join(parts), progress

            yield "\n\n".join(parts), 1.0

        except Exception as e:
            yield f"Ошибка: {e}", 1.0