mito-ai 0.1.55__py3-none-any.whl → 0.1.57__py3-none-any.whl

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -11,21 +11,35 @@ import json
 from typing import Final
 from mito_ai.utils.schema import MITO_FOLDER
 
-# Section headings used in prompts
-FILES_SECTION_HEADING = "Files in the current directory:"
-VARIABLES_SECTION_HEADING = "Defined Variables:"
-CODE_SECTION_HEADING = "Code in the active code cell:"
-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:"
+CHART_CONFIG_RULES = """
+When creating a matplotlib chart, you must use the `# === CHART CONFIG ===` and `# === END CONFIG ===` markers to indicate the start and end of the chart configuration section.
 
-# Placeholder text used when trimming content from messages
-CONTENT_REMOVED_PLACEHOLDER = "Content removed to save space" 
+The chart configuration section is a list of variables used to customize the chart. This includes the titles, labels, colors, and any variables that affect the chart's appearance.
 
-CITATION_RULES = """RULES FOR CITING YOUR WORK
+Rules:
+- All imports must appear at the top, before the chart configuration section.
+- Variables with multiple words should be underscore-separated.
+- All colors should be in hex format (e.g., "#3498db"). Use quotes around the hex string: COLOR = "#3498db" or COLOR = '#3498db'. Do NOT nest quotes.
+- Variables can only be strings, numbers, booleans, tuples, or lists. 
+- NEVER include comments on the same line as a variable assignment. Each variable assignment must be on its own line with no trailing comments.
+- For string values, use either single or double quotes (e.g., TITLE = "Sales by Product" or TITLE = 'Sales by Product'). Do not use nested quotes (e.g., do NOT use '"value"').
 
+Common Mistakes to Avoid:
+- WRONG: COLOR = '"#1877F2" # Meta Blue'  (nested quotes and inline comment)
+- WRONG: COLOR = "#1877F2" # Meta Blue  (inline comment)
+- WRONG: COLOR = '"#1877F2"'  (nested quotes)
+- CORRECT: COLOR = "#1877F2"  (simple hex string, no nested quotes, no inline comments)
+
+Example:
+# === CHART CONFIG ===
+TITLE = "Sales by Product"
+X_LABEL = "Product"
+Y_LABEL = "Sales"
+BAR_COLOR = "#000000"
+# === END CONFIG ===
+"""
+
+CITATION_RULES = """
 It is important that the user is able to verify any insights that you share with them about their data. To make this easy for the user, you must cite the lines of code that you are drawing the insight from. To provide a citation, use one of the following formats inline in your response:
 
 Single line citation:
@@ -36,19 +50,18 @@ Multiline citation (for citing a range of lines):
 
 Citation Rules:
 
-1. Every fact or statement derived from the user's notebook must include a citation. 
+1. Every fact or conclusion you draw based on the user's notebook must include a MITO_CITATION so that the user can verify your work.
 2. When choosing the citation, select the code that will most help the user validate the fact or statement that you shared with them.
 3. Place the citation immediately after the statement it supports. Do not explain the citation with phrases like "See", "Derived from", etc. Just provide the citation object.
 4. For the "line_number" field, use the line number within the cell that is most relevant to the citation. Important: The cell line number should be 0-indexed and should not skip comments.
 5. For multiline citations, use the "first_line-last_line" format when the insight spans multiple lines of code. Both line numbers should be 0-indexed.
 6. If you cannot find relevant information in the notebook to answer a question, clearly state this and do not provide a citation.
