mito-ai 0.1.33__py3-none-any.whl → 0.1.49__py3-none-any.whl

mito_ai/completions/handlers.py

@@ -14,6 +14,7 @@ import tornado.web
 from jupyter_core.utils import ensure_async
 from jupyter_server.base.handlers import JupyterHandler
 from tornado.websocket import WebSocketHandler
+from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.logger import get_logger
 from mito_ai.completions.models import (
@@ -48,11 +49,8 @@ from mito_ai.utils.telemetry_utils import identify
 
 FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
 
-# The GlobalMessageHistory is responsible for updating the message histories stored in the .mito/ai-chats directory.
-# We create one GlobalMessageHistory per backend server instance instead of one per websocket connection so that the
-# there is one manager of the locks for the .mito/ai-chats directory. This is my current understanding and it 
-# might be incorrect!
-message_history = GlobalMessageHistory()
+# The GlobalMessageHistory is now created in __init__.py and passed to handlers
+# to ensure there's only one instance managing the .mito/ai-chats directory locks
 
 # This handler is responsible for the mito_ai/completions endpoint.
 # It takes a message from the user, sends it to the OpenAI API, and returns the response.
@@ -61,12 +59,14 @@ message_history = GlobalMessageHistory()
 class CompletionHandler(JupyterHandler, WebSocketHandler):
     """Completion websocket handler."""
 
-    def initialize(self, llm: OpenAIProvider) -> None:
+    def initialize(self, llm: OpenAIProvider, message_history: GlobalMessageHistory) -> None:
         super().initialize()
         self.log.debug("Initializing websocket connection %s", self.request.path)
         self._llm = llm
+        self._message_history = message_history
         self.is_pro = is_pro()
         self._selected_model = FALLBACK_MODEL
+        self.is_electron = False
         identify(llm.key_type)
         
     @property
@@ -128,6 +128,18 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             parsed_message = json.loads(message)
             metadata_dict = parsed_message.get('metadata', {})
             type: MessageType = MessageType(parsed_message.get('type'))
+            
+            # Extract environment information from the message
+            environment = parsed_message.get('environment', {})
+            if environment:
+                is_electron = environment.get('isElectron', None)
+                if is_electron is not None:
+                    if is_electron != self.is_electron:
+                        # If the is_electron status is different, log it
+                        identify(key_type=self._llm.key_type, is_electron=is_electron)
+                    
+                    self.is_electron = is_electron
+
         except ValueError as e:
             self.log.error("Invalid completion request.", exc_info=e)
             return
@@ -136,7 +148,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
 
         # Clear history if the type is "start_new_chat"
         if type == MessageType.START_NEW_CHAT:
-            thread_id = message_history.create_new_thread()
+            thread_id = self._message_history.create_new_thread()
             
             reply = StartNewChatReply(
                 parent_id=parsed_message.get("message_id"),
@@ -147,7 +159,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
 
         # Handle get_threads: return list of chat threads
         if type == MessageType.GET_THREADS:
-            threads = message_history.get_threads()
+            threads = self._message_history.get_threads()
             reply = FetchThreadsReply(
                 parent_id=parsed_message.get("message_id"),
                 threads=threads
@@ -159,7 +171,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
         if type == MessageType.DELETE_THREAD:
             thread_id_to_delete = metadata_dict.get('thread_id')
             if thread_id_to_delete:
-                is_thread_deleted = message_history.delete_thread(thread_id_to_delete)
+                is_thread_deleted = self._message_history.delete_thread(thread_id_to_delete)
                 reply = DeleteThreadReply(
                     parent_id=parsed_message.get("message_id"),
                     success=is_thread_deleted
@@ -175,7 +187,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             
             # If a thread_id is provided, use that thread's history; otherwise, use newest.
             thread_id = metadata_dict.get('thread_id')
-            display_history = message_history.get_display_history(thread_id)
+            display_history = self._message_history.get_display_history(thread_id)
             
             reply = FetchHistoryReply(
                 parent_id=parsed_message.get('message_id'),
@@ -209,7 +221,32 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                 )
                 self.reply(reply)
             return
-        
+
+        if type == MessageType.STOP_AGENT:
+            thread_id_to_stop = metadata_dict.get('threadId')
+            if thread_id_to_stop:
+                self.log.info(f"Stopping agent, thread ID: {thread_id_to_stop}")
+                
+                ai_optimized_message: ChatCompletionMessageParam = {
+                    "role": "assistant",
+                    "content": "The user made the following request: Stop processing my last request. I want to change it. Please answer my future requests without going back and finising my previous request."
+                }
+                display_optimized_message: ChatCompletionMessageParam = {
+                    "role": "assistant",
+                    "content": "Agent interupted by user "
+                }
+                
+                await self._message_history.append_message(
+                    ai_optimized_message=ai_optimized_message,
+                    display_message=display_optimized_message,
+                    model=self._selected_model,
+                    llm_provider=self._llm,
+                    thread_id=thread_id_to_stop
+                )
+            else:
+                self.log.info("Trying to stop agent, but no thread ID available")
+            return
+
         try:
             # Get completion based on message type
             completion = None
@@ -227,7 +264,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_chat_completion(
                         chat_metadata, 
                         self._llm, 
-                        message_history, 
+                        self._message_history, 
                         message_id,
                         self.reply,
                         model
@@ -235,7 +272,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_chat_completion(chat_metadata, self._llm, message_history, model)
+                    completion = await get_chat_completion(chat_metadata, self._llm, self._message_history, model)
             elif type == MessageType.SMART_DEBUG:
                 smart_debug_metadata = SmartDebugMetadata(**metadata_dict)
                 # Handle streaming if requested and available
@@ -244,7 +281,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_smart_debug_completion(
                         smart_debug_metadata, 
                         self._llm, 
-                        message_history, 
+                        self._message_history, 
                         message_id,
                         self.reply,
                         model
@@ -252,7 +289,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, message_history, model)
+                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, self._message_history, model)
             elif type == MessageType.CODE_EXPLAIN:
                 code_explain_metadata = CodeExplainMetadata(**metadata_dict)
 
@@ -262,7 +299,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_code_explain_completion(
                         code_explain_metadata, 
                         self._llm, 
-                        message_history,
+                        self._message_history,
                         message_id,
                         self.reply,
                         model
@@ -270,16 +307,16 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, message_history, model)
+                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, self._message_history, model)
             elif type == MessageType.AGENT_EXECUTION:
                 agent_execution_metadata = AgentExecutionMetadata(**metadata_dict)
-                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, message_history, model)
+                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, self._message_history, model)
             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, message_history, model)
