Upload tools.py

src/tools.py

@@ -1,3 +1,6 @@
+import json
+
+from langchain_community.document_loaders import WikipediaLoader
 from langchain_core.tools import tool
 
 
@@ -60,3 +63,25 @@ def modulus(a: int, b: int) -> int:
         b: second int
     """
     return a % b
+
+
+@tool
+def wiki_search(query: str) -> str:
+    """
+    Search Wikipedia for a query and return maximum 2 results. The results
+    are formatted as a JSON string with the source and content of each
+    document.
+
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            {
+                "source": doc.metadata["source"],
+                "page": doc.metadata.get("page", ""),
+                "content": doc.page_content,
+            }
+            for doc in search_docs
+        ])
+    return json.dumps({"wiki_results": formatted_search_docs})