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

mito_ai/completions/completion_handlers/inline_completer_handler.py

@@ -5,7 +5,7 @@ from typing import List
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import InlineCompleterMetadata, MessageType
 from mito_ai.completions.prompt_builders.inline_completer_prompt import create_inline_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
 
@@ -17,9 +17,8 @@ class InlineCompleterHandler(CompletionHandler[InlineCompleterMetadata]):
     @staticmethod
     async def get_completion(
         metadata: InlineCompleterMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get an inline completion from the AI provider."""
         
@@ -37,9 +36,9 @@ class InlineCompleterHandler(CompletionHandler[InlineCompleterMetadata]):
         # Get the completion
         completion = await provider.request_completions(
             messages=messages, 
-            model=model,
             message_type=MessageType.INLINE_COMPLETION,
-            thread_id=None
+            thread_id=None,
+            use_fast_model=True
         )
         
         return completion

mito_ai/completions/completion_handlers/scratchpad_result_handler.py

@@ -5,7 +5,7 @@ from typing import List, Literal, Union
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import ScratchpadResultMetadata, MessageType, ResponseFormatInfo, AgentResponse
 from mito_ai.completions.prompt_builders.scratchpad_result_prompt import create_scratchpad_result_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 ScratchpadResultHandler(CompletionHandler[ScratchpadResultMetadata]):
     @staticmethod
     async def get_completion(
         metadata: ScratchpadResultMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get a scratchpad result completion from the AI provider."""
 
@@ -31,7 +30,7 @@ class ScratchpadResultHandler(CompletionHandler[ScratchpadResultMetadata]):
             )
 
         # Add the system message if it doesn't already exist
-        await append_agent_system_message(message_history, model, provider, metadata.threadId, True)
+        await append_agent_system_message(message_history, provider, metadata.threadId, True)
         
         # Create the prompt
         prompt = create_scratchpad_result_prompt(metadata)
@@ -41,12 +40,11 @@ class ScratchpadResultHandler(CompletionHandler[ScratchpadResultMetadata]):
         new_ai_optimized_message = create_ai_optimized_message(prompt, None, None)
         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 ScratchpadResultHandler(CompletionHandler[ScratchpadResultMetadata]):
         
         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/smart_debug_handler.py

@@ -15,7 +15,7 @@ from mito_ai.completions.prompt_builders.smart_debug_prompt import (
     create_error_prompt,
     remove_inner_thoughts_from_message,
 )
-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
@@ -29,9 +29,8 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
     @staticmethod
     async def get_completion(
         metadata: SmartDebugMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get a smart debug completion from the AI provider."""
 
@@ -43,7 +42,7 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
         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_error_prompt(error_message, active_cell_code, active_cell_id, variables, files)
@@ -53,13 +52,12 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
         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.SMART_DEBUG,
             user_input=error_message,
             thread_id=thread_id
@@ -78,7 +76,7 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
             "content": display_completion,
         }
         await message_history.append_message(
-            ai_response_message, display_response_message, model, provider, thread_id
+            ai_response_message, display_response_message, provider, thread_id
         )
 
         return display_completion
@@ -86,11 +84,10 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
     @staticmethod
     async def stream_completion(
         metadata: SmartDebugMetadata,
-        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 smart debug completions from the AI provider.
 
@@ -112,7 +109,7 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
         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_error_prompt(error_message, active_cell_code, active_cell_id, variables, files)
@@ -122,18 +119,17 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
         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.SMART_DEBUG,
             messages=message_history.get_ai_optimized_history(thread_id),
-            model=model,
             message_id=message_id,
+            thread_id=thread_id,
             reply_fn=reply_fn,
-            user_input=error_message,
-            thread_id=thread_id
+            user_input=error_message
         )
 
         # Process the completion to remove inner thoughts
@@ -149,7 +145,7 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
             "content": display_completion,
         }
         await message_history.append_message(
-            ai_response_message, display_response_message, model, provider, thread_id
+            ai_response_message, display_response_message, provider, thread_id
         )
 
         return display_completion

mito_ai/completions/completion_handlers/utils.py

@@ -5,7 +5,7 @@ import base64
 from typing import Optional, Union, List, Dict, Any, cast
 from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.completions.models import ThreadID
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.prompt_builders.chat_system_message import (
     create_chat_system_message_prompt,
@@ -17,8 +17,7 @@ from mito_ai.completions.prompt_builders.agent_system_message import (
 
 async def append_chat_system_message(
     message_history: GlobalMessageHistory,
-    model: str,
-    provider: OpenAIProvider,
+    provider: ProviderManager,
     thread_id: ThreadID,
 ) -> None:
 
@@ -39,7 +38,6 @@ async def append_chat_system_message(
     await message_history.append_message(
         ai_optimized_message=system_message,
         display_message=system_message,
-        model=model,
         llm_provider=provider,
         thread_id=thread_id,
     )
@@ -47,8 +45,7 @@ async def append_chat_system_message(
 
 async def append_agent_system_message(
     message_history: GlobalMessageHistory,
-    model: str,
-    provider: OpenAIProvider,
+    provider: ProviderManager,
     thread_id: ThreadID,
     isChromeBrowser: bool,
 ) -> None:
@@ -70,7 +67,6 @@ async def append_agent_system_message(
     await message_history.append_message(
         ai_optimized_message=system_message,
         display_message=system_message,
-        model=model,
         llm_provider=provider,
         thread_id=thread_id,
     )

mito_ai/completions/handlers.py

@@ -37,9 +37,10 @@ from mito_ai.completions.models import (
     ScratchpadResultMetadata,
     MessageType
 )
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.utils.create import initialize_user
 from mito_ai.utils.version_utils import is_pro
+from mito_ai.utils.model_utils import get_available_models
 from mito_ai.completions.completion_handlers.chat_completion_handler import get_chat_completion, stream_chat_completion
 from mito_ai.completions.completion_handlers.smart_debug_handler import get_smart_debug_completion, stream_smart_debug_completion
 from mito_ai.completions.completion_handlers.code_explain_handler import get_code_explain_completion, stream_code_explain_completion
@@ -49,8 +50,6 @@ from mito_ai.completions.completion_handlers.agent_auto_error_fixup_handler impo
 from mito_ai.completions.completion_handlers.scratchpad_result_handler import get_scratchpad_result_completion
 from mito_ai.utils.telemetry_utils import identify
 
-FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
-
 # The GlobalMessageHistory is now created in __init__.py and passed to handlers
 # to ensure there's only one instance managing the .mito/ai-chats directory locks
 
@@ -61,13 +60,12 @@ FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
 class CompletionHandler(JupyterHandler, WebSocketHandler):
     """Completion websocket handler."""
 
-    def initialize(self, llm: OpenAIProvider, message_history: GlobalMessageHistory) -> None:
+    def initialize(self, llm: ProviderManager, message_history: GlobalMessageHistory) -> None:
         super().initialize()
         self.log.debug("Initializing websocket connection %s", self.request.path)
         self._llm = llm
         self._message_history = message_history
         self.is_pro = is_pro()
-        self._selected_model = FALLBACK_MODEL
         self.is_electron = False
         identify(llm.key_type)
         
@@ -202,7 +200,25 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
         if type == MessageType.UPDATE_MODEL_CONFIG:
             model = metadata_dict.get('model')
             if model:
-                self._selected_model = model
+                # Validate model is in allowed list
+                available_models = get_available_models()
+                if model not in available_models:
+                    error = CompletionError(
+                        error_type="InvalidModelConfig",
+                        title="Invalid model configuration",
+                        traceback="",
+                        hint=f"Model '{model}' is not in the allowed model list. Available models: {', '.join(available_models)}"
+                    )
+                    reply = CompletionReply(
+                        items=[],
+                        error=error,
+                        parent_id=parsed_message.get('message_id')
+                    )
+                    self.reply(reply)
+                    return
+                
+                # Set the model in ProviderManager
+                self._llm.set_selected_model(model)
                 self.log.info(f"Model updated to: {model}")
                 reply = CompletionReply(
                     items=[CompletionItem(content=f"Model updated to {model}", isIncomplete=False)],
@@ -241,7 +257,6 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                 await self._message_history.append_message(
                     ai_optimized_message=ai_optimized_message,
                     display_message=display_optimized_message,
-                    model=self._selected_model,
                     llm_provider=self._llm,
                     thread_id=thread_id_to_stop
                 )
@@ -255,8 +270,6 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             message_id = parsed_message.get('message_id')
             stream = parsed_message.get('stream')
 
-            # When handling completions, always use the selected model
-            model = self._selected_model
             if type == MessageType.CHAT:
                 chat_metadata = ChatMessageMetadata(**metadata_dict)
                 
@@ -268,13 +281,12 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                         self._llm, 
                         self._message_history, 
                         message_id,
-                        self.reply,
-                        model
+                        self.reply
                     )
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_chat_completion(chat_metadata, self._llm, self._message_history, model)
+                    completion = await get_chat_completion(chat_metadata, self._llm, self._message_history)
             elif type == MessageType.SMART_DEBUG:
                 smart_debug_metadata = SmartDebugMetadata(**metadata_dict)
                 # Handle streaming if requested and available
@@ -285,13 +297,12 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                         self._llm, 
                         self._message_history, 
                         message_id,
-                        self.reply,
-                        model
+                        self.reply
                     )
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, self._message_history, model)
+                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, self._message_history)
             elif type == MessageType.CODE_EXPLAIN:
                 code_explain_metadata = CodeExplainMetadata(**metadata_dict)
 
@@ -303,25 +314,24 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                         self._llm, 
                         self._message_history,
                         message_id,
-                        self.reply,
-                        model
+                        self.reply
                     )
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, self._message_history, model)
+                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, self._message_history)
             elif type == MessageType.AGENT_EXECUTION:
                 agent_execution_metadata = AgentExecutionMetadata(**metadata_dict)
-                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, self._message_history, model)
+                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, self._message_history)
             elif type == MessageType.AGENT_AUTO_ERROR_FIXUP:
                 agent_auto_error_fixup_metadata = AgentSmartDebugMetadata(**metadata_dict)
-                completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_metadata, self._llm, self._message_history, model)
+                completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_metadata, self._llm, self._message_history)
             elif type == MessageType.AGENT_SCRATCHPAD_RESULT:
                 scratchpad_result_metadata = ScratchpadResultMetadata(**metadata_dict)
-                completion = await get_scratchpad_result_completion(scratchpad_result_metadata, self._llm, self._message_history, model)
+                completion = await get_scratchpad_result_completion(scratchpad_result_metadata, self._llm, self._message_history)
             elif type == MessageType.INLINE_COMPLETION:
                 inline_completer_metadata = InlineCompleterMetadata(**metadata_dict)
-                completion = await get_inline_completion(inline_completer_metadata, self._llm, self._message_history, model)
+                completion = await get_inline_completion(inline_completer_metadata, self._llm, self._message_history)
             else:
                 raise ValueError(f"Invalid message type: {type}")
             

mito_ai/completions/message_history.py

@@ -11,7 +11,7 @@ from typing import Dict, List, Optional
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import CompletionRequest, ChatThreadMetadata, MessageType, ThreadID
 from mito_ai.completions.prompt_builders.chat_name_prompt import create_chat_name_prompt
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.utils.schema import MITO_FOLDER
 from mito_ai.utils.message_history_utils import trim_old_messages
 
@@ -19,16 +19,15 @@ CHAT_HISTORY_VERSION = 2 # Increment this if the schema changes
 NEW_CHAT_NAME = "(New Chat)"
 NUMBER_OF_THREADS_CUT_OFF = 50
 
-async def generate_short_chat_name(user_message: str, assistant_message: str, model: str, llm_provider: OpenAIProvider) -> str:
+async def generate_short_chat_name(user_message: str, assistant_message: str, llm_provider: ProviderManager) -> str:
     prompt = create_chat_name_prompt(user_message, assistant_message)
 
     completion = await llm_provider.request_completions(
         messages=[{"role": "user", "content": prompt}], 
-        # We set the model so we can use the correct model provider, but request_completions will decide to 
-        # use the fast model from that provider to make the request.
-        model=model, 
+        # Use fast model from the selected provider for chat name generation
         message_type=MessageType.CHAT_NAME_GENERATION,
-        thread_id=None
+        thread_id=None,
+        use_fast_model=True
     )
     
     # Do a little cleanup of the completion. Gemini seems to return the string
@@ -135,7 +134,7 @@ class GlobalMessageHistory:
             Returns the AI-optimized history for the specified thread or newest thread.
         get_display_history(thread_id: Optional[ThreadID] = None) -> List[ChatCompletionMessageParam]:
             Returns the display-optimized history for the specified thread or newest thread.
-        append_message(ai_optimized_message: ChatCompletionMessageParam, display_message: ChatCompletionMessageParam, llm_provider: OpenAIProvider, thread_id: Optional[ThreadID] = None) -> None:
+        append_message(ai_optimized_message: ChatCompletionMessageParam, display_message: ChatCompletionMessageParam, llm_provider: ProviderManager, thread_id: Optional[ThreadID] = None) -> None:
             Appends messages to the specified thread (or newest thread) and generates a name if needed.
         truncate_histories(index: int, thread_id: Optional[ThreadID] = None) -> None:
             Truncates messages at the given index for the specified thread.
@@ -265,8 +264,7 @@ class GlobalMessageHistory:
         self, 
         ai_optimized_message: ChatCompletionMessageParam, 
         display_message: ChatCompletionMessageParam, 
-        model: str,
-        llm_provider: OpenAIProvider,
+        llm_provider: ProviderManager,
         thread_id: ThreadID
     ) -> None:
         """
@@ -305,7 +303,7 @@ class GlobalMessageHistory:
 
         # Outside the lock, await the name generation if needed
         if name_gen_input:
-            new_name = await generate_short_chat_name(str(name_gen_input[0]), str(name_gen_input[1]), model, llm_provider)
+            new_name = await generate_short_chat_name(str(name_gen_input[0]), str(name_gen_input[1]), llm_provider)
             with self._lock:
                 # Update the thread's name if still required
                 thread = self._chat_threads[thread_id]

mito_ai/completions/prompt_builders/chart_add_field_prompt.py

@@ -0,0 +1,35 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List
+from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
+from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
+from mito_ai.completions.prompt_builders.prompt_constants import CHART_CONFIG_RULES
+
+def create_chart_add_field_prompt(code: str, user_description: str, existing_variables: List[str]) -> str:
+    """
+    Create a prompt for adding a new field to the chart configuration.
+    
+    Args:
+        code: The current chart code
+        user_description: The user's description of what field they want to add
+        existing_variables: List of existing variable names in the config
+        
+    Returns:
+        A formatted prompt string
+    """
+    sections: List[PromptSection] = []
+
+    sections.append(SG.Generic("Instructions", "The user wants to add a new field to the chart configuration. You need to:\n1. Understand what field the user wants to add based on their description\n2. Add the appropriate variable to the chart configuration section\n3. Use the variable in the chart code where appropriate\n4. Return the complete updated code\n\nIMPORTANT: If you cannot add the requested field (e.g., the request is unclear, ambiguous, or not applicable to chart configuration), do NOT return any code block. Simply respond with a brief explanation without including any Python code blocks."))
+    
+    sections.append(SG.Generic("Chart Config Rules", CHART_CONFIG_RULES))
+    
+    existing_vars_text = ", ".join(existing_variables) if existing_variables else "none"
+    sections.append(SG.Generic("Existing Variables", f"The following variables already exist in the chart configuration: {existing_vars_text}"))
+    
+    sections.append(SG.Generic("User Request", f"The user wants to add a field for: {user_description}"))
+    
+    sections.append(SG.Generic("Current Code", f"```python\n{code}\n```"))
+    
+    prompt = Prompt(sections)
+    return str(prompt)

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -20,6 +20,8 @@ Rules:
 - All imports must appear at the top, before the chart configuration section.
 - Variables with multiple words should be underscore-separated.
 - All colors should be in hex format (e.g., "#3498db"). Use quotes around the hex string: COLOR = "#3498db" or COLOR = '#3498db'. Do NOT nest quotes.
+- Never use RGB/RGBA tuples/lists for colors (e.g. (0, 0.4, 0.8, 0.8) is forbidden).
+- If transparency is needed, store it separately as ALPHA = 0.8 and apply it in code (e.g. to_rgba(HEX_COLOR, ALPHA)).
 - Variables can only be strings, numbers, booleans, tuples, or lists. 
 - NEVER include comments on the same line as a variable assignment. Each variable assignment must be on its own line with no trailing comments.
 - For string values, use either single or double quotes (e.g., TITLE = "Sales by Product" or TITLE = 'Sales by Product'). Do not use nested quotes (e.g., do NOT use '"value"').

mito_ai/constants.py

@@ -2,10 +2,10 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import os
-from typing import Union
+from typing import Union, List
 
 # Claude
-CLAUDE_API_KEY = os.environ.get("CLAUDE_API_KEY")
+ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
 
 # Gemini
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
@@ -23,6 +23,35 @@ AZURE_OPENAI_API_VERSION = os.environ.get("AZURE_OPENAI_API_VERSION")
 AZURE_OPENAI_ENDPOINT = os.environ.get("AZURE_OPENAI_ENDPOINT")
 AZURE_OPENAI_MODEL = os.environ.get("AZURE_OPENAI_MODEL")
 
+def parse_comma_separated_models(models_str: str) -> List[str]:
+    """
+    Parse a comma-separated string of model names into a list.
+    Handles quoted and unquoted values, stripping whitespace and quotes.
+    
+    Args:
+        models_str: Comma-separated string of model names (e.g., "model1,model2" or '"model1","model2"')
+    
+    Returns:
+        List of model names with whitespace and quotes stripped
+    """
+    if not models_str:
+        return []
+    return [model.strip().strip('"\'') for model in models_str.split(",") if model.strip()]
+
+# LiteLLM Config (Enterprise mode only)
+LITELLM_BASE_URL = os.environ.get("LITELLM_BASE_URL")
+LITELLM_API_KEY = os.environ.get("LITELLM_API_KEY")
+LITELLM_MODELS_STR = os.environ.get("LITELLM_MODELS", "")
+# Parse comma-separated string into list, strip whitespace and quotes
+LITELLM_MODELS = parse_comma_separated_models(LITELLM_MODELS_STR)
+
+# Abacus AI Config (Enterprise mode only)
+ABACUS_BASE_URL = os.environ.get("ABACUS_BASE_URL")
+ABACUS_API_KEY = os.environ.get("ABACUS_API_KEY")
+ABACUS_MODELS_STR = os.environ.get("ABACUS_MODELS", "")
+# Parse comma-separated string into list, strip whitespace and quotes
+ABACUS_MODELS = parse_comma_separated_models(ABACUS_MODELS_STR)
+
 # Mito AI Base URLs and Endpoint Paths
 MITO_PROD_BASE_URL = "https://7eax4i53f5odkshhlry4gw23by0yvnuv.lambda-url.us-east-1.on.aws/v2"
 MITO_DEV_BASE_URL = "https://g5vwmogjg7gh7aktqezyrvcq6a0hyfnr.lambda-url.us-east-1.on.aws/v2"

mito_ai/enterprise/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) Saga Inc.
-# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+# Distributed under the terms of the Enterprise License at the root of this repository.