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

mito_ai/__init__.py

@@ -5,6 +5,7 @@ from typing import List, Dict
 from jupyter_server.utils import url_path_join
 from mito_ai.completions.handlers import CompletionHandler
 from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.app_deploy.handlers import AppDeployHandler
 from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
 from mito_ai.log.urls import get_log_urls
@@ -16,6 +17,8 @@ from mito_ai.auth.urls import get_auth_urls
 from mito_ai.streamlit_preview.urls import get_streamlit_preview_urls
 from mito_ai.app_manager.handlers import AppManagerHandler
 from mito_ai.file_uploads.urls import get_file_uploads_urls
+from mito_ai.user.urls import get_user_urls
+from mito_ai.chat_history.urls import get_chat_history_urls
 
 # Force Matplotlib to use the Jupyter inline backend.
 # Background: importing Streamlit sets os.environ["MPLBACKEND"] = "Agg" very early.
@@ -62,13 +65,17 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     base_url = web_app.settings["base_url"]
 
     open_ai_provider = OpenAIProvider(config=server_app.config)
+    
+    # Create a single GlobalMessageHistory instance for the entire server
+    # This ensures thread-safe access to the .mito/ai-chats directory
+    global_message_history = GlobalMessageHistory()
 
     # WebSocket handlers
     handlers = [
         (
             url_path_join(base_url, "mito-ai", "completions"),
             CompletionHandler,
-            {"llm": open_ai_provider},
+            {"llm": open_ai_provider, "message_history": global_message_history},
         ),
         (
             url_path_join(base_url, "mito-ai", "app-deploy"),
@@ -100,6 +107,8 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     handlers.extend(get_auth_urls(base_url))  # type: ignore
     handlers.extend(get_streamlit_preview_urls(base_url))  # type: ignore
     handlers.extend(get_file_uploads_urls(base_url)) # type: ignore
+    handlers.extend(get_user_urls(base_url)) # type: ignore
+    handlers.extend(get_chat_history_urls(base_url, global_message_history)) # type: ignore
 
     web_app.add_handlers(host_pattern, handlers)
     server_app.log.info("Loaded the mito_ai server extension")

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.45'
+__version__ = VERSION = '0.1.46'

mito_ai/anthropic_client.py

@@ -5,9 +5,9 @@ import json
 import anthropic
 from typing import Dict, Any, Optional, Tuple, Union, Callable, List, cast
 
-from anthropic.types import Message, MessageParam
-from mito_ai.completions.models import CompletionError, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
-from mito_ai.utils.mito_server_utils import ProviderCompletionException
+from anthropic.types import Message, MessageParam, TextBlockParam
+from mito_ai.completions.models import ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
+from mito_ai.constants import MESSAGE_HISTORY_TRIM_THRESHOLD
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_server, stream_anthropic_completion_from_mito_server, get_anthropic_completion_function_params
 
@@ -125,6 +125,90 @@ def get_anthropic_system_prompt_and_messages(messages: List[ChatCompletionMessag
     return system_prompt, anthropic_messages
 
 
+def add_cache_control_to_message(message: MessageParam) -> MessageParam:
+    """
+    Adds cache_control to a message's content.
+    Handles both string content and list of content blocks.
+    """
+    content = message.get("content")
+    
+    if isinstance(content, str):
+        # Simple string content - convert to list format with cache_control
+        return {
+            "role": message["role"],
+            "content": [
+                {
+                    "type": "text",
+                    "text": content,
+                    "cache_control": {"type": "ephemeral"}
+                }
+            ]
+        }
+    
+    elif isinstance(content, list) and len(content) > 0:
+        # List of content blocks - add cache_control to last block
+        content_blocks = content.copy()
+        last_block = content_blocks[-1].copy()
+        last_block["cache_control"] = {"type": "ephemeral"}
+        content_blocks[-1] = last_block
+        
+        return {
+            "role": message["role"],
+            "content": content_blocks
+        }
+    
+    else:
+        # Edge case: empty or malformed content
+        return message
+
+
+def get_anthropic_system_prompt_and_messages_with_caching(messages: List[ChatCompletionMessageParam]) -> Tuple[
+    Union[str, List[TextBlockParam], anthropic.Omit], List[MessageParam]]:
+    """
+    Convert a list of OpenAI messages to a list of Anthropic messages with caching applied.
+    
+    Caching Strategy:
+    1. System prompt (static) → Always cached
+    2. Stable conversation history → Cache at keep_recent boundary
+    3. Recent messages → Never cached (always fresh)
+    
+    The keep_recent parameter determines which messages are stable and won't be trimmed.
+    We cache at the keep_recent boundary because those messages are guaranteed to be stable.
+    """
+    
+    # Get the base system prompt and messages
+    system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+    
+    # 1. Cache the system prompt always
+    # If the system prompt is something like anthropic.Omit, we don't need to cache it
+    cached_system_prompt: Union[str, List[TextBlockParam], anthropic.Omit] = system_prompt
+    if isinstance(system_prompt, str):
+        cached_system_prompt = [{
+            "type": "text",
+            "text": system_prompt,
+            "cache_control": {"type": "ephemeral"}
+        }]
+    
+    # 2. Cache conversation history at the boundary where the messages are stable.
+    # Messages are stable after they are more than MESSAGE_HISTORY_TRIM_THRESHOLD old.
+    # At this point, the messages are not edited anymore, so they will not invalidate the cache.
+    # If we included the messages before the boundary in the cache, then every time we send a new
+    # message, we would invalidate the cache and we would never get a cache hit except for the system prompt.
+    messages_with_cache = []
+    
+    if len(anthropic_messages) > 0:
+        cache_boundary = len(anthropic_messages) - MESSAGE_HISTORY_TRIM_THRESHOLD - 1
+        
+        # Add all messages, but only add cache_control to the message at the boundary
+        for i, msg in enumerate(anthropic_messages):
+            if i == cache_boundary:
+                messages_with_cache.append(add_cache_control_to_message(msg))
+            else:
+                messages_with_cache.append(msg)
+    
+    return cached_system_prompt, messages_with_cache
+
+
 class AnthropicClient:
     """
     A client for interacting with the Anthropic API or the Mito server fallback.
@@ -149,7 +233,7 @@ class AnthropicClient:
         """
         Get a response from Claude or the Mito server that adheres to the AgentResponse format.
         """
-        anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+        anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
         
         provider_data = get_anthropic_completion_function_params(
             message_type=message_type,
@@ -166,6 +250,7 @@ class AnthropicClient:
             # Unpack provider_data for direct API call
             assert self.client is not None
             response = self.client.messages.create(**provider_data)
+            
             if provider_data.get("tool_choice") is not None:
                 result = extract_and_parse_anthropic_json_response(response)
                 return json.dumps(result) if not isinstance(result, str) else result
@@ -192,7 +277,7 @@ class AnthropicClient:
     async def stream_completions(self, messages: List[ChatCompletionMessageParam], model: str, message_id: str, message_type: MessageType,
                               reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]) -> str:
         try:
-            anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+            anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
             accumulated_response = ""
 
             if self.api_key:

mito_ai/chat_history/handlers.py

@@ -0,0 +1,63 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import tornado
+from typing import List, Any
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.completions.models import ChatThreadMetadata
+
+
+class ChatHistoryHandler(APIHandler):
+    """
+    Endpoints for working with chat history threads.
+    """
+
+    def initialize(self, message_history: GlobalMessageHistory) -> None:
+        """Initialize the handler with the global message history instance."""
+        super().initialize()
+        self._message_history = message_history
+
+    @tornado.web.authenticated
+    def get(self, *args: Any, **kwargs: Any) -> None:
+        """Get all chat history threads or a specific thread by ID."""
+        try:
+            # Check if a specific thread ID is provided in the URL
+            thread_id = kwargs.get("thread_id")
+
+            if thread_id:
+                # Get specific thread
+                if thread_id in self._message_history._chat_threads:
+                    thread = self._message_history._chat_threads[thread_id]
+                    thread_data = {
+                        "thread_id": thread.thread_id,
+                        "name": thread.name,
+                        "creation_ts": thread.creation_ts,
+                        "last_interaction_ts": thread.last_interaction_ts,
+                        "display_history": thread.display_history,
+                        "ai_optimized_history": thread.ai_optimized_history,
+                    }
+                    self.finish(thread_data)
+                else:
+                    self.set_status(404)
+                    self.finish({"error": f"Thread with ID {thread_id} not found"})
+            else:
+                # Get all threads
+                threads: List[ChatThreadMetadata] = self._message_history.get_threads()
+
+                # Convert to dict format for JSON serialization
+                threads_data = [
+                    {
+                        "thread_id": thread.thread_id,
+                        "name": thread.name,
+                        "creation_ts": thread.creation_ts,
+                        "last_interaction_ts": thread.last_interaction_ts,
+                    }
+                    for thread in threads
+                ]
+
+                self.finish({"threads": threads_data})
+
+        except Exception as e:
+            self.set_status(500)
+            self.finish({"error": str(e)})

mito_ai/chat_history/urls.py

@@ -0,0 +1,32 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Any
+from jupyter_server.utils import url_path_join
+from mito_ai.chat_history.handlers import ChatHistoryHandler
+from mito_ai.completions.message_history import GlobalMessageHistory
+
+
+def get_chat_history_urls(base_url: str, message_history: GlobalMessageHistory) -> List[Tuple[str, Any, dict]]:
+    """Get all chat history related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+        message_history: The global message history instance
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai/chat-history"
+    return [
+        (
+            url_path_join(BASE_URL, "threads"),
+            ChatHistoryHandler,
+            {"message_history": message_history},
+        ),
+        (
+            url_path_join(BASE_URL, "threads", "(?P<thread_id>[^/]+)"),
+            ChatHistoryHandler,
+            {"message_history": message_history},
+        ),
+    ]

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/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/streamlit_conversion/agent_utils.py

@@ -3,8 +3,14 @@
 
 from typing import List
 import re
+from anthropic.types import MessageParam
+from mito_ai.streamlit_conversion.streamlit_system_prompt import streamlit_system_prompt
+from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
 from unidiff import PatchSet
 from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER
+from mito_ai.completions.models import MessageType
+
+STREAMLIT_AI_MODEL = "claude-3-5-haiku-latest"
 
 def extract_todo_placeholders(agent_response: str) -> List[str]:
     """Extract TODO placeholders from the agent's response"""
@@ -23,6 +29,8 @@ def apply_patch_to_text(text: str, diff: str) -> str:
     diff : str
         A unified diff that transforms *text* into the desired output.
         The diff must reference exactly one file (the Streamlit app).
+        NOTE: This assumes a custom format where BOTH -X,Y and +X,Y
+        reference the original file line numbers.
 
     Returns
     -------
@@ -42,37 +50,46 @@ def apply_patch_to_text(text: str, diff: str) -> str:
     patch = PatchSet(diff.splitlines(keepends=True))
 
     # We expect a single-file patch (what the prompt asks the model to emit)
-    if len(patch) != 1:
+    if len(patch) == 0:
+        raise ValueError("No patches found in diff")
+    
+    # Check that all patches are for the same file
+    file_names = set(p.source_file for p in patch)
+    if len(file_names) > 1:
         raise ValueError(
-            f"Expected a patch for exactly one file, got {len(patch)} files."
+            f"Expected patches for exactly one file, got files: {file_names}"
         )
 
-    file_patch = patch[0]
-
+    # Apply all hunks from all patches (they should all be for the same file)
     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])
+    # Process all hunks from all patches
+    # We only expect one patch file, but it always returns as a list 
+    # so we just iterate over it
+    for file_patch in patch:
+        for hunk in file_patch:
+            # Since hunks reference the original file, just convert to 0-based
+            hunk_start = hunk.source_start - 1
+            
+            # Copy unchanged lines before this hunk
+            while cursor < hunk_start:
+                if cursor < len(original_lines):
+                    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)
+
+            # Apply hunk line-by-line
+            for line in hunk:
+                if line.is_context:
+                    # Use the line from the diff to preserve exact formatting
+                    result_lines.append(line.value)
+                    cursor += 1
+                elif line.is_removed:
+                    cursor += 1  # Skip this line from the original
+                elif line.is_added:
+                    # Use the line from the diff to preserve exact formatting
+                    result_lines.append(line.value)
 
     # Copy any remaining lines after the last hunk
     result_lines.extend(original_lines[cursor:])
@@ -80,11 +97,66 @@ def apply_patch_to_text(text: str, diff: str) -> str:
     return "".join(result_lines)
 
 
+def fix_context_lines(diff: str) -> str:
+    """
+    Fix context lines in unified diff to ensure they all start with a space character.
+    
+    In unified diffs, context lines (unchanged lines) must start with a single space ' ',
+    even if the line itself is empty. The AI sometimes generates diffs where empty
+    context lines are just blank lines without the leading space, which causes the
+    unidiff parser to fail.
+    
+    Args:
+        diff (str): The unified diff string
+        
+    Returns:
+        str: The corrected diff with proper context line formatting
+    """
+    lines = diff.split('\n')
+    corrected_lines = []
+    in_hunk = False
+    
+    for i, line in enumerate(lines):
+        # Check if we're entering a hunk
+        if line.startswith('@@'):
+            in_hunk = True
+            corrected_lines.append(line)
+            continue
+            
+        # Check if we're leaving a hunk (new file header)
+        if line.startswith('---') or line.startswith('+++'):
+            in_hunk = False
+            corrected_lines.append(line)
+            continue
+            
+        if in_hunk:
+            # We're inside a hunk
+            if line.startswith(' ') or line.startswith('-') or line.startswith('+'):
+                # Already has proper diff marker
+                corrected_lines.append(line)
+            elif line.strip() == '':
+                # Empty line should be a context line with leading space
+                corrected_lines.append(' ')
+            else:
+                # Line without diff marker - treat as context line
+                corrected_lines.append(' ' + line)
+        else:
+            # Outside hunk - keep as is
+            corrected_lines.append(line)
+    
+    return '\n'.join(corrected_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.
+    
+    If no header is provided at all, then there is nothing to fix.
     """
+    # First fix context lines to ensure they have proper leading spaces
+    diff = fix_context_lines(diff)
+    
     lines = diff.split('\n')
     
     for i, line in enumerate(lines):
@@ -99,18 +171,64 @@ def fix_diff_headers(diff: str) -> str:
                 old_count = 0
                 new_count = 0
                 
+                # Find the end of this hunk (next @@ line or end of file)
+                hunk_end = len(lines)
                 for j in range(i + 1, len(lines)):
-                    next_line = lines[j]
-                    if next_line.startswith('@@') or next_line.startswith('---'):
+                    if lines[j].startswith('@@'):
+                        hunk_end = j
                         break
-                    if next_line.startswith(' ') or next_line.startswith('-'):
+                
+                # Count lines in this hunk
+                for j in range(i + 1, hunk_end):
+                    hunk_line = lines[j]
+                    # Empty lines are treated as context lines
+                    if hunk_line == '' or hunk_line.startswith(' ') or hunk_line.startswith('-'):
                         old_count += 1
-                    if next_line.startswith(' ') or next_line.startswith('+'):
+                    if hunk_line == '' or hunk_line.startswith(' ') or hunk_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)
-
-
+    corrected_diff = '\n'.join(lines)
+    corrected_diff = corrected_diff.lstrip()
+    
+    # If there is no diff, just return it without fixing file headers
+    if len(corrected_diff) == 0:
+        return corrected_diff
+    
+    # Remove known problametic file component headers that the AI sometimes returns
+    problamatic_file_header_components = ['--- a/app.py +++ b/app.py']
+    for problamatic_file_header_component in problamatic_file_header_components:
+        corrected_diff = corrected_diff.removeprefix(problamatic_file_header_component).lstrip()
+        
+    
+    # If the diff is missing the file component of the header, add it
+    valid_header_component = """--- a/app.py
++++ b/app.py"""
+    if not corrected_diff.startswith(valid_header_component):
+        corrected_diff = valid_header_component + '\n' + corrected_diff
+    
+    return corrected_diff
+
+
+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 = 8192 # 64_000
+    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