Wikitool and agent model to gemini

.devcontainer/devcontainer.json

@@ -0,0 +1,7 @@
+{
+    "image": "pytorch/pytorch",
+    "runArgs": [
+        "--network=host"
+    ],
+    "postCreateCommand": "pip install -r /workspaces/qaolin_Final_Assignment/requirements.txt"
+}

.gitignore

@@ -0,0 +1,2 @@
+.env
+test*

agents/BasicAgent.py

@@ -1,7 +1,7 @@
 import os
 
-from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, HfApiModel
-
+from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, HfApiModel, LiteLLMModel, Tool
+from langchain_community.tools.wikidata.tool import WikidataAPIWrapper, WikidataQueryRun
 
 class BasicAgent(CodeAgent):
     def __init__(self):
@@ -16,19 +16,21 @@ class BasicAgent(CodeAgent):
         print(f"LLM Config: Model='{model_name}'")
 
         try:
-            llm = HfApiModel(
-                model_id=model_name,
-                temperature=0.1,
-            )
+            llm = LiteLLMModel(
+            model_id=model_name,
+            api_key=os.environ.get("GEMINI_API_KEY", "")
+        )
         except Exception as e:
             print(f"Error initializing LLM: {e}")
             raise
 
+        # TOOLS
         print("Initializing tools...")
         try:
             search_tool = DuckDuckGoSearchTool()
             visit_tool = VisitWebpageTool()
-            tools = [search_tool, visit_tool]
+            wikidata = Tool.from_langchain(WikidataQueryRun(api_wrapper=WikidataAPIWrapper())) 
+            tools = [search_tool, visit_tool, wikidata]
             print(f"Tools initialized: {[tool.name for tool in tools]}")
         except Exception as e:
             print(f"Error initializing tools: {e}")

requirements.txt

@@ -4,4 +4,8 @@ huggingface_hub
 python-dotenv
 requests
 smolagents
-smolagents[litellm]
+smolagents[litellm]
+langchain-community
+wikibase-rest-api-client
+mediawikiapi
+litellm==1.67.2