Update pipeline.py

pipeline.py

@@ -4,54 +4,41 @@ from weather_node import fetch_weather
 from rag_node import rag_answer
 from llm_node import refine_answer
 
-# Graph state definition
 class AgentState(TypedDict):
     query: str
     result: str
 
-# Weather node
-def weather_node(state: AgentState) -> AgentState:
-    city = state["query"].split()[-1]
-    state["result"] = fetch_weather(city)
+def weather_node_fn(state: AgentState):
+    state["result"] = fetch_weather(state["query"])
     return state
 
-# PDF node
-def rag_node_fn(state: AgentState) -> AgentState:
+def rag_node_fn(state: AgentState):
     state["result"] = rag_answer(state["query"])
     return state
 
-# Decision node
-def decision_node(state: AgentState) -> str:
+def decision_node(state: AgentState):
     keywords = ["weather", "temperature", "forecast", "climate"]
-    if any(word in state["query"].lower() for word in keywords):
-        return "weather"
-    return "rag"
+    return "weather" if any(k in state["query"].lower() for k in keywords) else "rag"
 
-# LLM node
-def llm_node_fn(state: AgentState) -> AgentState:
-    state["result"] = refine_answer(f"Summarize this: {state['result']}")
+def llm_node_fn(state: AgentState):
+    state["result"] = refine_answer(state["result"])
     return state
 
-# Build pipeline with LangGraph
 def build_pipeline():
-    workflow = StateGraph(AgentState)
+    graph = StateGraph(AgentState)
 
-    # Add nodes
-    workflow.add_node("weather", weather_node)
-    workflow.add_node("rag", rag_node_fn)
-    workflow.add_node("llm", llm_node_fn)
+    graph.add_node("weather", weather_node_fn)
+    graph.add_node("rag", rag_node_fn)
+    graph.add_node("llm", llm_node_fn)
 
-    # Start → Decision → Weather/PDF
-    workflow.add_conditional_edges(
+    graph.add_conditional_edges(
         START,
         decision_node,
-        {"weather": "weather", "rag": "rag"},
+        {"weather": "weather", "rag": "rag"}
     )
 
-    # After weather or rag → go to llm → end
-    workflow.add_edge("weather", "llm")
-    workflow.add_edge("rag", "llm")
-    workflow.add_edge("llm", END)
-
-    return workflow.compile()
+    graph.add_edge("weather", "llm")
+    graph.add_edge("rag", "llm")
+    graph.add_edge("llm", END)
 
+    return graph.compile()