mito-ai 0.1.50__py3-none-any.whl

mito_ai/completions/models.py

@@ -0,0 +1,404 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import traceback
+from dataclasses import dataclass, field
+from typing import List, Literal, Optional, NewType, Dict, Any
+from openai.types.chat import ChatCompletionMessageParam
+from enum import Enum
+from pydantic import BaseModel
+
+# The ThreadID is the unique identifier for the chat thread.
+ThreadID = NewType('ThreadID', str)
+  
+########################################################
+# Agent Response formats
+########################################################
+    
+class CellUpdate(BaseModel):
+    type: Literal['modification', 'new']
+    index: Optional[int]
+    id: Optional[str]
+    code: str
+    code_summary: str
+    cell_type: Optional[Literal['code', 'markdown']]
+
+
+# Using a discriminated Pydantic model doesn't work well with OpenAI's API, 
+# so instead we just combine all of the possible response types into a single class 
+# 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']
+    message: str
+    cell_update: Optional[CellUpdate]
+    get_cell_output_cell_id: Optional[str]
+    next_steps: Optional[List[str]]
+    analysis_assumptions: Optional[List[str]]
+    edit_streamlit_app_prompt: Optional[str]
+    
+    
+@dataclass(frozen=True)
+class ResponseFormatInfo():
+    name: str
+    # Use the type because we are actually just providing the type format, not an actual instance of the format
+    format: type[AgentResponse]
+
+########################################################
+# Message Types and Metadata
+########################################################
+
+class MessageType(Enum):
+    """
+    This is all of the different types of messages that we support through the on_message handler.
+    """
+    CHAT = "chat"
+    SMART_DEBUG = "smartDebug"
+    CODE_EXPLAIN = "codeExplain"
+    AGENT_EXECUTION = "agent:execution"
+    AGENT_AUTO_ERROR_FIXUP = "agent:autoErrorFixup"
+    INLINE_COMPLETION = "inline_completion"
+    CHAT_NAME_GENERATION = "chat_name_generation"
+    START_NEW_CHAT = "start_new_chat"
+    FETCH_HISTORY = "fetch_history"
+    GET_THREADS = "get_threads"
+    DELETE_THREAD = "delete_thread"
+    UPDATE_MODEL_CONFIG = "update_model_config"
+    STREAMLIT_CONVERSION = "streamlit_conversion"
+    STOP_AGENT = "stop_agent"
+    DEPLOY_APP = "deploy_app"
+
+    
+@dataclass(frozen=True)
+class AIOptimizedCell():
+  cell_type: str
+  id: str
+  code: str
+  
+
+@dataclass(frozen=True)
+class ChatMessageMetadata():
+    promptType: Literal['chat']
+    threadId: ThreadID
+    input: str
+    activeCellCode: str 
+    activeCellId: str
+    variables: Optional[List[str]] = None
+    files: Optional[List[str]] = None
+    base64EncodedActiveCellOutput: Optional[str] = None
+    index: Optional[int] = None
+    stream: bool = False
+    additionalContext: Optional[List[Dict[str, str]]] = None
+    
+    
+@dataclass(frozen=True)
+class AgentExecutionMetadata():
+    promptType: Literal['agent:execution']
+    threadId: ThreadID
+    input: str
+    aiOptimizedCells: List[AIOptimizedCell]
+    activeCellId: str
+    isChromeBrowser: bool
+    notebookPath: str
+    notebookID: str
+    base64EncodedActiveCellOutput: Optional[str] = None
+    variables: Optional[List[str]] = None
+    files: Optional[List[str]] = None
+    index: Optional[int] = None
+    additionalContext: Optional[List[Dict[str, str]]] = None
+    
+@dataclass(frozen=True)
+class AgentSmartDebugMetadata():
+    promptType: Literal['agent:autoErrorFixup']
+    threadId: ThreadID
+    aiOptimizedCells: List[AIOptimizedCell]
+    errorMessage: str
+    error_message_producing_code_cell_id: str
+    isChromeBrowser: bool
+    variables: Optional[List[str]] = None
+    files: Optional[List[str]] = None
+    
+@dataclass(frozen=True)
+class SmartDebugMetadata():
+    promptType: Literal['smartDebug']
+    threadId: ThreadID
+    errorMessage: str
+    activeCellCode: str 
+    activeCellId: str
+    variables: Optional[List[str]] = None
+    files: Optional[List[str]] = None
+    
+@dataclass(frozen=True)
+class CodeExplainMetadata():    
+    promptType: Literal['codeExplain']
+    threadId: ThreadID
+    variables: Optional[List[str]] = None
+    activeCellCode: Optional[str] = None
+    
+@dataclass(frozen=True)
+class InlineCompleterMetadata():
+    promptType: Literal['inline_completion']
+    prefix: str 
+    suffix: str
+    variables: Optional[List[str]] = None
+    files: Optional[List[str]] = None
+
+@dataclass(frozen=True)
+class CompletionRequest:
+    """
+    Message send by the client to request an AI chat response.
+    """
+
+    # Message type.
+    type: MessageType
+
+    # Message UID generated by the client.
+    message_id: str
+
+    # Chat messages.
+    messages: List[ChatCompletionMessageParam] = field(default_factory=list)
+
+    # Whether to stream the response (if supported by the model).
+    stream: bool = False
+    
+    # Environment information from the client
+    environment: Optional[Dict[str, Any]] = None
+    
+    
+@dataclass(frozen=True)
+class AICapabilities:
+    """
+    AI provider capabilities
+    """
+
+    # Configuration schema.
+    configuration: dict
+
+    # AI provider name.
+    provider: str
+
+    # Message type.
+    type: str = "ai_capabilities"
+
+
+@dataclass(frozen=True)
+class CompletionItemError:
+    """
+    Completion item error information.
+    """
+
+    # Error message.
+    message: Optional[str] = None
+
+
+@dataclass(frozen=True)
+class CompletionItem:
+    """
+    A completion suggestion.
+    """
+
+    # The completion.
+    content: str
+
+    # Whether the completion is incomplete or not.
+    isIncomplete: Optional[bool] = None
+    
+    # Unique token identifying the completion request in the frontend.
+    token: Optional[str] = None
+
+    # Error information for the completion item.
+    error: Optional[CompletionItemError] = None
+
+
+@dataclass(frozen=True)
+class CompletionError:
+    """
+    Completion error description.
+    """
+
+    # Error type.
+    error_type: str
+
+    # Error title.
+    title: str
+
+    # Error traceback.
+    traceback: str
+
+    # Hint to resolve the error.
+    hint: str = ""
+
+    @staticmethod
+    def from_exception(exception: BaseException, hint: str = "") -> "CompletionError":
+        """
+        Create a completion error from an exception.
+        
+        Note: OpenAI exceptions can include a 'body' attribute with detailed error information.
+        While mypy doesn't know about this attribute on BaseException, we need to handle it
+        to properly extract error messages from OpenAI API responses.
+        """
+        from mito_ai.utils.mito_server_utils import ProviderCompletionException
+
+        
+        # Handle ProviderCompletionException specially
+        if isinstance(exception, ProviderCompletionException):
+            return CompletionError(
+                error_type="LLM Provider Error", 
+                title=exception.user_friendly_title, 
+                traceback=traceback.format_exc(),
+                hint=exception.user_friendly_hint
+            )
+        
+        # Handle all other exceptions as before
+        error_type = type(exception)
+        error_module = getattr(error_type, "__module__", "")
+        
+        # Handle OpenAI exceptions that have a 'body' attribute
+        title = ""
+        body = getattr(exception, "body", None)
+        if body and hasattr(body, "get"):
+            title = body.get("message", "")
+        
+        if not title and exception.args:
+            title = exception.args[0]
+        
+        if not title:
+            title = "Exception"
+            
+        return CompletionError(
+            error_type=f"{error_module}.{error_type.__name__}"
+            if error_module
+            else error_type.__name__,
+            title=title,
+            traceback=traceback.format_exc(),
+            hint=hint,
+        )
+
+@dataclass(frozen=True)
+class ErrorMessage(CompletionError):
+    """
+    Error message.
+    """
+
+    # Message type.
+    type: Literal["error"] = "error"
+
+
+
+@dataclass(frozen=True)
+class CompletionReply:
+    """
+    Message sent from model to client with the completion suggestions.
+    """
+
+    # List of completion items.
+    items: List[CompletionItem]
+
+    # Parent message UID.
+    parent_id: str
+
+    # Message type.
+    type: Literal["reply"] = "reply"
+
+    # Completion error.
+    error: Optional[CompletionError] = None
+
+
+@dataclass(frozen=True)
+class CompletionStreamChunk:
+    """
+    Message sent from model to client with the infill suggestions
+    """
+
+    chunk: CompletionItem
+
+    # Parent message UID.
+    parent_id: str
+
+    # Whether the completion is done or not.
+    done: bool
+
+    # Message type.
+    type: Literal["chunk"] = "chunk"
+
+    # Completion error.
+    error: Optional[CompletionError] = None
+    """Completion error."""
+
+@dataclass(frozen=True)
+class FetchHistoryReply:
+    """
+    Message sent from model to client with the chat history.
+    """
+
+    # Message UID.
+    parent_id: str
+
+    # List of chat messages.
+    items: List[ChatCompletionMessageParam]
+
+    # Message type.
+    type: Literal["reply"] = "reply"
+
+@dataclass(frozen=True)
+class ChatThreadMetadata:
+    """
+    Chat thread item.
+    """
+
+    thread_id: ThreadID
+
+    name: str
+
+    creation_ts: float
+
+    last_interaction_ts: float
+
+@dataclass(frozen=True)
+class StartNewChatReply:
+    """
+    Message sent from model to client after starting a new chat thread.
+    """
+
+    # Message UID.
+    parent_id: str
+
+    # Chat thread item.
+    thread_id: ThreadID
+
+    # Message type.
+    type: Literal["reply"] = "reply"
+
+@dataclass(frozen=True)
+class FetchThreadsReply:
+    """
+    Message sent from model to client with the chat threads.
+    """
+
+    # Message UID.
+    parent_id: str
+
+    # List of chat threads.
+    threads: List[ChatThreadMetadata]
+
+    # Message type.
+    type: Literal["reply"] = "reply"
+
+@dataclass(frozen=True)
+class DeleteThreadReply:
+    """
+    Message sent from model to client after deleting a chat thread.
+    """
+
+    # Message UID.
+    parent_id: str
+
+    #Success message
+    success: bool
+
+    # Message type.
+    type: Literal["reply"] = "reply"
+
+@dataclass(frozen=True)
+class UpdateModelConfigMetadata:
+    model: str

mito_ai/completions/prompt_builders/__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/completions/prompt_builders/agent_execution_prompt.py

@@ -0,0 +1,57 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+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
+)
+
+
+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()

mito_ai/completions/prompt_builders/agent_smart_debug_prompt.py

@@ -0,0 +1,160 @@
+# 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 AgentSmartDebugMetadata
+from mito_ai.completions.prompt_builders.prompt_constants import (
+    FILES_SECTION_HEADING,
+    JUPYTER_NOTEBOOK_SECTION_HEADING,
+    VARIABLES_SECTION_HEADING
+)
+
+# 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'
+- 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 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:
+    {{
+        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]
+}})
+    }},
+    {{
+        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'
+
+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_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}
+
+Error Traceback:
+{md.errorMessage}
+"""