Timmy-time-dashboard/src/timmy/agentic_loop.py

"""Agentic loop — multi-step task execution with progress tracking.

Provides `run_agentic_loop()`, the engine behind the `plan_and_execute` tool.
When the model recognises a task needs 3+ sequential steps, it calls
`plan_and_execute(task)` which spawns this loop in the background.

Flow:
  1. Planning — ask the model to break the task into numbered steps
  2. Execution — run each step sequentially, feeding results forward
  3. Adaptation — on failure, ask the model to adapt the plan
  4. Summary — ask the model to summarise what was accomplished

Progress is broadcast via WebSocket so the dashboard can show live updates.
"""

from __future__ import annotations

import asyncio
import logging
import re
import threading
import time
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass
class AgenticStep:
    """Result of a single step in the agentic loop."""

    step_num: int
    description: str
    result: str
    status: str  # "completed" | "failed" | "adapted"
    duration_ms: int


@dataclass
class AgenticResult:
    """Final result of the entire agentic loop."""

    task_id: str
    task: str
    summary: str
    steps: list[AgenticStep] = field(default_factory=list)
    status: str = "completed"  # "completed" | "partial" | "failed"
    total_duration_ms: int = 0


# ---------------------------------------------------------------------------
# Agent factory
# ---------------------------------------------------------------------------

_loop_agent = None
_loop_agent_lock = threading.Lock()


def _get_loop_agent():
    """Create a fresh agent for the agentic loop.

    Returns the same type of agent as `create_timmy()` but with a
    dedicated session so it doesn't pollute the main chat history.
    """
    global _loop_agent
    if _loop_agent is None:
        with _loop_agent_lock:
            if _loop_agent is None:
                from timmy.agent import create_timmy

                _loop_agent = create_timmy()
    return _loop_agent


# ---------------------------------------------------------------------------
# Plan parser
# ---------------------------------------------------------------------------

_STEP_RE = re.compile(r"^\s*(\d+)[.)]\s*(.+)$", re.MULTILINE)


def _parse_steps(plan_text: str) -> list[str]:
    """Extract numbered steps from the model's planning output."""
    matches = _STEP_RE.findall(plan_text)
    if matches:
        return [desc.strip() for _, desc in matches]
    # Fallback: split on newlines, ignore blanks
    return [line.strip() for line in plan_text.strip().splitlines() if line.strip()]


# ---------------------------------------------------------------------------
# Extracted helpers
# ---------------------------------------------------------------------------


def _extract_content(run_result) -> str:
    """Extract text content from an agent run result."""
    return run_result.content if hasattr(run_result, "content") else str(run_result)


def _clean(text: str) -> str:
    """Clean a model response using session's response cleaner."""
    from timmy.session import _clean_response

    return _clean_response(text)


async def _plan_task(
    agent, task: str, session_id: str, max_steps: int
) -> tuple[list[str], bool] | str:
    """Run the planning phase — returns (steps, was_truncated) or error string."""
    plan_prompt = (
        f"Break this task into numbered steps (max {max_steps}). "
        f"Return ONLY a numbered list, nothing else.\n\n"
        f"Task: {task}"
    )
    try:
        plan_run = await asyncio.to_thread(
            agent.run, plan_prompt, stream=False, session_id=f"{session_id}_plan"
        )
        plan_text = _extract_content(plan_run)
    except Exception as exc:  # broad catch intentional: agent.run can raise any error
        logger.error("Agentic loop: planning failed: %s", exc)
        return f"Planning failed: {exc}"

    steps = _parse_steps(plan_text)
    if not steps:
        return "Planning produced no steps."

    planned_count = len(steps)
    steps = steps[:max_steps]
    return steps, planned_count > len(steps)


async def _execute_step(
    agent,
    task: str,
    step_desc: str,
    step_num: int,
    total_steps: int,
    recent_results: list[str],
    session_id: str,
) -> AgenticStep:
    """Execute a single step, returning an AgenticStep."""
    step_start = time.monotonic()
    context = (
        f"Task: {task}\n"
        f"Step {step_num}/{total_steps}: {step_desc}\n"
        f"Recent progress: {recent_results[-2:] if recent_results else []}\n\n"
        f"Execute this step and report what you did."
    )
    step_run = await asyncio.to_thread(
        agent.run, context, stream=False, session_id=f"{session_id}_step{step_num}"
    )
    step_result = _clean(_extract_content(step_run))
    return AgenticStep(
        step_num=step_num,
        description=step_desc,
        result=step_result,
        status="completed",
        duration_ms=int((time.monotonic() - step_start) * 1000),
    )


async def _adapt_step(
    agent,
    step_desc: str,
    step_num: int,
    error: Exception,
    step_start: float,
    session_id: str,
) -> AgenticStep:
    """Attempt adaptation after a step failure."""
    adapt_prompt = (
        f"Step {step_num} failed with error: {error}\n"
        f"Original step was: {step_desc}\n"
        f"Adapt the plan and try an alternative approach for this step."
    )
    adapt_run = await asyncio.to_thread(
        agent.run, adapt_prompt, stream=False, session_id=f"{session_id}_adapt{step_num}"
    )
    adapt_result = _clean(_extract_content(adapt_run))
    return AgenticStep(
        step_num=step_num,
        description=f"[Adapted] {step_desc}",
        result=adapt_result,
        status="adapted",
        duration_ms=int((time.monotonic() - step_start) * 1000),
    )


def _summarize(result: AgenticResult, total_steps: int, was_truncated: bool) -> None:
    """Fill in summary and final status on the result object (mutates in place)."""
    completed = sum(1 for s in result.steps if s.status == "completed")
    adapted = sum(1 for s in result.steps if s.status == "adapted")
    failed = sum(1 for s in result.steps if s.status == "failed")

    parts = [f"Completed {completed}/{total_steps} steps"]
    if adapted:
        parts.append(f"{adapted} adapted")
    if failed:
        parts.append(f"{failed} failed")
    result.summary = f"{result.task}: {', '.join(parts)}."

    if was_truncated or len(result.steps) < total_steps or failed:
        result.status = "partial"
    else:
        result.status = "completed"


# ---------------------------------------------------------------------------
# Core loop
# ---------------------------------------------------------------------------


async def run_agentic_loop(
    task: str,
    *,
    session_id: str = "agentic",
    max_steps: int = 0,
    on_progress: Callable | None = None,
) -> AgenticResult:
    """Execute a multi-step task with planning, execution, and adaptation.

    Args:
        task:        Full description of the task to execute.
        session_id:  Agno session_id for conversation continuity.
        max_steps:   Max steps to execute (0 = use config default).
        on_progress: Optional async callback(description, step_num, total_steps).

    Returns:
        AgenticResult with steps, summary, and status.
    """
    from config import settings

    if max_steps <= 0:
        max_steps = getattr(settings, "max_agent_steps", 10)

    task_id = str(uuid.uuid4())[:8]
    start_time = time.monotonic()
    agent = _get_loop_agent()
    result = AgenticResult(task_id=task_id, task=task, summary="")

    # Phase 1: Planning
    plan = await _plan_task(agent, task, session_id, max_steps)
    if isinstance(plan, str):
        result.status = "failed"
        result.summary = plan
        result.total_duration_ms = int((time.monotonic() - start_time) * 1000)
        return result

    steps, was_truncated = plan
    total_steps = len(steps)

    await _broadcast_progress(
        "agentic.plan_ready",
        {"task_id": task_id, "task": task, "steps": steps, "total": total_steps},
    )

    # Phase 2: Execution
    completed_results: list[str] = []
    for i, step_desc in enumerate(steps, 1):
        step_start = time.monotonic()
        try:
            step = await _execute_step(
                agent,
                task,
                step_desc,
                i,
                total_steps,
                completed_results,
                session_id,
            )
            result.steps.append(step)
            completed_results.append(f"Step {i}: {step.result[:200]}")
            await _broadcast_progress(
                "agentic.step_complete",
                {
                    "task_id": task_id,
                    "step": i,
                    "total": total_steps,
                    "description": step_desc,
                    "result": step.result[:200],
                },
            )
            if on_progress:
                await on_progress(step_desc, i, total_steps)

        except Exception as exc:  # broad catch intentional: agent.run can raise any error
            logger.warning("Agentic loop step %d failed: %s", i, exc)
            try:
                step = await _adapt_step(agent, step_desc, i, exc, step_start, session_id)
                result.steps.append(step)
                completed_results.append(f"Step {i} (adapted): {step.result[:200]}")
                await _broadcast_progress(
                    "agentic.step_adapted",
                    {
                        "task_id": task_id,
                        "step": i,
                        "total": total_steps,
                        "description": step_desc,
                        "error": str(exc),
                        "adaptation": step.result[:200],
                    },
                )
                if on_progress:
                    await on_progress(f"[Adapted] {step_desc}", i, total_steps)
            except Exception as adapt_exc:  # broad catch intentional
                logger.error("Agentic loop adaptation also failed: %s", adapt_exc)
                result.steps.append(
                    AgenticStep(
                        step_num=i,
                        description=step_desc,
                        result=f"Failed: {exc}; Adaptation also failed: {adapt_exc}",
                        status="failed",
                        duration_ms=int((time.monotonic() - step_start) * 1000),
                    )
                )
                completed_results.append(f"Step {i}: FAILED")

    # Phase 3: Summary
    _summarize(result, total_steps, was_truncated)
    result.total_duration_ms = int((time.monotonic() - start_time) * 1000)

    await _broadcast_progress(
        "agentic.task_complete",
        {
            "task_id": task_id,
            "status": result.status,
            "steps_completed": len(result.steps),
            "summary": result.summary[:300],
            "duration_ms": result.total_duration_ms,
        },
    )

    return result


# ---------------------------------------------------------------------------
# WebSocket broadcast helper
# ---------------------------------------------------------------------------


async def _broadcast_progress(event: str, data: dict) -> None:
    """Broadcast agentic loop progress via WebSocket (best-effort)."""
    try:
        from infrastructure.ws_manager.handler import ws_manager

        await ws_manager.broadcast(event, data)
    except (ImportError, AttributeError, ConnectionError, RuntimeError) as exc:
        logger.warning("Agentic loop broadcast failed: %s", exc)
        logger.debug("Agentic loop: WS broadcast failed for %s", event)