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

mito_ai/__init__.py

@@ -4,11 +4,11 @@
 from typing import List, Dict
 from jupyter_server.utils import url_path_join
 from mito_ai.completions.handlers import CompletionHandler
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.app_deploy.handlers import AppDeployHandler
-from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
 from mito_ai.log.urls import get_log_urls
+from mito_ai.utils.litellm_utils import is_litellm_configured
 from mito_ai.version_check import VersionCheckHandler
 from mito_ai.db.urls import get_db_urls
 from mito_ai.settings.urls import get_settings_urls
@@ -20,6 +20,8 @@ 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
+from mito_ai.utils.version_utils import is_enterprise
+from mito_ai import constants
 
 # Force Matplotlib to use the Jupyter inline backend.
 # Background: importing Streamlit sets os.environ["MPLBACKEND"] = "Agg" very early.
@@ -33,16 +35,6 @@ from mito_ai.chart_wizard.urls import get_chart_wizard_urls
 import os
 os.environ["MPLBACKEND"] = "module://matplotlib_inline.backend_inline"
 
-try:
-    from _version import __version__
-except ImportError:
-    # Fallback when using the package in dev mode without installing in editable mode with pip. It is highly recommended to install
-    # the package from a stable release or in editable mode: https://pip.pypa.io/en/stable/topics/local-project-installs/#editable-installs
-    import warnings
-    
-    warnings.warn("Importing 'mito_ai' outside a proper installation.")
-    __version__ = "dev"
-
 def _jupyter_labextension_paths() -> List[Dict[str, str]]:
     return [{"src": "labextension", "dest": "mito_ai"}]
 
@@ -65,7 +57,7 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     web_app = server_app.web_app
     base_url = web_app.settings["base_url"]
 
-    open_ai_provider = OpenAIProvider(config=server_app.config)
+    provider_manager = ProviderManager(config=server_app.config)
     
     # Create a single GlobalMessageHistory instance for the entire server
     # This ensures thread-safe access to the .mito/ai-chats directory
@@ -76,18 +68,13 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
         (
             url_path_join(base_url, "mito-ai", "completions"),
             CompletionHandler,
-            {"llm": open_ai_provider, "message_history": global_message_history},
+            {"llm": provider_manager, "message_history": global_message_history},
         ),
         (
             url_path_join(base_url, "mito-ai", "app-deploy"),
             AppDeployHandler,
             {}
         ),
