Merge pull request #1571 from mito-ds/fix/agent-error-fixuo

Fix/agent error fixup

Author: aarondr77

mito-ai/mito_ai/completion_handlers/agent_auto_error_fixup_handler.py

@@ -1,19 +1,21 @@
 from openai.types.chat import ChatCompletionMessageParam
-from mito_ai.models import SmartDebugMetadata, MessageType
-from mito_ai.prompt_builders.smart_debug_prompt import create_error_prompt, remove_inner_thoughts_from_message
+from mito_ai.models import AgentResponse, AgentSmartDebugMetadata, MessageType, ResponseFormatInfo
+from mito_ai.prompt_builders.agent_smart_debug_prompt import create_agent_smart_debug_prompt
 from mito_ai.providers import OpenAIProvider
 from mito_ai.message_history import GlobalMessageHistory
 from mito_ai.completion_handlers.completion_handler import CompletionHandler
 from mito_ai.completion_handlers.open_ai_models import MESSAGE_TYPE_TO_MODEL
 from mito_ai.completion_handlers.utils import append_agent_system_message
+
+
 __all__ = ["get_agent_auto_error_fixup_completion"]
 
-class AgentAutoErrorFixupHandler(CompletionHandler[SmartDebugMetadata]):
+class AgentAutoErrorFixupHandler(CompletionHandler[AgentSmartDebugMetadata]):
     """Handler for smart debug completions."""
 
     @staticmethod
     async def get_completion(
-        metadata: SmartDebugMetadata,
+        metadata: AgentSmartDebugMetadata,
         provider: OpenAIProvider,
         message_history: GlobalMessageHistory
     ) -> str:
@@ -22,40 +24,31 @@ async def get_completion(
         # Add the system message if it doens't alredy exist
         await append_agent_system_message(message_history, provider)
 
-        error_message = metadata.errorMessage
-        active_cell_code = metadata.activeCellCode or ''
-        variables = metadata.variables or []
-        files = metadata.files or []
-        
         # Create the prompt
-        prompt = create_error_prompt(
-            error_message, 
-            active_cell_code, 
-            variables,
-            files
-        )
+        prompt = create_agent_smart_debug_prompt(metadata)
 
         # Add the prompt to the message history
         new_ai_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": prompt}
-        new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": error_message}
+        new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": metadata.errorMessage}
         await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, provider)
 
         # Get the completion
         completion = await provider.request_completions(
             messages=message_history.ai_optimized_history, 
             model=MESSAGE_TYPE_TO_MODEL[MessageType.AGENT_AUTO_ERROR_FIXUP],
+            response_format_info=ResponseFormatInfo(
+                name='agent_response',
+                format=AgentResponse
+            ),
             message_type=MessageType.AGENT_AUTO_ERROR_FIXUP
         )
 
-        # Process the completion to remove inner thoughts
-        display_completion = remove_inner_thoughts_from_message(completion)
-        
         # Add the response to message history
         ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
-        display_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": display_completion}
+        display_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
         await message_history.append_message(ai_response_message, display_response_message, provider)
 
-        return display_completion
+        return completion
 
 # Use the static method directly
 get_agent_auto_error_fixup_completion = AgentAutoErrorFixupHandler.get_completion

mito-ai/mito_ai/completion_handlers/completion_handler.py

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

mito-ai/mito_ai/handlers.py

@@ -14,6 +14,7 @@
 from mito_ai.message_history import GlobalMessageHistory
 from mito_ai.logger import get_logger
 from mito_ai.models import (
+    AgentSmartDebugMetadata,
     CompletionError,
     CompletionItem,
     CompletionReply,
@@ -195,7 +196,7 @@ async def on_message(self, message: str) -> None: # type: ignore
                 agent_execution_metadata = AgentExecutionMetadata(**metadata_dict)
                 completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, message_history)
             elif type == MessageType.AGENT_AUTO_ERROR_FIXUP:
-                agent_auto_error_fixup_metadata = SmartDebugMetadata(**metadata_dict)
+                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)
             elif type == MessageType.INLINE_COMPLETION:
                 inline_completer_metadata = InlineCompleterMetadata(**metadata_dict)

mito-ai/mito_ai/models.py

@@ -9,7 +9,7 @@
 ThreadID = NewType('ThreadID', str)
 
 @dataclass(frozen=True)
-class AIOptimizedCells():
+class AIOptimizedCell():
   cell_type: str
   id: str
   code: str
@@ -66,7 +66,16 @@ class ChatMessageMetadata():
 class AgentExecutionMetadata():
     promptType: Literal['agent:execution']
     input: str
-    aiOptimizedCells: List[AIOptimizedCells]
+    aiOptimizedCells: List[AIOptimizedCell]
+    variables: Optional[List[str]] = None
+    files: Optional[List[str]] = None
+    
+@dataclass(frozen=True)
+class AgentSmartDebugMetadata():
+    promptType: Literal['agent:autoErrorFixup']
+    aiOptimizedCells: List[AIOptimizedCell]
+    errorMessage: str
+    error_message_producing_code_cell_id: str
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None
 

mito-ai/mito_ai/prompt_builders/agent_smart_debug_prompt.py

@@ -0,0 +1,143 @@
+from typing import List
+from mito_ai.models import AgentSmartDebugMetadata
+
+# TODO:
+# 1. In the future, it might make sense to pass the previous CELL_UPDATE to this prompt?
+# 2. In the future, we should let the agent fix up the error by updating a different cell. This is sometimes a better solution. 
+# However, to do this, we then need to know which code cells to run in order to validate the update is correct! If the error was 
+# produced by code cell 3, and the agent corrects the source of the error in code cell 2, we then need to run cell 2 and 3 to validate
+# the cell update worked properly. This could be many cells if there are intermediate cells. It might require something like a dependency 
+# 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 []])
+    
+    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.
+
+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. 
+
+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.
+
+<Instructions for each Phase />
+
+ERROR ANALYSIS:
+
+- Identify error type (Syntax, Runtime, Logic).
+- Use the defined variables and Jupyter Notebook to understand the error.
+- Consider kernel state and execution order
+
+INTENT PRESERVATION:
+
+- Try to understand the user's intent using the defined variables and the Jupyter Notebook
+
+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.
+- 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'
+
+<Example>
+
+<Input>
+
+Files in the current directory:
+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'])\"\"\"
+    }},
+]
+
+Defined Variables:
+{{
+    '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'
+
+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 >
+
+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.
+
+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. 
+
+</ 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))"
+    }}
+}}
+
+</Output>
+
+</Example>
+
+Files in the current directory:
+{files_str}
+
+Jupyter Notebook:
+{ai_optimized_cells_str}
+
+Defined Variables:
+{variables_str}
+
+Cell ID of the Error Producing Code Cell:
+{md.error_message_producing_code_cell_id}
+
+Error Traceback:
+{md.errorMessage}
+"""

mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx

@@ -2,7 +2,7 @@ import OpenAI from "openai";
 import { IContextManager } from "../ContextManager/ContextManagerPlugin";
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { getActiveCellCode, getActiveCellID, getAIOptimizedCells, getCellCodeByID } from "../../utils/notebook";
-import { AgentResponse, IAgentExecutionMetadata, IChatMessageMetadata, ICodeExplainMetadata, ISmartDebugMetadata } from "../../utils/websocket/models";
+import { AgentResponse, IAgentExecutionMetadata, IAgentSmartDebugMetadata, IChatMessageMetadata, ICodeExplainMetadata, ISmartDebugMetadata } from "../../utils/websocket/models";
 import { addMarkdownCodeFormatting } from "../../utils/strings";
 
 export type PromptType = 
@@ -145,7 +145,7 @@ export class ChatHistoryManager {
     }
 
 
-    addDebugErrorMessage(errorMessage: string, promptType: PromptType): ISmartDebugMetadata {
+    addSmartDebugMessage(errorMessage: string): ISmartDebugMetadata {
 
         const activeCellID = getActiveCellID(this.notebookTracker)
         const activeCellCode = getCellCodeByID(this.notebookTracker, activeCellID)
@@ -163,13 +163,39 @@ export class ChatHistoryManager {
                 message: getDisplayedOptimizedUserMessage(errorMessage, activeCellCode), 
                 type: 'openai message',
                 codeCellID: activeCellID,
-                promptType: promptType
+                promptType: 'smartDebug'
             }
         );
 
         return smartDebugMetadata
     }
 
+    addAgentSmartDebugMessage(errorMessage: string): IAgentSmartDebugMetadata {
+
+        const activeCellID = getActiveCellID(this.notebookTracker)
+        const activeCellCode = getActiveCellCode(this.notebookTracker)
+
+        const agentSmartDebugMetadata: IAgentSmartDebugMetadata = {
+            promptType: 'agent:autoErrorFixup',
+            aiOptimizedCells: getAIOptimizedCells(this.notebookTracker),
+            variables: this.contextManager.variables,
+            files: this.contextManager.files,
+            errorMessage: errorMessage,
+            error_message_producing_code_cell_id: activeCellID || ''
+        }
+
+        this.displayOptimizedChatHistory.push(
+            {
+                message: getDisplayedOptimizedUserMessage(errorMessage, activeCellCode), 
+                type: 'openai message',
+                codeCellID: activeCellID,
+                promptType: 'agent:autoErrorFixup'
+            }
+        );
+
+        return agentSmartDebugMetadata
+    }
+
     addExplainCodeMessage(): ICodeExplainMetadata {
 
         const activeCellID = getActiveCellID(this.notebookTracker)