Fixed streaming

mainV2.py

@@ -11,48 +11,51 @@ if __name__ == "__main__":
 
     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
+
+    def handle_undo(history, undo_data: gr.UndoData):
+        return history[:undo_data.index], history[undo_data.index]['content']
     
-    with gr.Blocks() as demo:
+    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"
+            type="messages",
+            show_copy_button=True,
+            editable="user",
+            scale=1
         )
-        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 = gr.Textbox(submit_btn=True, stop_btn=True, 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(
-            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],
+            inputs=[input_box, chatbot],
             outputs=[input_box, chatbot],
-            queue=False
+            queue=False,
         ).then(
             model_manager.run,  # Generate and stream response
-            inputs=[msg_store, chatbot],
-            outputs=chatbot
+            inputs=chatbot,
+            outputs=chatbot,
+            show_progress="full",
+            trigger_mode="always_last"
         )
+        input_box.submit(lambda: "", None, [input_box])
 
     demo.launch(share=True)

models/system3.prompt

@@ -103,7 +103,7 @@ Strictly follow the schema required for invoking the tools and agents. Do not de
 </Rule>
 
 <Rule>
-Once you have the answer, provide it to the user in a clear and concise manner ending with a "EOF" message.
+Once you have the answer, provide it to the user in a clear and concise manner.
 </Rule>
 
 <Rule>

src/manager.py

@@ -16,6 +16,7 @@ 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()
@@ -30,22 +31,24 @@ class GeminiManager:
 
     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(),
-                ),
-            )
-    
+            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}")
+            logger.info(
+                f"Function Name: {function_call.name}, Arguments: {function_call.args}")
             try:
-                toolResponse = self.toolsLoader.runTool(function_call.name, function_call.args)
+                toolResponse = self.toolsLoader.runTool(
+                    function_call.name, function_call.args)
             except Exception as e:
                 logger.warning(f"Error running tool: {e}")
                 toolResponse = {
@@ -55,53 +58,57 @@ class GeminiManager:
                 }
             logger.debug(f"Tool Response: {toolResponse}")
             tool_content = types.Part.from_function_response(
-                    name=function_call.name,
-                    response = {"result":toolResponse})
+                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'])
+                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)})
+                    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": "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", ""))]
+                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", "")))
+                        formatted_history.append(
+                            eval(message.get("content", "")))
                         continue
                     case "function_call":
                         role = "model"
-                        formatted_history.append(eval(message.get("content", "")))
+                        formatted_history.append(
+                            eval(message.get("content", "")))
                         continue
                     case _:
-                        role = "model" 
+                        role = "model"
                 formatted_history.append(types.Content(
                     role=role,
                     parts=parts
                 ))
         return formatted_history
-    
-    def run(self, message, messages):
+
+    def run(self, messages):
+        print("Messages: ", messages)
         chat_history = self.format_chat_history(messages)
         logger.debug(f"Chat history: {chat_history}")
         try:
@@ -109,18 +116,19 @@ class GeminiManager:
         except Exception as e:
             logger.debug(f"Error generating response: {e}")
             messages.append({
-                "role":"assistant",
-                "content":f"Error generating response: {e}"
+                "role": "assistant",
+                "content": f"Error generating response: {e}"
             })
             logger.error(f"Error generating response: {e}")
             return messages
         logger.debug(f"Response: {response}")
-        
+        print("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."}
+                "role": "assistant",
+                "content": "No response from the model.",
+                "metadata": {"title": "No response from the model."}
             })
 
         # Attach the llm response to the messages
@@ -130,22 +138,19 @@ class GeminiManager:
                 "content": response.text
             })
             yield messages
-            
+
         # 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",
+                "role": "function_call",
                 "content": repr(response.candidates[0].content),
             })
-            yield messages
-        
+
         # 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)
-            yield messages
-            for value in self.run(message, messages):
-                yield value
-            return
-        return messages
+            yield from self.run(messages)
+        print("Final messages: ", messages)
+        return messages