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

mito_ai/enterprise/litellm_client.py

@@ -0,0 +1,144 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the The Mito Enterprise license.
+
+from typing import Optional, List, Callable, Union, Dict, Any
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.models import (
+    MessageType,
+    ResponseFormatInfo,
+    CompletionReply,
+    CompletionStreamChunk,
+    CompletionItem,
+)
+from mito_ai.utils.litellm_utils import get_litellm_completion_function_params
+from mito_ai.utils.model_utils import strip_router_prefix
+import litellm
+
+class LiteLLMClient:
+    """
+    A client for interacting with LiteLLM server endpoints.
+    LiteLLM provides an OpenAI-compatible API, so we use the LiteLLM SDK directly.
+    """
+    
+    def __init__(self, api_key: Optional[str], base_url: str, timeout: int = 30, max_retries: int = 1):
+        self.api_key = api_key
+        self.base_url = base_url
+        self.timeout = timeout
+        self.max_retries = max_retries
+    
+    async def request_completions(
+        self,
+        messages: List[ChatCompletionMessageParam],
+        model: str,  # Should include provider prefix (e.g., "LiteLLM/openai/gpt-4o")
+        response_format_info: Optional[ResponseFormatInfo] = None,
+        message_type: MessageType = MessageType.CHAT
+    ) -> str:
+        """
+        Request completions from LiteLLM server.
+        
+        Args:
+            messages: List of chat messages
+            model: Model name with router and provider prefix (e.g., "LiteLLM/openai/gpt-4o")
+            response_format_info: Optional response format specification
+            message_type: Type of message (chat, agent execution, etc.)
+            
+        Returns:
+            The completion text response
+        """
+        # Strip router prefix if present (LiteLLM/ prefix)
+        model_for_litellm = strip_router_prefix(model)
+        
+        # Prepare parameters for LiteLLM
+        params = get_litellm_completion_function_params(
+            model=model_for_litellm,
+            messages=messages,
+            api_key=self.api_key,
+            api_base=self.base_url,
+            timeout=self.timeout,
+            stream=False,
+            response_format_info=response_format_info,
+        )
+        
+        try:
+            # Use LiteLLM's acompletion function
+            response = await litellm.acompletion(**params)
+            
+            # Extract content from response
+            if response and response.choices and len(response.choices) > 0:
+                content = response.choices[0].message.content
+                return content or ""
+            else:
+                return ""
+        except Exception as e:
+            raise Exception(f"LiteLLM completion error: {str(e)}")
+    
+    async def stream_completions(
+        self,
+        messages: List[ChatCompletionMessageParam],
+        model: str,
+        message_type: MessageType,
+        message_id: str,
+        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+        response_format_info: Optional[ResponseFormatInfo] = None
+    ) -> str:
+        """
+        Stream completions from LiteLLM server.
+        
+        Args:
+            messages: List of chat messages
+            model: Model name with router and provider prefix (e.g., "LiteLLM/openai/gpt-4o")
+            message_type: Type of message (chat, agent execution, etc.)
+            message_id: ID of the message being processed
+            reply_fn: Function to call with each chunk for streaming replies
+            response_format_info: Optional response format specification
+            
+        Returns:
+            The accumulated response string
+        """
+        accumulated_response = ""
+        
+        # Strip router prefix if present (LiteLLM/ prefix)
+        model_for_litellm = strip_router_prefix(model)
+        
+        # Prepare parameters for LiteLLM
+        params = get_litellm_completion_function_params(
+            model=model_for_litellm,
+            messages=messages,
+            api_key=self.api_key,
+            api_base=self.base_url,
+            timeout=self.timeout,
+            stream=True,
+            response_format_info=response_format_info,
+        )
+        
+        try:
+            # Use LiteLLM's acompletion with stream=True
+            # When stream=True, acompletion returns an async iterable after awaiting
+            stream = await litellm.acompletion(**params)
+            
+            # Process streaming chunks
+            async for chunk in stream:
+                if chunk and chunk.choices and len(chunk.choices) > 0:
+                    delta = chunk.choices[0].delta
+                    content = delta.content if delta and delta.content else ""
+                    
+                    if content:
+                        accumulated_response += content
+                        
+                        # Check if this is the final chunk
+                        is_finished = chunk.choices[0].finish_reason is not None
+                        
+                        # Send chunk to frontend
+                        reply_fn(CompletionStreamChunk(
+                            parent_id=message_id,
+                            chunk=CompletionItem(
+                                content=content,
+                                isIncomplete=not is_finished,
+                                token=message_id,
+                            ),
+                            done=is_finished,
+                        ))
+            
+            return accumulated_response
+        except Exception as e:
+            raise Exception(f"LiteLLM streaming error: {str(e)}")

mito_ai/enterprise/utils.py

