Customer_Support/workflow.py

"""
LangGraph workflow setup for the customer support system.
"""
from langgraph.graph import StateGraph, END
from typing import Dict
from .models import State
from .handlers import (
    categorize, analyze_sentiment, handle_technical,
    handle_billing, handle_general, escalate, route_query
)


def create_workflow() -> StateGraph:
    """Create and configure the LangGraph workflow"""
    # Create the graph
    workflow = StateGraph(State)

    # Add nodes
    workflow.add_node("categorize", categorize)
    workflow.add_node("analyze_sentiment", analyze_sentiment)
    workflow.add_node("handle_technical", handle_technical)
    workflow.add_node("handle_billing", handle_billing)
    workflow.add_node("handle_general", handle_general)
    workflow.add_node("escalate", escalate)

    # Add edges
    workflow.add_edge("categorize", "analyze_sentiment")
    workflow.add_conditional_edges(
        "analyze_sentiment",
        route_query,
        {
            "handle_technical": "handle_technical",
            "handle_billing": "handle_billing",
            "handle_general": "handle_general",
            "escalate": "escalate"
        }
    )
    workflow.add_edge("handle_technical", END)
    workflow.add_edge("handle_billing", END)
    workflow.add_edge("handle_general", END)
    workflow.add_edge("escalate", END)

    # Set entry point
    workflow.set_entry_point("categorize")

    return workflow.compile()


def run_customer_support(query: str) -> Dict[str, str]:
    """Process a customer query through the LangGraph workflow.

    Args:
        query (str): The customer's query

    Returns:
        Dict[str, str]: A dictionary containing the query's category, sentiment, and response
    """
    # Create a fresh instance of the workflow
    app = create_workflow()

    # Invoke the workflow with the query
    results = app.invoke({"query": query})

    return {
        "query": query,
        "category": results["category"],
        "sentiment": results["sentiment"],
        "response": results["response"]
    }