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

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -4,13 +4,13 @@
 from typing import List
 from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
 from mito_ai.completions.prompt_builders.prompt_constants import (
+    CHART_CONFIG_RULES,
     CITATION_RULES,
     CELL_REFERENCE_RULES,
     get_database_rules
 )
 from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
 
-
 def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
     
     # The GET_CELL_OUTPUT tool only works on Chrome based browsers. 
@@ -28,6 +28,8 @@ The user is going to ask you to guide them as they complete a task. You will hel
 You have access to a set of tools that you can use to accomplish the task you've been given. You can use one tool per message, and will receive the result of that tool use in the user's response. You use tools step-by-step to accomplish a given task, with each tool use informed by the result of the previous tool use.
 
 Each time you use a tool, except for the finished_task tool, the user will execute the tool and provide you with updated information about the notebook and variables defined in the kernel to help you decide what to do next."""))
+    
+    sections.append(SG.Generic("Chart Config Rules", CHART_CONFIG_RULES))
 
     sections.append(SG.Generic("TOOL: CELL_UPDATE", """
 
@@ -223,6 +225,85 @@ Important information:
 4. If running all cells results in an error, the system will automatically handle the error through the normal error fixing process.
 5. Do not use this tool repeatedly if it continues to produce errors - instead, focus on fixing the specific error that occurred."""))
     
+    # SCRATCHPAD tool
+    sections.append(SG.Generic("TOOL: SCRATCHPAD", """
+When you need to explore data, check the filesystem, analyze mappings, or look up values without leaving code in the notebook, use the SCRATCHPAD tool.
+
+Format:
+{{
+    type: 'scratchpad',
+    message: str,
+    scratchpad_code: str,
+    scratchpad_summary: str
+}}
+
+Important information:
+1. The scratchpad_code will execute silently against the same kernel as your notebook, so you have access to all variables and dataframes.
+2. Any variables you create in scratchpad code MUST be prefixed with "scratch_" (e.g., use "scratch_temp_df" not "temp_df", use "scratch_files" not "files"). This prevents them from appearing in the variable list and confusing future decisions.
+3. CRITICAL: Do NOT modify existing variables. If you need to explore or modify data, create a copy with the scratch_ prefix first. For example, use "scratch_df = df.copy()" and then modify scratch_df, rather than modifying df directly. This ensures existing variables remain unchanged.
+4. Structure your code to print the information you need. Use print() statements for output that will be captured.
+5. If you need structured data, consider using JSON: `import json; print(json.dumps(your_data))`
+6. The results (including any errors) will be included in your next message, so you can use them to inform your next action.
+7. If the code errors, the error message and traceback will be included in the results. You can then decide to fix the code and try again, ask the user a question, or take a different approach.
+8. Use scratchpad for exploration work that doesn't belong in the final notebook. Once you have the information, create clean CELL_UPDATES with hardcoded values.
+9. The scratchpad_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Checking files", "Exploring data", "Analyzing mappings", "Looking up values"). Avoid full sentences or explanations. This should read like a quick commit message or code label, not a description.
+
+Example:
+{{
+    type: 'scratchpad',
+    message: "I'll check what files are in the current directory to find the data file.",
+    scratchpad_code: "import os\\nscratch_files = os.listdir('.')\\nprint('Files:', scratch_files)\\nfor scratch_file in scratch_files:\\n    if scratch_file.endswith('.csv'):\\n        print(f'CSV file found: {scratch_file}')",
+    scratchpad_summary: "Checking files"
+}}
+
+Example with dataframe exploration:
+{{
+    type: 'scratchpad',
+    message: "I'll explore the dataframe structure to understand the columns.",
+    scratchpad_code: "scratch_df = df.copy()\\nprint('Columns:', scratch_df.columns.tolist())\\nprint('Shape:', scratch_df.shape)\\nprint('First few rows:')\\nprint(scratch_df.head())",
+    scratchpad_summary: "Exploring dataframe"
+}}
+"""))
+    
+    # ASK_USER_QUESTION tool 
+    sections.append(SG.Generic("TOOL: ASK_USER_QUESTION", f"""
+
+When you have a specific question that you the user to answer so that you can figure out how to proceed in your work, you can respond in this format:
+
+{{
+    type: 'ask_user_question',
+    message: str,
+    question: str,
+    answers: Optional[List[str]]
+}}
+
+Important information:
+1. Use this tool when you need clarification from the user on how to proceed. Common scenarios include:
+   - A file or resource doesn't exist and you need to know how to proceed
+   - There are multiple valid approaches and you want the user to choose
+   - You need clarification on ambiguous requirements
+   - You encounter an error that requires user input to resolve
+2. The message should be a short description of what you've tried so far and why you need to ask the user a question now. This helps the user understand the context. The message provides background information but should NOT contain the actual question.
+3. The question field is REQUIRED and must always be provided. It cannot be null or empty. The question should be a clear, direct question that ends with a question mark. It should be concise and direct - do NOT include instructions or explanations in the question itself, as the answer options should make it clear what information is needed. For example, instead of "Which companies would you like me to compare Meta's stock performance against? Please provide a list of company names or stock tickers", just ask "Which companies would you like me to compare Meta's stock performance against?" The answer options will make it clear that company names or tickers are expected.
+4. The message and question serve different purposes: the message provides context about what you've tried, while the question is the actual question the user needs to answer. If your message contains a question, extract that question into the question field. For example, if your message says "I need to understand how you'd like to access the tweets", your question should be something like "How would you like to access the tweets?"
+5. Use the optional list of answers to provide the user multiple choice options to choose from. If it is an open ended question that you cannot create helpful multiple choice answers for, leave it blank and the user will respond in the text input field. 
+6. When creating multiple choice answer options:
+   - Make each option distinct and meaningful - avoid options that differ only slightly from each other.
+   - If there are no obvious predefined answers, leave it blank and the user will respond in the text input field.
+7. After the user responds to your question, you will receive their response in the next message and can continue with the task based on their answer.
+8. Do not use this tool for trivial questions that you can infer from context. Only use it when you cannot proceed in the user's task without answering your specific question first.
+
+    <Example>
+    {{
+        type: 'ask_user_question',
+        message: "I tried importing apple_prices.csv and confirmed that it does not exist in the current working directory.",
+        question: "The file apple_prices.csv does not exist. How do you want to proceed?",
+        answers: ["Pull Apple Stock prices using yfinance API", "Create placeholder data", "Skip this step"]
+    }}
+    </Example>
+
+"""))
+    
     # CREATE_STREAMLIT_APP tool
     sections.append(SG.Generic("TOOL: CREATE_STREAMLIT_APP", """
 
@@ -316,18 +397,29 @@ Important information:
     # RULES section
     sections.append(SG.Generic("RULES", """
 - You are working in a Jupyter Lab environment in a .ipynb file.
-- In each message you can choose one of the tools to respond with. BUT YOU WILL GET TO SEND MULTIPLE MESSAGES TO THE USER TO ACCOMPLISH YOUR TASK SO DO NOT TRY TO ACCOMPLISH YOUR TASK IN A SINGLE MESSAGE.
+- In each message you can choose one of the tools to respond with. YOU WILL GET TO SEND MULTIPLE MESSAGES TO THE USER TO ACCOMPLISH YOUR TASK SO DO NOT TRY TO ACCOMPLISH YOUR TASK IN A SINGLE MESSAGE.
 - After you send a CELL_UPDATE, the user will send you a message with the updated variables, code, and files in the current directory. You will use this information to decide what to do next, so it is critical that you wait for the user's response after each CELL_UPDATE before deciding your next action.
-- When updating code, keep as much of the original code as possible and do not recreate variables that already exist.
 - When writing the message, do not explain to the user how to use the CELL_UPDATE or FINISHED_TASK response, they will already know how to use them. Just provide a summary of your thought process. Do not reference any Cell IDs in the message.
 - When writing the message, do not include leading words like "Explanation:" or "Thought process:". Just provide a summary of your thought process.
 - When writing the message, use tickmarks when referencing specific variable names. For example, write `sales_df` instead of "sales_df" or just sales_df."""))
     
     # CODE STYLE section
     sections.append(SG.Generic("CODE STYLE", """
-- Avoid using try/except blocks and other defensive programming patterns (like checking if files exist before reading them, verifying variables are defined before using them, etc.) unless there is a really good reason. In Jupyter notebooks, errors should surface immediately so users can identify and fix issues. When errors are caught and suppressed or when defensive checks hide problems, users continue running broken code without realizing it, and the agent's auto-error-fix loop cannot trigger. If a column doesn't exist, a file is missing, a variable isn't defined, or a module isn't installed, let it error. The user needs to know.
+- When updating code, keep as much of the original code as possible and do not recreate variables that already exist.
 - When you want to display a dataframe to the user, just write the dataframe on the last line of the code cell instead of writing print(<dataframe name>). Jupyter will automatically display the dataframe in the notebook.
-- When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter."""))
+- When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter.
+- Avoid adding try/except blocks unless there is a very good reason. Do not use them for things like: 
+    ```
+    try: 
+        df = pd.read_csv('my_data.csv')
+    except: 
+        print("File not found")
+    ```
+    Instead, just let the cell error and use the ask_user_question tool to figure out how to proceed.
+- Avoid defensive if statements like checking if a variable exists in the globals or verifying that a column exists. Instead, just let the code error and use the ask_user_question tool to figure out how to proceed.
+- Do not simulate the data without the user explicity asking you to do so.
+- Do not replace broken code with print statements that explain the issue. Instead, leave the broken code in the notebook and use the ask_user_question tools to communicate the issue to the user and figure out how to proceed.
+"""))
     
     # CITATION_RULES 
     sections.append(SG.Generic("Citation Rules", f"""{CITATION_RULES}

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/chart_conversion_prompt.py

@@ -0,0 +1,27 @@
+# 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_conversion_prompt(code: str) -> str:
+    """
+    Create a prompt for converting matplotlib chart code to be used with the Chart Wizard.
+    
+    Args:
+        code: The matplotlib chart code to convert
+        
+    Returns:
+        A formatted prompt string
+    """
+    sections: List[PromptSection] = []
+
+    sections.append(SG.Generic("Instructions", "The following code contains a matplotlib chart. However, the chart must be converted to a specific format for use in our tool. Below you will find the rules used to create an acceptable chart; use these rules to reformat the code."))
+    
+    sections.append(SG.Generic("Chart Config Rules", CHART_CONFIG_RULES))
+    sections.append(SG.Generic("Code to Convert", f"```python\n{code}\n```"))
+    
+    prompt = Prompt(sections)
+    return str(prompt)

mito_ai/completions/prompt_builders/chat_system_message.py

@@ -4,6 +4,7 @@
 from typing import List
 from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
 from mito_ai.completions.prompt_builders.prompt_constants import (
+    CHART_CONFIG_RULES,
     CHAT_CODE_FORMATTING_RULES,
     CITATION_RULES,
     CELL_REFERENCE_RULES,
@@ -31,6 +32,7 @@ Other useful information:
 2. If the user asks you to generate a dashboard, app, or streamlit app for them, you should tell them that they must use Agent mode to complete the task. You are not able to automatically switch the user to agent mode, but they can switch to it themselves by using the Chat/Agent mode toggle in the bottom left corner of the Ai taskpane.
 """))
 
+    sections.append(SG.Generic("Chart Config Rules", CHART_CONFIG_RULES))
     sections.append(SG.Generic("DatabaseRules", get_database_rules()))
     sections.append(SG.Generic("Citation Rules", CITATION_RULES))
     sections.append(SG.Generic("Cell Reference Rules", CELL_REFERENCE_RULES))

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -11,6 +11,34 @@ import json
 from typing import Final
 from mito_ai.utils.schema import MITO_FOLDER
 
+CHART_CONFIG_RULES = """
+When creating a matplotlib chart, you must use the `# === CHART CONFIG ===` and `# === END CONFIG ===` markers to indicate the start and end of the chart configuration section.
+
+The chart configuration section is a list of variables used to customize the chart. This includes the titles, labels, colors, and any variables that affect the chart's appearance.
+
+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.
+- 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"').
+
+Common Mistakes to Avoid:
+- WRONG: COLOR = '"#1877F2" # Meta Blue'  (nested quotes and inline comment)
+- WRONG: COLOR = "#1877F2" # Meta Blue  (inline comment)
+- WRONG: COLOR = '"#1877F2"'  (nested quotes)
+- CORRECT: COLOR = "#1877F2"  (simple hex string, no nested quotes, no inline comments)
+
+Example:
+# === CHART CONFIG ===
+TITLE = "Sales by Product"
+X_LABEL = "Product"
+Y_LABEL = "Sales"
+BAR_COLOR = "#000000"
+# === END CONFIG ===
+"""
+
 CITATION_RULES = """
 It is important that the user is able to verify any insights that you share with them about their data. To make this easy for the user, you must cite the lines of code that you are drawing the insight from. To provide a citation, use one of the following formats inline in your response:
 

mito_ai/completions/prompt_builders/scratchpad_result_prompt.py

@@ -0,0 +1,17 @@
+# 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.models import ScratchpadResultMetadata
+from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
+from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
+
+
+def create_scratchpad_result_prompt(md: ScratchpadResultMetadata) -> str:
+    sections: List[PromptSection] = [
+        SG.Generic("Reminder", "Continue working on your current task using the scratchpad results below."),
+        SG.Generic("Scratchpad Result", f"The result of your scratchpad is: {md.scratchpadResult}"),
+    ]
+
+    prompt = Prompt(sections)
+    return str(prompt)

mito_ai/constants.py

@@ -5,7 +5,7 @@ import os
 from typing import Union
 
 # 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,13 @@ 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")
 
+# 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
+LITELLM_MODELS = [model.strip() for model in LITELLM_MODELS_STR.split(",") if model.strip()] if LITELLM_MODELS_STR else []
+
 # 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.
 

mito_ai/enterprise/litellm_client.py

@@ -0,0 +1,137 @@
+# 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
+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., "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 provider prefix (e.g., "openai/gpt-4o")
+            response_format_info: Optional response format specification
+            message_type: Type of message (chat, agent execution, etc.)
+            
+        Returns:
+            The completion text response
+        """
+        # Prepare parameters for LiteLLM
+        params = get_litellm_completion_function_params(
+            model=model,
+            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 provider prefix (e.g., "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 = ""
+        
+        # Prepare parameters for LiteLLM
+        params = get_litellm_completion_function_params(
+            model=model,
+            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/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,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"])