src/nodes/response_generator.py

from langchain_core.messages import HumanMessage, AIMessage
from src.state import AgentState
from src.utils.llm_factory import get_llm, _invoke_with_backoff
from src.utils.prompt_templates import RESPONSE_SYSTEM
from datetime import datetime

# Maximum conversation history to prevent token bloat
MAX_HISTORY_TURNS = 5

def generate_response(state: AgentState) -> AgentState:
    """Generate final customer-facing response"""
    try:
        # Get recent conversation history for context (last N turns)
        recent_messages = state.get("messages", [])[-MAX_HISTORY_TURNS * 2:]
        
        # Build history string for LLM context
        history_str = ""
        if recent_messages:
            history_str = "\n\nPrior conversation:\n"
            for msg in recent_messages:
                role = "User" if isinstance(msg, HumanMessage) else "Agent"
                content = msg.content[:200] if len(msg.content) > 200 else msg.content
                history_str += f"{role}: {content}\n"
        
        # Prepare context
        sql_result = state.get("sql_result", {})
        rag_result = state.get("rag_result", "")
        compensation_offered = state.get("compensation_offered", False)
        compensation_tier = state.get("compensation_tier", "")
        frustration_score = state.get("frustration_score") or 0.0
        
        # Format SQL result for readability with special status handling
        sql_str = ""
        extra_context = ""
        
        if sql_result and sql_result.get("rows"):
            rows = sql_result["rows"]
            row = rows[0] if len(rows) == 1 else None
            
            # Check for special order statuses and delivery timing
            if row:
                order_status = row.get("order_status", "").lower()
                is_late = row.get("is_late", 0)
                days_overdue = row.get("days_overdue", 0)
                delivered_date = row.get("order_delivered_customer_date")
                estimated_date = row.get("order_estimated_delivery_date")
                
                if order_status in ["canceled", "cancelled"]:
                    extra_context = "\n\nIMPORTANT: This order has been CANCELLED. Clearly inform the customer and ask if they need help with a refund or placing a new order."
                elif is_late and days_overdue > 0:
                    extra_context = f"\n\nIMPORTANT: This order is LATE by {days_overdue} days. Sincerely apologize for the delay and mention the compensation offered."
                elif order_status == "delivered" and delivered_date and estimated_date:
                    # Check if delivered early
                    if delivered_date < estimated_date:
                        extra_context = "\n\nGOOD NEWS: This order was delivered EARLY (ahead of the estimated date). Celebrate this positive outcome with the customer."
                    else:
                        extra_context = "\n\nThis order was delivered on time. Provide a neutral, professional status update."
                elif order_status in ["shipped", "in_transit"]:
                    extra_context = "\n\nThis order is currently in transit. Provide tracking information and estimated delivery date."
            
            if len(rows) == 1:
                sql_str = str(rows[0])
            else:
                sql_str = f"{len(rows)} results found: {rows[:3]}"
        elif sql_result and sql_result.get("empty"):
            sql_str = "No records found"
        elif sql_result and sql_result.get("error"):
            # Don't expose internal errors to user
            sql_str = "Unable to retrieve data at this time"
        
        # Build prompt
        prompt = RESPONSE_SYSTEM.format(
            sql_result=sql_str,
            rag_result=rag_result or "No policy information retrieved",
            compensation_offered=compensation_offered,
            compensation_tier=compensation_tier,
            frustration_score=frustration_score
        )
        
        # Add conversation history for context
        prompt = history_str + "\n" + prompt
        
        # Add extra context for special cases
        prompt += extra_context
        
        # Add compensation code if offered
        if compensation_offered:
            code = state.get("session_context", {}).get("compensation_code", "")
            discount = state.get("session_context", {}).get("compensation_discount", "")
            prompt += f"\n\nCompensation code: {code} ({discount})"
        
        llm = get_llm(temperature=0.3)  # Slight variation for natural tone
        response = _invoke_with_backoff(llm, [HumanMessage(content=prompt + "\n\nUser question: " + state.get("user_input", ""))], provider="groq")
        
        state["final_response"] = response.content.strip()
        
        # Return only new messages (LangGraph's add_messages reducer will merge)
        # Don't mutate state["messages"] directly to avoid duplication
        return {
            **state,
            "messages": [
                HumanMessage(content=state.get("user_input", "")),
                AIMessage(content=state["final_response"])
            ]
        }
        
    except Exception as e:
        state["error_log"].append(f"[response_generator] LLM error: {str(e)[:100]}")
        
        # Template-based fallback response
        sql_result = state.get("sql_result", {})
        rag_result = state.get("rag_result", "")
        
        if sql_result and sql_result.get("rows"):
            row = sql_result["rows"][0]
            order_id = row.get("order_id", "your order")[:16]
            status = row.get("order_status", "unknown")
            
            state["final_response"] = (
                f"I'm experiencing technical difficulties with our AI system, "
                f"but I can see your order {order_id}... is currently {status}. "
                f"Would you like specific details about delivery date, price, or items?"
            )
        elif rag_result:
            # Use RAG result directly if available
            state["final_response"] = rag_result[:400] + "\n\nWould you like more information?"
        else:
            state["final_response"] = (
                "I'm experiencing technical difficulties. "
                "Please try again in a moment, or contact our support team for immediate assistance."
            )
    
    return state