Fix search tool

agent.py

@@ -27,9 +27,15 @@ def build_graph():
 
     def assistant(state: MessagesState):
         """Assistant node for invoking the LLM."""
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+        messages = state["messages"]
+        # Add system message if not present
+        if not any(isinstance(m, SystemMessage) for m in messages):
+            messages = [SystemMessage(content="You are a helpful AI assistant. Use the available tools to answer questions accurately. When providing your final answer, use the format: FINAL ANSWER: [your answer]")] + messages
+        response = llm_with_tools.invoke(messages)
+        return {"messages": [response]}
 
     builder = StateGraph(MessagesState)
+    # builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
     builder.add_edge(START, "assistant")
@@ -42,8 +48,9 @@ def build_graph():
     # Compile graph
     return builder.compile()
 
+
 if __name__ == "__main__":
-    question = "Give me a summary of the wikipedia page for photosynthesis"
+    question = "On June 6, 2023, an article by Carolyn Collins Petersen was published in Universe Today. This article mentions a team that produced a paper about their observations, linked at the bottom of the article. Find this paper. Under what NASA award number was the work performed by R. G. Arendt supported by?"
     # Build the graph
     graph = build_graph()
     # Run the graph

app.py

@@ -4,6 +4,7 @@ import requests
 import inspect
 import pandas as pd
 from langchain_core.messages import HumanMessage
+
 from agent import build_graph
 
 # (Keep Constants as is)

tools/search_tools.py

@@ -6,16 +6,21 @@ from langchain_community.document_loaders import WikipediaLoader
 
 @tool
 def web_search(query: str) -> str:
-    """Search the web via Tavily for a given query and return maximum 3 results.
-    
+    """Search Tavily for a query and return maximum 3 results.
     Args:
-        query (str): The search query
-
+        query: The search query.
+        
     Returns:
-        str: Content from the Tavily search
+        str: The result of the search query
     """
-    search_results = TavilySearchResults(max_results=3).invoke(query=query)
-    return {"web_results": search_results}
+    search_tool = TavilySearchResults(max_results=3)
+    search_docs = search_tool.invoke({"query": query})
+    # Convert search results to a single string
+    results_text = ""
+    for doc in search_docs:
+        results_text += f"{doc['title']}: {doc['content']}\n"
+    return results_text
+
 
 @tool
 def search_wikipedia(query: str) -> str:
@@ -35,8 +40,8 @@ def search_wikipedia(query: str) -> str:
         loader = WikipediaLoader(query=query, load_max_docs=1)
         docs = loader.load()
         if not docs:
-            return f"No Wikipedia page found for '{query}'"
+            raise ValueError(f"No Wikipedia page found for query: {query}")
         return docs[0].page_content
     except Exception as e:
-        return f"Error searching Wikipedia: {str(e)}"
+        raise ValueError(f"Error searching Wikipedia: {str(e)}")