Create SparrowV2Graph.py

src/graphs/SparrowV2Graph.py

@@ -0,0 +1,126 @@
+from langgraph.graph import StateGraph, START, END 
+from src.states.queryState import SparrowAgentState, SparrowInputState
+
+from src.nodes.queryNode import QueryNode
+from src.llms.groqllm import GroqLLM
+
+import logging 
+from src.graphs.masterGraph import master_graph
+from langchain_core.messages import AIMessage
+
+
+
+class SparrowV2GraphBuilder:
+    def __init__(self, llm):
+        self.llm = llm
+        self.graph = StateGraph(SparrowAgentState, input_schema=SparrowInputState)
+        self.logger = logging.getLogger(__name__)
+
+    def convert_sparrow_to_master(state: SparrowAgentState) -> dict:
+        """Convert SparrowAgentState to master graph input format"""
+        return {
+            "query_brief": state.get("query_brief", ""),
+            "execution_jobs": [],
+            "completed_jobs": [],
+            "worker_outputs": [],
+            "final_output": ''
+        }
+    
+    def update_sparrow_from_master(sparrow_state: SparrowAgentState, master_state:dict) -> SparrowAgentState:
+        """Update sparrow state with master results"""
+         
+        final_output = master_state.get("final_output", "")
+        if final_output:
+            sparrow_state["messages"] = sparrow_state.get("messages", []) + [AIMessage(content=final_output)]
+            sparrow_state["final_message"] = final_output
+
+        execution_jobs = master_state.get("execution_jobs", [])
+        completed_jobs = master_state.get("completed_jobs", [])
+
+        if execution_jobs:
+            sparrow_state["notes"] = sparrow_state.get("notes", []) + [f"Execution jobs: {', '.join(execution_jobs)}"]
+
+        if completed_jobs:
+            sparrow_state["notes"] = sparrow_state.get("notes", []) + [f"Completed: {', '.join(completed_jobs)}"]
+
+        return sparrow_state
+    
+   
+
+        
+ 
+    def route_after_clarification(self, state: SparrowAgentState) -> str:
+        """Route based on clarification status from queryNode response"""
+
+        if state.get("clarification_complete", False):
+            print("Routing: Clarification marked as complete")
+            return "master_subgraph"
+
+        if state.get("max_clarification_reached", False):
+            print("Routing: Max clarification attempts reached")
+            return "master_subgraph"
+        
+        print("Routing: Clarification still needed")
+        
+        return "clarify_with_user"
+    
+    def run_master_subgraph(self, state: SparrowAgentState) -> SparrowAgentState:
+        """Run the master subgraph - using sync version to avoid async issues with Send"""
+
+        try:
+            print("Running master subgraph..")
+            master_input = self.convert_sparrow_to_master(state)
+
+            master_result = master_graph.invoke(master_input)
+
+            return self.update_sparrow_from_master(state, master_result)
+
+        except Exception as e:
+            self.logger.error(f"Master subgraph failed: {e}")
+            return {**state, "error": f"Master subgraph failed: {str(e)}"}
+        
+
+    
+
+    def run_query_subgraph(self, state:SparrowAgentState) -> SparrowAgentState:
+        """Run the query subgraph - using sync version to avoid async issues with Send"""
+        try:
+            print("Running query subgraph..")
+
+        except Exception as e:
+            self.logger.error(f"Query subgraph failed: {e}")
+            return {**state, "error": f"Query subgraph failed: {str(e)}"}
+
+    
+    def build_query_graph(self):
+        """
+        Build a graph for customer query inquiry
+
+        """
+        self.query_node_obj= QueryNode(self.llm)
+        print(self.llm)
+
+        self.graph.add_node("clarify_with_user", self.query_node_obj.clarify_with_user)
+        self.graph.add_node("write_query_brief", self.query_node_obj.write_query_brief)
+        self.graph.add_node("master_subgraph", self.run_master_subgraph)
+
+        self.graph.add_edge(START, "clarify_with_user")
+        self.graph.add_edge("clarify_with_user", "write_query_brief")
+        self.graph.add_conditional_edges(
+            "write_query_brief",
+            self.route_after_clarification,
+            {
+                "clarify_with_user": END,
+                "master_subgraph": "master_subgraph"
+            }
+        )
+        
+        self.graph.add_edge("master_subgraph", END)
+
+        return self.graph
+    
+llm = GroqLLM().get_llm()
+
+graph_builder=SparrowV2GraphBuilder(llm)
+
+graph=graph_builder.build_query_graph().compile()