fix edges

agent.py

@@ -1,4 +1,3 @@
-# import libraries for langgraph, huggingface
 import os
 from dotenv import load_dotenv
 from typing import TypedDict, List, Dict, Any, Optional, Annotated
@@ -106,7 +105,7 @@ def search_web(query: str) -> Dict[str, str]:
     docs = SerpAPIWrapper()
     docs.run(query)
     formatted_result = f'<Document source="{docs.metadata["source"]}" page="{docs.metadata.get("page", "")}"/>\n{docs.page_content}\n</Document>'
-    return {"wiki_results": formatted_result}
+    return {"web_results": formatted_result}
 
 
 # ArXiv search tool
@@ -119,7 +118,7 @@ def search_arxiv(query: str) -> Dict[str, str]:
     arxiv = ArxivAPIWrapper()
     docs = arxiv.run(query)
     formatted_result = f'<Document source="{docs.metadata["source"]}" page="{docs.metadata.get("page", "")}"/>\n{docs.page_content}\n</Document>'
-    return {"wiki_results": formatted_result}
+    return {"arxiv_results": formatted_result}
 
 
 # build retriever
@@ -180,14 +179,16 @@ def build_graph():
     builder.add_edge(START, "assistant")
     builder.add_conditional_edges(
         "assistant",
-        # If the latest message requires a tool, route to tools
-        # Otherwise, provide a direct response
         tools_condition,
+        {
+            # If the latest message requires a tool, route to tools
+            "tools": "tools",
+            # Otherwise, provide a direct response
+            END: END,
+        }
     )
     builder.add_edge("tools", "assistant")
 
-    return builder.compile()
-
 
 if __name__ == "__main__":
     question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"

app.py

@@ -23,30 +23,9 @@ class BasicAgent:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         # Wrap the question in a HumanMessage from langchain_core
         messages = [HumanMessage(content=question)]
-        # The graph.invoke method returns the final state of the graph.
-        # The state is expected to contain a 'messages' key with a list of messages.
-        result = self.graph.invoke({"messages": messages})
-        # Extract the final answer message from the state
-        # The final message should be the content of the last AIMessage with content
-        submitted_answer = "Could not get answer from agent output." # Default if extraction fails
-
-        # Check if the result is a dictionary and contains a list of messages
-        if isinstance(result, dict) and "messages" in result and isinstance(result["messages"], list):
-            # Iterate through messages in reverse to find the last one from the AI
-            for msg in reversed(result["messages"]):
-                # Check if the message is an AIMessage and has content (non-empty string or non-empty list/dict representation)
-                # The 'content' attribute exists on AIMessage, so hasattr check is redundant.
-                if isinstance(msg, AIMessage) and msg.content is not None:
-                     # Ensure the content is treated as a string for submission.
-                     # LangChain AIMessage content can be str or list[str | dict].
-                     # str() handles conversion for both cases.
-                     submitted_answer = str(msg.content)
-                     break # Found the last AI message with content, stop searching
-
-        # If the loop finishes without finding an AIMessage with content,
-        # submitted_answer remains the default value.
-
-        return submitted_answer
+        messages = self.graph.invoke({"messages": messages})
+        answer = messages['messages'][-1].content
+        return answer[14:]
 
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):