History and multiple conversations work!

main.py

@@ -1,108 +1,78 @@
-from google.genai import types
-from src.manager import GeminiManager
-from src.tool_manager import ToolManager
 import gradio as gr
-import time
+
 import base64
+from src.manager import GeminiManager
+
+model_manager = GeminiManager(gemini_model="gemini-2.0-flash")
 
 _logo_bytes = open("HASHIRU_LOGO.png", "rb").read()
 _logo_b64 = base64.b64encode(_logo_bytes).decode()
 _header_html = f"""
 <div style="
     display: flex;
-    flex-direction: column;
+    flex-direction: row;
     align-items: center;
-    padding-right: 24px;
+    justify-content: flex-start;
+    width: 30%;
 ">
-  <img src="data:image/png;base64,{_logo_b64}" width="20" height="20" />
-  <span style="margin-top: 8px; font-size: 20px; font-weight: bold; color: white;">
+  <img src="data:image/png;base64,{_logo_b64}" width="40" class="logo"/>
+  <h1>
     HASHIRU AI
-  </span>
+  </h1>
 </div>
 """
+css = """
+    .logo { margin-right: 20px; }
+    """
 
-if __name__ == "__main__":
-    model_manager = GeminiManager(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+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)
+def run_model(message, history):
+    history.append({
+        "role": "user",
+        "content": message,
+    })
+    yield "", history
+    for messages in model_manager.run(history):
+        for message in messages:
+            if message.get("role") == "summary":
+                print(f"Summary: {message.get('content', '')}")
+        yield "", messages
+
+
+def update_model(model_name):
+    print(f"Model changed to: {model_name}")
+    pass
 
-    def update_model(model_name):
-        print(f"Model changed to: {model_name}")
-        pass
-    
-    css = """
-    #title-row { background: #2c2c2c; border-radius: 8px; padding: 8px; }
-    """
-    with gr.Blocks(css=css, fill_width=True, fill_height=True) as demo:
-        local_storage = gr.BrowserState(["", ""])
-        with gr.Column(scale=1):
-            with gr.Row(scale=0):
-                gr.HTML(_header_html)
-                model_dropdown = gr.Dropdown(
-                        choices=[
-                            "HASHIRU",
-                            "Static-HASHIRU",
-                            "Cloud-Only HASHIRU",
-                            "Local-Only HASHIRU",
-                            "No-Economy HASHIRU",
-                        ],
-                        value="HASHIRU",
-                        # label="HASHIRU",
-                        interactive=True,
-                )
 
-                model_dropdown.change(fn=update_model, inputs=model_dropdown, outputs=[])
-            with gr.Row(scale=1):
-                with gr.Sidebar(position="left"):
-                    buttons = []
-                    for i in range(1, 6):
-                        button = gr.Button(f"Button {i}", elem_id=f"button-{i}")
-                        button.click(fn=lambda x=i: print(f"Button {x} clicked"), inputs=[], outputs=[])
-                        buttons.append(button)
-                with gr.Column(scale=1):
-                    chatbot = gr.Chatbot(
-                        avatar_images=("HASHIRU_2.png", "HASHIRU.png"),
-                        type="messages",
-                        show_copy_button=True,
-                        editable="user",
-                        scale=1,
-                        render_markdown=True,
-                    )
-                    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.ask_llm,  # Generate and stream response
-                        inputs=chatbot,
-                        outputs=[chatbot, input_box],
-                        queue=True,
-                        show_progress="full",
-                        trigger_mode="always_last"
-                    )
-    
+with gr.Blocks(css=css, fill_width=True, fill_height=True) as demo:
+    with gr.Column(scale=1):
+        with gr.Row(scale=0):
+            gr.Markdown(_header_html)
+            model_dropdown = gr.Dropdown(
+                choices=[
+                    "HASHIRU",
+                    "Static-HASHIRU",
+                    "Cloud-Only HASHIRU",
+                    "Local-Only HASHIRU",
+                    "No-Economy HASHIRU",
+                ],
+                value="HASHIRU",
+                interactive=True,
+            )
+
+            model_dropdown.change(
+                fn=update_model, inputs=model_dropdown, outputs=[])
+        with gr.Row(scale=1):
+            chatbot = gr.Chatbot(
+                avatar_images=("HASHIRU_2.png", "HASHIRU.png"),
+                type="messages",
+                show_copy_button=True,
+                editable="user",
+                scale=1,
+                render_markdown=True,
+                placeholder="Type your message here...",
+            )
+            gr.ChatInterface(fn=run_model, type="messages", chatbot=chatbot,
+                             additional_outputs=[chatbot], save_history=True)
+if __name__ == "__main__":
     demo.launch()

