Updating tools and QOL changes

CEO/CEO.py

@@ -17,18 +17,27 @@ class GeminiManager:
             self.system_prompt = f.read()
 
     def request(self, messages):
-        response = self.client.models.generate_content(
-            #model='gemini-2.5-pro-preview-03-25',
-            model=self.model_name,
-            #model='gemini-2.5-pro-exp-03-25',
-            #model='gemini-2.0-flash',
-            contents=messages,
-            config=types.GenerateContentConfig(
-                system_instruction=self.system_prompt,
-                temperature=0.2,
-                tools=self.toolsLoader.getTools(),
-            ),
-        )
+        try:
+            response = self.client.models.generate_content(
+                #model='gemini-2.5-pro-preview-03-25',
+                model=self.model_name,
+                #model='gemini-2.5-pro-exp-03-25',
+                #model='gemini-2.0-flash',
+                contents=messages,
+                config=types.GenerateContentConfig(
+                    system_instruction=self.system_prompt,
+                    temperature=0.2,
+                    tools=self.toolsLoader.getTools(),
+                ),
+            )
+        except Exception as e:
+            print(f"Error: {e}")
+            shouldRetry = input("An error occurred. Do you want to retry? (y/n): ")
+            if shouldRetry.lower() == "y":
+                return self.request(messages)
+            else:
+                print("Ending the conversation.")
+                return messages
 
         print(f"Response: {response}")
         
@@ -65,7 +74,12 @@ class GeminiManager:
                 parts=parts
             )
             messages.append(function_response_content)
-            self.request(messages)
+            shouldContinue = input("Should I continue? (y/n): ")
+            if shouldContinue.lower() == "y":
+                return self.request(messages)
+            else:
+                print("Ending the conversation.")
+                return messages
         else:
             print("No tool calls found in the response.")
             return messages

CEO/tool_loader.py

@@ -4,6 +4,7 @@ import os
 import types
 import pip
 from google.genai import types
+import sys
 
 toolsImported = []
 
@@ -11,6 +12,8 @@ TOOLS_DIRECTORY = os.path.abspath("./tools")
 
 class Tool:
     def __init__(self, toolClass):
+        save_stdout = sys.stdout
+        sys.stdout = open('trash', 'w')
         self.tool = toolClass()
         self.inputSchema = self.tool.inputSchema
         self.name = self.inputSchema["name"]
@@ -24,6 +27,7 @@ class Tool:
                 if '==' in package:
                     package = package.split('==')[0]
                 pip.main(['install', package])
+        sys.stdout = save_stdout
 
     def run(self, query):
         return self.tool.run(**query)
@@ -35,6 +39,7 @@ class ToolLoader:
         pass
 
     def load_tools(self):
+        self.toolsImported = []
         for filename in os.listdir(TOOLS_DIRECTORY):
             if filename.endswith(".py") and filename != "__init__.py":
                 module_name = filename[:-3]

main.py

@@ -20,7 +20,6 @@ def web_search(website: str, query: str) -> List[str]:
     return results
 
 if __name__ == "__main__":
-    warnings.filterwarnings("ignore")
     # Define the tool metadata for orchestration.
     tools = [
         {
@@ -50,7 +49,7 @@ if __name__ == "__main__":
     # The prompt explicitly mentions that it can use the web_search tool if needed,
     # and that it is allowed to choose the website for the search.
     task_prompt = (
-        "Create a strategy for Ashton Hall to improve its online presence."
+        "Get me the current AAPL stock price and the latest news about it. "
     )
     
     # Request a CEO response with the prompt.

tools/agent_creater_tool.py

@@ -33,9 +33,6 @@ class AgentCreator():
             "required": ["agent_name", "base_model", "system_prompt", "description"],
         }
     }
-
-    def __init__(self):
-        pass
     
     def does_agent_exist(self, agent_name):
         ollama = importlib.import_module("ollama")

tools/arxiv_tool.py

@@ -14,7 +14,12 @@ class ArxivTool():
             "properties": {
                 "query": {
                     "type": "string",
-                    "description": "Search query for papers (e.g., 'superconductors gem5').",
+                    "description": "Search query for papers on arXiv.",
+                    "examples":[
+                        "superconductors gem5",
+                        "machine learning in healthcare",
+                        "quantum computing algorithms"
+                    ]
                 },
                 "max_results": {
                     "type": "integer",
@@ -26,9 +31,6 @@ class ArxivTool():
         }
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         query = kwargs.get("query")
         max_results = kwargs.get("max_results", 5)

tools/ask_agent_tool.py

@@ -27,9 +27,6 @@ class AskAgent():
         }
     }
 