@@ -5,11 +5,25 @@
 # Distributed under the terms of the The Mito Enterprise license.
 
 from mito_ai.utils.version_utils import is_enterprise, is_mitosheet_private
-from mito_ai.constants import AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL
+from mito_ai.constants import (
+    AZURE_OPENAI_API_KEY, 
+    AZURE_OPENAI_ENDPOINT, 
+    AZURE_OPENAI_API_VERSION, 
+    AZURE_OPENAI_MODEL,
+    ABACUS_BASE_URL,
+    ABACUS_MODELS
+)
 
 def is_azure_openai_configured() -> bool:
     """
     Azure OpenAI is only supported for Mito Enterprise users
     """
     is_allowed_to_use_azure = is_enterprise() or is_mitosheet_private()
-    return all([is_allowed_to_use_azure, AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL])
+    return all([is_allowed_to_use_azure, AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL])
+
+def is_abacus_configured() -> bool:
+    """
+    Abacus AI is only supported for Mito Enterprise users.
+    Checks if Abacus AI is configured with base URL and models.
+    """
+    return all([is_enterprise(), ABACUS_BASE_URL, ABACUS_MODELS])

mito_ai/log/handlers.py

@@ -32,7 +32,7 @@ class LogHandler(APIHandler):
         log_event = data['log_event']
         params = data.get('params', {})
         
-        key_type = MITO_SERVER_KEY if self.key_type == "mito_server_key" else USER_KEY
+        key_type = MITO_SERVER_KEY if self.key_type == MITO_SERVER_KEY else USER_KEY
         log(log_event, params, key_type=key_type)
 
 

mito_ai/openai_client.py

@@ -7,12 +7,13 @@ 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
-from mito_ai.enterprise.utils import is_azure_openai_configured
+from mito_ai.enterprise.utils import is_azure_openai_configured, is_abacus_configured
 from mito_ai.logger import get_logger
+from mito_ai.utils.model_utils import strip_router_prefix
 from mito_ai.completions.models import (
     AICapabilities,
     CompletionError,
@@ -24,28 +25,17 @@ from mito_ai.completions.models import (
     ResponseFormatInfo,
 )
 from mito_ai.utils.open_ai_utils import (
-    check_mito_server_quota,
     get_ai_completion_from_mito_server,
     get_open_ai_completion_function_params,
     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,
-)
+from mito_ai.utils.server_limits import update_mito_server_quota, check_mito_server_quota
 
 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,15 @@ 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 is_abacus_configured():
+            return AICapabilities(
+                configuration={
+                    "model": "<dynamic>"
+                },
+                provider="Abacus AI",
+            )
+
+        if constants.OLLAMA_MODEL:
             return AICapabilities(
                 configuration={
                     "model": constants.OLLAMA_MODEL
@@ -141,14 +84,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 +110,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 +129,16 @@ 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 is_abacus_configured():
+            base_url = constants.ABACUS_BASE_URL
+            llm_api_key = constants.ABACUS_API_KEY
+            self.log.debug(f"Using Abacus AI with base URL: {constants.ABACUS_BASE_URL}")
+        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")
@@ -221,17 +153,25 @@ This attribute is observed by the websocket provider to push the error to the cl
         )
         return client
 
-    def _adjust_model_for_azure_or_ollama(self, model: str) -> str:
+    def _adjust_model_for_provider(self, model: str) -> str:
         
         # If they have set an Azure OpenAI model, then we always use it
         if is_azure_openai_configured() and constants.AZURE_OPENAI_MODEL is not None:
             self.log.debug(f"Resolving to Azure OpenAI model: {constants.AZURE_OPENAI_MODEL}")
+            # TODO: We should update Azure so it works the way LiteLLM and Abacus do: 
+            # when configured, we only show models from Azure in the UI. 
             return constants.AZURE_OPENAI_MODEL
         
         # If they have set an Ollama model, then we use it
         if constants.OLLAMA_MODEL is not None:
             return constants.OLLAMA_MODEL
         
+        # If using Abacus, strip the "Abacus/" prefix from the model name
+        if is_abacus_configured() and model.lower().startswith('abacus/'):
+            stripped_model = strip_router_prefix(model)
+            self.log.debug(f"Stripping Abacus prefix: {model} -> {stripped_model}")
+            return stripped_model
+        
         # Otherwise, we use the model they provided
         return model
         
@@ -262,11 +202,11 @@ 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
-        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
+        completion_function_params["model"] = self._adjust_model_for_provider(completion_function_params["model"])
 
         if self._active_async_client is not None:
             response = await self._active_async_client.chat.completions.create(**completion_function_params)
@@ -313,10 +253,10 @@ 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"])
+        completion_function_params["model"] = self._adjust_model_for_provider(completion_function_params["model"])
 
         try:
             if self._active_async_client is not None: