Adjust Tools and introduce wiki search

agent.py

@@ -1,5 +1,7 @@
 import os
 
+from langchain_community.document_loaders import WikipediaLoader
+
 # --- Langchain / Langraph ---
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
@@ -18,7 +20,7 @@ from langgraph.prebuilt import ToolNode, tools_condition
 
 
 @tool
-def search_web_sources(query: str) -> str:
+def search_web_sources(query: str) -> dict:
     """
     Perform a web search using Tavily and return up to 3 relevant documents.
     This tool is useful for answering research-based queries that require
@@ -30,14 +32,35 @@ def search_web_sources(query: str) -> str:
     Returns:
         str: Formatted web search results with metadata and content.
     """
-    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    tavily_tool = TavilySearchResults(max_results=3)
+    search_docs = tavily_tool.invoke({"query": query})
     formatted_search_docs = "\n\n---\n\n".join(
         [
             f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
             for doc in search_docs
         ]
     )
-    return formatted_search_docs
+    return {"web_results": formatted_search_docs}
+
+
+@tool
+def search_wikipedia(query: str) -> dict:
+    """Search Wikipedia using LangChain's loader and return the first document summary."""
+    try:
+        loader = WikipediaLoader(query=query, lang="en", load_max_docs=2)
+        docs = loader.load()
+        if docs:
+            formatted_docs = "---".join(
+                [
+                    f'<WikipediaArticle title="{query}">{doc.page_content}</WikipediaArticle>'
+                    for doc in docs
+                ]
+            )
+            return {"wiki_results": formatted_docs}
+        else:
+            return {"wiki_results": "No content found."}
+    except Exception as e:
+        return {"wiki_results": f"Error fetching Wikipedia article: {e}"}
 
 
 @tool
@@ -74,7 +97,7 @@ def build_agent_graph(provider: str = "groq"):
     groq_api_key = os.environ.get("GROQ_API_KEY")
 
     # Define toolset
-    tools = [search_web_sources, run_python_code]
+    tools = [search_web_sources, search_wikipedia, run_python_code]
 
     # Instantiate LLM
     llm = ChatGroq(model="qwen-qwq-32b", temperature=0, api_key=groq_api_key)
@@ -83,10 +106,10 @@ def build_agent_graph(provider: str = "groq"):
     llm_with_tools = llm.bind_tools(tools)
 
     # Assistant: reasoning step that plans next action
-    def assistant_node(state: MessagesState):
+    def assistant_node(state: MessagesState) -> dict:
         messages = state["messages"]
         response = llm_with_tools.invoke(messages)
-        return {"messages": add_messages(state, [response])}
+        return {"messages": response}
 
     # Stubbed retriever node for future integration
     def retriever_node(state: MessagesState):