-    def __init__(self):
-        pass
-
     def does_agent_exist(self, ollama, agent_name):
         all_agents = [a.model for a in ollama.list().models]
         if agent_name in all_agents or f'{agent_name}:latest' in all_agents:

tools/ask_user.py

@@ -21,9 +21,6 @@ class AskUser():
         }
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         print("Running Ask User tool")
         question = kwargs.get("question")

tools/get_agents_tool.py

@@ -16,9 +16,6 @@ class GetAgents():
         },
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         with open("./models/models.json", "r", encoding="utf8") as f:
             models = f.read()

tools/get_website_tool.py

@@ -0,0 +1,63 @@
+import importlib
+
+__all__ = ['GetWebsiteTool']
+
+
+class GetWebsiteTool():
+    dependencies = ["requests"]
+
+    inputSchema = {
+        "name": "GetWebsiteTool",
+        "description": "Returns the content of a website based on a query string.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "url": {
+                    "type": "string",
+                    "description": "The URL of the website to fetch content from.",
+                },
+            },
+            "required": ["url"],
+        }
+    }
+
+    def run(self, **kwargs):
+        print("Running web search")
+
+        url = kwargs.get("url")
+
+        if not url:
+            return {
+                "status": "error",
+                "message": "Missing required parameters: 'url'",
+                "output": None
+            }
+        
+        output = None
+        requests = importlib.import_module("requests")
+        try:
+            response = requests.get(url)
+            if response.status_code == 200:
+                output = response.text
+            else:
+                return {
+                    "status": "error",
+                    "message": f"Failed to fetch content from {url}. Status code: {response.status_code}",
+                    "output": None
+                }
+            
+            # truncate the results to avoid excessive output
+            if len(output) > 1000:
+                output = output[:1000] + "... (truncated)"
+
+            return {
+                "status": "success",
+                "message": "Search completed successfully",
+                "output": output,
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e),
+                "output": None
+            }

tools/{web_search_tool.py → google_search_tool.py}

@@ -1,14 +1,14 @@
 import importlib
 
-__all__ = ['WebSearchTool']
+__all__ = ['GoogleSearchTool']
 
 
-class WebSearchTool():
+class GoogleSearchTool():
     dependencies = ["googlesearch-python==1.3.0"]
 
     inputSchema = {
-        "name": "WebSearchTool",
-        "description": "Searches a specific website for a given query using Google search.",
+        "name": "GoogleSearchTool",
+        "description": "Provides a list of URLs from google search results based on a query string.",
         "parameters": {
             "type": "object",
             "properties": {
@@ -21,27 +21,26 @@ class WebSearchTool():
                     "description": "The query string to search for on the website.",
                 }
             },
-            "required": ["website", "query"],
+            "required": ["query"],
         }
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         print("Running web search")
 
         website = kwargs.get("website")
         query = kwargs.get("query")
 
-        if not website or not query:
+        if not query:
             return {
                 "status": "error",
-                "message": "Missing required parameters: 'website' and 'query'",
+                "message": "Missing required parameters: 'query'",
                 "output": None
             }
-
-        search_query = f"site:{website} {query}"
+        search_query = query
+        if website:
+            search_query = f"site:{website} {query}"
+        
         results = []
 
         googlesearch = importlib.import_module("googlesearch")

tools/list_files.py

@@ -20,9 +20,6 @@ class ListFiles():
         }
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         print("Running List Files tool")
         directory = kwargs.get("directory")

tools/read_file.py

@@ -21,9 +21,6 @@ class ReadFile():
         }
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         print("Running Read File tool")
         file_path = kwargs.get("file_path")

tools/tool_creator.py

@@ -25,9 +25,6 @@ class ToolCreator():
         }
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         print("Running Tool Creator")
         name = kwargs.get("name")

tools/weather_tool.py

@@ -21,9 +21,6 @@ class WeatherApi():
         }
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         print("Running Weather API test tool")
         location = kwargs.get("location")

tools/wikipedia_tool.py

@@ -21,9 +21,6 @@ class WikipediaTool():
         }
     }
 
-    def __init__(self):
-        pass
-
     def run(self, **kwargs):
         question = kwargs.get("question")
         if not question: