mito-ai 0.1.45__py3-none-any.whl → 0.1.47__py3-none-any.whl

mito_ai/completions/handlers.py

@@ -49,11 +49,8 @@ from mito_ai.utils.telemetry_utils import identify
 
 FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
 
-# The GlobalMessageHistory is responsible for updating the message histories stored in the .mito/ai-chats directory.
-# We create one GlobalMessageHistory per backend server instance instead of one per websocket connection so that the
-# there is one manager of the locks for the .mito/ai-chats directory. This is my current understanding and it 
-# might be incorrect!
-message_history = GlobalMessageHistory()
+# The GlobalMessageHistory is now created in __init__.py and passed to handlers
+# to ensure there's only one instance managing the .mito/ai-chats directory locks
 
 # This handler is responsible for the mito_ai/completions endpoint.
 # It takes a message from the user, sends it to the OpenAI API, and returns the response.
@@ -62,10 +59,11 @@ message_history = GlobalMessageHistory()
 class CompletionHandler(JupyterHandler, WebSocketHandler):
     """Completion websocket handler."""
 
-    def initialize(self, llm: OpenAIProvider) -> None:
+    def initialize(self, llm: OpenAIProvider, message_history: GlobalMessageHistory) -> None:
         super().initialize()
         self.log.debug("Initializing websocket connection %s", self.request.path)
         self._llm = llm
+        self._message_history = message_history
         self.is_pro = is_pro()
         self._selected_model = FALLBACK_MODEL
         self.is_electron = False
@@ -150,7 +148,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
 
         # Clear history if the type is "start_new_chat"
         if type == MessageType.START_NEW_CHAT:
-            thread_id = message_history.create_new_thread()
+            thread_id = self._message_history.create_new_thread()
             
             reply = StartNewChatReply(
                 parent_id=parsed_message.get("message_id"),
@@ -161,7 +159,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
 
         # Handle get_threads: return list of chat threads
         if type == MessageType.GET_THREADS:
-            threads = message_history.get_threads()
+            threads = self._message_history.get_threads()
             reply = FetchThreadsReply(
                 parent_id=parsed_message.get("message_id"),
                 threads=threads
@@ -173,7 +171,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
         if type == MessageType.DELETE_THREAD:
             thread_id_to_delete = metadata_dict.get('thread_id')
             if thread_id_to_delete:
-                is_thread_deleted = message_history.delete_thread(thread_id_to_delete)
+                is_thread_deleted = self._message_history.delete_thread(thread_id_to_delete)
                 reply = DeleteThreadReply(
                     parent_id=parsed_message.get("message_id"),
                     success=is_thread_deleted
@@ -189,7 +187,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             
             # If a thread_id is provided, use that thread's history; otherwise, use newest.
             thread_id = metadata_dict.get('thread_id')
-            display_history = message_history.get_display_history(thread_id)
+            display_history = self._message_history.get_display_history(thread_id)
             
             reply = FetchHistoryReply(
                 parent_id=parsed_message.get('message_id'),
@@ -238,7 +236,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     "content": "Agent interupted by user "
                 }
                 
-                await message_history.append_message(
+                await self._message_history.append_message(
                     ai_optimized_message=ai_optimized_message,
                     display_message=display_optimized_message,
                     model=self._selected_model,
@@ -266,7 +264,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_chat_completion(
                         chat_metadata, 
                         self._llm, 
-                        message_history, 
+                        self._message_history, 
                         message_id,
                         self.reply,
                         model
@@ -274,7 +272,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_chat_completion(chat_metadata, self._llm, message_history, model)
+                    completion = await get_chat_completion(chat_metadata, self._llm, self._message_history, model)
             elif type == MessageType.SMART_DEBUG:
                 smart_debug_metadata = SmartDebugMetadata(**metadata_dict)
                 # Handle streaming if requested and available
@@ -283,7 +281,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_smart_debug_completion(
                         smart_debug_metadata, 
                         self._llm, 
-                        message_history, 
+                        self._message_history, 
                         message_id,
                         self.reply,
                         model
@@ -291,7 +289,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, message_history, model)
+                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, self._message_history, model)
             elif type == MessageType.CODE_EXPLAIN:
                 code_explain_metadata = CodeExplainMetadata(**metadata_dict)
 
@@ -301,7 +299,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_code_explain_completion(
                         code_explain_metadata, 
                         self._llm, 
-                        message_history,
+                        self._message_history,
                         message_id,
                         self.reply,
                         model
@@ -309,16 +307,16 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, message_history, model)
+                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, self._message_history, model)
             elif type == MessageType.AGENT_EXECUTION:
                 agent_execution_metadata = AgentExecutionMetadata(**metadata_dict)
-                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, message_history, model)
+                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, self._message_history, model)
             elif type == MessageType.AGENT_AUTO_ERROR_FIXUP:
                 agent_auto_error_fixup_metadata = AgentSmartDebugMetadata(**metadata_dict)
-                completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_metadata, self._llm, message_history, model)
+                completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_metadata, self._llm, self._message_history, model)
             elif type == MessageType.INLINE_COMPLETION:
                 inline_completer_metadata = InlineCompleterMetadata(**metadata_dict)
-                completion = await get_inline_completion(inline_completer_metadata, self._llm, message_history, model)
+                completion = await get_inline_completion(inline_completer_metadata, self._llm, self._message_history, model)
             else:
                 raise ValueError(f"Invalid message type: {type}")
             

mito_ai/completions/models.py

@@ -29,12 +29,13 @@ class CellUpdate(BaseModel):
 # for now and rely on the AI to respond with the correct types, following the format
 # that we show it in the system prompt.
 class AgentResponse(BaseModel):
-    type: Literal['cell_update', 'get_cell_output', 'run_all_cells', 'finished_task']
+    type: Literal['cell_update', 'get_cell_output', 'run_all_cells', 'finished_task', 'create_streamlit_app', 'edit_streamlit_app']
     message: str
     cell_update: Optional[CellUpdate]
     get_cell_output_cell_id: Optional[str]
     next_steps: Optional[List[str]]
     analysis_assumptions: Optional[List[str]]
+    edit_streamlit_app_prompt: Optional[str]
     
     
 @dataclass(frozen=True)
@@ -65,6 +66,7 @@ class MessageType(Enum):
     UPDATE_MODEL_CONFIG = "update_model_config"
     STREAMLIT_CONVERSION = "streamlit_conversion"
     STOP_AGENT = "stop_agent"
+    DEPLOY_APP = "deploy_app"
 
     
 @dataclass(frozen=True)
@@ -101,6 +103,7 @@ class AgentExecutionMetadata():
     files: Optional[List[str]] = None
     index: Optional[int] = None
     additionalContext: Optional[List[Dict[str, str]]] = None
+    streamlitAppIsOpen: Optional[bool] = None
     
 @dataclass(frozen=True)
 class AgentSmartDebugMetadata():

mito_ai/completions/prompt_builders/agent_execution_prompt.py

@@ -5,12 +5,14 @@ from mito_ai.completions.models import AgentExecutionMetadata
 from mito_ai.completions.prompt_builders.prompt_constants import (
     FILES_SECTION_HEADING,
     JUPYTER_NOTEBOOK_SECTION_HEADING,
+    STREAMLIT_APP_STATUS_SECTION_HEADING,
     VARIABLES_SECTION_HEADING,
     cell_update_output_str
 )
 from mito_ai.completions.prompt_builders.utils import (
     get_rules_str,
     get_selected_context_str,
+    get_streamlit_app_status_str,
 )
 
 
@@ -20,7 +22,7 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
     ai_optimized_cells_str = '\n'.join([f"{cell}" for cell in md.aiOptimizedCells or []])
     rules_str = get_rules_str(md.additionalContext)
     selected_context_str = get_selected_context_str(md.additionalContext)
-        
+    streamlit_status_str = get_streamlit_app_status_str(md.streamlitAppIsOpen)
     context_str = f"""Remember to choose the correct tool to respond with.
 
 {rules_str}
@@ -35,6 +37,9 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
 {FILES_SECTION_HEADING}
 {files_str}
 
+{STREAMLIT_APP_STATUS_SECTION_HEADING}
+{streamlit_status_str}
+
 {selected_context_str}
 
 {cell_update_output_str(md.base64EncodedActiveCellOutput is not None)}"""

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -225,6 +225,57 @@ Important information:
 5. Do not use this tool repeatedly if it continues to produce errors - instead, focus on fixing the specific error that occurred.
 ====
 
+TOOL: CREATE_STREAMLIT_APP
+
+When you want to create a new Streamlit app from the current notebook, respond with this format:
+
+{{
+    type: 'create_streamlit_app',
+    message: str
+}}
+
+Important information:
+1. The message is a short summary of why you're creating the Streamlit app.
+2. Only use this tool when the user explicitly asks to create or preview a Streamlit app AND no Streamlit app is currently open.
+3. This tool creates a new app from scratch - use EDIT_STREAMLIT_APP tool if the user is asking you to edit, update, or modify an app that already exists.
+4. Using this tool will automatically open the app so the user can see a preview of the app.
+5. When you use this tool, assume that it successfully created the Streamlit app unless the user explicitly tells you otherwise. The app will remain open so that the user can view it until the user decides to close it. You do not need to continually use the create_streamlit_app tool to keep the app open.
+
+<Example>
+
+Your task: Show me my notebook as an app.
+
+Output:
+{{
+    type: 'create_streamlit_app',
+    message: "I'll convert your notebook into an app."
+}}
+
+The user will see a preview of the app and because you fulfilled your task, you can next respond with a FINISHED_TASK tool message.
+
+<Example>
+
+====
+
+TOOL: EDIT_STREAMLIT_APP
+
+When you want to edit an existing Streamlit app, respond with this format:
+
+{{
+    type: 'edit_streamlit_app',
+    message: str,
+    edit_streamlit_app_prompt: str
+}}
+
+Important information:
+1. The message is a short summary of why you're editing the Streamlit app.
+2. The edit_streamlit_app_prompt is REQUIRED and must contain specific instructions for the edit (e.g., "Make the title text larger", "Change the chart colors to blue", "Add a sidebar with filters").
+3. Only use this tool when the user asks to edit, update, or modify a Streamlit app. 
+4. The app does not need to already be open for you to use the tool. Using this tool will automatically open the streamlit app after applying the changes so the user can view it. You do not need to call the create_streamlit_app tool first.
+5. When you use this tool, assume that it successfully edited the Streamlit app unless the user explicitly tells you otherwise. The app will remain open so that the user can view it until the user decides to close it. 
+
+====
+
 TOOL: FINISHED_TASK
 
 When you have completed the user's task, respond with a message in this format:
@@ -238,8 +289,8 @@ When you have completed the user's task, respond with a message in this format:
 Important information:
 1. The message is a short summary of the ALL the work that you've completed on this task. It should not just refer to the final message. It could be something like "I've completed the sales strategy analysis by exploring key relationships in the data and summarizing creating a report with three recommendations to boost sales.""
 2. The message should include citations for any insights that you shared with the user.
-3. The next_steps is an optional list of 2 or 3 suggested follow-up tasks or analyses that the user might want to perform next. These should be concise, actionable suggestions that build on the work you've just completed. For example: ["Visualize the results", "Export the cleaned data to CSV", "Perform statistical analysis on the key metrics"].
-4. The next_steps should be as relevant to the user's actual task as possible. Try your best not to make generic suggestions like "Analyze the data" or "Visualize the results". For example, if the user just asked you to calculate LTV of their customers, you might suggest the following next steps: ["Graph key LTV drivers: churn and average transaction value", "Visualize LTV per customer age group"].
+3. The next_steps is an optional list of 2 or 3 suggested follow-up tasks or analyses that the user might want to perform next. These should be concise, actionable suggestions that build on the work you've just completed. For example: ["Export the cleaned data to CSV", "Analyze revenue per customer", "Convert notebook into an app"].
+4. The next_steps should be as relevant to the user's actual task as possible. Try your best not to make generic suggestions like "Analyze the data" or "Visualize the results". For example, if the user just asked you to calculate LTV of their customers, you might suggest the following next steps: ["Graph key LTV drivers: churn and average transaction value", "Visualize LTV per age group"].
 5. If you are not sure what the user might want to do next, err on the side of suggesting next steps instead of making an assumption and using more CELL_UPDATES.
 6. If the user's task doesn't involve creating or modifying a code cell, you should respond with a FINISHED_TASK response. 
 7. If the user is just sending a friendly greeting (like "Hello", "Hi", "Hey", "How are you?", "What can you help me with?", etc.), you must respond with a FINISHED_TASK response message with a friendly message like this: "Hello! I'm Mito AI, your AI assistant for data analysis and Python programming in Jupyter notebooks. I can help you analyze datasets, create visualizations, clean data, and much more. What would you like to work on today?"
@@ -391,7 +442,9 @@ As you are guiding the user through the process of completing the task, send the
 
 The user is a beginning Python user, so you will need to be careful to send them only small steps to complete. Don't try to complete the task in a single response to the user. Instead, each message you send to the user should only contain a single, small step towards the end goal. When the user has completed the step, they will let you know that they are ready for the next step. 
 
-You will keep working in the following iterative format until you have decided that you have finished the user's request. When you decide that you have finished the user's request, respond with a FINISHED_TASK tool message. Otherwise, if you have not finished the user's request, respond with a CELL_UPDATE {OR_GET_CELL_OUTPUT} tool message. When you respond with a CELL_UPDATE, the user will apply the CELL_UPDATE to the notebook and run the new code cell. The user will then send you a message with an updated version of the variables defined in the kernel, code in the notebook, and files in the current directory. In addition, the user will check if the code you provided produced an errored when executed. If it did produce an error, the user will share the error message with you.
+You will keep working in the following iterative format until you have decided that you have finished the user's request. When you decide that you have finished the user's request, respond with a FINISHED_TASK tool message. Otherwise, if you have not finished the user's request, respond with one of your other tools. 
+
+When you respond with a CELL_UPDATE, the user will apply the CELL_UPDATE to the notebook and run the new code cell. The user will then send you a message with an updated version of the variables defined in the kernel, code in the notebook, and files in the current directory. In addition, the user will check if the code you provided produced an errored when executed. If it did produce an error, the user will share the error message with you.
 
 Whenever you get a message back from the user, you should:
 1. Ask yourself if the previous message you sent to the user was correct. You can answer this question by reviewing the updated code, variables, or output of the cell if you requested it.
@@ -410,4 +463,10 @@ REMEMBER, YOU ARE GOING TO COMPLETE THE USER'S TASK OVER THE COURSE OF THE ENTIR
 - If you are happy with the analysis, refer back to the original task provided by the user to decide your next steps. In this example, it is to graph the results, so you will send a CellAddition to construct the graph. 
 - Wait for the user to send you back the updated variables and notebook state.
 {'' if not isChromeBrowser else '- Send a GET_CELL_OUTPUT tool message to get the output of the cell you just created and check if you can improve the graph to make it more readable, informative, or professional.'}
-- If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK tool message."""
+- If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK tool message.
+
+====
+
+OTHER USEFUL INFORMATION:
+1. When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter
+"""

mito_ai/completions/prompt_builders/chat_system_message.py

@@ -21,6 +21,10 @@ There are three possible types of responses you might give:
 2. Explanation/Analysis: If the task does not require a code update, it might instead require you to provide an explanation of existing code or data, provide an analysis of the the data or chart.
 3. Friendly Response: If the user is just asking a question, saying hi, or you're just chatting, respond with a friendly response and do not return any code.
 
+Other useful information:
+1. The user has two types of modes that they can collaborate with you in: Chat Mode (this mode) and agent mode. Chat mode gives the user more control over the edits made to the notebook and only edits the active cell. Agent mode gives you more autonomy over completing the user's task across mulitple messages. In agent mode, you can edit or create new cells, see the entire notebook, automatically run the code you write, and more.
+2. If the user asks you to generate a dashboard, app, or streamlit app for them, you should tell them that they must use Agent mode to complete the task. You are not able to automatically switch the user to agent mode, but they can switch to it themselves by using the Chat/Agent mode toggle in the bottom left corner of the Ai taskpane.
+
 ====
 {CITATION_RULES}
 

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -19,6 +19,7 @@ ACTIVE_CELL_ID_SECTION_HEADING = "The ID of the active code cell:"
 ACTIVE_CELL_OUTPUT_SECTION_HEADING = "Output of the active code cell:"
 GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING = "Output of the code cell you just applied the CELL_UPDATE to:"
 JUPYTER_NOTEBOOK_SECTION_HEADING = "Jupyter Notebook:"
+STREAMLIT_APP_STATUS_SECTION_HEADING = "Streamlit App Status:"
 
 # Placeholder text used when trimming content from messages
 CONTENT_REMOVED_PLACEHOLDER = "Content removed to save space" 

mito_ai/completions/prompt_builders/utils.py

@@ -69,3 +69,17 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
 
     # STEP 3: Combine into a single string
     return "\n\n".join(context_parts)
+
+
+def get_streamlit_app_status_str(streamlit_app_is_open: Optional[bool]) -> str:
+    """
+    Get the streamlit app status string.
+    """
+    if streamlit_app_is_open is None:
+        return ""
+    
+    if streamlit_app_is_open:
+        return "A Streamlit app is currently open and running."
+    else:
+        return "No Streamlit app is currently open."
+

mito_ai/constants.py

@@ -58,3 +58,6 @@ COGNITO_CONFIG_DEV = {
 }
 
 ACTIVE_COGNITO_CONFIG = COGNITO_CONFIG_DEV # Change to COGNITO_CONFIG_DEV for dev
+
+
+MESSAGE_HISTORY_TRIM_THRESHOLD: int = 3

mito_ai/path_utils.py

@@ -0,0 +1,56 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import NewType
+import os
+from mito_ai.utils.error_classes import StreamlitPreviewError
+
+# Type definitions for better type safety
+AbsoluteNotebookPath = NewType('AbsoluteNotebookPath', str)
+AbsoluteNotebookDirPath = NewType('AbsoluteNotebookDirPath', str)
+AbsoluteAppPath = NewType('AbsoluteAppPath', str)
+
+def get_absolute_notebook_path(notebook_path: str) -> AbsoluteNotebookPath:
+    """
+    Convert any notebook path to an absolute path.
+    
+    Args:
+        notebook_path: Path to the notebook (can be relative or absolute)
+        
+    Returns:
+        AbsoluteNotebookPath: The absolute path to the notebook
+        
+    Raises:
+        ValueError: If the path is invalid or empty
+    """
+    if not notebook_path or not notebook_path.strip():
+        raise StreamlitPreviewError("Notebook path cannot be empty", 400)
+    
+    absolute_path = os.path.abspath(notebook_path)
+    return AbsoluteNotebookPath(absolute_path)
+    
+
+def get_absolute_notebook_dir_path(notebook_path: AbsoluteNotebookPath) -> AbsoluteNotebookDirPath:
+    """
+    Get the directory containing the notebook.
+    
+    Args:
+        notebook_path: Absolute path to the notebook
+        
+    Returns:
+        AbsoluteNotebookDirPath: The directory containing the notebook
+    """
+    return AbsoluteNotebookDirPath(os.path.dirname(notebook_path))
+
+def get_absolute_app_path(app_directory: AbsoluteNotebookDirPath) -> AbsoluteAppPath:
+    """
+    Check if the app.py file exists in the given directory.
+    """
+    return AbsoluteAppPath(os.path.join(app_directory, "app.py"))
+
+def does_app_path_exists(app_path: AbsoluteAppPath) -> bool:
+    """
+    Check if the app.py file exists in the given directory.
+    """
+    return os.path.exists(app_path)
+

mito_ai/streamlit_conversion/agent_utils.py

@@ -1,116 +1,37 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from typing import List
+from typing import List, Tuple
 import re
-from unidiff import PatchSet
+from anthropic.types import MessageParam
+from mito_ai.streamlit_conversion.prompts.streamlit_system_prompt import streamlit_system_prompt
+from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
 from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER
+from mito_ai.completions.models import MessageType
+
+STREAMLIT_AI_MODEL = "claude-sonnet-4-5-20250929"
 
 def extract_todo_placeholders(agent_response: str) -> List[str]:
     """Extract TODO placeholders from the agent's response"""
     return [line.strip() for line in agent_response.split('\n') if MITO_TODO_PLACEHOLDER in line]
 
-
-def apply_patch_to_text(text: str, diff: str) -> str:
-    """
-    Apply a *unified-diff* (git-style) patch to the given text and return
-    the updated contents.
-
-    Parameters
-    ----------
-    text : str
-        The original file contents.
-    diff : str
-        A unified diff that transforms *text* into the desired output.
-        The diff must reference exactly one file (the Streamlit app).
-
-    Returns
-    -------
-    str
-        The patched contents.
-
-    Raises
-    ------
-    ValueError
-        If the patch cannot be applied or references more than one file.
-    """
-    # Nothing to do
-    if not diff.strip():
-        return text
-
-    # Parse the patch
-    patch = PatchSet(diff.splitlines(keepends=True))
-
-    # We expect a single-file patch (what the prompt asks the model to emit)
-    if len(patch) != 1:
-        raise ValueError(
-            f"Expected a patch for exactly one file, got {len(patch)} files."
-        )
-
-    file_patch = patch[0]
-
-    original_lines = text.splitlines(keepends=True)
-    result_lines: List[str] = []
-
-    cursor = 0  # index in original_lines (0-based)
-
-    for hunk in file_patch:
-        # Copy unchanged lines before this hunk
-        while cursor < hunk.source_start - 1:
-            result_lines.append(original_lines[cursor])
-            cursor += 1
-
-        # Apply hunk line-by-line
-        for line in hunk:
-            if line.is_context:
-                result_lines.append(original_lines[cursor])
-                cursor += 1
-            elif line.is_removed:
-                cursor += 1  # Skip this line from the original
-            elif line.is_added:
-                # Ensure added line ends with newline for consistency
-                val = line.value
-                if not val.endswith("\n"):
-                    val += "\n"
-                result_lines.append(val)
-
-    # Copy any remaining lines after the last hunk
-    result_lines.extend(original_lines[cursor:])
-
-    return "".join(result_lines)
-
-
-def fix_diff_headers(diff: str) -> str:
-    """
-    The AI is generally not very good at counting the number of lines in the diff. If the hunk header has
-    an incorrect count, then the patch will fail. So instead we just calculate the counts ourselves, its deterministic.
-    """
-    lines = diff.split('\n')
-    
-    for i, line in enumerate(lines):
-        if line.startswith('@@'):
-            # Extract the starting line numbers
-            match = re.match(r'@@ -(\d+),\d+ \+(\d+),\d+ @@', line)
-            if match:
-                old_start = match.group(1)
-                new_start = match.group(2)
-                
-                # Count lines in this hunk
-                old_count = 0
-                new_count = 0
-                
-                for j in range(i + 1, len(lines)):
-                    next_line = lines[j]
-                    if next_line.startswith('@@') or next_line.startswith('---'):
-                        break
-                    if next_line.startswith(' ') or next_line.startswith('-'):
-                        old_count += 1
-                    if next_line.startswith(' ') or next_line.startswith('+'):
-                        new_count += 1
-                
-                # Replace the header with correct counts
-                lines[i] = f"@@ -{old_start},{old_count} +{new_start},{new_count} @@"
-    
-    return '\n'.join(lines)
-
-
+async def get_response_from_agent(message_to_agent: List[MessageParam]) -> str:
+    """Gets the streaming response from the agent using the mito server"""
+    model = STREAMLIT_AI_MODEL
+    max_tokens = 64000 # TODO: If we move to haiku, we must reset this to 8192
+    temperature = 0.2
+
+    accumulated_response = ""
+    async for stream_chunk in stream_anthropic_completion_from_mito_server(
+        model = model,
+        max_tokens = max_tokens,
+        temperature = temperature,
+        system = streamlit_system_prompt,
+        messages = message_to_agent,
+        stream=True,
+        message_type=MessageType.STREAMLIT_CONVERSION,
+        reply_fn=None,
+        message_id=""
+    ):
+        accumulated_response += stream_chunk
+    return accumulated_response