Switching to gemini for CEO

Co-authored-by: Kunal Pai <kunpai@users.noreply.github.com>
Co-authored-by: Harshil Patel <hpppatel@ucdavis.edu>

CEO/CEO.py

@@ -1,144 +1,65 @@
-from enum import Enum
-from pydantic import BaseModel, Field
-from typing import List, Dict, Optional
-from pathlib import Path
-import ollama
+from google import genai
+from google.genai import types
+import os
+from dotenv import load_dotenv
 
 from CEO.tool_loader import ToolLoader
 
-# Enum for Model Types
-class ModelType(Enum):
-    LM_3B = "LM-3B"
-    LM_5B = "LM-5B"
-    LM_7B = "LM-7B"
-    LLM = "LLM"
-
-# Enum for AI Companies
-class AICompany(Enum):
-    OPENAI = "OpenAI"
-    GOOGLE = "Google"
-    META = "Meta"
-    CLAUDE = "Claude"
-    MISTRAL = "Mistral"
-
-# Enum for Agent Specializations
-class Specialization(Enum):
-    NLP = "Natural Language Processing"
-    CV = "Computer Vision"
-    RL = "Reinforcement Learning"
-    ML = "Machine Learning"
-    DATA_SCIENCE = "Data Science"
-
-# Enum for Model Parameters (Temperature, num_ctx, etc.)
-class ModelParameters(Enum):
-    NUM_CTX = 4096
-    TEMPERATURE = 0.2  # A typical temperature value for model responses
-    TOP_K = 50         # Number of top tokens to consider during generation
-
-class Subtask(BaseModel):
-    subtask_id: str = Field(..., description="Unique identifier for the subtask")
-    description: str = Field(..., description="Description of the subtask")
-    assigned_to: str = Field(..., description="ID of the agent or API handling the subtask")
-
-class Agent(BaseModel):
-    agent_id: str = Field(..., description="Unique identifier for the hired agent")
-    model_type: ModelType = Field(..., description="Parameters of model used: 3 billion, 5 billion, 7 billion, LLM")
-    company: AICompany = Field(..., description="Company name of the agent: OpenAI, Google, Meta, Claude, Mistral")
-    specialization: Specialization = Field(..., description="Task specialization of the agent")
-    cost: float = Field(..., description="Cost of hiring the agent")
-
-class APIUtilization(BaseModel):
-    api_name: str = Field(..., description="Name of the external API used")
-    endpoint: str = Field(..., description="API endpoint URL")
-    parameters: Dict[str, str] = Field(..., description="Input parameters and their types")
-    reasoning: str = Field(..., description="Explanation for using this API")
-
-class AgentManagement(BaseModel):
-    hired: List[Agent] = Field(default=[], description="List of hired agents")
-
-class ToolCall(BaseModel):
-    function: str = Field(..., description="Name of the function to be called")
-    arguments: Dict[str, str] = Field(..., description="Arguments for the function call")
-
-class CEOResponse(BaseModel):
-    # decision: str = Field(..., description="Decision made by the CEO: Hire or Assign_API")
-    # task_breakdown: List[Subtask] = Field(..., description="List of decomposed subtasks")
-    # agent_management: AgentManagement = Field(..., description="Details of agent hiring")
-    # api_utilization: Optional[List[APIUtilization]] = Field(default=None, description="List of utilized APIs, if any")
-    tools: List[ToolCall] = Field(default=None, description="List of tool or agent calls made by the model")
-    message: str = Field(default=None, description="Message content from the model")
-
-class OllamaModelManager:
-    def __init__(self, toolsLoader: ToolLoader, model_name="HASHIRU-CEO", system_prompt_file="./models/system2.prompt"):
-        self.model_name = model_name
-        # Get the directory of the current script and construct the path to system.prompt
-        script_dir = Path(__file__).parent
-        self.system_prompt_file = system_prompt_file
+class GeminiManager:
+    def __init__(self, toolsLoader: ToolLoader, system_prompt_file="./models/system2.prompt"):
+        load_dotenv()
+        self.API_KEY = os.getenv("GEMINI_KEY")
+        self.client = genai.Client(api_key=self.API_KEY)
         self.toolsLoader = toolsLoader
         self.toolsLoader.load_tools()
