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

mito_ai/{completions/providers.py → provider_manager.py}

@@ -6,7 +6,7 @@ import asyncio
 from typing import Any, Callable, Dict, List, Optional, Union, cast
 from mito_ai import constants
 from openai.types.chat import ChatCompletionMessageParam
-from traitlets import Instance, Unicode, default, validate
+from traitlets import Instance
 from traitlets.config import LoggingConfigurable
 from openai.types.chat import ChatCompletionMessageParam
 
@@ -24,32 +24,23 @@ from mito_ai.completions.models import (
     CompletionReply,
     CompletionStreamChunk,
     MessageType,
-    ResponseFormatInfo, CompletionItemError,
+    ResponseFormatInfo,
 )
+from mito_ai.utils.litellm_utils import is_litellm_configured
 from mito_ai.utils.telemetry_utils import (
-    KEY_TYPE_PARAM,
-    MITO_AI_COMPLETION_ERROR,
-    MITO_AI_COMPLETION_RETRY,
     MITO_SERVER_KEY,
     USER_KEY,
-    log,
     log_ai_completion_error,
     log_ai_completion_retry,
     log_ai_completion_success,
 )
 from mito_ai.utils.provider_utils import get_model_provider
-from mito_ai.utils.mito_server_utils import ProviderCompletionException
+from mito_ai.utils.model_utils import get_available_models, get_fast_model_for_selected_model, get_smartest_model_for_selected_model
 
-__all__ = ["OpenAIProvider"]
+__all__ = ["ProviderManager"]
 