+                completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_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, message_history, model)
+                completion = await get_inline_completion(inline_completer_metadata, self._llm, self._message_history, model)
             else:
                 raise ValueError(f"Invalid message type: {type}")
             

mito_ai/completions/message_history.py

@@ -251,7 +251,15 @@ class GlobalMessageHistory:
         with self._lock:
             if thread_id not in self._chat_threads:
                 return []
-            return self._chat_threads[thread_id].display_history
+            
+            thread = self._chat_threads[thread_id]
+            display_history = thread.display_history
+            
+            # When we get a thread, update it's last interaction time so that if the 
+            # user refreshes their browser, this chat will re-appear as the last opened chat.
+            self._update_last_interaction(thread)
+            self._save_thread_to_disk(thread)
+            return display_history
 
     async def append_message(
         self, 

mito_ai/completions/models.py

@@ -3,10 +3,10 @@
 
 import traceback
 from dataclasses import dataclass, field
-from typing import Annotated, List, Literal, Optional, Type, Union, NewType
+from typing import List, Literal, Optional, NewType, Dict, Any
 from openai.types.chat import ChatCompletionMessageParam
 from enum import Enum
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
 # The ThreadID is the unique identifier for the chat thread.
 ThreadID = NewType('ThreadID', str)
@@ -20,6 +20,7 @@ class CellUpdate(BaseModel):
     index: Optional[int]
     id: Optional[str]
     code: str
+    code_summary: str
     cell_type: Optional[Literal['code', 'markdown']]
 
 
@@ -28,11 +29,13 @@ 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', 'finished_task']
+    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)
@@ -61,6 +64,9 @@ class MessageType(Enum):
     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)
@@ -82,7 +88,7 @@ class ChatMessageMetadata():
     base64EncodedActiveCellOutput: Optional[str] = None
     index: Optional[int] = None
     stream: bool = False
-    selectedRules: Optional[List[str]] = None
+    additionalContext: Optional[List[Dict[str, str]]] = None
     
     
 @dataclass(frozen=True)
@@ -91,12 +97,15 @@ class AgentExecutionMetadata():
     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
-    selectedRules: Optional[List[str]] = None
+    additionalContext: Optional[List[Dict[str, str]]] = None
     
 @dataclass(frozen=True)
 class AgentSmartDebugMetadata():
@@ -133,7 +142,7 @@ class InlineCompleterMetadata():
     suffix: str
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None
-    
+
 @dataclass(frozen=True)
 class CompletionRequest:
     """
@@ -152,6 +161,9 @@ class CompletionRequest:
     # 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:
@@ -225,6 +237,19 @@ class CompletionError:
         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__", "")
         
@@ -249,7 +274,6 @@ class CompletionError:
             hint=hint,
         )
 
-
 @dataclass(frozen=True)
 class ErrorMessage(CompletionError):
     """

mito_ai/completions/prompt_builders/agent_execution_prompt.py

@@ -3,18 +3,29 @@
 
 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
+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.selectedRules)
+    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.
 
@@ -30,6 +41,14 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> 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: 

mito_ai/completions/prompt_builders/agent_smart_debug_prompt.py

@@ -54,6 +54,14 @@ ERROR CORRECTION:
 - 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>