-        self.create_model()
-
-    def is_model_loaded(self, model):
-        loaded_models = [m.model for m in ollama.list().models]
-        return model in loaded_models or f'{model}:latest' in loaded_models
-
-    def create_model(self):
-        with open(self.system_prompt_file, 'r', encoding="utf8") as f:
-            system = f.read()
-            # system += "Tools\n"+str(self.toolsLoader.getTools())
-
-        if not self.is_model_loaded(self.model_name):
-            print(f"Creating model {self.model_name}")
-            ollama.create(
-                model=self.model_name,
-                from_='mistral-nemo',
-                system=system,
-                # parameters={"temperature": ModelParameters.TEMPERATURE.value}
-            )
+        with open(system_prompt_file, 'r', encoding="utf8") as f:
+            self.system_prompt = f.read()
 
     def request(self, messages):
-        print(f"messages: {messages}")
-        response = ollama.chat(
-            model=self.model_name,
-            messages=messages,
-            # format=CEOResponse.model_json_schema(),
-            tools=self.toolsLoader.getTools(),
+        response = self.client.models.generate_content(
+            model='gemini-2.0-flash-001',
+            contents=messages,
+            config=types.GenerateContentConfig(
+                system_instruction=self.system_prompt,
+                temperature=0.0,
+                tools=self.toolsLoader.getTools(),
+            ),
         )
-        print(f"Response: {response}")
-        # response = CEOResponse.model_validate_json(response['message']['content'])
-        if "EOF" in response.message.content:
-            return messages
-        if response.message.tool_calls:
-            for tool_call in response.message.tool_calls:
-                print(f"Tool Name: {tool_call.function.name}, Arguments: {tool_call.function.arguments}")
-                toolResponse = self.toolsLoader.runTool(tool_call.function.name, tool_call.function.arguments)
+
+        if response.text is not None:
+            print(f"Response: {response.text}")
+            assistant_content = types.Content(
+                role='assistant',
+                parts=[types.Part.from_text(text=response.text)],
+            )
+            messages.append(assistant_content)
+            if "EOF" in response.text:
+                return messages
+        if response.candidates[0].content:
+            messages.append(response.candidates[0].content)
+        if response.function_calls:
+            parts = []
+            for function_call in response.function_calls:
+                toolResponse = None
+                print(f"Function Name: {function_call.name}, Arguments: {function_call.args}")
+                toolResponse = self.toolsLoader.runTool(function_call.name, function_call.args)
                 print(f"Tool Response: {toolResponse}")
-                role = "tool"
-                if "role" in toolResponse:
-                    role = toolResponse["role"]
-                messages.append({"role": role, "content": str(toolResponse)})
+                tool_content = types.Part.from_function_response(
+                        name=function_call.name,
+                        response = {"result":toolResponse})
                 try:
                     self.toolsLoader.load_tools()
                 except Exception as e:
                     print(f"Error loading tools: {e}")
-                    messages.append({"role": "assistant", "content": "Error loading new tools."})
+                    tool_content = types.Part.from_function_response(
+                            name=function_call.name,
+                            response={"result":"Error loading new tools."+str(e)})
+                parts.append(tool_content)
+            function_response_content = types.Content(
+                role='tool', parts=parts
+            )
+            messages.append(function_response_content)
             self.request(messages)
-        # if response.tools:
-        #     for tool_call in response.tools:
-        #         print(f"Tool Name: {tool_call.function}, Arguments: {tool_call.arguments}")
-        #         toolResponse = self.toolsLoader.runTool(tool_call.function, tool_call.arguments)
-        #         print(f"Tool Response: {toolResponse}")
-        #         role = "tool"
-        #         if "role" in toolResponse:
-        #             role = toolResponse["role"]
-        #         messages.append({"role": role, "content": str(toolResponse)})
-        #     self.request(messages)
         else:
             print("No tool calls found in the response.")
-            messages.append({"role": "assistant", "content": response.message})
-            print(f"Messages: {messages}")
-            # ask_user_tool = AskUser()
-            # ask_user_response = ask_user_tool.run(prompt=response.message.content)
-            # messages.append({"role": "user", "content": ask_user_response})
-            # self.request(messages)
             return messages

CEO/tool_loader.py

@@ -1,7 +1,9 @@
 import importlib
 import importlib.util
 import os
+import types
 import pip
+from google.genai import types
 
 toolsImported = []
 
@@ -57,10 +59,27 @@ class ToolLoader:
     def getTools(self):
         toolsList = []
         for tool in self.toolsImported:
-            toolsList.append({
-                "type": "function",
-                "function": tool.inputSchema
-            })
+            parameters = types.Schema()
+            parameters.type = tool.inputSchema["parameters"]["type"]
+            properties = {}
+            for prop, value in tool.inputSchema["parameters"]["properties"].items():
+                properties[prop] = types.Schema(
+                    type=value["type"],
+                    description=value["description"]
+                )
+            parameters.properties = properties
+            parameters.required = tool.inputSchema["parameters"].get("required", [])
+            function = types.FunctionDeclaration(
+                name=tool.inputSchema["name"],
+                description=tool.inputSchema["description"],
+                parameters=parameters,
+            )
+            toolType = types.Tool(function_declarations=[function])
+            toolsList.append(toolType)
+            # toolsList.append({
+            #     "type": "function",
+            #     "function": tool.inputSchema
+            # })
         return toolsList
 
 toolLoader = ToolLoader()

main.py

@@ -1,6 +1,7 @@
+from google.genai import types
 from typing import List
 from googlesearch import search
-from CEO.CEO import OllamaModelManager
+from CEO.CEO import GeminiManager
 from CEO.tool_loader import ToolLoader
 
 # Define the web search tool function.
@@ -40,18 +41,20 @@ if __name__ == "__main__":
 
     # Load the tools using the ToolLoader class.
     tool_loader = ToolLoader()
-    # tools.extend(tool_loader.getTools())
 
-    # Create the Ollama model manager and ensure the model is set up.
-    model_manager = OllamaModelManager(toolsLoader=tool_loader)
+    model_manager = GeminiManager(toolsLoader=tool_loader)
     
     # Example prompt instructing the CEO model to create a strategy for Ashton Hall.
     # 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 tool to get the current system time and invoke it to get the current time."
+        "Get me the current time here"
     )
     
     # Request a CEO response with the prompt.
-    response = model_manager.request([{"role": "user", "content": task_prompt}])
+    user_prompt_content = types.Content(
+        role='user',
+        parts=[types.Part.from_text(text=task_prompt)],
+    )
+    response = model_manager.request([user_prompt_content])
     print("\nCEO Response:", response)

models/system2.prompt

@@ -26,7 +26,7 @@ Always invoke GetAgents tool to get the list of available agents and their capab
 </Rule>
 
 <Rule>
-If an agent isn't already available, invoke the CreateAgent tool to create a new agent with the required capabilities. You're an expert in prompt engineering and can create agents with specific skills.
+If an agent isn't already available, invoke the AgentCreator tool to create a new agent with the required capabilities. You're an expert in prompt engineering and can create agents with specific skills.
 </Rule>
 
 <Rule>
@@ -38,7 +38,7 @@ If the agent is not able to answer the question, invoke the AskUser tool to get
 </Rule>
 
 <Rule>
-In order to execute tasks on real time data, math calculations, or any other operations, invoke the CreateTool tool to create a new tool with the required capabilities. The tools are created in Python and must follow this strict schema:
+In order to execute tasks on real time data, math calculations, or any other operations, invoke the ToolCreator tool to create a new tool with the required capabilities. The tools are created in Python and must follow this strict schema:
 import importlib
 
 __all__ = ['WeatherApi']

models/system_.prompt

@@ -13,9 +13,11 @@ General Guidelines:
 3. When more details are needed, ask the user for clarification by invoking the AskUser (ASK_USR) command.
 4. Prior to invoking any other tools, always invoke the GetAgents (GET) command to retrieve the list of available agents and their capabilities.
 5. If a needed agent isn't available, create one using the CREATE_LM (or CreateAgent) command with the required capabilities.
-6. When interacting with an agent, use the AskAgent tool to forward the query.
+6. When interacting with an agent, use the CALL_TOOL tool to forward the query.
 7. If the agent cannot answer the query, then use ASK_USR to request additional clarification.
 8. For real-time data, computations, or specialized operations, create a new tool using CREATE_TOOL following the strict schema provided.