-        (
-            url_path_join(base_url, "mito-ai", "streamlit-preview"),
-            StreamlitPreviewHandler,
-            {}
-        ),
         (
             url_path_join(base_url, "mito-ai", "version-check"),
             VersionCheckHandler,
@@ -104,13 +91,20 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     handlers.extend(get_db_urls(base_url))  # type: ignore
     handlers.extend(get_settings_urls(base_url))  # type: ignore
     handlers.extend(get_rules_urls(base_url))  # type: ignore
-    handlers.extend(get_log_urls(base_url, open_ai_provider.key_type))  # type: ignore
+    handlers.extend(get_log_urls(base_url, provider_manager.key_type))  # type: ignore
     handlers.extend(get_auth_urls(base_url))  # type: ignore
-    handlers.extend(get_streamlit_preview_urls(base_url))  # type: ignore
+    handlers.extend(get_streamlit_preview_urls(base_url, provider_manager))  # 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
+    handlers.extend(get_chart_wizard_urls(base_url, provider_manager)) # type: ignore
 
     web_app.add_handlers(host_pattern, handlers)
+    
+    # Log enterprise mode status and LiteLLM configuration
+    if is_enterprise():
+        server_app.log.info("Enterprise mode enabled")
+        if is_litellm_configured():
+            server_app.log.info(f"LiteLLM configured: endpoint={constants.LITELLM_BASE_URL}, models={constants.LITELLM_MODELS}")
+    
     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.57'
+__version__ = VERSION = '0.1.58'

mito_ai/anthropic_client.py

@@ -9,7 +9,7 @@ from anthropic.types import Message, MessageParam, TextBlockParam
 from mito_ai.completions.models import ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
 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
+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, LARGE_CONTEXT_MODEL, EXTENDED_CONTEXT_BETA
 
 # Max tokens is a required parameter for the Anthropic API. 
 # We set it to a high number so that we can edit large code cells
@@ -220,7 +220,10 @@ class AnthropicClient:
         self.max_retries = max_retries
         self.client: Optional[anthropic.Anthropic]
         if api_key:
-            self.client = anthropic.Anthropic(api_key=api_key)
+            # Use a higher timeout to avoid the 10-minute streaming requirement for long requests
+            # The default SDK timeout is 600s (10 minutes), but we set it higher for agent mode
+            # TODO: We should update agent mode to use streaming like anthropic suggests
+            self.client = anthropic.Anthropic(api_key=api_key, timeout=1200.0)  # 20 minutes
         else:
             self.client = None
 
@@ -249,7 +252,8 @@ class AnthropicClient:
         if self.api_key:
             # Unpack provider_data for direct API call
             assert self.client is not None
-            response = self.client.messages.create(**provider_data)
+            # Beta API accepts MessageParam (compatible at runtime with BetaMessageParam)
+            response = self.client.beta.messages.create(**provider_data)  # type: ignore[arg-type]
             
             if provider_data.get("tool_choice") is not None:
                 result = extract_and_parse_anthropic_json_response(response)
@@ -284,21 +288,27 @@ class AnthropicClient:
 
             if self.api_key:
                 assert self.client is not None
-                stream = self.client.messages.create(
-                    model=model,
-                    max_tokens=MAX_TOKENS,
-                    temperature=0,
-                    system=anthropic_system_prompt,
-                    messages=anthropic_messages,
-                    stream=True
-                )
+                # Beta API accepts MessageParam (compatible at runtime with BetaMessageParam)
+                # Enable extended context beta when using LARGE_CONTEXT_MODEL
+                create_params = {
+                    "model": model,
+                    "max_tokens": MAX_TOKENS,
+                    "temperature": 0,
+                    "system": anthropic_system_prompt,
+                    "messages": anthropic_messages,  # type: ignore[arg-type]
+                    "stream": True
+                }
+                if model == LARGE_CONTEXT_MODEL:
+                    create_params["betas"] = [EXTENDED_CONTEXT_BETA]
+                stream = self.client.beta.messages.create(**create_params)  # type: ignore[call-overload]
 
                 for chunk in stream:
-                    if chunk.type == "content_block_delta" and chunk.delta.type == "text_delta":
-                        content = chunk.delta.text
+                    # Type checking for beta API streaming chunks (runtime type checking, types are compatible)
+                    if chunk.type == "content_block_delta" and chunk.delta.type == "text_delta":  # type: ignore[union-attr]
+                        content = chunk.delta.text  # type: ignore[union-attr]
                         accumulated_response += content
 
-                        is_finished = chunk.type == "message_stop"
+                        is_finished = chunk.type == "message_stop"  # type: ignore[union-attr]
 
                         reply_fn(CompletionStreamChunk(
                             parent_id=message_id,

mito_ai/chart_wizard/handlers.py

@@ -6,13 +6,18 @@ 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.provider_manager import ProviderManager
 from mito_ai.completions.models import MessageType
-from mito_ai.completions.prompt_builders.chart_conversion_prompt import create_chart_conversion_prompt
+from mito_ai.completions.prompt_builders.chart_conversion_prompt import (
+    create_chart_conversion_prompt,
+)
+from mito_ai.completions.prompt_builders.chart_add_field_prompt import (
+    create_chart_add_field_prompt,
+)
 
-class ChartWizardHandler(APIHandler):
-    def initialize(self, llm: OpenAIProvider) -> None:
+
+class ConvertChartHandler(APIHandler):
+    def initialize(self, llm: ProviderManager) -> None:
         """Initialize the handler with the LLM provider."""
         super().initialize()
         self._llm = llm
@@ -21,26 +26,82 @@ class ChartWizardHandler(APIHandler):
     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', '')
-            
+            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}]
+            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
+                thread_id=None,
+                use_fast_model=True,
             )
-            
+
             # Return the converted code
