Spaces:
Sleeping
Sleeping
File size: 6,495 Bytes
cf796c5 2bfa9ba cf796c5 b1c949c d2c5868 cf796c5 2bfa9ba cf796c5 b6ae869 cf796c5 2dc2e18 cf796c5 2bfa9ba cf796c5 b6ae869 d2c5868 8622dab 2bfa9ba b6ae869 cf796c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from src.state import AgentState
from src.nodes.sanitizer import sanitize_input
from src.nodes.entity_extractor import extract_entities
from src.nodes.clarification import check_clarification_needed
from src.nodes.intent_classifier import classify_intent
from src.nodes.sql_generator import generate_sql
from src.nodes.sql_safety import check_sql_safety
from src.nodes.sql_executor import execute_sql
from src.nodes.rag_retriever import retrieve_rag
from src.nodes.sentiment_analyzer import analyze_sentiment
from src.nodes.proactive_logic import check_proactive_logic
from src.nodes.compensation import offer_compensation
from src.nodes.escalation import handle_escalation
from src.nodes.response_generator import generate_response
# Routing Functions
def route_after_sanitizer(state: AgentState) -> str:
if not state.get("sanitized"):
return "END"
# If escalation detected in sanitizer, skip to escalation
if state.get("escalation_required"):
return "escalation"
return "entity_extractor"
def route_after_clarification(state: AgentState) -> str:
# If clarification is needed, skip to END with clarification message
if state.get("clarification_needed"):
return "END"
return "intent_classifier"
def route_by_intent(state: AgentState) -> str:
intent = state.get("intent", "hybrid")
return {
"transactional": "sql_generator",
"informational": "rag_retriever",
"sentimental": "sentiment_analyzer",
"hybrid": "sql_generator", # Simplified: hybrid starts with SQL
"out_of_scope": "END",
}.get(intent, "sql_generator")
def route_after_sql_executor(state: AgentState) -> str:
sql = state.get("sql_result", {})
retry = state.get("retry_count", 0)
# Retry if error and not maxed out
if sql.get("error") and sql["error"] not in ("unsafe_query", "max_retries_exceeded", "no_query") and retry < 3:
return "sql_generator"
return "proactive_logic"
def route_after_proactive(state: AgentState) -> str:
if state.get("escalation_required"):
return "escalation"
if state.get("is_late_delivery") and not state.get("compensation_offered"):
return "compensation"
# For hybrid intent, also retrieve RAG if not done yet
if state.get("intent") == "hybrid" and state.get("rag_result") is None:
return "rag_retriever"
return "response_generator"
def route_after_sentiment(state: AgentState) -> str:
if state.get("escalation_required"):
return "escalation"
return "response_generator"
def route_after_rag(state: AgentState) -> str:
# Check if escalation was set by the node
if state.get("escalation_required"):
return "escalation"
return "response_generator"
def route_after_compensation(state: AgentState) -> str:
# After compensation, check if we need RAG for hybrid
if state.get("intent") == "hybrid" and state.get("rag_result") is None:
return "rag_retriever"
return "response_generator"
# Graph Construction
def build_graph() -> StateGraph:
builder = StateGraph(AgentState)
# Add all nodes
builder.add_node("sanitizer", sanitize_input)
builder.add_node("entity_extractor", extract_entities)
builder.add_node("clarification", check_clarification_needed)
builder.add_node("intent_classifier", classify_intent)
builder.add_node("sql_generator", generate_sql)
builder.add_node("sql_safety", check_sql_safety)
builder.add_node("sql_executor", execute_sql)
builder.add_node("rag_retriever", retrieve_rag)
builder.add_node("sentiment_analyzer", analyze_sentiment)
builder.add_node("proactive_logic", check_proactive_logic)
builder.add_node("compensation", offer_compensation)
builder.add_node("escalation", handle_escalation)
builder.add_node("response_generator", generate_response)
# Add edges
builder.add_edge(START, "sanitizer")
builder.add_conditional_edges("sanitizer", route_after_sanitizer,
{"entity_extractor": "entity_extractor",
"escalation": "escalation",
"END": END})
builder.add_edge("entity_extractor", "clarification")
builder.add_conditional_edges("clarification", route_after_clarification,
{"intent_classifier": "intent_classifier", "END": END})
builder.add_conditional_edges("intent_classifier", route_by_intent,
{"sql_generator": "sql_generator",
"rag_retriever": "rag_retriever",
"sentiment_analyzer": "sentiment_analyzer",
"END": END})
# Transactional path
builder.add_edge("sql_generator", "sql_safety")
builder.add_edge("sql_safety", "sql_executor")
builder.add_conditional_edges("sql_executor", route_after_sql_executor,
{"sql_generator": "sql_generator",
"proactive_logic": "proactive_logic"})
builder.add_conditional_edges("proactive_logic", route_after_proactive,
{"escalation": "escalation",
"compensation": "compensation",
"rag_retriever": "rag_retriever",
"response_generator": "response_generator"})
# Informational path
builder.add_conditional_edges("rag_retriever", route_after_rag,
{"escalation": "escalation",
"response_generator": "response_generator"})
# Sentimental path
builder.add_conditional_edges("sentiment_analyzer", route_after_sentiment,
{"escalation": "escalation",
"response_generator": "response_generator"})
# Terminal nodes
builder.add_conditional_edges("compensation", route_after_compensation,
{"rag_retriever": "rag_retriever",
"response_generator": "response_generator"})
builder.add_edge("escalation", END)
builder.add_edge("response_generator", END)
checkpointer = MemorySaver()
return builder.compile(checkpointer=checkpointer)
# Singleton
graph = build_graph()
|