+9. Make sure all the code you write for tools is production-ready with no dummy configs or defaults.
+It should be ready-to-use out of the box.
 
 Command Set:
 ------------
@@ -65,7 +67,7 @@ Tools: External programs used to perform specific tasks. They can be created, in
 </Info>
 
 <Info>
-Agents: Special entities that handle complex tasks or queries. They are invoked via tools using AskAgent. Create agents with specific capabilities when necessary.
+Agents: Special entities that handle complex tasks or queries. They are invoked via tools using CALL_TOOL. Create agents with specific capabilities when necessary.
 </Info>
 
 <Rule>
@@ -74,7 +76,7 @@ Agents: Special entities that handle complex tasks or queries. They are invoked
 3. If more information is required, invoke ASK_USR to request clarification.
 4. Always invoke GET to list available agents before using any tool.
 5. If an agent is missing, use CREATE_LM (or a similar CreateAgent command) to create one.
-6. Use AskAgent to interact with an agent once created.
+6. Use CALL_TOOL to interact with an agent once created.
 7. If the agent cannot resolve the query, then request more details from the user using ASK_USR.
 8. For tasks involving real-time data, computations, or similar operations, create a tool with CREATE_TOOL. The tool’s code must strictly adhere to the provided schema. Be aware of recognizing if a task requires a tool.
 9. In case of API Keys, ask the user to provide one.

tools/GetCurrentTime.py

@@ -1,4 +0,0 @@
-import datetime
-
-def get_current_time():
-    return datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')

tools/TimeApi.py

@@ -0,0 +1,43 @@
+import datetime
+
+__all__ = ['TimeApi']
+
+class TimeApi():
+    dependencies = []
+
+    inputSchema = {
+        "name": "TimeApi",
+        "description": "Returns the current time for a given location",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "location": {
+                    "type": "string",
+                    "description": "The location for which to get the current time",
+                },
+            },
+            "required": ["location"],
+        }
+    }
+
+    def __init__(self):
+        pass
+
+    def run(self, **kwargs):
+        location = kwargs.get("location")
+        try:
+            # This will only work if the timezone is configured correctly on the server
+            now = datetime.datetime.now(datetime.timezone.utc)
+            return {
+                "status": "success",
+                "message": f"Current time in {location} is {now.strftime('%Y-%m-%d %H:%M:%S %Z%z')}",
+                "error": None,
+                "output": now.strftime('%Y-%m-%d %H:%M:%S %Z%z')
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": f"Could not get current time for {location}",
+                "error": str(e),
+                "output": None
+            }

tools/agent_creater_tool.py

@@ -29,7 +29,8 @@ class AgentCreator():
                     "type": "string",
                     "description": "Description of the agent. This is a string that describes the agent and its capabilities. It should be a single line description.",
                 },
-            }
+            },
+            "required": ["agent_name", "base_model", "system_prompt", "description"],
         }
     }
 

tools/ask_agent_tool.py

@@ -9,7 +9,7 @@ class AskAgent():
                     "pydantic_core==2.33.0"]
 
     inputSchema = {
-        "name": "AgentCreator",
+        "name": "AskAgent",
         "description": "Asks an AI agent a question and gets a response. The agent must be created using the AgentCreator tool before using this tool.",
         "parameters": {
             "type": "object",
@@ -22,7 +22,8 @@ class AskAgent():
                     "type": "string",
                     "description": "This is the prompt that will be used to ask the agent a question. It should be a string that describes the question to be asked.",
                 }
-            }
+            },
+            "required": ["agent_name", "prompt"],
         }
     }
 

tools/get_agents_tool.py

@@ -11,8 +11,9 @@ class GetAgents():
         "description": "Retrieves a list of available AI agents. This tool is used to get the list of available models that can be invoked using the AskAgent tool.",
         "parameters": {
             "type": "object",
-            "properties": {}
-        }
+            "properties": {},
+            "required": [],
+        },
     }
 
     def __init__(self):

tools/tool_creator.py

@@ -21,7 +21,7 @@ class ToolCreator():
                     "description": "The code of the tool to create",
                 },
             },
-            "required": ["name", "content"],
+            "required": ["name", "tool_code"],
         }
     }