Cleanup tasks

Co-authored-by: Kunal Pai <kunpai@users.noreply.github.com>

CEO/CEO.py

@@ -4,6 +4,7 @@ import os
 from dotenv import load_dotenv
 
 from CEO.tool_loader import ToolLoader
+from CEO.utils.supress_outputs import supress_output
 
 class GeminiManager:
     def __init__(self, toolsLoader: ToolLoader, system_prompt_file="./models/system2.prompt", gemini_model="gemini-2.5-pro-exp-03-25"):
@@ -18,7 +19,7 @@ class GeminiManager:
 
     def request(self, messages):
         try:
-            response = self.client.models.generate_content(
+            response = supress_output(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',

CEO/tool_loader.py

@@ -6,14 +6,17 @@ import pip
 from google.genai import types
 import sys
 
+from CEO.utils.supress_outputs import supress_output
+
 toolsImported = []
 
 TOOLS_DIRECTORY = os.path.abspath("./tools")
 
 class Tool:
     def __init__(self, toolClass):
-        save_stdout = sys.stdout
-        sys.stdout = open('trash', 'w')
+        supress_output(self.load_tool)(toolClass)
+
+    def load_tool(self, toolClass):
         self.tool = toolClass()
         self.inputSchema = self.tool.inputSchema
         self.name = self.inputSchema["name"]
@@ -27,7 +30,6 @@ 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)

CEO/utils/supress_outputs.py

@@ -0,0 +1,22 @@
+from contextlib import redirect_stderr, redirect_stdout
+
+
+def supress_output(func):
+    """
+    Decorator to suppress output of a function.
+    """
+    def wrapper(*args, **kwargs):
+        import sys
+        import os
+        sys.stdout = open(os.devnull, 'w')
+        sys.stderr = open(os.devnull, 'w')
+        
+        try:
+            result = func(*args, **kwargs)
+        finally:
+            sys.stdout = sys.__stdout__
+            sys.stderr = sys.__stderr__
+        return result
+
+
+    return wrapper

README.md

@@ -1,4 +1,4 @@
-# HASHIRU: Hierarchical Agent System for Hybrid Intelligent Resource Utilization 
+# HASHIRU: Hierarchical Agent System for Hybrid Intelligent Resource Utilization
 (For AgentX competition)
 
 ![HASHIRU_ARCH](HASHIRU_ARCH.png)
@@ -6,19 +6,77 @@
 ## Overview
 HASHIRU is an agent-based framework designed to dynamically allocate and manage large language models (LLMs) and external APIs through a CEO model. The CEO model acts as a central manager, capable of hiring, firing, and directing multiple specialized agents (employees) over a given budget. It can also create and utilize external APIs as needed, making it highly flexible and scalable.
 
+## High-Level Overview
+
+1. Loads available tools using ToolLoader
+2. Instantiates a Gemini-powered CEO using GeminiManager
+3. Wraps the user prompt into a structured Content object
+4. Calls Gemini with the prompt
+5. Executes external tool calls if needed
+6. Returns a full response to the user
+7. Can ask the user for clarification if needed
+
+After every step in the request, there is a checkpoint so that we can check what happened in that step (i.e., which tool was called, what was the response, etc.). This is useful for debugging and understanding the flow of the program.
+
 ## Features
 - **Cost-Benefit Matrix**:  
   Selects the best LLM model (LLaMA, Mixtral, Gemini, DeepSeek, etc.) for any task using Ollama, based on latency, size, cost, quality, and speed.
-## Usage:
+- **Dynamic Agent Management**:  
+  The CEO agent dynamically hires and fires specialized agents based on task requirements and budget constraints.
+- **API Integration**:  
+  Seamlessly integrates external APIs for enhanced functionality and scalability.
+
+## How to Run
+
+Clone the repo and install the required dependencies:
+
+```bash
+git clone <repository-url>
+cd HASHIRU
+pip install -r requirements.txt
+```
+
+Run the main script:
+
+```bash
+python main.py
+```
+
+## Usage
+
+### Cost-Benefit Matrix
+
+The `cost_benefit.py` tool allows you to select the best LLM model for a given task based on customizable weights:
 
 ```bash
 python tools/cost_benefit.py \
   --prompt "Best places to visit in Davis" \
   --latency 4 --size 2 --cost 5 --speed 3
 ```
+
 Each weight is on a scale of **1** (least important) to **5** (most important):
 
 - `--latency`: Prefer faster responses (lower time to answer)
 - `--size`: Prefer smaller models (use less memory/resources)
 - `--cost`: Prefer cheaper responses (fewer tokens, lower token price)
 - `--speed`: Prefer models that generate tokens quickly (tokens/sec)
+
+### Example Output
+
+```plaintext
+Selected Model: Gemini
+Reason: Optimal balance of cost, speed, and latency for the given weights.
+```
+
+## Contributing
+
+Contributions are welcome! Please follow these steps:
+
+1. Fork the repository.
+2. Create a new branch for your feature or bug fix.
+3. Submit a pull request with a detailed description of your changes.
+
+## License
+
+This project is licensed under the MIT License. See the `LICENSE` file for details.
+

main.py

@@ -1,45 +1,9 @@
 from google.genai import types
-from typing import List
-from googlesearch import search
 from CEO.CEO import GeminiManager
 from CEO.tool_loader import ToolLoader
-import warnings
-
-# Define the web search tool function.
-def web_search(website: str, query: str) -> List[str]:
-    """
-    Searches the specified website for the given query.
-    The search query is formed by combining the website domain and the query string.
-    """
-    search_query = f"site:{website} {query}"
-    results = []
-    for result in search(search_query, num_results=10):
-        # Filter out irrelevant search pages
-        if "/search?num=" not in result:
-            results.append(result)
-    return results
 
 if __name__ == "__main__":
     # Define the tool metadata for orchestration.
-    tools = [
-        {
-            'type': 'function',
-            'function': {
-                'name': 'web_search',
-                'description': 'Search for results on a specified website using a query string. '
-                               'The CEO model should define which website to search from and the query to use.',
-                'parameters': {
-                    'type': 'object',
-                    'required': ['website', 'query'],
-                    'properties': {
-                        'website': {'type': 'string', 'description': 'The website domain to search from (e.g., huggingface.co)'},
-                        'query': {'type': 'string', 'description': 'The search query to use on the specified website'},
-                    },
-                },
-            },
-        }
-    ]
-
     # Load the tools using the ToolLoader class.
     tool_loader = ToolLoader()
 

requirements.txt

@@ -1,40 +1,7 @@
-annotated-types==0.7.0
-anyio==4.9.0
-beautifulsoup4==4.13.3
-cachetools==5.5.2
-certifi==2025.1.31
-charset-normalizer==3.4.1
-google-ai-generativelanguage==0.6.15
-google-api-core==2.24.2
-google-api-python-client==2.166.0
-google-auth==2.38.0
-google-auth-httplib2==0.2.0
-google-genai==1.9.0
-googleapis-common-protos==1.69.2
-grpcio==1.71.0
-grpcio-status==1.71.0
-h11==0.14.0
-httpcore==1.0.7
-httplib2==0.22.0
-httpx==0.28.1
-idna==3.10
+mistralai==1.6.0
 ollama==0.4.7
-pathlib==1.0.1
-proto-plus==1.26.1
+pip==25.0.1
 protobuf==5.29.4
-pyasn1==0.6.1
-pyasn1_modules==0.4.2
-pydantic==2.11.1
-pydantic_core==2.33.0
-pyparsing==3.2.3
+pydantic==2.11.2
 python-dotenv==1.1.0
-requests==2.32.3
-rsa==4.9
-sniffio==1.3.1
-soupsieve==2.6
-tqdm==4.67.1
-typing-inspection==0.4.0
-typing_extensions==4.13.0
-uritemplate==4.1.1
-urllib3==2.3.0
-websockets==15.0.1
+google-genai==1.9.0

tools/ask_user.py

@@ -1,4 +1,3 @@
-import importlib
 
 __all__ = ['AskUser']
 
@@ -24,8 +23,7 @@ class AskUser():
     def run(self, **kwargs):
         print("Running Ask User tool")
         question = kwargs.get("question")
-        print(f"Question: {question}")
-        output = input(question)
+        output = input(f"{question}: ")
         return {
             "status": "success",
             "message": "Ask User tool executed successfully",

tools/get_agents_tool.py

@@ -1,4 +1,3 @@
-import importlib
 import json
 
 __all__ = ['GetAgents']

tools/list_files.py

@@ -1,7 +1,6 @@
 
 __all__ = ['ListFiles']
 
-
 class ListFiles():
     dependencies = []
 

tools/read_file.py

@@ -1,4 +1,3 @@
-import importlib
 
 __all__ = ['ReadFile']
 

tools/tool_creator.py

@@ -1,4 +1,3 @@
-import importlib
 
 __all__ = ['ToolCreator']