-7. You ONLY need to provide a citation when sharing an insight from the data in the message part of the response. If all you are doing is writing/updating code, then there is no need to provide a citation.
+7. You only need to provide a citation when sharing an insight with the user. For example you should use a MITO_CITATION when you writing insights like any of the following: "The highest trading volume day was on January 15th, 2025", "The MSE was 6.8", "Apple's COGS are represented by the Orange bar in the graph". If you are not writing a quantitative insight and instead are just referring to a block of code like "I updated Cell 5 to include a graph", then you should instead use a MITO_CELL_REF.
 8. Do not include the citation in the code block as a comment. ONLY include the citation in the message field of your response.
 """
 
-CELL_REFERENCE_RULES = """RULES FOR REFERENCING CELLS
-
-When referring to specific cells in the notebook in your messages, use cell references so the user can easily navigate to the cell you're talking about. The user sees cells numbered as "Cell 1", "Cell 2", etc., but internally cells are identified by their unique IDs.
+CELL_REFERENCE_RULES = """
+When referring to specific cells in the notebook in your messages, use cell references so the user can easily navigate to the cell you're talking about. Cell references are displayed to the user  "Cell 1", "Cell 2", etc., but internally cells are identified by their unique IDs.
 
 To reference a cell, use this format inline in your message:
 [MITO_CELL_REF:cell_id]
@@ -58,35 +71,12 @@ This will be displayed to the user as a clickable "Cell N" link that navigates t
 Cell Reference Rules:
 
 1. Use cell references when discussing specific cells you've created or modified (e.g., "I've added the data cleaning code in [MITO_CELL_REF:abc123]").
-2. Use cell references when referring to cells the user mentioned or that contain relevant context.
-3. The cell_id must be an actual cell ID from the notebook - do not make up IDs.
-4. Place the reference inline where it makes sense in your message, similar to how you would write "Cell 3" in natural language.
-5. Do not use cell references in code - only in the message field of your responses.
-6. Cell references are different from citations. Use citations for specific line-level insights; use cell references for general cell-level navigation.
-
-Example:
-"I've loaded the sales data in [MITO_CELL_REF:c68fdf19-db8c-46dd-926f-d90ad35bb3bc] and will now calculate the monthly totals."
+2. The cell_id must be an actual cell ID from the notebook - do not make up IDs.
+3. Place the reference inline where it makes sense in your message, similar to how you would write "Cell 3" in natural language.
+4. Do not use cell references in code - only in the message field of your responses.
+5. You only need to provide a cell reference when you want to make it easy for the user to navigate to a specific cell in the notebook. For example you should use a MITO_CELL_REF when you are stating things like: "I've loaded the sales data in [MITO_CELL_REF:c68fdf19-db8c-46dd-926f-d90ad35bb3bc]" or "[MITO_CELL_REF:a91fde20-cc7f-g6ee-146g-e10bc34abdbh] creates the graph showing the total highest closing stock price for each company". If you are not referencing an entire code block and instead of providing justification for a specific conclucions that you drew like "The most common used car in the lot is a 2005 Honda CRV", then you should instead use a MITO_CITATION.
 """
 
-def get_active_cell_output_str(has_active_cell_output: bool) -> str:
-    """
-    Used to tell the AI about the output of the active code cell. 
-    We use this in the chat prompt.
-    """
-    if has_active_cell_output:
-        return f"{ACTIVE_CELL_OUTPUT_SECTION_HEADING}\nAttatched is an image of the output of the active code cell for your context."
-    else:
-        return ""
-    
-def cell_update_output_str(has_cell_update_output: bool) -> str:
-    """
-    Used to respond to the GET_CELL_OUTPUT tool, telling the agent the output of the cell it requested
-    """
-    if has_cell_update_output:
-        return f"{GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING}\nAttatched is an image of code cell output that you requested."
-    else:
-        return ""
-
 def redact_sensitive_info(connections: dict) -> dict:
     """
     Redacts sensitive information from connections data.
@@ -184,7 +174,7 @@ Here is the schema:
     return DATABASE_RULES
 
 
-CHAT_CODE_FORMATTING_RULES = """CRITICAL CODE UPDATE RULES:
+CHAT_CODE_FORMATTING_RULES = """
 - COMPLETE REPLACEMENT: Your code will COMPLETELY REPLACE the entire contents of the active code cell. 
 - INCLUDE ALL CODE: You MUST return the COMPLETE, FULL contents of the entire code cell - including ALL existing code that should remain plus your modifications.
 - NEVER PARTIAL CODE: NEVER return only a portion, snippet, or subset of the code cell. Partial responses will break the user's notebook by deleting important code.

mito_ai/completions/prompt_builders/prompt_section_registry/__init__.py

@@ -0,0 +1,70 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
+
+
+from typing import List, Type
+from .files import FilesSection
+from .variables import VariablesSection
+from .active_cell_code import ActiveCellCodeSection
+from .notebook import NotebookSection
+from .active_cell_id import ActiveCellIdSection
+from .active_cell_output import ActiveCellOutputSection
+from .get_cell_output_tool_response import GetCellOutputToolResponseSection
+from .streamlit_app_status import StreamlitAppStatusSection
+from .selected_context import SelectedContextSection
+from .rules import RulesSection
+from .task import TaskSection
+from .error_traceback import ErrorTracebackSection
+from .example import ExampleSection
+from .generic import GenericSection
+from .base import Prompt, PromptSection
+
+
+class SectionRegistry:
+    """Namespace for easy section construction."""
+    Files = FilesSection
+    Variables = VariablesSection
+    ActiveCellCode = ActiveCellCodeSection
+    Notebook = NotebookSection
+    ActiveCellId = ActiveCellIdSection
+    ActiveCellOutput = ActiveCellOutputSection
+    GetCellOutputToolResponse = GetCellOutputToolResponseSection
+    StreamlitAppStatus = StreamlitAppStatusSection
+    SelectedContext = SelectedContextSection
+    Rules = RulesSection
+    Task = TaskSection
+    ErrorTraceback = ErrorTracebackSection
+    Example = ExampleSection
+    Generic = GenericSection
+
+
+# Export as SG (Section Generator) for easy usage
+SG = SectionRegistry()
+
+# Also export function to get all section classes for trimming
+def get_all_section_classes() -> List[Type[PromptSection]]:
+    """Returns all section classes for building trimming mapping."""
+    prompt_section_classes: List[Type[PromptSection]] = PromptSection.__subclasses__()
+    return prompt_section_classes
+
+
+def get_max_trim_after_messages() -> int:
+    """
+    Returns the maximum trim_after_messages value from all section classes.
+    
+    This is used for cache boundary calculation - we cache messages that are
+    older than the max trim threshold, since those messages are stable and
+    won't be edited anymore.
+    
+    Returns 0 if all sections have trim_after_messages = None.
+    """
+    section_classes = get_all_section_classes()
+    max_value = 0
+    for section_class in section_classes:
+        trim_value = getattr(section_class, 'trim_after_messages', None)
+        if trim_value is not None and trim_value > max_value:
+            max_value = trim_value
+    return max_value
+

mito_ai/completions/prompt_builders/prompt_section_registry/active_cell_code.py

@@ -0,0 +1,15 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+
+
+class ActiveCellCodeSection(PromptSection):
+    """Section for code in the active code cell."""
+    trim_after_messages: int = 3
+    
+    def __init__(self, code: str):
+        self.code = code
+        self.content = f"```python\n{code}\n```"
+        self.name = "ActiveCellCode"
+

mito_ai/completions/prompt_builders/prompt_section_registry/active_cell_id.py

@@ -0,0 +1,10 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+
+
+class ActiveCellIdSection(PromptSection):
+    """Section for the ID of the active code cell."""
+    trim_after_messages = None
+

mito_ai/completions/prompt_builders/prompt_section_registry/active_cell_output.py

@@ -0,0 +1,20 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+
+
+class ActiveCellOutputSection(PromptSection):
+    """Section for output of the active code cell."""
+    trim_after_messages: int = 3
+    exclude_if_empty: bool = True
+    
+    def __init__(self, has_active_cell_output: bool):
+        self.name = "ActiveCellOutput"
+        self.has_active_cell_output = has_active_cell_output
+        self.content = ""
+        if has_active_cell_output:
+            # The actual image is attatched to the message, its not part of the text content
+            self.content = f"Attatched is an image of the output of the active code cell."
+
+

mito_ai/completions/prompt_builders/prompt_section_registry/base.py

@@ -0,0 +1,37 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from abc import ABC
+from typing import Optional, List
+
+
+class PromptSection(ABC):
+    """Abstract base class for all prompt sections."""
+    
+    # Class variable: trimming threshold (None = never trim)
+    trim_after_messages: Optional[int] = 3
+    
+    # Class variable: if True, exclude XML tags when content is empty
+    exclude_if_empty: bool = False
+    
+    def __init__(self, content: str):
+        self.content = content
+        self.name = self.__class__.__name__.replace("Section", "")
+    
+    def __str__(self) -> str:
+        # If exclude_if_empty is True and content is empty, return empty string
+        if self.exclude_if_empty and (self.content is None or self.content.strip() == ""):
+            return ""
+        return f"<{self.name}>\n{self.content}\n</{self.name}>"
+
+
+class Prompt:
+    """Container for multiple prompt sections."""
+    def __init__(self, sections: List[PromptSection]):
+        self.sections = sections
+    
+    def __str__(self) -> str:
+        # Filter out empty strings to exclude sections that return "" when they have no content
+        section_strings = [str(section) for section in self.sections]
+        return "\n\n".join(s for s in section_strings if s)
+

mito_ai/completions/prompt_builders/prompt_section_registry/error_traceback.py

@@ -0,0 +1,17 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+from typing import Optional
+
+
+class ErrorTracebackSection(PromptSection):
+    """Section for error traceback - never trimmed."""
+    trim_after_messages: Optional[int] = None
+    
+    def __init__(self, code_cell_id: str, traceback: str):
+        self.code_cell_id = code_cell_id
+        self.traceback = traceback
+        self.content = f"Cell ID: {code_cell_id}\n\n{traceback}"
+        self.name = "Error"
+

mito_ai/completions/prompt_builders/prompt_section_registry/example.py

@@ -0,0 +1,19 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+
+
+class ExampleSection(PromptSection):
+    """Section for examples - contains nested section XML tags."""
+    trim_after_messages: int = 3
+    
+    def __init__(self, name: str, content: str):
+        # content can contain nested XML like: "<Files>...</Files><Variables>...</Variables>"
+        super().__init__(content)
+        self.example_name = name  # e.g., "Example 1", "Cell Modification Example"
+    
+    def __str__(self) -> str:
+        # Render as: <Example name="Example 1">...</Example>
+        return f'<Example name="{self.example_name}">{self.content}</Example>'
+

mito_ai/completions/prompt_builders/prompt_section_registry/files.py

@@ -0,0 +1,17 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Optional
+from .base import PromptSection
+
+
+class FilesSection(PromptSection):
+    """Section for files in the current directory."""
+    trim_after_messages: int = 3
+    exclude_if_empty: bool = True
+    
+    def __init__(self, files: Optional[List[str]]):
+        self.files = files
+        self.content = '\n'.join([f"file_name: {file}" for file in files or []])
+        self.name = "Files"
+

mito_ai/completions/prompt_builders/prompt_section_registry/generic.py

@@ -0,0 +1,15 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+from typing import Optional
+
+
+class GenericSection(PromptSection):
+    """Generic section that can be used with a custom name and content."""
+    trim_after_messages: Optional[int] = None
+    
+    def __init__(self, name: str, content: str):
+        self.content = content
+        self.name = name
+

mito_ai/completions/prompt_builders/prompt_section_registry/get_cell_output_tool_response.py

@@ -0,0 +1,21 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Optional
+from .base import PromptSection
+
+
+class GetCellOutputToolResponseSection(PromptSection):
+    """Section for output of the code cell after applying CELL_UPDATE."""
+    trim_after_messages: int = 3
+    exclude_if_empty: bool = True
+    
+    def __init__(self, base64EncodedCellOutput: Optional[str]):
+        self.name = "GetCellOutputToolResponse"
+        self.base64EncodedCellOutput = base64EncodedCellOutput
+        
+        self.content = ""
+        if base64EncodedCellOutput is not None and base64EncodedCellOutput != '':
+            # The actual image is attatched to the message, its not part of the text content
+            self.content = f"Attatched is an image of code cell output that you requested."
+        

mito_ai/completions/prompt_builders/prompt_section_registry/notebook.py

@@ -0,0 +1,19 @@
+# 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 AIOptimizedCell
+from .base import PromptSection
+
+
+class NotebookSection(PromptSection):
+    """Section for Jupyter notebook content."""
+    trim_after_messages: int = 6
+    exclude_if_empty: bool = False
+    
+    def __init__(self, cells: List[AIOptimizedCell]):
+        self.cells = cells
+        self.content = '\n'.join([f"{cell}" for cell in cells or []])
+        self.name = "Notebook"
+

mito_ai/completions/prompt_builders/prompt_section_registry/rules.py

@@ -0,0 +1,39 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+from typing import Optional, List, Dict
+from mito_ai.rules.utils import get_rule
+
+def get_rules_str(additional_context: Optional[List[Dict[str, str]]]) -> str:
+    """
+    Extract the rules from the additional context array, and retrieve the rule content.
+    """
+    if not additional_context:
+        return ""
+
+    selected_rules = [context["value"] for context in additional_context if context.get("type") == "rule"]
+    if len(selected_rules) == 0:
+        return ""
+
+    rules_content = []
+    for rule in selected_rules:
+        rule_content = get_rule(rule)
+        if rule_content is None or rule_content == "":
+            continue
+
+        rules_content.append(f"{rule}:\n\n{rule_content}")
+
+    return '\n'.join(rules_content)
+
+
+class RulesSection(PromptSection):
+    """Section for rules - never trimmed."""
+    trim_after_messages: Optional[int] = None
+    exclude_if_empty: bool = True
+    
+    def __init__(self, additional_context: Optional[List[Dict[str, str]]]):
+        self.additional_context = additional_context
+        self.content = get_rules_str(additional_context)
+        self.name = "Rules"
+

mito_ai/completions/prompt_builders/{utils.py → prompt_section_registry/selected_context.py}

@@ -1,91 +1,100 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from typing import List, Optional, Dict
-from mito_ai.rules.utils import get_rule
-
-
-def get_rules_str(additional_context: Optional[List[Dict[str, str]]]) -> str:
-    """
-    Extract the rules from the additional context array, and retrieve the rule content.
-    """
-    if not additional_context:
-        return ""
-
-    selected_rules = [context["value"] for context in additional_context if context.get("type") == "rule"]
-    if len(selected_rules) == 0:
-        return ""
-
-    rules_str = ""
-    for rule in selected_rules:
-        rule_content = get_rule(rule)
-        if rule_content is None or rule_content == "":
-            continue
-
-        rules_str += f"===========\n\nCustom Instructions Provided by User: {rule}\n\n{rule_content}\n\n==========="
-
-    return rules_str
-
+from .base import PromptSection
+from typing import Optional, List, Dict
+import json
 
 def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]]) -> str:
     """
-    Get the selected context from the additional context array.
+    Render the selected context from the additional context array.
     """
+    
     if not additional_context:
         return ""
-
+    
     # STEP 1: Extract each context type into a separate list
+    # Filter out any non-dict items and ensure we only process dictionaries
     selected_variables = [context["value"] for context in additional_context if context.get("type") == "variable"]
     selected_files = [context["value"] for context in additional_context if context.get("type") == "file"]
     selected_db_connections = [context["value"] for context in additional_context if context.get("type") == "db"]
     selected_images = [context["value"] for context in additional_context if context.get("type", "").startswith("image/")]
     selected_cells = [context["value"] for context in additional_context if context.get("type") == "cell"]
+    selected_line_selections = [context["value"] for context in additional_context if context.get("type") == "line_selection"]
 
     # STEP 2: Create a list of strings (instructions) for each context type
     context_parts = []
-
+    
     if len(selected_variables) > 0:
         context_parts.append(
             "The following variables have been selected by the user to be used in the task:\n"
             + "\n".join(selected_variables)
         )
-
+    
     if len(selected_files) > 0:
         context_parts.append(
             "The following files have been selected by the user to be used in the task:\n"
             + "\n".join(selected_files)
         )
-
+    
     if len(selected_db_connections) > 0:
         context_parts.append(
             "The following database connections have been selected by the user to be used in the task:\n"
             + "\n".join(selected_db_connections)
         )
-
+    
     if len(selected_images) > 0:
         context_parts.append(
             "The following images have been selected by the user to be used in the task:\n"
             + "\n".join(selected_images)
         )
-        
+
     if len(selected_cells) > 0:
         context_parts.append(
             "The following cells have been selected by the user to be used in the task:\n"
             + "\n".join(selected_cells)
         )
 
+    if len(selected_line_selections) > 0:
+        # Parse the line selection JSON values and format them for the prompt
+        line_selection_strs = []
+        for line_selection_json in selected_line_selections:
+            try:
+                selection_info = json.loads(line_selection_json)
+                cell_id = selection_info.get("cellId", "")
+                start_line = selection_info.get("startLine", 0)
+                end_line = selection_info.get("endLine", 0)
+                selected_code = selection_info.get("selectedCode", "")
+
+                # Format: Cell {cell_id} lines X-Y (0 indexed)\n[selected code]
+                if start_line == end_line:
+                    line_info = f"Cell {cell_id} line {start_line} (0 indexed)"
+                else:
+                    line_info = f"Cell {cell_id} lines {start_line}-{end_line} (0 indexed)"
+
+                line_selection_strs.append(f"{line_info}\n```python\n{selected_code}\n```")
+            except (json.JSONDecodeError, KeyError):
+                continue
+
+        if line_selection_strs:
+            context_parts.append(
+                "The user has selected the following lines of code to focus on:\n"
+                + "\n\n".join(line_selection_strs)
+            )
+
     # STEP 3: Combine into a single string
     return "\n\n".join(context_parts)
 
 
-def get_streamlit_app_status_str(notebook_id: str, notebook_path: str) -> str:
-    """
-    Get the streamlit app status string.
-    """
-    from mito_ai.path_utils import does_notebook_id_have_corresponding_app
-    if does_notebook_id_have_corresponding_app(notebook_id, notebook_path):
-        return "The notebook has an existing Streamlit app that you can edit"
-    return "The notebook does not have an existing Streamlit app. If you want to show an app to the user, you must create a new one."
-
-
+class SelectedContextSection(PromptSection):
+    """Section for selected context - never trimmed."""
+    trim_after_messages: Optional[int] = None
+    exclude_if_empty: bool = True
+    
+    def __init__(self, additional_context: Optional[List[Dict[str, str]]]):
+        self.additional_context = additional_context
+        self.content = get_selected_context_str(additional_context)
+        self.name = "SelectedContext"
+    
+    
 

mito_ai/completions/prompt_builders/prompt_section_registry/streamlit_app_status.py

@@ -0,0 +1,25 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+
+def get_streamlit_app_status_str(notebook_id: str, notebook_path: str) -> str:
+    """
+    Get the streamlit app status string.
+    """
+    from mito_ai.path_utils import does_notebook_id_have_corresponding_app
+    if does_notebook_id_have_corresponding_app(notebook_id, notebook_path):
+        return "The notebook has an existing Streamlit app that you can edit"
+    return "The notebook does not have an existing Streamlit app. If you want to show an app to the user, you must create a new one."
+
+
+class StreamlitAppStatusSection(PromptSection):
+    """Section for Streamlit app status."""
+    trim_after_messages: int = 3
+    
+    def __init__(self, notebook_id: str, notebook_path: str):
+        self.notebook_id = notebook_id
+        self.notebook_path = notebook_path
+        self.content = get_streamlit_app_status_str(notebook_id, notebook_path)
+        self.name = "StreamlitAppStatus"
+

mito_ai/completions/prompt_builders/prompt_section_registry/task.py

@@ -0,0 +1,12 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .base import PromptSection
+from typing import Optional
+
+
+class TaskSection(PromptSection):
+    """Section for user task - never trimmed."""
+    trim_after_messages: Optional[int] = None
+    exclude_if_empty: bool = False
+

mito_ai/completions/prompt_builders/prompt_section_registry/variables.py

@@ -0,0 +1,18 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Optional
+from .base import PromptSection
+
+
+class VariablesSection(PromptSection):
+    """Section for defined variables."""
+    trim_after_messages: int = 6
+    exclude_if_empty: bool = False
+    
+    def __init__(self, variables: Optional[List[str]]):
+        self.variables = variables
+        self.content = '\n'.join([f"{variable}" for variable in variables or []])
+        self.name = "Variables"
+
+