agent/graph.py

import json
import os
import re
from typing import Any, Dict, List, Optional, TypedDict

from groq import Groq
from langgraph.graph import END, StateGraph

from agent.audit import log_step
from agent.tools import TOOL_MAP
from agent.validator import validate_tool_output

_groq = Groq(api_key=os.getenv("GROQ_API_KEY", ""))
_MODEL = "llama-3.3-70b-versatile"
MAX_STEPS = 8


class AgentState(TypedDict):
    user_message: str
    session_id: str
    intent: str
    pending_tools: List[Dict[str, Any]]
    tool_results: List[Dict[str, Any]]
    last_tool_name: Optional[str]
    last_tool_input: Optional[Dict[str, Any]]
    last_tool_output: Optional[Any]
    step_count: int
    final_response: str


def _groq_chat(system: str, user: str, temperature: float = 0.1) -> str:
    resp = _groq.chat.completions.create(
        model=_MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=temperature,
    )
    return resp.choices[0].message.content.strip()


def _extract_json(text: str) -> Dict:
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    match = re.search(r"\{[\s\S]*\}", text)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass
    return {}


def intent_router(state: AgentState) -> Dict:
    system = """You are a customer support AI. Analyze the customer message and plan the exact tool calls needed.

Available tools:
- order_lookup: Look up order details. Args: {"order_id": "ORD-XXX"}
- refund_trigger: Initiate a refund. Args: {"order_id": "ORD-XXX", "reason": "explanation"}
- ticket_create: Create a support ticket. Args: {"issue": "full description", "priority": "high|medium|low"}
- faq_retriever: Search policy/FAQ docs. Args: {"query": "search terms"}

Planning rules:
- Wrong item received → [order_lookup, refund_trigger, ticket_create] (high priority ticket)
- Damaged item → [order_lookup, refund_trigger, ticket_create] (high priority)
- Refund request → [order_lookup, refund_trigger]
- Order status/tracking → [order_lookup]
- Return policy / warranty / general questions → [faq_retriever]
- Explicit ticket request → [ticket_create]
- Extract order IDs exactly as written (e.g., ORD-003, ORD-001)

Respond ONLY with valid JSON:
{
  "intent": "order_lookup|refund|ticket|faq|unknown",
  "reasoning": "one sentence explaining the plan",
  "tools": [
    {"name": "tool_name", "args": {}}
  ]
}"""

    raw = _groq_chat(system, state["user_message"])
    parsed = _extract_json(raw)

    intent = parsed.get("intent", "unknown")
    tools = parsed.get("tools", [])
    reasoning = parsed.get("reasoning", "Classified user intent and planned tool calls.")

    log_step(
        session_id=state["session_id"],
        node_name="intent_router",
        reasoning=reasoning,
        tool_called=None,
        tool_input={"user_message": state["user_message"]},
        tool_output={"intent": intent, "planned_tools": [t.get("name") for t in tools]},
    )

    return {
        "intent": intent,
        "pending_tools": tools,
        "tool_results": [],
        "last_tool_name": None,
        "last_tool_input": None,
        "last_tool_output": None,
        "step_count": state["step_count"] + 1,
    }


def tool_executor(state: AgentState) -> Dict:
    if not state["pending_tools"]:
        return {"step_count": state["step_count"] + 1}

    current = state["pending_tools"][0]
    remaining = state["pending_tools"][1:]
    tool_name = current.get("name", "")
    tool_args = current.get("args", {})

    tool_fn = TOOL_MAP.get(tool_name)
    raw_result: Dict[str, Any] = {}

    if tool_fn:
        try:
            raw = tool_fn.invoke(tool_args)
            raw_result = raw if isinstance(raw, dict) else {"result": str(raw)}
        except Exception as exc:
            raw_result = {"error": str(exc), "tool": tool_name}
    else:
        raw_result = {"error": f"Unknown tool: {tool_name}"}

    return {
        "pending_tools": remaining,
        "last_tool_name": tool_name,
        "last_tool_input": tool_args,
        "last_tool_output": raw_result,
        "step_count": state["step_count"] + 1,
    }


def action_validator(state: AgentState) -> Dict:
    tool_name = state.get("last_tool_name") or ""
    raw_output = state.get("last_tool_output") or {}

    validated = validate_tool_output(tool_name, raw_output)
    validated_dict = (
        validated.model_dump() if hasattr(validated, "model_dump") else raw_output
    )

    updated_results = state["tool_results"] + [
        {"tool": tool_name, "result": validated_dict}
    ]

    return {
        "last_tool_output": validated_dict,
        "tool_results": updated_results,
        "step_count": state["step_count"] + 1,
    }


def audit_logger(state: AgentState) -> Dict:
    tool_name = state.get("last_tool_name")
    log_step(
        session_id=state["session_id"],
        node_name="audit_logger",
        reasoning=f"Completed tool call: {tool_name}" if tool_name else "Step complete",
        tool_called=tool_name,
        tool_input=state.get("last_tool_input"),
        tool_output=state.get("last_tool_output"),
    )
    return {"step_count": state["step_count"] + 1}


def responder(state: AgentState) -> Dict:
    results_text = json.dumps(state["tool_results"], indent=2, default=str)
    system = """You are a helpful customer support agent. Write a clear, friendly, concise response based on the tool results.
Be specific: mention order IDs, ticket IDs, refund timelines, and next steps.
Keep under 150 words. Do not mention internal tool names or system details."""

    user = f"Customer message: {state['user_message']}\n\nTool results:\n{results_text}"
    response = _groq_chat(system, user, temperature=0.3)

    log_step(
        session_id=state["session_id"],
        node_name="responder",
        reasoning="Generating final customer-facing response from all tool results",
        tool_called=None,
        tool_input={"num_tool_results": len(state["tool_results"])},
        tool_output={"response_preview": response[:120] + "..." if len(response) > 120 else response},
    )

    return {"final_response": response, "step_count": state["step_count"] + 1}


def _route_after_audit(state: AgentState) -> str:
    if state["step_count"] >= MAX_STEPS:
        return "respond"
    if state["pending_tools"]:
        return "execute_tool"
    return "respond"


def build_graph():
    graph = StateGraph(AgentState)

    graph.add_node("intent_router", intent_router)
    graph.add_node("tool_executor", tool_executor)
    graph.add_node("action_validator", action_validator)
    graph.add_node("audit_logger", audit_logger)
    graph.add_node("responder", responder)

    graph.set_entry_point("intent_router")
    graph.add_edge("intent_router", "tool_executor")
    graph.add_edge("tool_executor", "action_validator")
    graph.add_edge("action_validator", "audit_logger")
    graph.add_conditional_edges(
        "audit_logger",
        _route_after_audit,
        {"execute_tool": "tool_executor", "respond": "responder"},
    )
    graph.add_edge("responder", END)

    return graph.compile()


_compiled_graph = None


def get_graph():
    global _compiled_graph
    if _compiled_graph is None:
        _compiled_graph = build_graph()
    return _compiled_graph


def run_agent(session_id: str, message: str) -> str:
    graph = get_graph()
    initial_state: AgentState = {
        "user_message": message,
        "session_id": session_id,
        "intent": "",
        "pending_tools": [],
        "tool_results": [],
        "last_tool_name": None,
        "last_tool_input": None,
        "last_tool_output": None,
        "step_count": 0,
        "final_response": "",
    }
    result = graph.invoke(initial_state)
    return result.get("final_response", "I'm sorry, I wasn't able to process your request.")