src/agent.py

"""
P08 · SRE Agent — LangGraph implementation
ReAct pattern: Reason → Act → Observe → repeat until done.
Human-in-the-loop: tools marked requires_approval=True need explicit user confirmation.
"""

import json
import re
from typing import Any

from .tools import TOOLS, ToolResult


# ── Agent state ───────────────────────────────────────────────────────────────
class AgentState:
    def __init__(self, query: str):
        self.query = query
        self.steps: list[dict] = []          # full trace
        self.pending_approval: dict | None = None   # tool waiting for human
        self.final_answer: str = ""
        self.done: bool = False
        self.error: str = ""

    def add_step(self, step_type: str, content: Any):
        self.steps.append({"type": step_type, "content": content})

    def to_trace(self) -> str:
        lines = []
        for s in self.steps:
            if s["type"] == "thought":
                lines.append(f"🤔 **Thought:** {s['content']}")
            elif s["type"] == "tool_call":
                lines.append(f"🔧 **Tool:** `{s['content']['tool']}` → `{json.dumps(s['content']['args'])}`")
            elif s["type"] == "tool_result":
                r: ToolResult = s["content"]
                if r.success:
                    lines.append(f"✅ **Result ({r.latency_ms}ms):** ```json\n{json.dumps(r.data, indent=2)}\n```")
                else:
                    lines.append(f"❌ **Error:** {r.error}")
            elif s["type"] == "approval_needed":
                lines.append(f"⚠️ **Approval needed for:** `{s['content']}`")
            elif s["type"] == "approved":
                lines.append(f"✅ **Approved:** `{s['content']}`")
            elif s["type"] == "rejected":
                lines.append(f"🚫 **Rejected:** `{s['content']}`")
        return "\n\n".join(lines)


# ── Prompt builder ────────────────────────────────────────────────────────────
def build_system_prompt() -> str:
    tool_descriptions = "\n".join(
        f"- {name}: {info['description']}"
        for name, info in TOOLS.items()
    )
    return f"""You are an SRE on-call assistant. Answer questions about service health,
SLOs, alerts, and runbooks using the available tools.

Available tools:
{tool_descriptions}

To use a tool, respond with EXACTLY this format (nothing else on that line):
TOOL: tool_name
ARGS: {{"arg1": "value1", "arg2": "value2"}}

To give a final answer, respond with:
ANSWER: your complete answer here

Rules:
- Always use tools to get real data before answering
- Use at most 5 tool calls per query
- If a tool fails, try an alternative approach
- Be concise and actionable in your final answer
- Always include the service name and severity in your answer"""


def build_user_prompt(state: AgentState) -> str:
    if not state.steps:
        return f"User query: {state.query}"

    history = []
    for step in state.steps:
        if step["type"] == "thought":
            history.append(f"Thought: {step['content']}")
        elif step["type"] == "tool_call":
            tc = step["content"]
            history.append(f"Tool called: {tc['tool']} with {tc['args']}")
        elif step["type"] == "tool_result":
            r: ToolResult = step["content"]
            if r.success:
                history.append(f"Tool result: {json.dumps(r.data)}")
            else:
                history.append(f"Tool error: {r.error}")

    history_str = "\n".join(history)
    return f"User query: {state.query}\n\nHistory:\n{history_str}\n\nContinue:"


# ── Tool call parser ──────────────────────────────────────────────────────────
def parse_tool_call(text: str) -> dict | None:
    """Parse TOOL/ARGS block from model output."""
    tool_match = re.search(r"TOOL:\s*(\w+)", text)
    args_match = re.search(r"ARGS:\s*(\{.*?\})", text, re.DOTALL)
    answer_match = re.search(r"ANSWER:\s*(.+)", text, re.DOTALL)

    if answer_match:
        return {"type": "answer", "content": answer_match.group(1).strip()}

    if tool_match:
        tool_name = tool_match.group(1).strip()
        args = {}
        if args_match:
            try:
                args = json.loads(args_match.group(1))
            except json.JSONDecodeError:
                pass
        return {"type": "tool_call", "tool": tool_name, "args": args}

    return None


# ── Execute tool ──────────────────────────────────────────────────────────────
def execute_tool(tool_name: str, args: dict) -> ToolResult:
    if tool_name not in TOOLS:
        return ToolResult(
            tool=tool_name,
            success=False,
            data=None,
            error=f"Unknown tool '{tool_name}'. Available: {list(TOOLS.keys())}",
        )
    fn = TOOLS[tool_name]["fn"]
    try:
        return fn(**args)
    except TypeError as e:
        return ToolResult(
            tool=tool_name,
            success=False,
            data=None,
            error=f"Invalid arguments for {tool_name}: {e}",
        )
    except Exception as e:
        return ToolResult(
            tool=tool_name,
            success=False,
            data=None,
            error=f"Tool execution failed: {e}",
        )