-class OpenAIProvider(LoggingConfigurable):
-    """Provide AI feature through OpenAI services."""
-
-    api_key = Unicode(
-        config=True,
-        allow_none=True,
-        help="OpenAI API key. Default value is read from the OPENAI_API_KEY environment variable.",
-    )
+class ProviderManager(LoggingConfigurable):
+    """Manage AI providers (Claude, Gemini, OpenAI) and route requests to the appropriate client."""
 
     last_error = Instance(
         CompletionError,
@@ -61,29 +52,57 @@ This attribute is observed by the websocket provider to push the error to the cl
 
     def __init__(self, **kwargs: Dict[str, Any]) -> None:
         config = kwargs.get('config', {})
-        if 'api_key' in kwargs:
-            config['OpenAIClient'] = {'api_key': kwargs['api_key']}
         kwargs['config'] = config
 
         super().__init__(log=get_logger(), **kwargs)
         self.last_error = None
         self._openai_client: Optional[OpenAIClient] = OpenAIClient(**config)
+        # Initialize with the first available model to ensure it's always valid
+        # This respects LiteLLM configuration: if LiteLLM is configured, uses first LiteLLM model
+        # Otherwise, uses first standard model
+        available_models = get_available_models()
+        self._selected_model: str = available_models[0] if available_models else "gpt-4.1"
+    
+    def get_selected_model(self) -> str:
+        """Get the currently selected model."""
+        return self._selected_model
+    
+    def set_selected_model(self, model: str) -> None:
+        """Set the selected model."""
+        self._selected_model = model
 
     @property
     def capabilities(self) -> AICapabilities:
         """
         Returns the capabilities of the AI provider.
         """
-        if constants.CLAUDE_API_KEY and not self.api_key:
+        # TODO: We should validate that these keys are actually valid for the provider
+        # otherwise it will look like we are using the user_key when actually falling back 
+        # to the mito server because the key is invalid. 
+        if is_litellm_configured():
+            return AICapabilities(
+                configuration={"model": "<dynamic>"},
+                provider="LiteLLM",
+            )
+            
+        if constants.OPENAI_API_KEY:
+            return AICapabilities(
+                configuration={"model": "<dynamic>"},
+                provider="OpenAI",
+            )
+            
+        if constants.ANTHROPIC_API_KEY:
             return AICapabilities(
                 configuration={"model": "<dynamic>"},
                 provider="Claude",
             )
-        if constants.GEMINI_API_KEY and not self.api_key:
+            
+        if constants.GEMINI_API_KEY:
             return AICapabilities(
                 configuration={"model": "<dynamic>"},
                 provider="Gemini",
             )
+            
         if self._openai_client:
             return self._openai_client.capabilities
         
@@ -94,65 +113,106 @@ This attribute is observed by the websocket provider to push the error to the cl
 
     @property
     def key_type(self) -> str:
-        if constants.CLAUDE_API_KEY and not self.api_key:
-            return "claude"
-        if constants.GEMINI_API_KEY and not self.api_key:
-            return "gemini"
-        if self._openai_client:
-            return self._openai_client.key_type
+        # TODO: We should validate that these keys are actually valid for the provider
+        # otherwise it will look like we are using the user_key when actually falling back 
+        # to the mito server because the key is invalid. 
+        if is_litellm_configured():
+            return USER_KEY
+        
+        if constants.ANTHROPIC_API_KEY or constants.GEMINI_API_KEY or constants.OPENAI_API_KEY or constants.OLLAMA_MODEL:  
+            return USER_KEY
+        
         return MITO_SERVER_KEY
 
     async def request_completions(
         self,
         message_type: MessageType,
         messages: List[ChatCompletionMessageParam],
-        model: str,
         response_format_info: Optional[ResponseFormatInfo] = None,
         user_input: Optional[str] = None,
         thread_id: Optional[str] = None,
-        max_retries: int = 3
+        max_retries: int = 3,
+        use_fast_model: bool = False,
+        use_smartest_model: bool = False
     ) -> str:
         """
         Request completions from the AI provider.
+        
+        Args:
+            message_type: Type of message
+            messages: List of chat messages
+            response_format_info: Optional response format specification
+            user_input: Optional user input for logging
+            thread_id: Optional thread ID for logging
+            max_retries: Maximum number of retries
+            use_fast_model: If True, use the fastest model from the selected provider
+            use_smartest_model: If True, use the smartest model from the selected provider
         """
         self.last_error = None
         completion = None
         last_message_content = str(messages[-1].get('content', '')) if messages else ""
-        model_type = get_model_provider(model)
+        
+        # Get the model to use (selected model, fast model, or smartest model if requested)
+        selected_model = self.get_selected_model()
+        if use_smartest_model:
+            resolved_model = get_smartest_model_for_selected_model(selected_model)
+        elif use_fast_model:
+            resolved_model = get_fast_model_for_selected_model(selected_model)
+        else:
+            resolved_model = selected_model
+        
+        # Validate model is in allowed list (uses same function as endpoint)
+        available_models = get_available_models()
+        if resolved_model not in available_models:
+            raise ValueError(f"Model {resolved_model} is not in the allowed model list: {available_models}")
+        
+        # Get model provider type
+        model_type = get_model_provider(resolved_model)
         
         # Retry loop
         for attempt in range(max_retries + 1):
             try:
-                if model_type == "claude":
-                    api_key = constants.CLAUDE_API_KEY
+                if model_type == "litellm":
+                    from mito_ai.enterprise.litellm_client import LiteLLMClient
+                    if not constants.LITELLM_BASE_URL:
+                        raise ValueError("LITELLM_BASE_URL is required for LiteLLM models")
+                    litellm_client = LiteLLMClient(api_key=constants.LITELLM_API_KEY, base_url=constants.LITELLM_BASE_URL)
+                    completion = await litellm_client.request_completions(
+                        messages=messages,
+                        model=resolved_model,
+                        response_format_info=response_format_info,
+                        message_type=message_type
+                    )
+                elif model_type == "claude":
+                    api_key = constants.ANTHROPIC_API_KEY
                     anthropic_client = AnthropicClient(api_key=api_key)
-                    completion = await anthropic_client.request_completions(messages, model, response_format_info, message_type)
+                    completion = await anthropic_client.request_completions(messages, resolved_model, response_format_info, message_type)
                 elif model_type == "gemini":
                     api_key = constants.GEMINI_API_KEY
                     gemini_client = GeminiClient(api_key=api_key)
                     messages_for_gemini = [dict(m) for m in messages]
-                    completion = await gemini_client.request_completions(messages_for_gemini, model, response_format_info, message_type)
+                    completion = await gemini_client.request_completions(messages_for_gemini, resolved_model, response_format_info, message_type)
                 elif model_type == "openai":
                     if not self._openai_client:
                         raise RuntimeError("OpenAI client is not initialized.")
                     completion = await self._openai_client.request_completions(
                         message_type=message_type,
                         messages=messages,
-                        model=model,
+                        model=resolved_model,
                         response_format_info=response_format_info
                     )
                 else:
-                    raise ValueError(f"No AI provider configured for model: {model}")
+                    raise ValueError(f"No AI provider configured for model: {resolved_model}")
                 
                 # Success! Log and return
                 log_ai_completion_success(
-                    key_type=USER_KEY if self.key_type == "user" else MITO_SERVER_KEY,
+                    key_type=USER_KEY if self.key_type == USER_KEY else MITO_SERVER_KEY,
                     message_type=message_type,
                     last_message_content=last_message_content,
                     response={"completion": completion},
                     user_input=user_input or "",
                     thread_id=thread_id or "",
-                    model=model
+                    model=resolved_model
                 )
                 return completion # type: ignore
             
@@ -160,7 +220,7 @@ This attribute is observed by the websocket provider to push the error to the cl
                 # If we hit a free tier limit, then raise an exception right away without retrying.
                 self.log.exception(f"Error during request_completions: {e}")
                 self.last_error = CompletionError.from_exception(e)
-                log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
+                log_ai_completion_error(USER_KEY if self.key_type != MITO_SERVER_KEY else MITO_SERVER_KEY, thread_id or "", message_type, e)
                 raise
 
             except BaseException as e:
@@ -169,14 +229,14 @@ This attribute is observed by the websocket provider to push the error to the cl
                     # Exponential backoff: wait 2^attempt seconds
                     wait_time = 2 ** attempt
                     self.log.info(f"Retrying request_completions after {wait_time}s (attempt {attempt + 1}/{max_retries + 1}): {str(e)}")
-                    log_ai_completion_retry('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
+                    log_ai_completion_retry(USER_KEY if self.key_type != MITO_SERVER_KEY else MITO_SERVER_KEY, thread_id or "", message_type, e)
                     await asyncio.sleep(wait_time)
                     continue
                 else:
                     # Final failure after all retries - set error state and raise
                     self.log.exception(f"Error during request_completions after {attempt + 1} attempts: {e}")
                     self.last_error = CompletionError.from_exception(e)
-                    log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
+                    log_ai_completion_error(USER_KEY if self.key_type != MITO_SERVER_KEY else MITO_SERVER_KEY, thread_id or "", message_type, e)
                     raise
         
         # This should never be reached due to the raise in the except block,
@@ -187,21 +247,50 @@ This attribute is observed by the websocket provider to push the error to the cl
         self,
         message_type: MessageType,
         messages: List[ChatCompletionMessageParam],
-        model: str,
         message_id: str,
         thread_id: str,
         reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
         user_input: Optional[str] = None,
-        response_format_info: Optional[ResponseFormatInfo] = None
+        response_format_info: Optional[ResponseFormatInfo] = None,
+        use_fast_model: bool = False,
+        use_smartest_model: bool = False
     ) -> str:
         """
         Stream completions from the AI provider and return the accumulated response.
+        
+        Args:
+            message_type: Type of message
+            messages: List of chat messages
+            message_id: ID of the message being processed
+            thread_id: Thread ID for logging
+            reply_fn: Function to call with each chunk for streaming replies
+            user_input: Optional user input for logging
+            response_format_info: Optional response format specification
+            use_fast_model: If True, use the fastest model from the selected provider
+            use_smartest_model: If True, use the smartest model from the selected provider
+            
         Returns: The accumulated response string.
         """
         self.last_error = None
         accumulated_response = ""
         last_message_content = str(messages[-1].get('content', '')) if messages else ""
-        model_type = get_model_provider(model)
+        
+        # Get the model to use (selected model, fast model, or smartest model if requested)
+        selected_model = self.get_selected_model()
+        if use_smartest_model:
+            resolved_model = get_smartest_model_for_selected_model(selected_model)
+        elif use_fast_model:
+            resolved_model = get_fast_model_for_selected_model(selected_model)
+        else:
+            resolved_model = selected_model
+                    
+        # Validate model is in allowed list (uses same function as endpoint)
+        available_models = get_available_models()
+        if resolved_model not in available_models:
+            raise ValueError(f"Model {resolved_model} is not in the allowed model list: {available_models}")
+        
+        # Get model provider type
+        model_type = get_model_provider(resolved_model)
         reply_fn(CompletionReply(
             items=[
                 CompletionItem(content="", isIncomplete=True, token=message_id)
@@ -210,12 +299,28 @@ This attribute is observed by the websocket provider to push the error to the cl
         ))
 
         try:
-            if model_type == "claude":
-                api_key = constants.CLAUDE_API_KEY
+            if model_type == "litellm":
+                from mito_ai.enterprise.litellm_client import LiteLLMClient
+                if not constants.LITELLM_BASE_URL:
+                    raise ValueError("LITELLM_BASE_URL is required for LiteLLM models")
+                litellm_client = LiteLLMClient(
+                    api_key=constants.LITELLM_API_KEY,
+                    base_url=constants.LITELLM_BASE_URL
+                )
+                accumulated_response = await litellm_client.stream_completions(
+                    messages=messages,
+                    model=resolved_model,
+                    message_type=message_type,
+                    message_id=message_id,
+                    reply_fn=reply_fn,
+                    response_format_info=response_format_info
+                )
+            elif model_type == "claude":
+                api_key = constants.ANTHROPIC_API_KEY
                 anthropic_client = AnthropicClient(api_key=api_key)
                 accumulated_response = await anthropic_client.stream_completions(
                     messages=messages,
-                    model=model,
+                    model=resolved_model,
                     message_type=message_type,
                     message_id=message_id,
                     reply_fn=reply_fn
@@ -228,7 +333,7 @@ This attribute is observed by the websocket provider to push the error to the cl
                 messages_for_gemini = [dict(m) for m in messages]
                 accumulated_response = await gemini_client.stream_completions(
                     messages=messages_for_gemini,
-                    model=model,
+                    model=resolved_model,
                     message_id=message_id,
                     reply_fn=reply_fn,
                     message_type=message_type
@@ -239,7 +344,7 @@ This attribute is observed by the websocket provider to push the error to the cl
                 accumulated_response = await self._openai_client.stream_completions(
                     message_type=message_type,
                     messages=messages,
-                    model=model,
+                    model=resolved_model,
                     message_id=message_id,
                     thread_id=thread_id,
                     reply_fn=reply_fn,
@@ -247,24 +352,24 @@ This attribute is observed by the websocket provider to push the error to the cl
                     response_format_info=response_format_info
                 )
             else:
-                raise ValueError(f"No AI provider configured for model: {model}")
+                raise ValueError(f"No AI provider configured for model: {resolved_model}")
 
             # Log the successful completion
             log_ai_completion_success(
-                key_type=USER_KEY if self.key_type == "user" else MITO_SERVER_KEY,
+                key_type=USER_KEY if self.key_type == USER_KEY else MITO_SERVER_KEY,
                 message_type=message_type,
                 last_message_content=last_message_content,
                 response={"completion": accumulated_response},
                 user_input=user_input or "",
                 thread_id=thread_id,
-                model=model
+                model=resolved_model
             )
             return accumulated_response
 
         except BaseException as e:
             self.log.exception(f"Error during stream_completions: {e}")
             self.last_error = CompletionError.from_exception(e)
-            log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id, message_type, e)
+            log_ai_completion_error(USER_KEY if self.key_type != MITO_SERVER_KEY else MITO_SERVER_KEY, thread_id, message_type, e)
 
             # Send error message to client before raising
             reply_fn(CompletionStreamChunk(
@@ -281,4 +386,3 @@ This attribute is observed by the websocket provider to push the error to the cl
                 error=CompletionError.from_exception(e),
             ))
             raise
-

mito_ai/settings/enterprise_handler.py

@@ -0,0 +1,26 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import tornado
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.utils.model_utils import get_available_models
+
+
+class AvailableModelsHandler(APIHandler):
+    """REST handler for returning available models to the frontend."""
+    
+    @tornado.web.authenticated
+    async def get(self) -> None:
+        """GET endpoint that returns the list of available models."""
+        try:
+            available_models = get_available_models()
+            
+            self.write({
+                "models": available_models
+            })
+            self.finish()
+        except Exception as e:
+            self.set_status(500)
+            self.write({"error": str(e)})
+            self.finish()

mito_ai/settings/urls.py

@@ -4,6 +4,7 @@
 from typing import Any, List, Tuple
 from jupyter_server.utils import url_path_join
 from mito_ai.settings.handlers import SettingsHandler
+from mito_ai.settings.enterprise_handler import AvailableModelsHandler
 
 def get_settings_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
     """Get all settings related URL patterns.
@@ -17,4 +18,5 @@ def get_settings_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
     BASE_URL = base_url + "/mito-ai"
     return [
         (url_path_join(BASE_URL, "settings/(.*)"), SettingsHandler, {}),
+        (url_path_join(BASE_URL, "available-models"), AvailableModelsHandler, {}),
     ]

mito_ai/streamlit_conversion/agent_utils.py

@@ -1,37 +1,9 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from typing import List, Tuple
-import re
-from anthropic.types import MessageParam
-from mito_ai.streamlit_conversion.prompts.streamlit_system_prompt import streamlit_system_prompt
-from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
+from typing import List
 from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER
-from mito_ai.completions.models import MessageType
-
-STREAMLIT_AI_MODEL = "claude-sonnet-4-5-20250929"
 
 def extract_todo_placeholders(agent_response: str) -> List[str]:
     """Extract TODO placeholders from the agent's response"""
-    return [line.strip() for line in agent_response.split('\n') if MITO_TODO_PLACEHOLDER in line]
-
-async def get_response_from_agent(message_to_agent: List[MessageParam]) -> str:
-    """Gets the streaming response from the agent using the mito server"""
-    model = STREAMLIT_AI_MODEL
-    max_tokens = 64000 # TODO: If we move to haiku, we must reset this to 8192
-    temperature = 0.2
-
-    accumulated_response = ""
-    async for stream_chunk in stream_anthropic_completion_from_mito_server(
-        model = model,
-        max_tokens = max_tokens,
-        temperature = temperature,
-        system = streamlit_system_prompt,
-        messages = message_to_agent,
-        stream=True,
-        message_type=MessageType.STREAMLIT_CONVERSION,
-        reply_fn=None,
-        message_id=""
-    ):
-        accumulated_response += stream_chunk
-    return accumulated_response
+    return [line.strip() for line in agent_response.split('\n') if MITO_TODO_PLACEHOLDER in line]

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -1,9 +1,10 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from anthropic.types import MessageParam
-from typing import List, cast
-from mito_ai.streamlit_conversion.agent_utils import extract_todo_placeholders, get_response_from_agent
+from typing import List
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.streamlit_conversion.agent_utils import extract_todo_placeholders
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.streamlit_conversion.prompts.streamlit_app_creation_prompt import get_streamlit_app_creation_prompt
 from mito_ai.streamlit_conversion.prompts.streamlit_error_correction_prompt import get_streamlit_error_correction_prompt
 from mito_ai.streamlit_conversion.prompts.streamlit_finish_todo_prompt import get_finish_todo_prompt
@@ -15,22 +16,23 @@ from mito_ai.completions.models import MessageType
 from mito_ai.utils.error_classes import StreamlitConversionError
 from mito_ai.utils.telemetry_utils import log_streamlit_app_validation_retry, log_streamlit_app_conversion_success
 from mito_ai.path_utils import AbsoluteNotebookPath, AppFileName, get_absolute_notebook_dir_path, get_absolute_app_path, get_app_file_name
+from mito_ai.streamlit_conversion.prompts.streamlit_system_prompt import streamlit_system_prompt
 
-async def generate_new_streamlit_code(notebook: List[dict], streamlit_app_prompt: str) -> str:
+async def generate_new_streamlit_code(notebook: List[dict], streamlit_app_prompt: str, provider: ProviderManager) -> str:
     """Send a query to the agent, get its response and parse the code"""
     
     prompt_text = get_streamlit_app_creation_prompt(notebook, streamlit_app_prompt)
     
-    messages: List[MessageParam] = [
-        cast(MessageParam, {
-            "role": "user",
-            "content": [{
-                "type": "text",
-                "text": prompt_text
-            }]
-        })
+    messages: List[ChatCompletionMessageParam] = [
+        {"role": "system", "content": streamlit_system_prompt},
+        {"role": "user", "content": prompt_text}
     ]
-    agent_response = await get_response_from_agent(messages)
+    agent_response = await provider.request_completions(
+        message_type=MessageType.STREAMLIT_CONVERSION,
+        messages=messages,
+        use_smartest_model=True,
+        thread_id=None
+    )
     converted_code = extract_code_blocks(agent_response)
     
     # Extract the TODOs from the agent's response
@@ -39,16 +41,16 @@ async def generate_new_streamlit_code(notebook: List[dict], streamlit_app_prompt
     for todo_placeholder in todo_placeholders:
         print(f"Processing AI TODO: {todo_placeholder}")
         todo_prompt = get_finish_todo_prompt(notebook, converted_code, todo_placeholder)
-        todo_messages: List[MessageParam] = [
-            cast(MessageParam, {
-                "role": "user",
-                "content": [{
-                    "type": "text",
-                    "text": todo_prompt
-                }]
-            })
+        todo_messages: List[ChatCompletionMessageParam] = [
+            {"role": "system", "content": streamlit_system_prompt},
+            {"role": "user", "content": todo_prompt}
         ]
-        todo_response = await get_response_from_agent(todo_messages)
+        todo_response = await provider.request_completions(
+            message_type=MessageType.STREAMLIT_CONVERSION,
+            messages=todo_messages,
+            use_smartest_model=True,
+            thread_id=None
+        )
         
         # Apply the search/replace to the streamlit app
         search_replace_pairs = extract_search_replace_blocks(todo_response)
@@ -57,21 +59,21 @@ async def generate_new_streamlit_code(notebook: List[dict], streamlit_app_prompt
     return converted_code
 
 
-async def update_existing_streamlit_code(notebook: List[dict], streamlit_app_code: str, edit_prompt: str) -> str:
+async def update_existing_streamlit_code(notebook: List[dict], streamlit_app_code: str, edit_prompt: str, provider: ProviderManager) -> str:
     """Send a query to the agent, get its response and parse the code"""
     prompt_text = get_update_existing_app_prompt(notebook, streamlit_app_code, edit_prompt)
     
-    messages: List[MessageParam] = [
-        cast(MessageParam, {
-            "role": "user",
-            "content": [{
-                "type": "text",
-                "text": prompt_text
-            }]
-        })
+    messages: List[ChatCompletionMessageParam] = [
+        {"role": "system", "content": streamlit_system_prompt},
+        {"role": "user", "content": prompt_text}
     ]
     
-    agent_response = await get_response_from_agent(messages)
+    agent_response = await provider.request_completions(
+        message_type=MessageType.STREAMLIT_CONVERSION,
+        messages=messages,
+        use_smartest_model=True,
+        thread_id=None
+    )
     print(f"[Mito AI Search/Replace Tool]:\n {agent_response}")
 
     # Apply the search/replace to the streamlit app
@@ -81,18 +83,18 @@ async def update_existing_streamlit_code(notebook: List[dict], streamlit_app_cod
     return converted_code
 
 
-async def correct_error_in_generation(error: str, streamlit_app_code: str) -> str:
+async def correct_error_in_generation(error: str, streamlit_app_code: str, provider: ProviderManager) -> str:
     """If errors are present, send it back to the agent to get corrections in code"""
-    messages: List[MessageParam] = [
-        cast(MessageParam, {
-            "role": "user",
-            "content": [{
-                "type": "text",
-                "text": get_streamlit_error_correction_prompt(error, streamlit_app_code)
-            }]
-        })
+    messages: List[ChatCompletionMessageParam] = [
+        {"role": "system", "content": streamlit_system_prompt},
+        {"role": "user", "content": get_streamlit_error_correction_prompt(error, streamlit_app_code)}
     ]
-    agent_response = await get_response_from_agent(messages)
+    agent_response = await provider.request_completions(
+        message_type=MessageType.STREAMLIT_CONVERSION,
+        messages=messages,
+        use_smartest_model=True,
+        thread_id=None
+    )
     
     # Apply the search/replace to the streamlit app
     search_replace_pairs = extract_search_replace_blocks(agent_response)
@@ -100,7 +102,7 @@ async def correct_error_in_generation(error: str, streamlit_app_code: str) -> st
 
     return streamlit_app_code
 
-async def streamlit_handler(create_new_app: bool, notebook_path: AbsoluteNotebookPath, app_file_name: AppFileName, streamlit_app_prompt: str = "") -> None:
+async def streamlit_handler(create_new_app: bool, notebook_path: AbsoluteNotebookPath, app_file_name: AppFileName, streamlit_app_prompt: str, provider: ProviderManager) -> None:
     """Handler function for streamlit code generation and validation"""
 
     # Convert to absolute path for consistent handling
@@ -110,7 +112,7 @@ async def streamlit_handler(create_new_app: bool, notebook_path: AbsoluteNoteboo
     
     if create_new_app:
         # Otherwise generate a new streamlit app
-        streamlit_code = await generate_new_streamlit_code(notebook_code, streamlit_app_prompt)
+        streamlit_code = await generate_new_streamlit_code(notebook_code, streamlit_app_prompt, provider)
     else:
         # If the user is editing an existing streamlit app, use the update function
         existing_streamlit_code = get_app_code_from_file(app_path)
@@ -118,14 +120,14 @@ async def streamlit_handler(create_new_app: bool, notebook_path: AbsoluteNoteboo
         if existing_streamlit_code is None:
             raise StreamlitConversionError("Error updating existing streamlit app because app.py file was not found.", 404)
         
-        streamlit_code = await update_existing_streamlit_code(notebook_code, existing_streamlit_code, streamlit_app_prompt) 
+        streamlit_code = await update_existing_streamlit_code(notebook_code, existing_streamlit_code, streamlit_app_prompt, provider) 
        
     # Then, after creating/updating the app, validate that the new code runs 
     errors = validate_app(streamlit_code, notebook_path)
     tries = 0
     while len(errors) > 0 and tries < 5:
         for error in errors:
-            streamlit_code = await correct_error_in_generation(error, streamlit_code)
+            streamlit_code = await correct_error_in_generation(error, streamlit_code, provider)
         
         errors = validate_app(streamlit_code, notebook_path)