Update agent.py

agent.py

@@ -21,6 +21,7 @@ def build_graph():
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
+    # Node: retriever
     def retriever_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
         if query in qa_dict:
@@ -29,40 +30,48 @@ def build_graph():
         print("🔍 No match found. Falling back to LLM.")
         return {"messages": state["messages"]}
 
-def assistant_node(state: MessagesState):
-    query = state["messages"][-1].content.strip()
-
-    # Instrukcja — jak ma wyglądać odpowiedź
-    system_prompt = (
-        "You are a helpful AI assistant taking part in the GAIA evaluation benchmark.\n"
-        "You must return only the final answer to the user's question.\n"
-        "- No explanations.\n"
-        "- No formatting like 'Final answer:' or similar.\n"
-        "- If the answer is a list, return comma-separated values.\n"
-        "- If the answer is unknown, return 'Unknown'."
-    )
+    # Node: assistant (LLM with fallback prompt logic)
+    def assistant_node(state: MessagesState):
+        query = state["messages"][-1].content.strip()
 
-    try:
-        # Preferowany tryb: chat format
-        chat_input = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": query}
-        ]
-        response = llm.invoke(chat_input)
-        print("✅ Used chat-style prompt.")
-    except Exception as e:
-        print(f"⚠️ Chat-style failed: {e}")
-        # Fallback: classic prompt-based
-        fallback_prompt = (
-            f"{system_prompt}\n\n"
-            f"Question: {query}\n"
-            f"Answer:"
+        system_prompt = (
+            "You are a helpful AI assistant taking part in the GAIA evaluation benchmark.\n"
+            "You must return only the final answer to the user's question.\n"
+            "- No explanations.\n"
+            "- No formatting like 'Final answer:' or similar.\n"
+            "- If the answer is a list, return comma-separated values.\n"
+            "- If the answer is unknown, return 'Unknown'."
         )
-        response = llm.invoke(fallback_prompt)
-        print("🔁 Used fallback prompt format.")
 
-    return {"messages": [AIMessage(content=response.strip())]}
+        try:
+            chat_input = [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": query}
+            ]
+            response = llm.invoke(chat_input)
+            print("✅ Used chat-style prompt.")
+        except Exception as e:
+            print(f"⚠️ Chat-style failed: {e}")
+            fallback_prompt = (
+                f"{system_prompt}\n\n"
+                f"Question: {query}\n"
+                f"Answer:"
+            )
+            response = llm.invoke(fallback_prompt)
+            print("🔁 Used fallback prompt format.")
+
+        return {"messages": [AIMessage(content=response.strip())]}
+
+    # Build the LangGraph
+    builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever_node)
+    builder.add_node("assistant", assistant_node)
+
+    builder.set_entry_point("retriever")
+    builder.add_edge("retriever", "assistant")
+    builder.set_finish_point("assistant")
 
+    return builder.compile()
 
 
 # --- Agent class wrapper for app.py ---