Update app.py

app.py

@@ -11,7 +11,7 @@ from typing import TypedDict, List, Dict, Any, Optional,Callable
 from langgraph.graph import StateGraph, END
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
-import wikipedia
+from langchain_community.document_loaders import WikipediaLoader
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -101,29 +101,17 @@ def web_search(query: str) -> str:
     return f"Search results for '{query}' : {result}."
 
 @tool
-def wikipedia_tool(query: str, llm: Optional[Callable[[str], str]] =openrouter_inference) -> str:
-    """
-    Searches Wikipedia for the given query and uses the LLM to extract a direct answer.
-    Args:
-        query (str): The search term or question.
-        llm (callable, optional): The language model to extract the answer.
-    Returns:
-        str: A direct answer or summary from Wikipedia.
-    """
+def wikipedia_tool(query: str) -> str:
+    """Search Wikipedia for a query and return up to 2 results as formatted text."""
     try:
-        summary = wikipedia.summary(query, sentences=5)
-        if llm:
-            # Ask the LLM to extract a direct answer from the summary
-            prompt = (
-                f"Given the following Wikipedia summary, answer the question as directly as possible:\n"
-                f"Question: {query}\n"
-                f"Summary: {summary}\n"
-                f"Answer:"
-            )
-            answer = llm(prompt)
-            return answer.strip()
-        else:
-            return summary
+        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+        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
     except Exception as e:
         return f"Wikipedia search error: {e}"