mito-ai 0.1.35__py3-none-any.whl → 0.1.37__py3-none-any.whl

mito_ai/utils/mito_server_utils.py

@@ -1,12 +1,13 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from mito_ai.completions.models import MessageType
+import asyncio
+import json
+import time
+from typing import Any, Dict, Optional, Callable, Union, AsyncGenerator
+from mito_ai.completions.models import MessageType, CompletionReply, CompletionStreamChunk, CompletionItem
 from mito_ai.utils.server_limits import check_mito_server_quota, update_mito_server_quota
 from tornado.httpclient import HTTPResponse
-import time
-import json
-from typing import Any, Dict, Optional
 from mito_ai.constants import MITO_GEMINI_URL
 from mito_ai.utils.utils import _create_http_client
 
@@ -88,5 +89,147 @@ async def get_response_from_mito_server(
             pass
         
         http_client.close()
+
+
+async def stream_response_from_mito_server(
+    url: str,
+    headers: Dict[str, str],
+    data: Dict[str, Any],
+    timeout: int,
+    max_retries: int,
+    message_type: MessageType,
+    reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+    message_id: str,
+    chunk_processor: Optional[Callable[[str], str]] = None,
+    provider_name: str = "Mito Server",
+) -> AsyncGenerator[str, None]:
+    """
+    Stream responses from the Mito server.
+    
+    This is a unified streaming function that can be used by all providers (OpenAI, Anthropic, Gemini).
+    
+    Args:
+        url: The Mito server URL to stream from
+        headers: Request headers
+        data: Request data
+        timeout: Request timeout in seconds
+        max_retries: Maximum number of retries
+        message_type: The message type for quota tracking
+        provider_name: Name of the provider for error messages
+        reply_fn: Optional function to call with each chunk for streaming replies
+        message_id: The message ID to track the request
+        chunk_processor: Optional function to process chunks before yielding (e.g., for Gemini's special processing)
+        
+    Yields:
+        Chunks of text from the streaming response
+    """
+    # Check the mito server quota
+    check_mito_server_quota(message_type)
+    
+    # Create HTTP client with appropriate timeout settings
+    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
+    
+    # Set up streaming infrastructure
+    start_time = time.time()
+    chunk_queue: asyncio.Queue[str] = asyncio.Queue()
+    fetch_complete = False
+    
+    # Define a callback to process chunks and add them to the queue
+    def chunk_callback(chunk: bytes) -> None:
+        try:
+            chunk_str = chunk.decode('utf-8')
+            asyncio.create_task(chunk_queue.put(chunk_str))
+        except Exception as e:
+            print(f"Error processing {provider_name} streaming chunk: {str(e)}")
+    
+    # Execute the streaming request
+    fetch_future = None
+    try:
+        fetch_future = http_client.fetch(
+            url, 
+            method="POST", 
+            headers=headers, 
+            body=json.dumps(data), 
+            request_timeout=http_client_timeout,
+            streaming_callback=chunk_callback
+        )
+        
+        # Create a task to wait for the fetch to complete
+        async def wait_for_fetch() -> None:
+            try:
+                await fetch_future
+                nonlocal fetch_complete
+                fetch_complete = True
+                print(f"{provider_name} fetch completed")
+            except Exception as e:
+                print(f"Error in {provider_name} fetch: {str(e)}")
+                raise
+        
+        # Start the task to wait for fetch completion
+        fetch_task = asyncio.create_task(wait_for_fetch())
+        
+        # Yield chunks as they arrive
+        while not (fetch_complete and chunk_queue.empty()):
+            try:
+                # Wait for a chunk with a timeout to prevent deadlocks
+                chunk = await asyncio.wait_for(chunk_queue.get(), timeout=0.1)
+                
+                # Process chunk if processor is provided
+                processed_chunk = chunk
+                if chunk_processor:
+                    processed_chunk = chunk_processor(chunk)
+
+                if reply_fn is not None and message_id is not None:
+                    # Send the chunk directly to the frontend
+                    reply_fn(CompletionStreamChunk(
+                        parent_id=message_id,
+                        chunk=CompletionItem(
+                            content=processed_chunk,
+                            isIncomplete=True,
+                            token=message_id,
+                        ),
+                        done=False,
+                    ))
+                
+                yield chunk
+            except asyncio.TimeoutError:
+                # No chunk available within timeout, check if fetch is complete
+                if fetch_complete and chunk_queue.empty():
+                    break
+                
+                # Otherwise continue waiting for chunks
+                continue
+                
+        print(f"\n{provider_name} stream completed in {time.time() - start_time:.2f} seconds")
+        
+        if reply_fn is not None and message_id is not None:
+            # Send a final chunk to indicate completion
+            reply_fn(CompletionStreamChunk(
+                parent_id=message_id,
+                chunk=CompletionItem(
+                    content="",
+                    isIncomplete=False,
+                    token=message_id,
+                ),
+                done=True,
+            ))
+    except Exception as e:
+        print(f"\n{provider_name} stream failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
+        # If an exception occurred, ensure the fetch future is awaited to properly clean up
+        if fetch_future:
+            try:
+                await fetch_future
+            except Exception:
+                pass
+        raise
+    finally:
+        # Clean up resources
+        try:
+            # We always update the quota, even if there is an error
+            update_mito_server_quota(message_type)
+        except Exception as e:
+            pass
+        
+        http_client.close()
         
     

mito_ai/utils/open_ai_utils.py

@@ -10,7 +10,7 @@ import asyncio
 import json
 import time
 from typing import Any, Dict, List, Optional, Final, Union, AsyncGenerator, Tuple, Callable
-from mito_ai.utils.mito_server_utils import get_response_from_mito_server
+from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
 from mito_ai.utils.provider_utils import does_message_require_fast_model
 from tornado.httpclient import AsyncHTTPClient
 from openai.types.chat import ChatCompletionMessageParam
@@ -106,8 +106,8 @@ async def stream_ai_completion_from_mito_server(
     timeout: int,
     max_retries: int,
     message_type: MessageType,
-    reply_fn: Optional[Callable[[Union[CompletionReply, CompletionStreamChunk]], None]] = None,
-    message_id: Optional[str] = None,
+    reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+    message_id: str,
 ) -> AsyncGenerator[str, None]:
     """
     Stream AI completions from the Mito server.
@@ -127,7 +127,7 @@ async def stream_ai_completion_from_mito_server(
     Yields:
         Chunks of text from the streaming response
     """
-    # ===== STEP 1: Prepare request data and headers =====
+    # Prepare request data and headers
     data, headers = _prepare_request_data_and_headers(
         last_message_content, 
         ai_completion_data, 
@@ -136,109 +136,20 @@ async def stream_ai_completion_from_mito_server(
         message_type
     )
     
-    # ===== STEP 2: Create HTTP client with appropriate timeout settings =====
-    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
-    
-    # ===== STEP 3: Set up streaming infrastructure =====
-    start_time = time.time()
-    chunk_queue: asyncio.Queue[str] = asyncio.Queue()
-    fetch_complete = False
-    
-    # Define a callback to process chunks and add them to the queue
-    def chunk_callback(chunk: bytes) -> None:
-        try:
-            chunk_str = chunk.decode('utf-8')
-            asyncio.create_task(chunk_queue.put(chunk_str))
-        except Exception as e:
-            print(f"Error processing streaming chunk: {str(e)}")
-    
-    # ===== STEP 4: Execute the streaming request =====
-    fetch_future = None
-    try:
-        # Use fetch with streaming_callback to handle streaming responses.
-        # The streaming_callback is not sent as part of the POST request.
-        # It's a parameter for the Tornado AsyncHTTPClient.fetch() method that specifies
-        # how to handle incoming data chunks as they arrive from the server.
-        # When the server sends data in chunks, this callback function is called each time
-        # a new chunk arrives, allowing for immediate processing without waiting for the
-        # entire response to complete.
-        fetch_future = http_client.fetch(
-            MITO_OPENAI_URL, 
-            method="POST", 
-            headers=headers, 
-            body=json.dumps(data), 
-            request_timeout=http_client_timeout,
-            streaming_callback=chunk_callback
-        )
-        
-        # Create a task to wait for the fetch to complete
-        async def wait_for_fetch() -> None:
-            try:
-                await fetch_future
-                nonlocal fetch_complete
-                fetch_complete = True
-                print("Fetch completed")
-            except Exception as e:
-                print(f"Error in fetch: {str(e)}")
-                raise
-        
-        # Start the task to wait for fetch completion
-        fetch_task = asyncio.create_task(wait_for_fetch())
-        
-        # ===== STEP 5: Yield chunks as they arrive =====
-        while not (fetch_complete and chunk_queue.empty()):
-            try:
-                # Wait for a chunk with a timeout. By setting the timeout, we 1. prevent deadlocks
-                # which could happen if fetch_complete has not been set to true yet, and 2. it enables
-                # periodic checking if the queue has a new chunk.
-                chunk = await asyncio.wait_for(chunk_queue.get(), timeout=0.1)
-                
-                # If reply_fn is provided, send the chunk directly to the frontend
-                if reply_fn and message_id:
-                    reply_fn(CompletionStreamChunk(
-                        parent_id=message_id,
-                        chunk=CompletionItem(
-                            content=chunk,
-                            isIncomplete=True,
-                            token=message_id,
-                        ),
-                        done=False,
-                    ))
-                
-                yield chunk
-            except asyncio.TimeoutError:
-                # No chunk available within timeout, check if fetch is complete
-                if fetch_complete and chunk_queue.empty():
-                    break
-                
-                # Otherwise continue waiting for chunks
-                continue
-                
-        print(f"\nStream completed in {time.time() - start_time:.2f} seconds")
-        
-        # Send a final chunk to indicate completion if reply_fn is provided
-        if reply_fn and message_id:
-            reply_fn(CompletionStreamChunk(
-                parent_id=message_id,
-                chunk=CompletionItem(
-                    content="",
-                    isIncomplete=False,
-                    token=message_id,
-                ),
-                done=True,
-            ))
-    except Exception as e:
-        print(f"\nStream failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
-        # If an exception occurred, ensure the fetch future is awaited to properly clean up
-        if fetch_future:
-            try:
-                await fetch_future
-            except Exception:
-                pass
-        raise
-    finally:
-        # ===== STEP 6: Clean up resources =====
-        http_client.close()
+    # Use the unified streaming function
+    async for chunk in stream_response_from_mito_server(
+        url=MITO_OPENAI_URL,
+        headers=headers,
+        data=data,
+        timeout=timeout,
+        max_retries=max_retries,
+        message_type=message_type,
+        reply_fn=reply_fn,
+        message_id=message_id,
+        chunk_processor=None,
+        provider_name="OpenAI",
+    ):
+        yield chunk
 
 
 def get_open_ai_completion_function_params(