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

mito_ai/__init__.py

@@ -19,6 +19,7 @@ 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
+from mito_ai.chart_wizard.urls import get_chart_wizard_urls
 
 # Force Matplotlib to use the Jupyter inline backend.
 # Background: importing Streamlit sets os.environ["MPLBACKEND"] = "Agg" very early.
@@ -109,6 +110,7 @@ def _load_jupyter_server_extension(server_app) -> None: # 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
+    handlers.extend(get_chart_wizard_urls(base_url, open_ai_provider)) # 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.55'
+__version__ = VERSION = '0.1.57'

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/chart_wizard/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+

mito_ai/chart_wizard/handlers.py

@@ -0,0 +1,52 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import tornado
+from typing import List
+from jupyter_server.base.handlers import APIHandler
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.utils.anthropic_utils import FAST_ANTHROPIC_MODEL
+from mito_ai.completions.models import MessageType
+from mito_ai.completions.prompt_builders.chart_conversion_prompt import create_chart_conversion_prompt
+
+class ChartWizardHandler(APIHandler):
+    def initialize(self, llm: OpenAIProvider) -> None:
+        """Initialize the handler with the LLM provider."""
+        super().initialize()
+        self._llm = llm
+
+    @tornado.web.authenticated
+    async def post(self) -> None:
+        """POST endpoint that receives code from the frontend and sends it to LLM."""
+        try:
+            data = json.loads(self.request.body.decode('utf-8'))
+            code = data.get('code', '')
+            
+            # Create prompt using the prompt builder
+            prompt = create_chart_conversion_prompt(code)
+            
+            # Call LLM
+            messages: List[ChatCompletionMessageParam] = [{"role": "user", "content": prompt}]
+            converted_code = await self._llm.request_completions(
+                messages=messages,
+                model=FAST_ANTHROPIC_MODEL,
+                message_type=MessageType.CHAT,
+                thread_id=None
+            )
+            
+            # Return the converted code
+            self.write({
+                "message": "Code converted successfully",
+                "converted_code": converted_code
+            })
+            self.finish()
+        except json.JSONDecodeError:
+            self.set_status(400)
+            self.write({"error": "Invalid JSON in request body"})
+            self.finish()
+        except Exception as e:
+            self.set_status(500)
+            self.write({"error": str(e)})
+            self.finish()

mito_ai/chart_wizard/urls.py

@@ -0,0 +1,23 @@
+# 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.chart_wizard.handlers import ChartWizardHandler
+from mito_ai.completions.providers import OpenAIProvider
+
+
+def get_chart_wizard_urls(base_url: str, llm: OpenAIProvider) -> List[Tuple[str, Any, dict]]:
+    """Get all chart wizard related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+        llm: The OpenAI provider instance
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    return [
+        (url_path_join(BASE_URL, "chart-wizard"), ChartWizardHandler, {"llm": llm}),
+    ]

mito_ai/completions/completion_handlers/completion_handler.py

@@ -3,11 +3,20 @@
 
 from typing import Protocol, TypeVar
 from abc import abstractmethod, ABCMeta
-from mito_ai.completions.models import ChatMessageMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, InlineCompleterMetadata, AgentSmartDebugMetadata
+from mito_ai.completions.models import ChatMessageMetadata, ScratchpadResultMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, InlineCompleterMetadata, AgentSmartDebugMetadata
 from mito_ai.completions.providers import OpenAIProvider
 from mito_ai.completions.message_history import GlobalMessageHistory
 
-T = TypeVar('T', ChatMessageMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, AgentSmartDebugMetadata, InlineCompleterMetadata, contravariant=True)
+T = TypeVar('T', 
+            ChatMessageMetadata, 
+            SmartDebugMetadata, 
+            CodeExplainMetadata, 
+            AgentExecutionMetadata, 
+            AgentSmartDebugMetadata, 
+            InlineCompleterMetadata,
+            ScratchpadResultMetadata,
+            contravariant=True
+        )
 
 class CompletionHandler(Protocol[T], metaclass=ABCMeta):
     """Protocol defining the interface for completion handlers.

mito_ai/completions/completion_handlers/scratchpad_result_handler.py

@@ -0,0 +1,66 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Literal, Union
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.models import ScratchpadResultMetadata, MessageType, ResponseFormatInfo, AgentResponse
+from mito_ai.completions.prompt_builders.scratchpad_result_prompt import create_scratchpad_result_prompt
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.completions.completion_handlers.completion_handler import CompletionHandler
+from mito_ai.completions.completion_handlers.utils import append_agent_system_message, create_ai_optimized_message
+
+__all__ = ["get_scratchpad_result_completion"]
+
+class ScratchpadResultHandler(CompletionHandler[ScratchpadResultMetadata]):
+    """Handler for scratchpad result completions."""
+    
+    @staticmethod
+    async def get_completion(
+        metadata: ScratchpadResultMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get a scratchpad result completion from the AI provider."""
+
+        if metadata.index is not None:
+            message_history.truncate_histories(
+                thread_id=metadata.threadId,
+                index=metadata.index
+            )
+
+        # Add the system message if it doesn't already exist
+        await append_agent_system_message(message_history, model, provider, metadata.threadId, True)
+        
+        # Create the prompt
+        prompt = create_scratchpad_result_prompt(metadata)
+        display_prompt = ""
+        
+        # Add the prompt to the message history
+        new_ai_optimized_message = create_ai_optimized_message(prompt, None, None)
+        new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": display_prompt}
+        
+        await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, model, provider, metadata.threadId)
+        
+        # Get the completion
+        completion = await provider.request_completions(
+            messages=message_history.get_ai_optimized_history(metadata.threadId), 
+            model=model,
+            response_format_info=ResponseFormatInfo(
+                name='agent_response',
+                format=AgentResponse
+            ),
+            message_type=MessageType.AGENT_SCRATCHPAD_RESULT,
+            user_input="",
+            thread_id=metadata.threadId
+        )
+        
+        ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
+        
+        await message_history.append_message(ai_response_message, ai_response_message, model, provider, metadata.threadId)
+        
+        return completion
+
+# Use the static method directly
+get_scratchpad_result_completion = ScratchpadResultHandler.get_completion

mito_ai/completions/handlers.py

@@ -34,6 +34,7 @@ from mito_ai.completions.models import (
     CodeExplainMetadata,
     AgentExecutionMetadata,
     InlineCompleterMetadata,
+    ScratchpadResultMetadata,
     MessageType
 )
 from mito_ai.completions.providers import OpenAIProvider
@@ -45,6 +46,7 @@ from mito_ai.completions.completion_handlers.code_explain_handler import get_cod
 from mito_ai.completions.completion_handlers.inline_completer_handler import get_inline_completion
 from mito_ai.completions.completion_handlers.agent_execution_handler import get_agent_execution_completion
 from mito_ai.completions.completion_handlers.agent_auto_error_fixup_handler import get_agent_auto_error_fixup_completion
+from mito_ai.completions.completion_handlers.scratchpad_result_handler import get_scratchpad_result_completion
 from mito_ai.utils.telemetry_utils import identify
 
 FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
@@ -314,6 +316,9 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             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, self._message_history, model)
+            elif type == MessageType.AGENT_SCRATCHPAD_RESULT:
+                scratchpad_result_metadata = ScratchpadResultMetadata(**metadata_dict)
+                completion = await get_scratchpad_result_completion(scratchpad_result_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, self._message_history, model)

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
@@ -29,13 +29,26 @@ 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', 'create_streamlit_app', 'edit_streamlit_app']
+    type: Literal[
+        'cell_update', 
+        'get_cell_output', 
+        'run_all_cells', 
+        'finished_task', 
+        'create_streamlit_app', 
+        'edit_streamlit_app', 
+        'ask_user_question', 
+        'scratchpad',
+    ]
     message: str
     cell_update: Optional[CellUpdate]
     get_cell_output_cell_id: Optional[str]
     next_steps: Optional[List[str]]
     analysis_assumptions: Optional[List[str]]
     streamlit_app_prompt: Optional[str]
+    question: Optional[str]
+    answers: Optional[List[str]]
+    scratchpad_code: Optional[str]
+    scratchpad_summary: Optional[str]
     
     
 @dataclass(frozen=True)
@@ -67,6 +80,7 @@ class MessageType(Enum):
     STREAMLIT_CONVERSION = "streamlit_conversion"
     STOP_AGENT = "stop_agent"
     DEPLOY_APP = "deploy_app"
+    AGENT_SCRATCHPAD_RESULT = "agent:scratchpad-result"
 
     
 @dataclass(frozen=True)
@@ -136,13 +150,20 @@ class CodeExplainMetadata():
     activeCellCode: Optional[str] = None
     
 @dataclass(frozen=True)
-class InlineCompleterMetadata():
+class InlineCompleterMetadata():    
     promptType: Literal['inline_completion']
     prefix: str 
     suffix: str
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None
 
+@dataclass(frozen=True)
+class ScratchpadResultMetadata():
+    promptType: Literal['agent:scratchpad-result']
+    threadId: ThreadID
+    scratchpadResult: str
+    index: Optional[int] = None
+
 @dataclass(frozen=True)
 class CompletionRequest:
     """

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,22 +15,21 @@ 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.
 
-It is very important that When fixing this error, you do not change the original intent of the code cell. 
+It is very important that when fixing this error, you do not change the original intent of the code cell. 
 
 To fix this error, take the following approach: 
 Step 1: ERROR ANALYSIS: Analyze the error message to identify why the code cell errored.
 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.
+Step 3: ERROR CORRECTION: Respond with a new CELL_UPDATE that is applied to the same cell as the erroring CELL_UPDATE or use the ASK_USER_QUESTION tool to get more information about how to proceed.
 
-<Instructions for each Phase />
+INSTRUCTIONS FOR EACH PHASE
 
 ERROR ANALYSIS:
 
@@ -47,13 +43,15 @@ INTENT PRESERVATION:
 
 ERROR CORRECTION:
 
-- Return the full, updated version of cell {md.error_message_producing_code_cell_id} with the error fixed and a short explanation of the error.
+- Use one of your tools to correct the error or get more information from the user on how to proceed.
+- If you use the CELL_UPDATE tool, you must reutn the full updated version of cell {md.error_message_producing_code_cell_id} with the error fixed and a short explanation of the error.
 - You can only update code in {md.error_message_producing_code_cell_id}. You are unable to edit the code in any other cell when resolving this error.
 - Propose a solution that fixes the error and does not change the user's intent.
 - Make the solution as simple as possible.
 - Reuse as much of the existing code as possible.
 - DO NOT ADD TEMPORARY COMMENTS like '# Fixed the typo here' or '# Added this line to fix the error'
 - If you encounter a ModuleNotFoundError, you can install the package by adding the the following line to the top of the code cell: `!pip install <package_name> --quiet`.
+- If the error is not resolvable without getting more information from the user, you can respond with a ASK_USER_QUESTION tool call.
 - If you encounter a NameError, you can use the RUN_ALL_CELLS tool to run all cells from the top of the notebook to the bottom to bring the variable into scope.
     RUN_ALL_CELLS:
     When you want to execute all cells in the notebook from top to bottom, respond with this format:
@@ -61,100 +59,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)