models/system3.prompt

@@ -33,6 +33,10 @@ There is a strict resource constraint (budget) you need to follow. You start wit
 If you're over this budget, you can no longer create new tools. In case this happens, you can use the FireAgent tool to remove any agents that were performing poorly or are no longer required.
 </Info>
 
+<Info>
+Once the you're finished with a conversation, you can use the DigestConversation to clear the conversation history and create a summary of the interaction. This summary can be used to improve the system's performance and user experience in future interactions.
+</Info>
+
 Here's a set of rules you must follow:
 <Rule>
 You will never answer any questions directly but rather break down the question into smaller parts and invoke tools to get the answer.
@@ -116,4 +120,8 @@ If you think there are multiple paths to proceed, ask the user on which path to
 
 <Rule>
 When you go over the resource budget, you must carefully evaluate which agent is least likely to be useful going forward before firing it. Only fire an agent to create a new one when absolutely necessary.
+</Rule>
+
+<Rule>
+When using the DigestConversation tool, ensure that the summary is clear and concise, capturing the main points of the conversation without unnecessary details. Include any important information such as how errors were solved, user's preferences, and any other relevant context that could help in future interactions.
 </Rule>

src/manager.py

@@ -40,13 +40,31 @@ class GeminiManager:
         self.messages = []
         
     def generate_response(self, messages):
+        tools = self.toolsLoader.getTools()
+        function = types.FunctionDeclaration(
+                name="DigestConversation",
+                description="Digest the conversation and store the summary provided.",
+                parameters=types.Schema(
+                    type = "object",
+                    properties={
+                        # string that summarizes the conversation
+                        "summary": types.Schema(
+                            type="string",
+                            description="A summary of the conversation including all the important points.",
+                        ),
+                    },
+                    required=["summary"],
+                ),
+            )
+        toolType = types.Tool(function_declarations=[function])
+        tools.append(toolType)
         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(),
+                tools=tools,
             ),
         )
 
@@ -56,6 +74,13 @@ class GeminiManager:
             toolResponse = None
             logger.info(
                 f"Function Name: {function_call.name}, Arguments: {function_call.args}")
+            if function_call.name == "DigestConversation":
+                logger.info("Digesting conversation...")
+                summary = function_call.args["summary"]
+                return {
+                    "role": "summary",
+                    "content": f"{summary}",
+                }
             try:
                 toolResponse = self.toolsLoader.runTool(
                     function_call.name, function_call.args)
@@ -116,9 +141,6 @@ class GeminiManager:
                     parts=parts
                 ))
         return formatted_history
-    
-    def ask_llm(self, messages):
-        yield from self.run(messages)
         
     def run(self, messages):
         chat_history = self.format_chat_history(messages)
@@ -132,7 +154,7 @@ class GeminiManager:
                 "content": f"Error generating response: {e}"
             })
             logger.error(f"Error generating response: {e}")
-            yield messages, gr.update(interactive=True)
+            yield messages
             return
         logger.debug(f"Response: {response}")
 
@@ -149,7 +171,7 @@ class GeminiManager:
                 "role": "assistant",
                 "content": response.text
             })
-            yield messages, gr.update(interactive=False,)
+            yield messages
 
         # Attach the function call response to the messages
         if response.candidates[0].content and response.candidates[0].content.parts:
@@ -165,4 +187,4 @@ class GeminiManager:
             messages.append(calls)
             yield from self.run(messages)
             return
-        yield messages, gr.update(interactive=True)
+        yield messages