# ── Main agent class ──────────────────────────────────────────────────────────
class SREAgent:
    """
    LangGraph-style ReAct agent with human-in-the-loop approval.
    Uses local transformers model — no external API calls.
    """

    def __init__(self, pipe):
        """pipe: a transformers text-generation pipeline."""
        self.pipe = pipe
        self.max_steps = 5

    def _call_llm(self, state: AgentState) -> str:
        system = build_system_prompt()
        user = build_user_prompt(state)

        prompt = (
            f"<|im_start|>system\n{system}<|im_end|>\n"
            f"<|im_start|>user\n{user}<|im_end|>\n"
            f"<|im_start|>assistant\n"
        )

        output = self.pipe(prompt, return_full_text=False)[0]["generated_text"]
        # Stop at next turn marker
        return output.split("<|im_end|>")[0].strip()

    def run(self, query: str) -> AgentState:
        """
        Run the agent synchronously.
        Tools requiring approval will pause and set state.pending_approval.
        Call run_with_approval() to continue after user approves.
        """
        state = AgentState(query)
        step_count = 0

        while not state.done and step_count < self.max_steps:
            step_count += 1

            # Get LLM decision
            llm_output = self._call_llm(state)
            parsed = parse_tool_call(llm_output)

            if parsed is None:
                # No structured output — treat as final answer
                state.final_answer = llm_output
                state.done = True
                break

            if parsed["type"] == "answer":
                state.final_answer = parsed["content"]
                state.done = True
                break

            if parsed["type"] == "tool_call":
                tool_name = parsed["tool"]
                args = parsed["args"]

                state.add_step("tool_call", {"tool": tool_name, "args": args})

                # Check if tool requires human approval
                tool_info = TOOLS.get(tool_name, {})
                if tool_info.get("requires_approval", False):
                    state.add_step("approval_needed", f"{tool_name}({args})")
                    state.pending_approval = {"tool": tool_name, "args": args}
                    break  # Pause for human input

                # Execute tool
                result = execute_tool(tool_name, args)
                state.add_step("tool_result", result)

        if step_count >= self.max_steps and not state.done:
            state.final_answer = "Reached maximum steps. Based on gathered data, please review the tool results above."
            state.done = True

        return state

    def approve_and_continue(self, state: AgentState) -> AgentState:
        """Continue after human approves a pending tool call."""
        if not state.pending_approval:
            return state

        tool_name = state.pending_approval["tool"]
        args = state.pending_approval["args"]

        state.add_step("approved", f"{tool_name}({args})")
        result = execute_tool(tool_name, args)
        state.add_step("tool_result", result)
        state.pending_approval = None

        # Continue the agent loop
        return self.run_continue(state)

    def reject_and_continue(self, state: AgentState) -> AgentState:
        """Continue after human rejects a pending tool call."""
        if not state.pending_approval:
            return state

        tool_name = state.pending_approval["tool"]
        state.add_step("rejected", tool_name)
        state.add_step(
            "tool_result",
            ToolResult(
                tool=tool_name,
                success=False,
                data=None,
                error="Tool call rejected by user",
            ),
        )
        state.pending_approval = None
        return self.run_continue(state)

    def run_continue(self, state: AgentState) -> AgentState:
        """Continue running from an existing state."""
        step_count = len(state.steps)

        while not state.done and step_count < self.max_steps:
            step_count += 1
            llm_output = self._call_llm(state)
            parsed = parse_tool_call(llm_output)

            if parsed is None or parsed["type"] == "answer":
                state.final_answer = (
                    parsed["content"] if parsed and parsed["type"] == "answer"
                    else llm_output
                )
                state.done = True
                break

            if parsed["type"] == "tool_call":
                tool_name = parsed["tool"]
                args = parsed["args"]
                state.add_step("tool_call", {"tool": tool_name, "args": args})

                tool_info = TOOLS.get(tool_name, {})
                if tool_info.get("requires_approval", False):
                    state.add_step("approval_needed", f"{tool_name}({args})")
                    state.pending_approval = {"tool": tool_name, "args": args}
                    break

                result = execute_tool(tool_name, args)
                state.add_step("tool_result", result)

        if step_count >= self.max_steps and not state.done:
            state.final_answer = "Reached maximum steps. Review tool results above."
            state.done = True

        return state