mito-ai 0.1.54__py3-none-any.whl → 0.1.56__py3-none-any.whl

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.54'
+__version__ = VERSION = '0.1.56'

mito_ai/anthropic_client.py

@@ -7,7 +7,7 @@ from typing import Dict, Any, Optional, Tuple, Union, Callable, List, cast
 
 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 mito_ai.completions.prompt_builders.prompt_section_registry import get_max_trim_after_messages
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_server, select_correct_model, stream_anthropic_completion_from_mito_server, get_anthropic_completion_function_params
 
@@ -15,6 +15,10 @@ from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_ser
 # We set it to a high number so that we can edit large code cells
 MAX_TOKENS = 64_000
 
+# Calculate the max trim threshold once at module load time.
+# This is used for cache boundary calculation - messages older than this threshold are stable.
+MAX_TRIM_THRESHOLD = get_max_trim_after_messages()
+
 def extract_and_parse_anthropic_json_response(response: Message) -> Union[object, Any]:
     """
     Extracts and parses the JSON response from the Claude API.
@@ -170,9 +174,6 @@ def get_anthropic_system_prompt_and_messages_with_caching(messages: List[ChatCom
     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
@@ -189,14 +190,14 @@ def get_anthropic_system_prompt_and_messages_with_caching(messages: List[ChatCom
         }]
     
     # 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.
+    # Messages are stable after they are older than the max trim_after_messages threshold.
     # 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
+        cache_boundary = len(anthropic_messages) - MAX_TRIM_THRESHOLD - 1
         
         # Add all messages, but only add cache_control to the message at the boundary
         for i, msg in enumerate(anthropic_messages):

mito_ai/completions/models.py

@@ -17,7 +17,7 @@ ThreadID = NewType('ThreadID', str)
     
 class CellUpdate(BaseModel):
     type: Literal['modification', 'new']
-    index: Optional[int]
+    after_cell_id: Optional[str]
     id: Optional[str]
     code: str
     code_summary: str

mito_ai/completions/prompt_builders/agent_execution_prompt.py

@@ -1,57 +1,25 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
+from typing import List
 from mito_ai.completions.models import AgentExecutionMetadata
-from mito_ai.completions.prompt_builders.prompt_constants import (
-    ACTIVE_CELL_ID_SECTION_HEADING,
-    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
-)
+from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
+from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
 
 
 def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
-    variables_str = '\n'.join([f"{variable}" for variable in md.variables or []])
-    files_str = '\n'.join([f"{file}" for file in md.files or []])
-    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.notebookID, md.notebookPath)    
-    
-    context_str = f"""Remember to choose the correct tool to respond with.
-
-{rules_str}
-
-
-{JUPYTER_NOTEBOOK_SECTION_HEADING}
-{ai_optimized_cells_str}
-
-{VARIABLES_SECTION_HEADING}
-{variables_str}
-
-{FILES_SECTION_HEADING}
-{files_str}
-
-{STREAMLIT_APP_STATUS_SECTION_HEADING}
-{streamlit_status_str}
-
-{ACTIVE_CELL_ID_SECTION_HEADING}
-{md.activeCellId}
-
-{selected_context_str}
-
-{cell_update_output_str(md.base64EncodedActiveCellOutput is not None)}"""
-
-    task_str = '' if md.input == '' else f"""Your task: 
-{md.input}"""
-
-    return '\n\n'.join([context_str, task_str]).strip()
+    sections: List[PromptSection] = [
+        SG.Generic("Reminder", "Remember to choose the correct tool to respond with."),
+        SG.Rules(md.additionalContext),
+        SG.StreamlitAppStatus(md.notebookID, md.notebookPath),
+        SG.Files(md.files),
+        SG.Variables(md.variables),
+        SG.SelectedContext(md.additionalContext),
+        SG.ActiveCellId(md.activeCellId),
+        SG.Notebook(md.aiOptimizedCells),
+        SG.GetCellOutputToolResponse(md.base64EncodedActiveCellOutput),
+        SG.Task(f"{md.input}"),
+    ]
+
+    prompt = Prompt(sections)
+    return str(prompt)

mito_ai/completions/prompt_builders/agent_smart_debug_prompt.py

@@ -3,11 +3,8 @@
 
 from typing import List
 from mito_ai.completions.models import AgentSmartDebugMetadata
-from mito_ai.completions.prompt_builders.prompt_constants import (
-    FILES_SECTION_HEADING,
-    JUPYTER_NOTEBOOK_SECTION_HEADING,
-    VARIABLES_SECTION_HEADING
-)
+from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
+from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
 
 # TODO:
 # 1. In the future, it might make sense to pass the previous CELL_UPDATE to this prompt?
@@ -18,11 +15,10 @@ from mito_ai.completions.prompt_builders.prompt_constants import (
 # graph of cells that we calculate ourselves, not relying on the AI. 
 
 def create_agent_smart_debug_prompt(md: AgentSmartDebugMetadata) -> str:
-    variables_str = '\n'.join([f"{variable}" for variable in md.variables or []])
-    files_str = '\n'.join([f"{file}" for file in md.files or []])
-    ai_optimized_cells_str = '\n'.join([f"{cell}" for cell in md.aiOptimizedCells or []])
+    sections: List[PromptSection] = []
     
-    return f"""I just applied and executed the CELL_UPDATE that you just shared with me, but it errored. Below I am sharing with you a strategy for how I want you to resolve this error and information about the actual error that occured.
+    # Add intro text
+    sections.append(SG.Generic("Instructions", f"""I just applied and executed the CELL_UPDATE that you just shared with me, but it errored. Below I am sharing with you a strategy for how I want you to resolve this error and information about the actual error that occured.
 
 Use this strategy for this message only. After this message, continue using the original set of instructions that I provided you.
 
@@ -33,7 +29,7 @@ Step 1: ERROR ANALYSIS: Analyze the error message to identify why the code cell
 Step 2: INTENT PRESERVATION: Make sure you understand the intent of the CELL_UPDATE so that you can be sure to preserve it when you create a new CELL_UPDATE
 Step 3: ERROR CORRECTION: Respond with a new CELL_UPDATE that is applied to the same cell as the erroring CELL_UPDATE.
 
-<Instructions for each Phase />
+INSTRUCTIONS FOR EACH PHASE
 
 ERROR ANALYSIS:
 
@@ -61,100 +57,89 @@ ERROR CORRECTION:
         type: 'run_all_cells',
         message: str
     }}
-    Note that if the name error persists even after using run_all_cells, it means that the variable is not defined in the notebook and you should not reuse this tool. Additionally, this tool could also be used to refresh the notebook state.
-
-<Example>
-
-<Input>
-
-{FILES_SECTION_HEADING}
-file_name: sales.csv
-
-Jupyter Notebook:
-[
-    {{
-        cell_type: 'markdown'
-        id: '9e38c62b-38f8-457d-bb8d-28bfc52edf2c'
-        code: \"\"\"# Transaction Analysis \"\"\"
-    }},
-    {{
-        cell_type: 'code'
-        id: 'adslkaf-jf73-l8xn-92j7-kjd8kdcnd2kso'
-        code: \"\"\" 'df' = pd.DataFrame({{
-    'order_id': [1, 2, 3, 4],
-    'date': ['Mar 7, 2025', 'Sep 24, 2024', '25 June, 2024', 'June 29, 2024'],
-    'amount': [100, 150, 299, 99]
-}})
-    }},
+    Note that if the name error persists even after using run_all_cells, it means that the variable is not defined in the notebook and you should not reuse this tool. Additionally, this tool could also be used to refresh the notebook state."""))
+    
+    # Add example
+    example_content = f"""
+    <Input>
+    
+    Files:
+    "file_name: sales.csv"
+
+    Jupyter Notebook:
+    [
+        {{
+            cell_type: 'markdown'
+            id: '9e38c62b-38f8-457d-bb8d-28bfc52edf2c'
+            code: \"\"\"# Transaction Analysis \"\"\"
+        }},
+        {{
+            cell_type: 'code'
+            id: 'adslkaf-jf73-l8xn-92j7-kjd8kdcnd2kso'
+            code: \"\"\" 'df' = pd.DataFrame({{
+                'order_id': [1, 2, 3, 4],
+                'date': ['Mar 7, 2025', 'Sep 24, 2024', '25 June, 2024', 'June 29, 2024'],
+                'amount': [100, 150, 299, 99]
+            }})
+        }},
+        {{
+            cell_type: 'code'
+            id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc'
+            code: \"\"\"df['date'] = pd.to_datetime(df['date'])\"\"\"
+        }},
+    ]
+    
+    Variables:
     {{
-        cell_type: 'code'
-        id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc'
-        code: \"\"\"df['date'] = pd.to_datetime(df['date'])\"\"\"
-    }},
-]
-
-{VARIABLES_SECTION_HEADING}
-{{
-    'df': pd.DataFrame({{
-        'order_id': [1, 2, 3, 4],
-        'date': ['Mar 7, 2025', 'Sep 24, 2024', '25 June, 2024', 'June 29, 2024'],
-        'amount': [100, 150, 299, 99]
-    }})
-}}
-
-Cell ID of the Error Producing Code Cell:
-'c68fdf19-db8c-46dd-926f-d90ad35bb3bc'
+        'df': pd.DataFrame({{
+            'order_id': [1, 2, 3, 4],
+            'date': ['Mar 7, 2025', 'Sep 24, 2024', '25 June, 2024', 'June 29, 2024'],
+            'amount': [100, 150, 299, 99]
+        }})
+    }}
 
-Error Traceback:
-Cell In[27], line 1
-----> 1 df['date'] = pd.to_datetime(df['date'])
+    Cell ID of the Error Producing Code Cell:
+    'c68fdf19-db8c-46dd-926f-d90ad35bb3bc'
 
-ValueError: time data "25 June, 2024" doesn't match format "%b %d, %Y", at position 2. You might want to try:
-    - passing `format` if your strings have a consistent format;
-    - passing `format='ISO8601'` if your strings are all ISO8601 but not necessarily in exactly the same format;
-    - passing `format='mixed'`, and the format will be inferred for each element individually. You might want to use `dayfirst` alongside this.
+    Error Traceback:
+    Cell In[27], line 1
+    ----> 1 df['date'] = pd.to_datetime(df['date'])
 
+    ValueError: time data "25 June, 2024" doesn't match format "%b %d, %Y", at position 2. You might want to try:
+        - passing `format` if your strings have a consistent format;
+        - passing `format='ISO8601'` if your strings are all ISO8601 but not necessarily in exactly the same format;
+        - passing `format='mixed'`, and the format will be inferred for each element individually. You might want to use `dayfirst` alongside this.
 
-</ Input>
 
-< Your Thinking >
+    </ Input>
 
-ERROR ANALYSIS
-This is a ValueError caused by applying the wrong format to a specific date string. Because it was triggered at position 2, the first date string must have successfully converted. By looking at the defined variables, I can see that first date string is in the format "Mar 7, 2025", but the third date string is in the format "25 June, 2024". Those dates are not in the same format, so the conversion failed.
+    < Your Thinking >
 
-INTENT PRESERVATION:
-User is trying to convert the date column to a datetime object even though the dates are not in the same starting format. 
+    ERROR ANALYSIS
+    This is a ValueError caused by applying the wrong format to a specific date string. Because it was triggered at position 2, the first date string must have successfully converted. By looking at the defined variables, I can see that first date string is in the format "Mar 7, 2025", but the third date string is in the format "25 June, 2024". Those dates are not in the same format, so the conversion failed.
 
-</ Your Thinking >
+    INTENT PRESERVATION:
+    User is trying to convert the date column to a datetime object even though the dates are not in the same starting format. 
 
-<Output>
+    </ Your Thinking >
 
+    <Output>
 
-{{
-    is_finished: false, 
-    cell_update: {{
-        type: 'modification'
-        id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc'
-        code: "def parse_date(date_str):\n    formats = ['%b %d, %Y', '%d %B, %Y']\n\n    for fmt in formats:\n        try:\n            return pd.to_datetime(date_str, format=fmt)\n        except ValueError:\n            # Try next format\n            continue\n\n    # If not format worked, return Not a Time\n    return pd.NaT\n\ndf['date'] = df['date'].apply(lambda x: parse_date(x))"
+    {{
+        is_finished: false, 
+        cell_update: {{
+            type: 'modification'
+            id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc'
+            code: "def parse_date(date_str):\n    formats = ['%b %d, %Y', '%d %B, %Y']\n\n    for fmt in formats:\n        try:\n            return pd.to_datetime(date_str, format=fmt)\n        except ValueError:\n            # Try next format\n            continue\n\n    # If not format worked, return Not a Time\n    return pd.NaT\n\ndf['date'] = df['date'].apply(lambda x: parse_date(x))"
+        }}
     }}
-}}
-
-</Output>
-
-</Example>
-
-{FILES_SECTION_HEADING}
-{files_str}
-
-{JUPYTER_NOTEBOOK_SECTION_HEADING}
-{ai_optimized_cells_str}
-
-{VARIABLES_SECTION_HEADING}
-{variables_str}
 
-Cell ID of the Error Producing Code Cell:
-{md.error_message_producing_code_cell_id}
+    </Output>"""
+    sections.append(SG.Example("Example", example_content))
+    sections.append(SG.Files(md.files))
+    sections.append(SG.Notebook(md.aiOptimizedCells))
+    sections.append(SG.Variables(md.variables))
+    sections.append(SG.ErrorTraceback(md.error_message_producing_code_cell_id, md.errorMessage))
 
-Error Traceback:
-{md.errorMessage}
-"""
+    prompt = Prompt(sections)
+    return str(prompt)