mito-ai 0.1.28__py3-none-any.whl → 0.1.30__py3-none-any.whl

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.28'
+__version__ = VERSION = '0.1.30'

mito_ai/openai_client.py

@@ -197,7 +197,7 @@ This attribute is observed by the websocket provider to push the error to the cl
             return openai.AsyncAzureOpenAI(
                 api_key=constants.AZURE_OPENAI_API_KEY,
                 api_version=constants.AZURE_OPENAI_API_VERSION,
-                azure_endpoint=constants.AZURE_OPENAI_ENDPOINT or OPENAI_MODEL_FALLBACK,
+                azure_endpoint=constants.AZURE_OPENAI_ENDPOINT, # type: ignore
                 max_retries=self.max_retries,
                 timeout=self.timeout,
             )
@@ -222,13 +222,25 @@ This attribute is observed by the websocket provider to push the error to the cl
         )
         return client
 
-    def _resolve_model(self, model: Optional[str] = None) -> str:
+    def _resolve_model(self, model: Optional[str] = None, response_format_info: Optional[ResponseFormatInfo] = None) -> 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}")
             return constants.AZURE_OPENAI_MODEL
-        elif constants.OLLAMA_MODEL is not None:
+        
+        # Otherwise, we use the fast model for anything other than the agent mode
+        if response_format_info:
+            return OPENAI_FAST_MODEL
+        
+        # If they have set an Ollama model, then we use it
+        if constants.OLLAMA_MODEL is not None:
             return constants.OLLAMA_MODEL
-        elif model:
+        
+        # If they have set a model, then we use it
+        if model:
             return model
+        
         return OPENAI_MODEL_FALLBACK
 
     async def request_completions(
@@ -253,7 +265,11 @@ This attribute is observed by the websocket provider to push the error to the cl
         completion = None
 
         try:
-            model = self._resolve_model(model) if response_format_info else OPENAI_FAST_MODEL
+            
+            # Make sure we are using the correct model
+            # TODO: If we bring back inline completions or another action that needs to 
+            # respond fast, we must require the user to configure a fast model with Azure as well. 
+            model = self._resolve_model(model, response_format_info)
 
             # Handle other providers as before
             completion_function_params = get_open_ai_completion_function_params(
@@ -301,7 +317,7 @@ This attribute is observed by the websocket provider to push the error to the cl
         accumulated_response = ""
         
         # Validate that the model is supported.
-        model = self._resolve_model(model)
+        model = self._resolve_model(model, response_format_info)
             
         # Send initial acknowledgment
         reply_fn(CompletionReply(