mito-ai 0.1.33__py3-none-any.whl → 0.1.49__py3-none-any.whl

mito_ai/completions/providers.py

@@ -2,6 +2,7 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 from __future__ import annotations
+import asyncio
 from typing import Any, Callable, Dict, List, Optional, Union, cast
 from mito_ai import constants
 from openai.types.chat import ChatCompletionMessageParam
@@ -28,12 +29,16 @@ from mito_ai.completions.models import (
 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.constants import get_model_provider
+from mito_ai.utils.provider_utils import get_model_provider
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
 
 __all__ = ["OpenAIProvider"]
 
@@ -66,6 +71,9 @@ This attribute is observed by the websocket provider to push the error to the cl
 
     @property
     def capabilities(self) -> AICapabilities:
+        """
+        Returns the capabilities of the AI provider.
+        """
         if constants.CLAUDE_API_KEY and not self.api_key:
             return AICapabilities(
                 configuration={"model": "<dynamic>"},
@@ -78,6 +86,7 @@ This attribute is observed by the websocket provider to push the error to the cl
             )
         if self._openai_client:
             return self._openai_client.capabilities
+        
         return AICapabilities(
             configuration={"model": "<dynamic>"},
             provider="Mito server",
@@ -100,7 +109,8 @@ This attribute is observed by the websocket provider to push the error to the cl
         model: str,
         response_format_info: Optional[ResponseFormatInfo] = None,
         user_input: Optional[str] = None,
-        thread_id: Optional[str] = None
+        thread_id: Optional[str] = None,
+        max_retries: int = 3
     ) -> str:
         """
         Request completions from the AI provider.
@@ -109,43 +119,69 @@ This attribute is observed by the websocket provider to push the error to the cl
         completion = None
         last_message_content = str(messages[-1].get('content', '')) if messages else ""
         model_type = get_model_provider(model)
-        try:
-            if model_type == "claude":
-                api_key = constants.CLAUDE_API_KEY
-                anthropic_client = AnthropicClient(api_key=api_key, model=model)
-                completion = await anthropic_client.request_completions(messages, response_format_info, message_type)
-            elif model_type == "gemini":
-                api_key = constants.GEMINI_API_KEY
-                gemini_client = GeminiClient(api_key=api_key, model=model)
-                messages_for_gemini = [dict(m) for m in messages]
-                completion = await gemini_client.request_completions(messages_for_gemini, 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(
+        
+        # Retry loop
+        for attempt in range(max_retries + 1):
+            try:
+                if model_type == "claude":
+                    api_key = constants.CLAUDE_API_KEY
+                    anthropic_client = AnthropicClient(api_key=api_key)
+                    completion = await anthropic_client.request_completions(messages, 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)
+                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,
+                        response_format_info=response_format_info
+                    )
+                else:
+                    raise ValueError(f"No AI provider configured for model: {model}")
+                
+                # Success! Log and return
+                log_ai_completion_success(
+                    key_type=USER_KEY if self.key_type == "user" else MITO_SERVER_KEY,
                     message_type=message_type,
-                    messages=messages,
-                    model=model,
-                    response_format_info=response_format_info
+                    last_message_content=last_message_content,
+                    response={"completion": completion},
+                    user_input=user_input or "",
+                    thread_id=thread_id or "",
+                    model=model
                 )
-            else:
-                raise ValueError(f"No AI provider configured for model: {model}")
-            log_ai_completion_success(
-                key_type=USER_KEY if self.key_type == "user" 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
-            )
-            return completion
+                return completion # type: ignore
+            
+            except PermissionError as e:
+                # 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)
+                raise
 
-        except BaseException as e:
-            self.log.exception(f"Error during request_completions: {e}")
-            self.last_error = CompletionError.from_exception(e)
-            log(MITO_AI_COMPLETION_ERROR, params={KEY_TYPE_PARAM: self.key_type}, error=e)
-            raise
+            except BaseException as e:
+                # Check if we should retry (not on the last attempt)
+                if attempt < max_retries:
+                    # 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)
+                    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)
+                    raise
+        
+        # This should never be reached due to the raise in the except block,
+        # but added to satisfy the linter
+        raise RuntimeError("Unexpected code path in request_completions")
 
     async def stream_completions(
         self,
@@ -176,19 +212,23 @@ 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
-                anthropic_client = AnthropicClient(api_key=api_key, model=model)
-                accumulated_response = await anthropic_client.stream_response(
+                anthropic_client = AnthropicClient(api_key=api_key)
+                accumulated_response = await anthropic_client.stream_completions(
                     messages=messages,
+                    model=model,
                     message_type=message_type,
                     message_id=message_id,
                     reply_fn=reply_fn
                 )
             elif model_type == "gemini":
                 api_key = constants.GEMINI_API_KEY
-                gemini_client = GeminiClient(api_key=api_key, model=model)
+                gemini_client = GeminiClient(api_key=api_key)
+                # TODO: We shouldn't need to do this because the messages should already be dictionaries... 
+                # but if we do have to do some pre-processing, we should do it in the gemini_client instead.
                 messages_for_gemini = [dict(m) for m in messages]
                 accumulated_response = await gemini_client.stream_completions(
                     messages=messages_for_gemini,
+                    model=model,
                     message_id=message_id,
                     reply_fn=reply_fn,
                     message_type=message_type
@@ -224,14 +264,8 @@ This attribute is observed by the websocket provider to push the error to the cl
         except BaseException as e:
             self.log.exception(f"Error during stream_completions: {e}")
             self.last_error = CompletionError.from_exception(e)
-            log(
-                MITO_AI_COMPLETION_ERROR,
-                params={
-                    KEY_TYPE_PARAM: self.key_type,
-                    'message_type': message_type.value,
-                },
-                error=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(
                 parent_id=message_id,

mito_ai/constants.py

@@ -23,30 +23,10 @@ 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 get_model_provider(model: str) -> Union[str, None]:
-    """
-    Determine the model type based on the model name prefix
-    """
-    if not model:
-        return None
-
-    model_lower = model.lower()
-
-    if model_lower.startswith('claude'):
-        return 'claude'
-    elif model_lower.startswith('gemini'):
-        return 'gemini'
-    elif model_lower.startswith('ollama'):
-        return 'ollama'
-    elif model_lower.startswith('gpt'):
-        return 'openai'
-
-    return None
-
-
 # Mito AI Base URLs and Endpoint Paths
-MITO_PROD_BASE_URL = "https://yxwyadgaznhavqvgnbfuo2k6ca0jboku.lambda-url.us-east-1.on.aws"
-MITO_DEV_BASE_URL = "https://x3rafympznv4abp7phos44gzgu0clbui.lambda-url.us-east-1.on.aws"
+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_LOCAL_BASE_URL = "http://127.0.0.1:8000/v2" # When you are running the mito completion server locally
 
 # Set ACTIVE_BASE_URL manually
 ACTIVE_BASE_URL = MITO_PROD_BASE_URL  # Change to MITO_DEV_BASE_URL for dev
@@ -59,4 +39,25 @@ OPENAI_PATH = "openai/completions"
 # Full URLs (always use ACTIVE_BASE_URL)
 MITO_ANTHROPIC_URL = f"{ACTIVE_BASE_URL}/{ANTHROPIC_PATH}"
 MITO_GEMINI_URL = f"{ACTIVE_BASE_URL}/{GEMINI_PATH}"
-MITO_OPENAI_URL = f"{ACTIVE_BASE_URL}/{OPENAI_PATH}"
+MITO_OPENAI_URL = f"{ACTIVE_BASE_URL}/{OPENAI_PATH}"
+
+# Streamlit conversion endpoints
+MITO_STREAMLIT_DEV_BASE_URL = "https://fr12uvtfy5.execute-api.us-east-1.amazonaws.com"
+MITO_STREAMLIT_TEST_BASE_URL = "https://iyual08t6d.execute-api.us-east-1.amazonaws.com"
+
+# Set ACTIVE_BASE_URL manually
+# TODO: Modify to PROD url before release
+ACTIVE_STREAMLIT_BASE_URL = MITO_STREAMLIT_DEV_BASE_URL  # Change to MITO_STREAMLIT_DEV_BASE_URL for dev
+
+# AWS Cognito configuration
+COGNITO_CONFIG_DEV = {
+    'TOKEN_ENDPOINT': 'https://mito-app-auth.auth.us-east-1.amazoncognito.com/oauth2/token',
+    'CLIENT_ID': '6ara3u3l8sss738hrhbq1qtiqf',
+    'CLIENT_SECRET': '',
+    'REDIRECT_URI': 'http://localhost:8888/lab'
+}
+
+ACTIVE_COGNITO_CONFIG = COGNITO_CONFIG_DEV # Change to COGNITO_CONFIG_DEV for dev
+
+
+MESSAGE_HISTORY_TRIM_THRESHOLD: int = 3

mito_ai/file_uploads/__init__.py

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

mito_ai/file_uploads/handlers.py

@@ -0,0 +1,248 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import tempfile
+import tornado
+from typing import Dict, Any
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.utils.telemetry_utils import (
+    log_file_upload_attempt,
+    log_file_upload_failure,
+)
+
+MAX_IMAGE_SIZE_MB = 3
+
+
+def _is_image_file(filename: str) -> bool:
+    image_extensions = {
+        ".jpg",
+        ".jpeg",
+        ".png",
+        ".gif",
+        ".bmp",
+        ".tiff",
+        ".tif",
+        ".webp",
+        ".svg",
+    }
+    file_extension = os.path.splitext(filename)[1].lower()
+    return file_extension in image_extensions
+
+
+def _check_image_size_limit(file_data: bytes, filename: str) -> None:
+    if not _is_image_file(filename):
+        return
+
+    file_size_mb = len(file_data) / (1024 * 1024)  # Convert bytes to MB
+
+    if file_size_mb > MAX_IMAGE_SIZE_MB:
+        raise ValueError(f"Image exceeded {MAX_IMAGE_SIZE_MB}MB limit.")
+
+
+class FileUploadHandler(APIHandler):
+    # Class-level dictionary to store temporary directories for each file upload
+    # This persists across handler instances since Tornado recreates handlers per request
+    # Key: filename, Value: dict with temp_dir, total_chunks, received_chunks, logged_upload
+    _temp_dirs: Dict[str, Dict[str, Any]] = {}
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+    @tornado.web.authenticated
+    def post(self) -> None:
+        """Handle file upload with multipart form data."""
+        try:
+            # Validate request has file
+            if not self._validate_file_upload():
+                return
+
+            uploaded_file = self.request.files["file"][0]
+            filename = uploaded_file["filename"]
+            file_data = uploaded_file["body"]
+
+            # Get notebook directory from request
+            notebook_dir = self.get_argument("notebook_dir", ".")
+
+            # Check if this is a chunked upload
+            chunk_number = self.get_argument("chunk_number", None)
+            total_chunks = self.get_argument("total_chunks", None)
+
+            if chunk_number and total_chunks:
+                self._handle_chunked_upload(
+                    filename, file_data, chunk_number, total_chunks, notebook_dir
+                )
+            else:
+                # Log the file upload attempt for regular (non-chunked) uploads
+                file_extension = filename.split(".")[-1].lower()
+                log_file_upload_attempt(filename, file_extension, False, 0)
+                self._handle_regular_upload(filename, file_data, notebook_dir)
+
+            self.finish()
+
+        except Exception as e:
+            self._handle_error(str(e))
+
+    def _validate_file_upload(self) -> bool:
+        """Validate that a file was uploaded in the request."""
+        if "file" not in self.request.files:
+            self._handle_error("No file uploaded", status_code=400)
+            return False
+        return True
+
+    def _handle_chunked_upload(
+        self,
+        filename: str,
+        file_data: bytes,
+        chunk_number: str,
+        total_chunks: str,
+        notebook_dir: str,
+    ) -> None:
+        """Handle chunked file upload."""
+        chunk_num = int(chunk_number)
+        total_chunks_num = int(total_chunks)
+
+        # Log the file upload attempt only for the first chunk
+        if chunk_num == 1:
+            file_extension = filename.split(".")[-1].lower()
+            log_file_upload_attempt(filename, file_extension, True, total_chunks_num)
+
+        # Save chunk to temporary file
+        self._save_chunk(filename, file_data, chunk_num, total_chunks_num)
+
+        # Check if all chunks are received and reconstruct if complete
+        if self._are_all_chunks_received(filename, total_chunks_num):
+            self._reconstruct_file(filename, total_chunks_num, notebook_dir)
+            self._send_chunk_complete_response(filename, notebook_dir)
+        else:
+            self._send_chunk_received_response(chunk_num, total_chunks_num)
+
+    def _handle_regular_upload(
+        self, filename: str, file_data: bytes, notebook_dir: str
+    ) -> None:
+        """Handle regular (non-chunked) file upload."""
+        # Check image file size limit before saving
+        _check_image_size_limit(file_data, filename)
+
+        file_path = os.path.join(notebook_dir, filename)
+        with open(file_path, "wb") as f:
+            f.write(file_data)
+
+        self.write({"success": True, "filename": filename, "path": file_path})
+
+    def _save_chunk(
+        self, filename: str, file_data: bytes, chunk_number: int, total_chunks: int
+    ) -> None:
+        """Save a chunk to a temporary file."""
+        # Initialize temporary directory for this file if it doesn't exist
+        if filename not in self._temp_dirs:
+            temp_dir = tempfile.mkdtemp(prefix=f"mito_upload_{filename}_")
+            self._temp_dirs[filename] = {
+                "temp_dir": temp_dir,
+                "total_chunks": total_chunks,
+                "received_chunks": set(),
+            }
+
+        # Save the chunk to the temporary directory
+        chunk_filename = os.path.join(
+            self._temp_dirs[filename]["temp_dir"], f"chunk_{chunk_number}"
+        )
+        with open(chunk_filename, "wb") as f:
+            f.write(file_data)
+
+        # Mark this chunk as received
+        self._temp_dirs[filename]["received_chunks"].add(chunk_number)
+
+    def _are_all_chunks_received(self, filename: str, total_chunks: int) -> bool:
+        """Check if all chunks for a file have been received."""
+        if filename not in self._temp_dirs:
+            return False
+
+        received_chunks = self._temp_dirs[filename]["received_chunks"]
+        is_complete = len(received_chunks) == total_chunks
+        return is_complete
+
+    def _reconstruct_file(
+        self, filename: str, total_chunks: int, notebook_dir: str
+    ) -> None:
+        """Reconstruct the final file from all chunks and clean up temporary directory."""
+
+        if filename not in self._temp_dirs:
+            raise ValueError(f"No temporary directory found for file: {filename}")
+
+        temp_dir = self._temp_dirs[filename]["temp_dir"]
+        file_path = os.path.join(notebook_dir, filename)
+
+        try:
+            # First, read all chunks to check total file size for images
+            all_file_data = b""
+            for i in range(1, total_chunks + 1):
+                chunk_filename = os.path.join(temp_dir, f"chunk_{i}")
+                with open(chunk_filename, "rb") as chunk_file:
+                    chunk_data = chunk_file.read()
+                    all_file_data += chunk_data
+
+            # Check image file size limit before saving
+            _check_image_size_limit(all_file_data, filename)
+
+            # Write the complete file
+            with open(file_path, "wb") as final_file:
+                final_file.write(all_file_data)
+        finally:
+            # Clean up the temporary directory
+            self._cleanup_temp_dir(filename)
+
+    def _cleanup_temp_dir(self, filename: str) -> None:
+        """Clean up the temporary directory for a file."""
+        if filename in self._temp_dirs:
+            temp_dir = self._temp_dirs[filename]["temp_dir"]
+            try:
+                import shutil
+
+                shutil.rmtree(temp_dir)
+            except Exception as e:
+                # Log the error but don't fail the upload
+                print(
+                    f"Warning: Failed to clean up temporary directory {temp_dir}: {e}"
+                )
+            finally:
+                # Remove from tracking dictionary
+                del self._temp_dirs[filename]
+
+    def _send_chunk_complete_response(self, filename: str, notebook_dir: str) -> None:
+        """Send response indicating all chunks have been processed and file is complete."""
+        file_path = os.path.join(notebook_dir, filename)
+        self.write(
+            {
+                "success": True,
+                "filename": filename,
+                "path": file_path,
+                "chunk_complete": True,
+            }
+        )
+
+    def _send_chunk_received_response(
+        self, chunk_number: int, total_chunks: int
+    ) -> None:
+        """Send response indicating a chunk was received but file is not yet complete."""
+        self.write(
+            {
+                "success": True,
+                "chunk_received": True,
+                "chunk_number": chunk_number,
+                "total_chunks": total_chunks,
+            }
+        )
+
+    def _handle_error(self, error_message: str, status_code: int = 500) -> None:
+        """Handle errors and send appropriate error response."""
+        log_file_upload_failure(error_message)
+        self.set_status(status_code)
+        self.write({"error": error_message})
+        self.finish()
+
+    def on_finish(self) -> None:
+        """Clean up any remaining temporary directories when the handler is finished."""
+        super().on_finish()
+        # Note: We don't clean up here anymore since we want to preserve state across requests
+        # The cleanup happens when the file is fully reconstructed

mito_ai/file_uploads/urls.py

@@ -0,0 +1,21 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Any
+from jupyter_server.utils import url_path_join
+from mito_ai.file_uploads.handlers import FileUploadHandler
+
+
+def get_file_uploads_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all file uploads related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    return [
+        (url_path_join(BASE_URL, "upload"), FileUploadHandler, {}),
+    ]

mito_ai/gemini_client.py

@@ -5,10 +5,9 @@ from typing import Any, Callable, Dict, List, Optional, Union, Tuple
 from google import genai
 from google.genai import types
 from google.genai.types import GenerateContentConfig, Part, Content, GenerateContentResponse
-from mito_ai.completions.models import CompletionItem, CompletionReply, CompletionStreamChunk, MessageType, ResponseFormatInfo
+from mito_ai.completions.models import CompletionError, CompletionItem, CompletionReply, CompletionStreamChunk, MessageType, ResponseFormatInfo
 from mito_ai.utils.gemini_utils import get_gemini_completion_from_mito_server, stream_gemini_completion_from_mito_server, get_gemini_completion_function_params
-
-GEMINI_FAST_MODEL = "gemini-2.0-flash-lite"
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
 
 def extract_and_parse_gemini_json_response(response: GenerateContentResponse) -> Optional[str]:
     """
@@ -100,65 +99,62 @@ def get_gemini_system_prompt_and_messages(messages: List[Dict[str, Any]]) -> Tup
 
 
 class GeminiClient:
-    def __init__(self, api_key: Optional[str], model: str):
+    def __init__(self, api_key: Optional[str]):
         self.api_key = api_key
-        self.model = model
         if api_key:
             self.client = genai.Client(api_key=api_key)
 
     async def request_completions(
         self,
         messages: List[Dict[str, Any]],
+        model: str,
         response_format_info: Optional[ResponseFormatInfo] = None,
         message_type: MessageType = MessageType.CHAT
     ) -> str:
-        try:
-            # Extract system instructions and contents
-            system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
-
-            # Get provider data for Gemini completion
-            provider_data = get_gemini_completion_function_params(
-                model=self.model if response_format_info else GEMINI_FAST_MODEL,
-                contents=contents,
+        # Extract system instructions and contents
+        system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
+
+        # Get provider data for Gemini completion
+        provider_data = get_gemini_completion_function_params(
+            model=model,
+            contents=contents,
+            message_type=message_type,
+            response_format_info=response_format_info
+        )
+
+        if self.api_key:
+            # Generate content using the Gemini client
+            response_config = GenerateContentConfig(
+                system_instruction=system_instructions,
+                response_mime_type=provider_data.get("config", {}).get("response_mime_type"),
+                response_schema=provider_data.get("config", {}).get("response_schema")
+            )
+            response = self.client.models.generate_content(
+                model=provider_data["model"],
+                contents=contents,  # type: ignore
+                config=response_config
+            )
+            
+            result = extract_and_parse_gemini_json_response(response)
+            
+            if not result:
+                return "No response received from Gemini API"
+
+            return result
+        else:
+            # Fallback to Mito server for completion
+            return await get_gemini_completion_from_mito_server(
+                model=provider_data["model"],
+                contents=messages, # Use the extracted contents instead of converted messages to avoid serialization issues
                 message_type=message_type,
-                response_format_info=response_format_info
+                config=provider_data.get("config", None),
+                response_format_info=response_format_info,
             )
 
-            if self.api_key:
-                # Generate content using the Gemini client
-                response_config = GenerateContentConfig(
-                    system_instruction=system_instructions,
-                    response_mime_type=provider_data.get("config", {}).get("response_mime_type"),
-                    response_schema=provider_data.get("config", {}).get("response_schema")
-                )
-                response = self.client.models.generate_content(
-                    model=provider_data["model"],
-                    contents=contents,  # type: ignore
-                    config=response_config
-                )
-                
-                result = extract_and_parse_gemini_json_response(response)
-                
-                if not result:
-                    return "No response received from Gemini API"
-
-                return result
-            else:
-                # Fallback to Mito server for completion
-                return await get_gemini_completion_from_mito_server(
-                    model=provider_data["model"],
-                    contents=messages, # Use the extracted contents instead of converted messages to avoid serialization issues
-                    message_type=message_type,
-                    config=provider_data.get("config", None),
-                    response_format_info=response_format_info,
-                )
-
-        except Exception as e:
-            return f"Error generating content: {str(e)}"
-
     async def stream_completions(
             self,
             messages: List[Dict[str, Any]],
+            model: str,
             message_id: str,
             reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
             message_type: MessageType = MessageType.CHAT
@@ -169,7 +165,7 @@ class GeminiClient:
             system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
             if self.api_key:
                 for chunk in self.client.models.generate_content_stream(
-                        model=self.model,
+                        model=model,
                         contents=contents,  # type: ignore
                         config=GenerateContentConfig(
                             system_instruction=system_instructions
@@ -208,7 +204,7 @@ class GeminiClient:
                 return accumulated_response
             else:
                 async for chunk_text in stream_gemini_completion_from_mito_server(
-                        model=self.model,
+                        model=model,
                         contents=messages,  # Use the extracted contents instead of converted messages to avoid serialization issues
                         message_type=message_type,
                         message_id=message_id,