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

mito_ai/openai_client.py

@@ -7,7 +7,7 @@ from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Union
 from mito_ai.utils.mito_server_utils import ProviderCompletionException
 import openai
 from openai.types.chat import ChatCompletionMessageParam
-from traitlets import Instance, Unicode, default, validate
+from traitlets import Instance, default, validate
 from traitlets.config import LoggingConfigurable
 
 from mito_ai import constants
@@ -30,22 +30,12 @@ from mito_ai.utils.open_ai_utils import (
     stream_ai_completion_from_mito_server,
 )
 from mito_ai.utils.server_limits import update_mito_server_quota
-from mito_ai.utils.telemetry_utils import (
-    MITO_SERVER_KEY,
-    USER_KEY,
-)
 
 OPENAI_MODEL_FALLBACK = "gpt-4.1"
 
 class OpenAIClient(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.",
-    )
-
     last_error = Instance(
         CompletionError,
         allow_none=True,
@@ -65,61 +55,6 @@ This attribute is observed by the websocket provider to push the error to the cl
         super().__init__(log=get_logger(), **kwargs)
         self.last_error = None
         self._async_client: Optional[openai.AsyncOpenAI] = None
-
-    @default("api_key")
-    def _api_key_default(self) -> Optional[str]:
-        default_key = constants.OPENAI_API_KEY
-        return self._validate_api_key(default_key)
-
-    @validate("api_key")
-    def _validate_api_key(self, api_key: Optional[str]) -> Optional[str]:
-        if not api_key:
-            self.log.debug(
-                "No OpenAI API key provided; following back to Mito server API."
-            )
-            return None
-
-        client = openai.OpenAI(api_key=api_key)
-        try:
-            # Make an http request to OpenAI to make sure it works
-            client.models.list()
-        except openai.AuthenticationError as e:
-            self.log.warning(
-                "Invalid OpenAI API key provided.",
-                exc_info=e,
-            )
-            self.last_error = CompletionError.from_exception(
-                e,
-                hint="You're missing the OPENAI_API_KEY environment variable. Run the following code in your terminal to set the environment variable and then relaunch the jupyter server `export OPENAI_API_KEY=<your-api-key>`",
-            )
-            return None
-        except openai.PermissionDeniedError as e:
-            self.log.warning(
-                "Invalid OpenAI API key provided.",
-                exc_info=e,
-            )
-            self.last_error = CompletionError.from_exception(e)
-            return None
-        except openai.InternalServerError as e:
-            self.log.debug(
-                "Unable to get OpenAI models due to OpenAI error.", exc_info=e
-            )
-            return api_key
-        except openai.RateLimitError as e:
-            self.log.debug(
-                "Unable to get OpenAI models due to rate limit error.", exc_info=e
-            )
-            return api_key
-        except openai.APIConnectionError as e:
-            self.log.warning(
-                "Unable to connect to OpenAI API.",
-                exec_info=e,
-            )
-            self.last_error = CompletionError.from_exception(e)
-            return None
-        else:
-            self.log.debug("User OpenAI API key validated.")
-            return api_key
         
     @property
     def capabilities(self) -> AICapabilities:
@@ -133,7 +68,7 @@ This attribute is observed by the websocket provider to push the error to the cl
                 provider="Azure OpenAI",
             )
 
-        if constants.OLLAMA_MODEL and not self.api_key:
+        if constants.OLLAMA_MODEL:
             return AICapabilities(
                 configuration={
                     "model": constants.OLLAMA_MODEL
@@ -141,14 +76,12 @@ This attribute is observed by the websocket provider to push the error to the cl
                 provider="Ollama",
             )
 
-        if self.api_key:
-            self._validate_api_key(self.api_key)
-
+        if constants.OPENAI_API_KEY:
             return AICapabilities(
                 configuration={
-                    "model": OPENAI_MODEL_FALLBACK,
+                    "model": "<dynamic>"
                 },
-                provider="OpenAI (user key)",
+                provider="OpenAI",
             )
 
         try:
@@ -169,19 +102,6 @@ This attribute is observed by the websocket provider to push the error to the cl
         if not self._async_client or self._async_client.is_closed():
             self._async_client = self._build_openai_client()
         return self._async_client
-    
-    
-    @property
-    def key_type(self) -> str:
-        """Returns the authentication key type being used."""
-
-        if self.api_key:
-            return USER_KEY
-
-        if constants.OLLAMA_MODEL:
-            return "ollama"
-
-        return MITO_SERVER_KEY
 
     def _build_openai_client(self) -> Optional[Union[openai.AsyncOpenAI, openai.AsyncAzureOpenAI]]:
         base_url = None
@@ -201,12 +121,12 @@ This attribute is observed by the websocket provider to push the error to the cl
                 timeout=self.timeout,
             )
         
-        elif constants.OLLAMA_MODEL and not self.api_key:
+        elif constants.OLLAMA_MODEL:
             base_url = constants.OLLAMA_BASE_URL
             llm_api_key = "ollama"
             self.log.debug(f"Using Ollama with model: {constants.OLLAMA_MODEL}")
-        elif self.api_key:
-            llm_api_key = self.api_key
+        elif constants.OPENAI_API_KEY:
+            llm_api_key = constants.OPENAI_API_KEY
             self.log.debug("Using OpenAI with user-provided API key")
         else:
             self.log.warning("No valid API key or model configuration provided")
@@ -262,7 +182,7 @@ This attribute is observed by the websocket provider to push the error to the cl
 
         # Handle other providers as before
         completion_function_params = get_open_ai_completion_function_params(
-            message_type, model, messages, False, response_format_info
+            model, messages, False, response_format_info
         )
         
         # If they have set an Azure OpenAI or Ollama model, then we use it
@@ -313,7 +233,7 @@ This attribute is observed by the websocket provider to push the error to the cl
 
         # Handle other providers as before
         completion_function_params = get_open_ai_completion_function_params(
-            message_type, model, messages, True, response_format_info
+            model, messages, True, response_format_info
         )
         
         completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])

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]