-            self.write({
-                "message": "Code converted successfully",
-                "converted_code": 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()
+
+
+class AddFieldHandler(APIHandler):
+    def initialize(self, llm: ProviderManager) -> None:
+        """Initialize the handler with the LLM provider."""
+        super().initialize()
+        self._llm = llm
+
+    @tornado.web.authenticated
+    async def post(self) -> None:
+        """POST endpoint that adds a new field to the chart configuration."""
+        try:
+            data = json.loads(self.request.body.decode("utf-8"))
+            code = data.get("code", "")
+            user_description = data.get("user_description", "")
+            existing_variables = data.get("existing_variables", [])
+
+            if not user_description:
+                self.set_status(400)
+                self.write({"error": "user_description is required"})
+                self.finish()
+                return
+
+            # Create prompt using the prompt builder
+            prompt = create_chart_add_field_prompt(
+                code, user_description, existing_variables
+            )
+
+            # Call LLM
+            messages: List[ChatCompletionMessageParam] = [
+                {"role": "user", "content": prompt}
+            ]
+            updated_code = await self._llm.request_completions(
+                messages=messages,
+                message_type=MessageType.CHAT,
+                thread_id=None,
+                use_fast_model=True,
+            )
+
+            # Return the updated code
+            self.write(
+                {"message": "Field added successfully", "updated_code": updated_code}
+            )
             self.finish()
         except json.JSONDecodeError:
             self.set_status(400)

mito_ai/chart_wizard/urls.py

@@ -3,11 +3,13 @@
 
 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
+from mito_ai.chart_wizard.handlers import ConvertChartHandler, AddFieldHandler
+from mito_ai.provider_manager import ProviderManager
 
 
-def get_chart_wizard_urls(base_url: str, llm: OpenAIProvider) -> List[Tuple[str, Any, dict]]:
+def get_chart_wizard_urls(
+    base_url: str, llm: ProviderManager
+) -> List[Tuple[str, Any, dict]]:
     """Get all chart wizard related URL patterns.
 
     Args:
@@ -17,7 +19,8 @@ def get_chart_wizard_urls(base_url: str, llm: OpenAIProvider) -> List[Tuple[str,
     Returns:
         List of (url_pattern, handler_class, handler_kwargs) tuples
     """
-    BASE_URL = base_url + "/mito-ai"
+    BASE_URL = base_url + "/mito-ai/chart-wizard"
     return [
-        (url_path_join(BASE_URL, "chart-wizard"), ChartWizardHandler, {"llm": llm}),
+        (url_path_join(BASE_URL, "convert"), ConvertChartHandler, {"llm": llm}),
+        (url_path_join(BASE_URL, "add-field"), AddFieldHandler, {"llm": llm}),
     ]

mito_ai/completions/completion_handlers/agent_auto_error_fixup_handler.py

@@ -4,7 +4,7 @@
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import AgentResponse, AgentSmartDebugMetadata, MessageType, ResponseFormatInfo
 from mito_ai.completions.prompt_builders.agent_smart_debug_prompt import create_agent_smart_debug_prompt
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 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
@@ -17,14 +17,13 @@ class AgentAutoErrorFixupHandler(CompletionHandler[AgentSmartDebugMetadata]):
     @staticmethod
     async def get_completion(
         metadata: AgentSmartDebugMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get an agent auto error fixup completion from the AI provider."""
         
         # Add the system message if it doesn't alredy exist
-        await append_agent_system_message(message_history, model, provider, metadata.threadId, metadata.isChromeBrowser)
+        await append_agent_system_message(message_history, provider, metadata.threadId, metadata.isChromeBrowser)
         
         # Create the prompt
         prompt = create_agent_smart_debug_prompt(metadata)
@@ -34,12 +33,11 @@ class AgentAutoErrorFixupHandler(CompletionHandler[AgentSmartDebugMetadata]):
         new_ai_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": prompt}
         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)
+        await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, 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
@@ -51,7 +49,7 @@ class AgentAutoErrorFixupHandler(CompletionHandler[AgentSmartDebugMetadata]):
         
         ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
         
-        await message_history.append_message(ai_response_message, ai_response_message, model, provider, metadata.threadId)
+        await message_history.append_message(ai_response_message, ai_response_message, provider, metadata.threadId)
         
         return completion
 

mito_ai/completions/completion_handlers/agent_execution_handler.py

@@ -5,7 +5,7 @@ from typing import List, Literal, Union
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import AgentExecutionMetadata, MessageType, ResponseFormatInfo, AgentResponse
 from mito_ai.completions.prompt_builders.agent_execution_prompt import create_agent_execution_prompt
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 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
@@ -18,9 +18,8 @@ class AgentExecutionHandler(CompletionHandler[AgentExecutionMetadata]):
     @staticmethod
     async def get_completion(
         metadata: AgentExecutionMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get an agent execution completion from the AI provider."""
 
@@ -31,7 +30,7 @@ class AgentExecutionHandler(CompletionHandler[AgentExecutionMetadata]):
             )
 
         # Add the system message if it doesn't alredy exist
-        await append_agent_system_message(message_history, model, provider, metadata.threadId, metadata.isChromeBrowser)
+        await append_agent_system_message(message_history, provider, metadata.threadId, metadata.isChromeBrowser)
         
         # Create the prompt
         prompt = create_agent_execution_prompt(metadata)
@@ -41,12 +40,11 @@ class AgentExecutionHandler(CompletionHandler[AgentExecutionMetadata]):
         new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
         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)
+        await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, 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
@@ -58,7 +56,7 @@ class AgentExecutionHandler(CompletionHandler[AgentExecutionMetadata]):
         
         ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
         
-        await message_history.append_message(ai_response_message, ai_response_message, model, provider, metadata.threadId)
+        await message_history.append_message(ai_response_message, ai_response_message, provider, metadata.threadId)
         
         return completion
 

mito_ai/completions/completion_handlers/chat_completion_handler.py

@@ -6,7 +6,7 @@ from typing import List, Union, AsyncGenerator, Callable
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import ChatMessageMetadata, MessageType, CompletionRequest, CompletionStreamChunk, CompletionReply
 from mito_ai.completions.prompt_builders.chat_prompt import create_chat_prompt
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 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_chat_system_message, create_ai_optimized_message
@@ -19,9 +19,8 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
     @staticmethod
     async def get_completion(
         metadata: ChatMessageMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get a chat completion from the AI provider."""
 
@@ -32,7 +31,7 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
             )
 
         # Add the system message if it doesn't alredy exist
-        await append_chat_system_message(message_history, model, provider, metadata.threadId)
+        await append_chat_system_message(message_history, provider, metadata.threadId)
         
         # Create the prompt
         prompt = create_chat_prompt(
@@ -49,30 +48,28 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
         # Add the prompt to the message history
         new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
         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)
+        await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, provider, metadata.threadId)
         
         # Get the completion (non-streaming)
         completion = await provider.request_completions(
             messages=message_history.get_ai_optimized_history(metadata.threadId), 
-            model=model,
             message_type=MessageType.CHAT,
             user_input=metadata.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)
+        await message_history.append_message(ai_response_message, ai_response_message, provider, metadata.threadId)
 
         return completion
     
     @staticmethod
     async def stream_completion(
         metadata: ChatMessageMetadata,
-        provider: OpenAIProvider,
+        provider: ProviderManager,
         message_history: GlobalMessageHistory,
         message_id: str,
-        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
-        model: str
+        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]
     ) -> str:
         """Stream chat completions from the AI provider.
         
@@ -95,7 +92,7 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
             )
         
         # Add the system message if it doesn't already exist
-        await append_chat_system_message(message_history, model, provider, metadata.threadId)
+        await append_chat_system_message(message_history, provider, metadata.threadId)
         
         # Create the prompt
         prompt = create_chat_prompt(
@@ -112,17 +109,16 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
         # Add the prompt to the message history
         new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
         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)
+        await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, provider, metadata.threadId)
         
         # Stream the completions using the provider's stream method
         accumulated_response = await provider.stream_completions(
             message_type=MessageType.CHAT,
             messages=message_history.get_ai_optimized_history(metadata.threadId),
-            model=model,
             message_id=message_id,
+            thread_id=metadata.threadId,
             reply_fn=reply_fn,
-            user_input=metadata.input,
-            thread_id=metadata.threadId
+            user_input=metadata.input
         )
 
         # Save the accumulated response to message history
