leocrm/app/ai/agent_loop.py

"""Core ReAct (Reasoning + Acting) loop for AI agents.

Implements a true ReAct loop that alternates between LLM reasoning and tool
execution. Each step records the thought (LLM content), action (tool name),
action_input (tool arguments), and observation (tool result).

The loop terminates when:
- The LLM returns a final response without tool calls (completed)
- max_steps is reached (stopped_max_steps)
- timeout is exceeded (stopped_timeout)
- A permanent error occurs (stopped_error)

Error handling uses ``ErrorCategory`` from ``app.core.error_codes``:
- TRANSIENT → retry the LLM call (up to 3 retries per step)
- PERMANENT → stop the loop immediately
- PARTIAL → continue with partial results

Usage::

    from app.ai.agent_loop import run_react_loop

    result = await run_react_loop(
        agent_definition=agent,
        messages=[{"role": "user", "content": "Summarize recent emails"}],
        tools=tool_schemas,
        tool_registry=registry,
        db=db_session,
        tenant_id=tenant_id,
        user_id=user_id,
    )
    print(result.final_content, result.total_cost_usd, result.steps_taken)
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

from app.ai.llm_client import llm_complete
from app.core.error_codes import ErrorCategory, classify_exception

if TYPE_CHECKING:
    from collections.abc import Callable

    from sqlalchemy.ext.asyncio import AsyncSession

    from app.plugins.builtins.ai_assistant.tool_registry import ToolRegistry

logger = logging.getLogger(__name__)

# Maximum retries for transient errors per LLM step
_MAX_TRANSIENT_RETRIES = 3


# ──────────────────────────────────────────────────────────────────────────
# Data structures
# ──────────────────────────────────────────────────────────────────────────


@dataclass
class ReActStep:
    """A single step in the ReAct loop (Thought → Action → Observation)."""

    step_number: int
    thought: str  # LLM content before tool calls
    action: str | None  # Tool name (None if final response)
    action_input: dict[str, Any] | None  # Tool arguments
    observation: str | None  # Tool result
    cost_usd: float
    timestamp: str  # ISO format


@dataclass
class ReActResult:
    """Final result of the ReAct loop."""

    final_content: str
    steps: list[ReActStep] = field(default_factory=list)
    total_cost_usd: float = 0.0
    steps_taken: int = 0
    status: str = "completed"  # completed | stopped_max_steps | stopped_timeout | stopped_error
    error: str | None = None


# ──────────────────────────────────────────────────────────────────────────
# Core loop
# ──────────────────────────────────────────────────────────────────────────


def _extract_tool_calls(raw_response: Any) -> list[dict[str, Any]]:
    """Extract tool calls from a LiteLLM raw response.

    Returns a list of dicts with keys: ``id``, ``name``, ``arguments``.
    """
    tool_calls: list[dict[str, Any]] = []
    try:
        msg = raw_response.choices[0].message
        if hasattr(msg, "tool_calls") and msg.tool_calls:
            for tc in msg.tool_calls:
                tool_calls.append({
                    "id": tc.id or "",
                    "name": tc.function.name if tc.function else "",
                    "arguments": tc.function.arguments if tc.function and tc.function.arguments else "{}",
                })
    except (AttributeError, IndexError, TypeError) as exc:
        logger.debug("Failed to extract tool calls from response: %s", exc)
    return tool_calls


async def _execute_tool(
    tool_registry: ToolRegistry,
    tool_name: str,
    arguments: dict[str, Any],
    context: dict[str, Any],
) -> str:
    """Execute a single tool call via the registry.

    Returns the tool result as a string, or an error message.
    """
    tool = tool_registry.get(tool_name)
    if tool is None:
        return f"Error: Tool '{tool_name}' not found"

    try:
        result = await tool.handler(arguments=arguments, context=context)
        return result if isinstance(result, str) else json.dumps(result)
    except Exception as exc:
        logger.exception("Tool '%s' execution failed", tool_name)
        return f"Error: {exc}"


async def run_react_loop(
    agent_definition: Any,  # AgentDefinition from automation models
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
    tool_registry: ToolRegistry,
    db: AsyncSession,
    tenant_id: uuid.UUID,
    user_id: uuid.UUID,
    agent_run_id: uuid.UUID | None = None,
    max_steps: int = 20,
    timeout_seconds: int = 300,
    trace_id: str | None = None,
    on_step: Callable | None = None,
) -> ReActResult:
    """Execute a ReAct loop: LLM reasoning → tool execution → repeat.

    Args:
        agent_definition: AgentDefinition with llm_model, system_prompt, etc.
        messages: Initial chat messages (without system prompt).
        tools: OpenAI-format tool schemas for function calling.
        tool_registry: ToolRegistry instance for tool execution.
        db: Async DB session.
        tenant_id: Tenant ID for multi-tenancy.
        user_id: User ID for permission context.
        agent_run_id: Optional AgentRun ID for step persistence.
        max_steps: Maximum loop iterations (default 20).
        timeout_seconds: Overall timeout (default 300).
        trace_id: Optional trace ID for correlation.
        on_step: Optional async callback fired after each step.

    Returns:
        ReActResult with final content, steps, cost, and status.
    """
    from app.core.hooks import do_action

    result = ReActResult(final_content="", status="completed")
    start_time = time.monotonic()

    # Build LLM parameters from agent definition
    litellm_model = getattr(agent_definition, "llm_model", None) or "gpt-4o"
    system_prompt = getattr(agent_definition, "system_prompt", "") or "You are a helpful AI assistant."
    api_key = getattr(agent_definition, "api_key", None)
    api_base = getattr(agent_definition, "api_base", None)
    provider = getattr(agent_definition, "provider", None)
    max_tokens = getattr(agent_definition, "max_tokens", None) or 1000

    # Build the full message list with system prompt prepended
    full_messages: list[dict[str, Any]] = [
        {"role": "system", "content": system_prompt},
        *messages,
    ]

    tool_context: dict[str, Any] = {
        "tenant_id": str(tenant_id),
        "user_id": str(user_id),
        "db": db,
    }

    for step_num in range(1, max_steps + 1):
        # ── Timeout check ──
        elapsed = time.monotonic() - start_time
        if elapsed >= timeout_seconds:
            result.status = "stopped_timeout"
            result.error = f"Timeout after {elapsed:.1f}s (limit {timeout_seconds}s)"
            logger.warning("ReAct loop timed out at step %d: %s", step_num, result.error)
            break

        # ── LLM call with transient retry ──
        llm_result: dict[str, Any] | None = None
        last_error: str | None = None

        for retry in range(_MAX_TRANSIENT_RETRIES + 1):
            try:
                llm_result = await llm_complete(
                    model=litellm_model,
                    messages=full_messages,
                    tools=tools if tools else None,
                    temperature=0.3,
                    max_tokens=max_tokens,
                    api_key=api_key,
                    api_base=api_base,
                    provider=provider,
                    trace_id=trace_id,
                    tenant_id=tenant_id,
                    db=db,
                )
                break
            except Exception as exc:
                last_error = str(exc)
                category = classify_exception(exc)

                if category == ErrorCategory.PERMANENT:
                    result.status = "stopped_error"
                    result.error = f"Permanent error at step {step_num}: {exc}"
                    logger.error("ReAct loop permanent error: %s", result.error)
                    return result

                if category == ErrorCategory.TRANSIENT and retry < _MAX_TRANSIENT_RETRIES:
                    backoff = 2 ** retry
                    logger.warning(
                        "Transient error at step %d (retry %d/%d): %s — retrying in %ds",
                        step_num, retry + 1, _MAX_TRANSIENT_RETRIES, exc, backoff,
                    )
                    await asyncio.sleep(backoff)
                    continue

                # PARTIAL or exhausted retries
                if category == ErrorCategory.PARTIAL:
                    logger.warning("Partial error at step %d: %s — continuing", step_num, exc)
                    last_error = str(exc)
                    break

                # Exhausted transient retries
                result.status = "stopped_error"
                result.error = f"Error after {retry + 1} retries at step {step_num}: {exc}"
                logger.error("ReAct loop error: %s", result.error)
                return result

        if llm_result is None:
            result.status = "stopped_error"
            result.error = f"LLM call failed at step {step_num}: {last_error}"
            return result

        # ── Extract response data ──
        content = llm_result.get("content", "")
        cost_usd = llm_result.get("cost_usd", 0.0)
        result.total_cost_usd += cost_usd

        tool_calls = _extract_tool_calls(llm_result.get("raw_response"))

        # ── No tool calls → final response ──
        if not tool_calls:
            step = ReActStep(
                step_number=step_num,
                thought=content,
                action=None,
                action_input=None,
                observation=None,
                cost_usd=cost_usd,
                timestamp=datetime.now(UTC).isoformat(),
            )
            result.steps.append(step)
            result.final_content = content
            result.steps_taken = step_num

            # Fire hook
            await do_action(
                "agent.step",
                agent_id=str(getattr(agent_definition, "id", "")),
                step_number=step_num,
                thought=content,
                action=None,
                observation=None,
                cost_usd=cost_usd,
                agent_run_id=str(agent_run_id) if agent_run_id else None,
                trace_id=trace_id,
            )

            # Callback
            if on_step:
                try:
                    await on_step(step)
                except Exception:
                    logger.debug("on_step callback failed", exc_info=True)

            break

        # ── Execute tool calls ──
        # Append assistant message with tool calls to conversation
        full_messages.append({
            "role": "assistant",
            "content": content,
            "tool_calls": [
                {
                    "id": tc["id"],
                    "type": "function",
                    "function": {"name": tc["name"], "arguments": tc["arguments"]},
                }
                for tc in tool_calls
            ],
        })

        # Execute each tool call and collect observations
        observations: list[str] = []
        for tc in tool_calls:
            tool_name = tc["name"]
            try:
                args = json.loads(tc["arguments"]) if tc["arguments"] else {}
            except json.JSONDecodeError:
                args = {}
                logger.warning("Invalid JSON arguments for tool '%s': %s", tool_name, tc["arguments"])

            observation = await _execute_tool(tool_registry, tool_name, args, tool_context)
            observations.append(observation)

            # Feed tool result back into conversation
            full_messages.append({
                "role": "tool",
                "tool_call_id": tc["id"],
                "content": observation,
            })

            # Record step
            step = ReActStep(
                step_number=step_num,
                thought=content,
                action=tool_name,
                action_input=args,
                observation=observation,
                cost_usd=cost_usd / len(tool_calls) if tool_calls else cost_usd,
                timestamp=datetime.now(UTC).isoformat(),
            )
            result.steps.append(step)

            # Fire hook
            await do_action(
                "agent.step",
                agent_id=str(getattr(agent_definition, "id", "")),
                step_number=step_num,
                thought=content,
                action=tool_name,
                observation=observation,
                cost_usd=cost_usd,
                agent_run_id=str(agent_run_id) if agent_run_id else None,
                trace_id=trace_id,
            )

            # Callback
            if on_step:
                try:
                    await on_step(step)
                except Exception:
                    logger.debug("on_step callback failed", exc_info=True)

        result.steps_taken = step_num

        # If this was the last allowed step, stop gracefully
        if step_num >= max_steps:
            result.status = "stopped_max_steps"
            result.error = f"Reached max_steps limit ({max_steps})"
            result.final_content = content
            logger.warning("ReAct loop stopped at max_steps=%d", max_steps)
            break

    # If loop completed without a final response (e.g. all steps had tool calls)
    if not result.final_content and result.steps:
        result.final_content = result.steps[-1].thought or ""

    if result.status == "completed" and not result.final_content:
        result.final_content = ""

    return result