Update agent.py

agent.py

@@ -29,20 +29,40 @@ 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()
-        response = llm.invoke(query)
-        return {"messages": [AIMessage(content=response)]}
+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'."
+    )
 
-    builder = StateGraph(MessagesState)
-    builder.add_node("retriever", retriever_node)
-    builder.add_node("assistant", assistant_node)
+    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:"
+        )
+        response = llm.invoke(fallback_prompt)
+        print("🔁 Used fallback prompt format.")
 
-    builder.set_entry_point("retriever")
-    builder.add_edge("retriever", "assistant")
-    builder.set_finish_point("assistant")
+    return {"messages": [AIMessage(content=response.strip())]}
 
-    return builder.compile()
 
 
 # --- Agent class wrapper for app.py ---