Updated to use pro and fixed retry and edit

main.py

@@ -1,24 +1,63 @@
 from google.genai import types
-from src.CEO import GeminiManager
+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")
+    model_manager = GeminiManager(toolsLoader=tool_loader, gemini_model="gemini-2.5-pro-preview-03-25")
+    
+    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
+
+    def handle_undo(history, undo_data: gr.UndoData):
+        return history[:undo_data.index], history[undo_data.index]['content']
     
-    task_prompt = (
-        "What is the peak price of trump coin in the last 30 days? "
-        "Please provide the price in USD. "
-    )
+    def handle_retry(history, retry_data: gr.RetryData):
+        new_history = history[:retry_data.index+1]
+        # yield new_history, gr.update(interactive=False,)
+        yield from model_manager.run(new_history)
+        
+    def handle_edit(history, edit_data: gr.EditData):
+        new_history = history[:edit_data.index+1]
+        new_history[-1]['content'] = edit_data.value
+        # yield new_history, gr.update(interactive=False,)
+        yield from model_manager.run(new_history)
     
-    # Request a CEO response with the prompt.
-    # user_prompt_content = types.Content(
-    #     role='user',
-    #     parts=[types.Part.from_text(text=task_prompt)],
-    # )
-    # response = model_manager.request([user_prompt_content])
-    response = model_manager.start_conversation()
+    with gr.Blocks(fill_width=True, fill_height=True) as demo:
+        gr.Markdown("# Hashiru AI")
+        
+        chatbot = gr.Chatbot(
+            avatar_images=("HASHIRU_2.png", "HASHIRU.png"),
+            type="messages",
+            show_copy_button=True,
+            editable="user",
+            scale=1
+        )
+        input_box = gr.Textbox(label="Chat Message", scale=0, interactive=True, submit_btn=True)
+        
+        chatbot.undo(handle_undo, chatbot, [chatbot, input_box])
+        chatbot.retry(handle_retry, chatbot, [chatbot, input_box])
+        chatbot.edit(handle_edit, chatbot, [chatbot, input_box])
+        
+        input_box.submit(
+            user_message,  # Add user message to chat
+            inputs=[input_box, chatbot],
+            outputs=[input_box, chatbot],
+            queue=False,
+        ).then(
+            model_manager.run,  # Generate and stream response
+            inputs=chatbot,
+            outputs=[chatbot, input_box],
+            queue=True,
+            show_progress="full",
+            trigger_mode="always_last"
+        )
+
+    demo.launch(share=True)

mainV2.py

@@ -1,61 +0,0 @@
-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 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
-
-    def handle_undo(history, undo_data: gr.UndoData):
-        return history[:undo_data.index], history[undo_data.index]['content']
-    
-    def handle_retry(history, retry_data: gr.RetryData):
-        new_history = history[:retry_data.index]
-        yield from model_manager.run(new_history)
-        
-    def handle_edit(history, edit_data: gr.EditData):
-        new_history = history[:edit_data.index]
-        new_history[-1]['content'] = edit_data.value
-        return new_history
-    
-    with gr.Blocks(fill_width=True, fill_height=True) as demo:
-        gr.Markdown("# Hashiru AI")
-        
-        chatbot = gr.Chatbot(
-            avatar_images=("HASHIRU_2.png", "HASHIRU.png"),
-            type="messages",
-            show_copy_button=True,
-            editable="user",
-            scale=1
-        )
-        input_box = gr.Textbox(max_lines=5, label="Chat Message", scale=0)
-        
-        chatbot.undo(handle_undo, chatbot, [chatbot, input_box])
-        chatbot.retry(handle_retry, chatbot, chatbot)
-        chatbot.edit(handle_edit, chatbot, chatbot)
-        
-        input_box.submit(
-            user_message,  # Add user message to chat
-            inputs=[input_box, chatbot],
-            outputs=[input_box, chatbot],
-            queue=False,
-        ).then(
-            model_manager.run,  # Generate and stream response
-            inputs=chatbot,
-            outputs=[chatbot, input_box],
-            queue=True,
-            show_progress="full",
-            trigger_mode="always_last"
-        )
-
-    demo.launch(share=True)

src/CEO.py

@@ -1,134 +0,0 @@
-from google import genai
-from google.genai import types
-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 src.utils.streamlit_interface import get_user_message, output_assistant_response
-
-logger = logging.getLogger(__name__)
-handler = logging.StreamHandler(sys.stdout)
-handler.setLevel(logging.INFO)
-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='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(),
-                ),
-            )
-    
-    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 types.Content(
-            role='model' if self.model_name == "gemini-2.5-pro-exp-03-25" else 'tool',
-            parts=parts
-        )
-    
-    def run(self, messages):
-        try:
-            response = suppress_output(self.generate_response)(messages)
-        except Exception as e:
-            logger.debug(f"Error generating response: {e}")
-            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:
-                output_assistant_response("Ending the conversation.")
-                return messages
-
-        logger.debug(f"Response: {response}")
-        
-        if (not response.text and not response.function_calls):
-            output_assistant_response("No response from the model.")
-        
-        # Attach the llm response to the messages
-        if response.text is not None:
-            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)],
-            )
-            messages.append(assistant_content)
-            
-        # 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)
-            
-        # 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 = get_user_message("Should I continue? (y/n): ")
-            if shouldContinue.lower() == "y":
-                return self.run(messages)
-            else:
-                output_assistant_response("Ending the conversation.")
-                return messages
-        else:
-            logger.debug("No tool calls found in the response.")
-            # Start the loop again
-            return self.start_conversation(messages)
-    
-    def start_conversation(self, messages=[]):
-        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',
-            parts=[types.Part.from_text(text=question)],
-        )
-        messages.append(user_content)
-        
-        # Start the conversation loop
-        return self.run(messages)