@@ -131,7 +127,7 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
             "content": accumulated_response,
         }
         await message_history.append_message(
-            ai_response_message, ai_response_message, model, provider, metadata.threadId
+            ai_response_message, ai_response_message, provider, metadata.threadId
         )
 
         return accumulated_response

mito_ai/completions/completion_handlers/code_explain_handler.py

@@ -5,7 +5,7 @@ from typing import List, Union, AsyncGenerator, Callable
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import CodeExplainMetadata, MessageType, CompletionRequest, CompletionStreamChunk, CompletionReply
 from mito_ai.completions.prompt_builders.explain_code_prompt import create_explain_code_prompt
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 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_chat_system_message
@@ -18,16 +18,15 @@ class CodeExplainHandler(CompletionHandler[CodeExplainMetadata]):
     @staticmethod
     async def get_completion(
         metadata: CodeExplainMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get a code explain completion from the AI provider."""
         active_cell_code = metadata.activeCellCode or ''
         thread_id = metadata.threadId
 
         # Add the system message if it doesn't already exist
-        await append_chat_system_message(message_history, model, provider, thread_id)
+        await append_chat_system_message(message_history, provider, thread_id)
 
         # Create the prompt
         prompt = create_explain_code_prompt(active_cell_code)
@@ -37,31 +36,29 @@ class CodeExplainHandler(CompletionHandler[CodeExplainMetadata]):
         new_ai_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": prompt}
         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, thread_id
+            new_ai_optimized_message, new_display_optimized_message, provider, thread_id
         )
 
         # Get the completion
         completion = await provider.request_completions(
             messages=message_history.get_ai_optimized_history(thread_id),
-            model=model,
             message_type=MessageType.CODE_EXPLAIN,
             thread_id=thread_id
         )
 
         # Add the response to message history
         ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
-        await message_history.append_message(ai_response_message, ai_response_message, model, provider, thread_id)
+        await message_history.append_message(ai_response_message, ai_response_message, provider, thread_id)
 
         return completion
 
     @staticmethod
     async def stream_completion(
         metadata: CodeExplainMetadata,
-        provider: OpenAIProvider,
+        provider: ProviderManager,
         message_history: GlobalMessageHistory,
         message_id: str,
-        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
-        model: str
+        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]
     ) -> str:
         """Stream code explain completions from the AI provider.
 
@@ -79,7 +76,7 @@ class CodeExplainHandler(CompletionHandler[CodeExplainMetadata]):
         thread_id = metadata.threadId
 
         # Add the system message if it doesn't already exist
-        await append_chat_system_message(message_history, model, provider, thread_id)
+        await append_chat_system_message(message_history, provider, thread_id)
 
         # Create the prompt
         prompt = create_explain_code_prompt(active_cell_code)
@@ -89,22 +86,21 @@ class CodeExplainHandler(CompletionHandler[CodeExplainMetadata]):
         new_ai_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": prompt}
         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, thread_id
+            new_ai_optimized_message, new_display_optimized_message, provider, thread_id
         )
 
         # Stream the completions using the provider's stream method
         accumulated_response = await provider.stream_completions(
             message_type=MessageType.CODE_EXPLAIN,
             messages=message_history.get_ai_optimized_history(thread_id),
-            model=model,
             message_id=message_id,
-            reply_fn=reply_fn,
-            thread_id=thread_id
+            thread_id=thread_id,
+            reply_fn=reply_fn
         )
 
         # Add the response to message history
         ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": accumulated_response}
-        await message_history.append_message(ai_response_message, ai_response_message, model, provider, thread_id)
+        await message_history.append_message(ai_response_message, ai_response_message, provider, thread_id)
 
         return accumulated_response
 

mito_ai/completions/completion_handlers/completion_handler.py

@@ -4,7 +4,7 @@
 from typing import Protocol, TypeVar
 from abc import abstractmethod, ABCMeta
 from mito_ai.completions.models import ChatMessageMetadata, ScratchpadResultMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, InlineCompleterMetadata, AgentSmartDebugMetadata
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.completions.message_history import GlobalMessageHistory
 
 T = TypeVar('T', 
@@ -29,9 +29,8 @@ class CompletionHandler(Protocol[T], metaclass=ABCMeta):
     @abstractmethod
     async def get_completion(
         metadata: T,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get a completion from the AI provider.
         
@@ -39,7 +38,6 @@ class CompletionHandler(Protocol[T], metaclass=ABCMeta):
             metadata: Metadata about the completion request
             provider: The AI provider to use
             message_history: The history of messages in the conversation
-            model: The model to use for the completion
             
         Returns:
             The completion string from the AI