add test

app.py

@@ -11,6 +11,10 @@ from langchain_community.retrievers import WikipediaRetriever
 from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
 from langchain_community.utilities import WikipediaAPIWrapper
 
+import wikipedia
+import wikipediaapi
+from bs4 import BeautifulSoup
+
 from langchain_core import messages
 from langchain_core.messages import (
     BaseMessage,
@@ -173,17 +177,29 @@ class BasicAgent:
 
     def wiki_search(self, state: AnswerState) -> Dict[str, Any]:
         print("[wiki_search] Searching for: " + str(state["wiki_request"]))
-        import wikipedia
+
         wikipedia.set_lang("en")
         page_titles = wikipedia.search(
             state["wiki_request"], results=1
         )
-        print(f"[miki] Wiki page_titles: {page_titles}")
-        res = wikipedia.page(title=page_titles[0], auto_suggest=False)
-        print(f"[miki] Wiki res: {res}")
-
+        print("[wiki_search] Fetching the whole page")
+        wiki = wikipediaapi.Wikipedia(language='en', user_agent='gaia_example')
+        page = wiki.page(page_titles[0])
+        if page.exists():
+            html_url = f"https://en.wikipedia.org/wiki/{page.title.replace(' ', '_')}"
+            html = requests.get(html_url).text
+            soup = BeautifulSoup(html, 'html.parser')
+            state["messages"].append(
+                    ChatMessage(
+                        role="user",
+                        content="WikipediaAgent tool results are: " + str(soup)
+                    )
+                )
+            return {"messages": state["messages"]}
 
+        print("[wiki_search] Fetching the page via wikiRetriever")
         wiki_tool = WikipediaRetriever(
+            wiki_client = wikipedia,
             top_k_results=1,
             doc_content_chars_max=100000,
             load_all_available_meta=True,
@@ -198,20 +214,23 @@ class BasicAgent:
                     + results[0].page_content,
                 )
             )
-        else:
-            wiki_tool = WikipediaAPIWrapper(
-                top_k_results=1,
-                doc_content_chars_max=100000,
-                load_all_available_meta=True,
-            )
-            results = wiki_tool.run(str(state["wiki_request"]))
-            print(f"Wiki results: {results}")
-            state["messages"].append(
-                ChatMessage(
-                    role="user",
-                    content="WikipediaAgent tool results are: " + results,
-                )
+            return {"messages": state["messages"]}
+
+        print("[wiki_search] Fetching the page via WikipediaAPIWrapper")
+        wiki_tool = WikipediaAPIWrapper(
+            wiki_client = wikipedia,
+            top_k_results=1,
+            doc_content_chars_max=100000,
+            load_all_available_meta=True,
+        )
+        results = wiki_tool.run(str(state["wiki_request"]))
+        print(f"Wiki results: {results}")
+        state["messages"].append(
+            ChatMessage(
+                role="user",
+                content="WikipediaAgent tool results are: " + results,
             )
+        )
         return {"messages": state["messages"]}
 
     def python_exec(self, state: AnswerState) -> Dict[str, Any]: