Updating to use webUI

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

deleteAgents.py

@@ -1,14 +1,15 @@
 import json
 import ollama
+from src.utils.streamlit_interface import output_assistant_response
 
 with open("./models/models.json", "r", encoding="utf8") as f:
     models = f.read()
 models = json.loads(models)
 for agent in models:
-    print(f"Deleting agent: {agent}")
+    output_assistant_response(f"Deleting agent: {agent}")
     try:
         ollama.delete(agent)
     except Exception as e:
-        print(f"Error deleting agent {agent}: {e}")
+        output_assistant_response(f"Error deleting agent {agent}: {e}")
     with open("./models/models.json", "w", encoding="utf8") as f:
         f.write(json.dumps({}, indent=4))

main.py

@@ -2,6 +2,7 @@ from google.genai import types
 from src.CEO import GeminiManager
 from src.tool_loader import ToolLoader
 
+
 if __name__ == "__main__":
     # Define the tool metadata for orchestration.
     # Load the tools using the ToolLoader class.

mainV2.py

@@ -0,0 +1,55 @@
+from google.genai import types
+from src.manager import GeminiManager
+from src.tool_loader import ToolLoader
+import gradio as gr
+import time
+
+if __name__ == "__main__":
+    # Define the tool metadata for orchestration.
+    # Load the tools using the ToolLoader class.
+    tool_loader = ToolLoader()
+
+    model_manager = GeminiManager(toolsLoader=tool_loader, gemini_model="gemini-2.0-flash")
+    
+    def respond(message, chat_history):
+        return model_manager.ask(message, chat_history)
+    
+    def user_message(msg: str, history: list) -> tuple[str, list]:
+        """Adds user message to chat history"""
+        history.append(gr.ChatMessage(role="user", content=msg))
+        return "", history
+    
+    with gr.Blocks() as demo:
+        chatbot = gr.Chatbot(type="messages")
+        input_box = gr.Textbox()
+        clear = gr.ClearButton([input_box, chatbot])
+
+        def respond(message, chat_history):
+            
+            chat_history.append({
+                "role":"user",
+                "content":message
+            })
+            print("Chat history:", chat_history)
+            chat_history = model_manager.run(chat_history)
+            return "", chat_history
+
+        msg_store = gr.State("")
+        
+        input_box.submit(
+            lambda msg: (msg, msg, ""),  # Store message and clear input
+            inputs=[input_box],
+            outputs=[msg_store, input_box, input_box],
+            queue=False
+        ).then(
+            user_message,  # Add user message to chat
+            inputs=[msg_store, chatbot],
+            outputs=[input_box, chatbot],
+            queue=False
+        ).then(
+            model_manager.run,  # Generate and stream response
+            inputs=[msg_store, chatbot],
+            outputs=chatbot
+        )
+
+    demo.launch()

models/system3.prompt

@@ -43,7 +43,7 @@ Never answer any questions yourself, instead use tools. Only exception to this r
 </Rule>
 
 <Rule>
-If you need more information to answer the question, ask the user for clarification or additional details by invoking the AskUser tool.
+If you need more information to answer the question, ask the user for clarification or additional details.
 </Rule>
 
 <Rule>
@@ -59,7 +59,7 @@ Once an Agent is created, use the AskAgent tool to ask the agent the question or
 </Rule>
 
 <Rule>
-If the agent is not able to answer the question, invoke the AskUser tool to get more information or clarify the question.
+If the agent is not able to answer the question, ask the user to get more information or clarify the question.
 </Rule>
 
 <Rule>

requirements.txt

@@ -1,7 +1,105 @@
+aiofiles==24.1.0
+altair==5.5.0
+annotated-types==0.7.0
+anyio==4.9.0
+appdirs==1.4.4
+arxiv==2.2.0
+attrs==25.3.0
+beautifulsoup4==4.13.3
+blinker==1.9.0
+cachetools==5.5.2
+certifi==2025.1.31
+charset-normalizer==3.4.1
+click==8.1.8
+colorama==0.4.6
+eval_type_backport==0.2.2
+fastapi==0.115.12
+feedparser==6.0.11
+ffmpy==0.5.0
+filelock==3.18.0
+frozendict==2.4.6
+fsspec==2025.3.2
+gitdb==4.0.12
+GitPython==3.1.44
+google-ai-generativelanguage==0.6.15
+google-api-core==2.24.2
+google-auth==2.38.0
+google-genai==1.9.0
+googleapis-common-protos==1.69.2
+googlesearch-python==1.3.0
+gradio==5.27.0
+gradio_client==1.9.0
+groovy==0.1.2
+grpcio==1.71.0
+grpcio-status==1.71.0
+h11==0.14.0
+httpcore==1.0.7
+httpx==0.28.1
+huggingface-hub==0.30.2
+idna==3.10
+importlib_metadata==8.6.1
+Jinja2==3.1.6
+jsonschema==4.23.0
+jsonschema-specifications==2025.4.1
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+mdurl==0.1.2
 mistralai==1.6.0
+multitasking==0.0.11
+narwhals==1.36.0
+numpy==2.2.5
 ollama==0.4.7
-pip==25.0.1
+orjson==3.10.16
+packaging==24.2
+pandas==2.2.3
+peewee==3.17.9
+pillow==11.2.1
+platformdirs==4.3.7
+proto-plus==1.26.1
 protobuf==5.29.4
+pyarrow==19.0.1
+pyasn1==0.6.1
+pyasn1_modules==0.4.2
 pydantic==2.11.2
+pydantic_core==2.33.1
+pydeck==0.9.1
+pydub==0.25.1
+pyee==11.1.1
+Pygments==2.19.1
+pyppeteer==2.0.0
+python-dateutil==2.9.0.post0
 python-dotenv==1.1.0
-google-genai==1.9.0
+python-multipart==0.0.20
+pytz==2025.2
+PyYAML==6.0.2
+referencing==0.36.2
+requests==2.32.3
+rich==14.0.0
+rpds-py==0.24.0
+rsa==4.9
+ruff==0.11.7
+safehttpx==0.1.6
+semantic-version==2.10.0
+sgmllib3k==1.0.0
+shellingham==1.5.4
+six==1.17.0
+smmap==5.0.2
+sniffio==1.3.1
+soupsieve==2.6
+starlette==0.46.2
+streamlit==1.44.1
+tenacity==9.1.2
+toml==0.10.2
+tomlkit==0.13.2
+tornado==6.4.2
+tqdm==4.67.1
+typer==0.15.2
+typing-inspection==0.4.0
+typing_extensions==4.13.1
+tzdata==2025.2
+urllib3==1.26.20
+uvicorn==0.34.2
+watchdog==6.0.0
+websockets==10.4
+yfinance==0.2.56
+zipp==3.21.0

src/CEO.py

@@ -7,6 +7,8 @@ from src.tool_loader import ToolLoader
 from src.utils.suppress_outputs import suppress_output
 import logging
 
+from src.utils.streamlit_interface import get_user_message, output_assistant_response
+
 logger = logging.getLogger(__name__)
 handler = logging.StreamHandler(sys.stdout)
 handler.setLevel(logging.INFO)
@@ -22,6 +24,7 @@ class GeminiManager:
         self.model_name = gemini_model
         with open(system_prompt_file, 'r', encoding="utf8") as f:
             self.system_prompt = f.read()
+        self.messages = []
 
     def generate_response(self, messages):
         return self.client.models.generate_content(
@@ -75,21 +78,22 @@ class GeminiManager:
             response = suppress_output(self.generate_response)(messages)
         except Exception as e:
             logger.debug(f"Error generating response: {e}")
-            shouldRetry = input("An error occurred. Do you want to retry? (y/n): ")
-            if shouldRetry.lower() == "y":
+            shouldRetry = get_user_message("An error occurred. Do you want to retry? (y/n): ")
+            if shouldRetry and shouldRetry.lower() == "y":
                 return self.run(messages)
             else:
-                print("Ending the conversation.")
+                output_assistant_response("Ending the conversation.")
                 return messages
 
         logger.debug(f"Response: {response}")
         
         if (not response.text and not response.function_calls):
-            print("No response from the model.")
+            output_assistant_response("No response from the model.")
         
         # Attach the llm response to the messages
         if response.text is not None:
-            print("CEO:", response.text)
+            output_assistant_response("CEO: " + response.text)
+            # print("CEO:", response.text)
             assistant_content = types.Content(
                 role='model' if self.model_name == "gemini-2.5-pro-exp-03-25" else 'assistant',
                 parts=[types.Part.from_text(text=response.text)],
@@ -103,11 +107,11 @@ class GeminiManager:
         # Invoke the function calls if any and attach the response to the messages
         if response.function_calls:
             messages.append(self.handle_tool_calls(response))
-            shouldContinue = input("Should I continue? (y/n): ")
+            shouldContinue = get_user_message("Should I continue? (y/n): ")
             if shouldContinue.lower() == "y":
                 return self.run(messages)
             else:
-                print("Ending the conversation.")
+                output_assistant_response("Ending the conversation.")
                 return messages
         else:
             logger.debug("No tool calls found in the response.")
@@ -115,9 +119,10 @@ class GeminiManager:
             return self.start_conversation(messages)
     
     def start_conversation(self, messages=[]):
-        question = input("User: ")
-        if ("exit" in question.lower() or "quit" in question.lower()):
-            print("Ending the conversation.")
+        question = get_user_message("User: ")
+        # question = input("User: ")
+        if question and ("exit" in question.lower() or "quit" in question.lower()):
+            output_assistant_response("Ending the conversation.")
             return messages
         user_content = types.Content(
             role='user',

src/agent_manager.py

@@ -3,7 +3,8 @@ from typing import Dict, Type, Any, Optional
 import os
 import json
 import ollama
-from src.singleton import singleton
+from src.utils.singleton import singleton
+from src.utils.streamlit_interface import output_assistant_response
 
 class Agent(ABC):
     
@@ -47,12 +48,12 @@ class OllamaAgent(Agent):
         )
     
     def ask_agent(self, prompt):
-        print(f"Asked Agent {self.agent_name} a question")
+        output_assistant_response(f"Asked Agent {self.agent_name} a question")
         agent_response = ollama.chat(
             model=self.agent_name,
             messages=[{"role": "user", "content": prompt}],
         )
-        print(f"Agent {self.agent_name} answered with {agent_response.message.content}")
+        output_assistant_response(f"Agent {self.agent_name} answered with {agent_response.message.content}")
         return agent_response.message.content
     
     def delete_agent(self):
@@ -122,7 +123,7 @@ class AgentManager():
             else:
                 return {}
         except Exception as e:
-            print(f"Error listing agents: {e}")
+            output_assistant_response(f"Error listing agents: {e}")
             return {}
     
     def delete_agent(self, agent_name: str) -> None:
@@ -140,7 +141,7 @@ class AgentManager():
                 with open("./models/models.json", "w", encoding="utf8") as f:
                     f.write(json.dumps(models, indent=4))
         except Exception as e:
-            print(f"Error deleting agent: {e}")
+            output_assistant_response(f"Error deleting agent: {e}")
     
     
     def _save_agent(self, agent_name: str, base_model: str, system_prompt: str, 
@@ -176,7 +177,7 @@ class AgentManager():
                 f.write(json.dumps(models, indent=4))
                 
         except Exception as e:
-            print(f"Error saving agent {agent_name}: {e}")
+            output_assistant_response(f"Error saving agent {agent_name}: {e}")
 
     def _get_agent_type(self, base_model)->str:
 
@@ -216,4 +217,4 @@ class AgentManager():
                         invoke_cost
                     )
         except Exception as e:
-            print(f"Error loading agents: {e}")
+            output_assistant_response(f"Error loading agents: {e}")

src/budget_manager.py

@@ -1,4 +1,4 @@
-from src.singleton import singleton
+from src.utils.singleton import singleton
 
 @singleton
 class BudgetManager():

src/llm_models.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from google import genai
 from google.genai import types
 from mistralai import Mistral
+from src.utils.streamlit_interface import output_assistant_response
 
 
 class AbstractModelManager(ABC):
@@ -39,7 +40,7 @@ class OllamaModelManager(AbstractModelManager):
             system = f.read()
         
         if not self.is_model_loaded(self.model_name):
-            print(f"Creating model {self.model_name}")
+            output_assistant_response(f"Creating model {self.model_name}")
             ollama.create(
                 model=self.model_name,
                 from_=base_model,
@@ -60,10 +61,10 @@ class OllamaModelManager(AbstractModelManager):
     
     def delete(self):
         if self.is_model_loaded("C2Rust:latest"):
-            print(f"Deleting model {self.model_name}")
+            output_assistant_response(f"Deleting model {self.model_name}")
             ollama.delete("C2Rust:latest")
         else:
-            print(f"Model {self.model_name} not found, skipping deletion.")
+            output_assistant_response(f"Model {self.model_name} not found, skipping deletion.")
 
 class GeminiModelManager(AbstractModelManager):
     def __init__(self, api_key):

src/manager.py

@@ -0,0 +1,147 @@
+from google import genai
+from google.genai import types
+from google.genai.types import *
+import os
+from dotenv import load_dotenv
+import sys
+from src.tool_loader import ToolLoader
+from src.utils.suppress_outputs import suppress_output
+import logging
+from gradio import ChatMessage
+
+from src.utils.streamlit_interface import get_user_message, output_assistant_response
+
+logger = logging.getLogger(__name__)
+handler = logging.StreamHandler(sys.stdout)
+# handler.setLevel(logging.DEBUG)
+logger.addHandler(handler)
+
+class GeminiManager:
+    def __init__(self, toolsLoader: ToolLoader, system_prompt_file="./models/system3.prompt", gemini_model="gemini-2.5-pro-exp-03-25"):
+        load_dotenv()
+        self.API_KEY = os.getenv("GEMINI_KEY")
+        self.client = genai.Client(api_key=self.API_KEY)
+        self.toolsLoader: ToolLoader = toolsLoader
+        self.toolsLoader.load_tools()
+        self.model_name = gemini_model
+        with open(system_prompt_file, 'r', encoding="utf8") as f:
+            self.system_prompt = f.read()
+        self.messages = []
+
+    def generate_response(self, messages):
+        return self.client.models.generate_content(
+                model=self.model_name,
+                contents=messages,
+                config=types.GenerateContentConfig(
+                    system_instruction=self.system_prompt,
+                    temperature=0.2,
+                    tools=self.toolsLoader.getTools(),
+                ),
+            )
+    
+    def handle_tool_calls(self, response):
+        parts = []
+        for function_call in response.function_calls:
+            toolResponse = None
+            logger.info(f"Function Name: {function_call.name}, Arguments: {function_call.args}")
+            try:
+                toolResponse = self.toolsLoader.runTool(function_call.name, function_call.args)
+            except Exception as e:
+                logger.warning(f"Error running tool: {e}")
+                toolResponse = {
+                    "status": "error",
+                    "message": f"Tool {function_call.name} failed to run.",
+                    "output": str(e),
+                }
+            logger.debug(f"Tool Response: {toolResponse}")
+            tool_content = types.Part.from_function_response(
+                    name=function_call.name,
+                    response = {"result":toolResponse})
+            try:
+                self.toolsLoader.load_tools()
+            except Exception as e:
+                logger.info(f"Error loading tools: {e}. Deleting the tool.")
+                # delete the created tool
+                self.toolsLoader.delete_tool(toolResponse['output']['tool_name'], toolResponse['output']['tool_file_path'])
+                tool_content = types.Part.from_function_response(
+                        name=function_call.name,
+                        response={"result":f"{function_call.name} with {function_call.args} doesn't follow the required format, please read the other tool implementations for reference." + str(e)})
+            parts.append(tool_content)
+        return {
+                "role": "tool",
+                "content": repr(types.Content(
+                    role='model' if self.model_name == "gemini-2.5-pro-exp-03-25" else 'tool',
+                    parts=parts
+                ))
+            }
+    
+    def format_chat_history(self, messages=[]):
+        formatted_history = []
+        for message in messages:
+            # Skip thinking messages (messages with metadata)
+            if not (message.get("role") == "assistant" and "metadata" in message):
+                role = "model"
+                parts=[types.Part.from_text(text=message.get("content", ""))]
+                match message.get("role"):
+                    case "user":
+                        role = "user"
+                    case "tool":
+                        role = "tool"
+                        formatted_history.append(eval(message.get("content", "")))
+                        continue
+                    case "function_call":
+                        role = "model"
+                        formatted_history.append(eval(message.get("content", "")))
+                        continue
+                    case _:
+                        role = "model" 
+                formatted_history.append(types.Content(
+                    role=role,
+                    parts=parts
+                ))
+        return formatted_history
+    
+    def run(self, message, messages):
+        chat_history = self.format_chat_history(messages)
+        logger.debug(f"Chat history: {chat_history}")
+        try:
+            response = suppress_output(self.generate_response)(chat_history)
+        except Exception as e:
+            logger.debug(f"Error generating response: {e}")
+            messages.append({
+                "role":"assistant",
+                "content":f"Error generating response: {e}"
+            })
+            logger.error(f"Error generating response: {e}")
+            return messages
+        logger.debug(f"Response: {response}")
+        
+        if (not response.text and not response.function_calls):
+            messages.append({
+                "role":"assistant",
+                "content":"No response from the model.",
+                "metadata":{"title":"No response from the model."}
+            })
+
+        # Attach the llm response to the messages
+        if response.text is not None and response.text != "":
+            messages.append({
+                "role": "assistant",
+                "content": response.text
+            })
+            
+        # Attach the function call response to the messages
+        if response.candidates[0].content and response.candidates[0].content.parts:
+            # messages.append(response.candidates[0].content)
+            messages.append({
+                "role":"function_call",
+                "content": repr(response.candidates[0].content),
+            })
+            pass
+        
+        # Invoke the function calls if any and attach the response to the messages
+        if response.function_calls:
+            calls = self.handle_tool_calls(response)
+            messages.append(calls)
+            return self.run(message, messages)
+        return messages

src/tool_loader.py

@@ -7,10 +7,11 @@ from google.genai import types
 import sys
 
 from src.budget_manager import BudgetManager
-from src.singleton import singleton
+from src.utils.singleton import singleton
 from src.utils.suppress_outputs import suppress_output
 from tools.get_agents_tool import GetAgents
 from tools.tool_deletor import ToolDeletor
+from src.utils.streamlit_interface import output_assistant_response
 
 toolsImported = []
 
@@ -59,7 +60,7 @@ class ToolLoader:
                 if tool.inputSchema["create_cost"] is not None:
                     self.budget_manager.add_to_expense(tool.inputSchema["create_cost"])
 
-        print(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
+        output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
 
     def load_tools(self):
         newToolsImported = []
@@ -76,12 +77,12 @@ class ToolLoader:
         self.toolsImported = newToolsImported
 
     def runTool(self, toolName, query):
-        print(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
+        output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
         for tool in self.toolsImported:
             if tool.name == toolName:
                 self.update_budget(query, tool.inputSchema)
                 return tool.run(query)
-        print(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
+        output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
         return {
             "status": "error",
             "message": f"Tool {toolName} not found",
@@ -152,7 +153,7 @@ class ToolLoader:
                 "output": None
             }
 
-toolLoader = ToolLoader()
+# toolLoader = ToolLoader()
 
 # Example usage
 # print(toolLoader.getTools())

src/utils/streamlit_interface.py

@@ -0,0 +1,8 @@
+from src.utils.singleton import singleton
+
+def get_user_message(message):
+    user_input = input(message)
+    return user_input
+
+def output_assistant_response(response):
+    print(response)

tools/ask_user.py

@@ -1,32 +0,0 @@
-
-__all__ = ['AskUser']
-
-
-class AskUser():
-    dependencies = []
-
-    inputSchema = {
-        "name": "AskUser",
-        "description": "Asks a question to the user and gets a response. Only use this when you need more clarification from the user on the question.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "question": {
-                    "type": "string",
-                    "description": "The question to ask the user",
-                },
-            },
-            "required": ["question"],
-        }
-    }
-
-    def run(self, **kwargs):
-        print("Running Ask User tool")
-        question = kwargs.get("question")
-        output = input(f"{question}: ")
-        return {
-            "status": "success",
-            "message": "Ask User tool executed successfully",
-            "output": output,
-            "role": "user",
-        }