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.")