Update main.py

main.py

@@ -10,7 +10,7 @@ import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, Union, AsyncGenerator
 
-from aiohttp import ClientSession, ClientResponseError
+from aiohttp import ClientSession, ClientResponseError, ClientTimeout
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
@@ -315,7 +315,7 @@ class Blackbox:
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        async with ClientSession(headers=common_headers) as session:
+        async with ClientSession(headers=common_headers, timeout=ClientTimeout(total=60)) as session:
             try:
                 async with session.post(
                     cls.api_endpoint,
@@ -324,9 +324,25 @@ class Blackbox:
                     proxy=proxy
                 ) as response_api_chat:
                     response_api_chat.raise_for_status()
-                    text = await response_api_chat.text()
-                    cleaned_response = cls.clean_response(text)
-                    return cleaned_response
+                    # Instead of reading the entire response, iterate over chunks
+                    async for data in response_api_chat.content.iter_chunked(1024):
+                        decoded_data = data.decode('utf-8')
+                        cleaned_response = cls.clean_response(decoded_data)
+                        if model in cls.image_models:
+                            match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
+                            if match:
+                                image_url = match.group(1)
+                                image_response = ImageResponseModel(images=image_url, alt="Generated Image")
+                                return image_response.dict()
+                            else:
+                                return cleaned_response
+                        else:
+                            if '$~~~$' in cleaned_response:
+                                final_response = cleaned_response.split('$~~~$')[0].strip()
+                            else:
+                                final_response = cleaned_response
+
+                            return final_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -347,7 +363,7 @@ class Blackbox:
         proxy: Optional[str] = None,
         websearch: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponseModel], None]:
+    ) -> AsyncGenerator[str, None]:
         """
         Creates an asynchronous generator for streaming responses from Blackbox AI.
 
@@ -359,7 +375,7 @@ class Blackbox:
             **kwargs: Additional keyword arguments.
 
         Yields:
-            Union[str, ImageResponseModel]: Segments of the generated response or ImageResponse objects.
+            str: Segments of the generated response.
         """
         model = cls.get_model(model)
 
@@ -439,19 +455,7 @@ class Blackbox:
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        headers_chat = {
-            'Accept': 'text/x-component',
-            'Content-Type': 'text/plain;charset=UTF-8',
-            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
-            'next-action': next_action,
-            'next-router-state-tree': next_router_state_tree,
-            'next-url': '/'
-        }
-        headers_chat_combined = {**common_headers, **headers_chat}
-
-        data_chat = '[]'
-
-        async with ClientSession(headers=common_headers) as session:
+        async with ClientSession(headers=common_headers, timeout=ClientTimeout(total=60)) as session:
             try:
                 async with session.post(
                     cls.api_endpoint,
@@ -460,43 +464,25 @@ class Blackbox:
                     proxy=proxy
                 ) as response_api_chat:
                     response_api_chat.raise_for_status()
-                    text = await response_api_chat.text()
-                    cleaned_response = cls.clean_response(text)
-
-                    if model in cls.image_models:
-                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
-                        if match:
-                            image_url = match.group(1)
-                            image_response = ImageResponseModel(images=image_url, alt="Generated Image")
-                            yield image_response
-                        else:
-                            yield cleaned_response
-                    else:
-                        if websearch:
-                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
+                    # Iterate over the response in chunks
+                    async for data in response_api_chat.content.iter_any():
+                        decoded_data = data.decode('utf-8')
+                        cleaned_response = cls.clean_response(decoded_data)
+                        if model in cls.image_models:
+                            match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
                             if match:
-                                source_part = match.group(1).strip()
-                                answer_part = cleaned_response[match.end():].strip()
-                                try:
-                                    sources = json.loads(source_part)
-                                    source_formatted = "**Source:**\n"
-                                    for item in sources:
-                                        title = item.get('title', 'No Title')
-                                        link = item.get('link', '#')
-                                        position = item.get('position', '')
-                                        source_formatted += f"{position}. [{title}]({link})\n"
-                                    final_response = f"{answer_part}\n\n{source_formatted}"
-                                except json.JSONDecodeError:
-                                    final_response = f"{answer_part}\n\nSource information is unavailable."
+                                image_url = match.group(1)
+                                image_response = ImageResponseModel(images=image_url, alt="Generated Image")
+                                yield f"Image URL: {image_response.images}\n"
                             else:
-                                final_response = cleaned_response
+                                yield cleaned_response
                         else:
                             if '$~~~$' in cleaned_response:
                                 final_response = cleaned_response.split('$~~~$')[0].strip()
                             else:
                                 final_response = cleaned_response
 
-                        yield final_response
+                            yield f"{final_response}\n"
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -509,29 +495,6 @@ class Blackbox:
             except Exception as e:
                 yield f"Unexpected error during /api/chat request: {str(e)}"
 
-            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
-
-            try:
-                async with session.post(
-                    chat_url,
-                    headers=headers_chat_combined,
-                    data=data_chat,
-                    proxy=proxy
-                ) as response_chat:
-                    response_chat.raise_for_status()
-                    pass
-            except ClientResponseError as e:
-                error_text = f"Error {e.status}: {e.message}"
-                try:
-                    error_response = await e.response.text()
-                    cleaned_error = cls.clean_response(error_response)
-                    error_text += f" - {cleaned_error}"
-                except Exception:
-                    pass
-                yield error_text
-            except Exception as e:
-                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
-
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
     def __init__(self, model: str):
@@ -654,15 +617,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     proxy=None,  # Add proxy if needed
                     websearch=False  # Modify if websearch is needed
                 ):
-                    if isinstance(chunk, ImageResponseModel):
-                        # Yield image URLs as plain text
-                        yield f"Image URL: {chunk.images}\n"
-                    else:
-                        # Ensure chunk is a string and yield as plain text
-                        if isinstance(chunk, str):
-                            yield f"{chunk}\n"
-                        else:
-                            yield f"{str(chunk)}\n"
+                    yield chunk
 
             logger.info(f"Initiating streaming response for API key: {api_key} | IP: {client_ip}")
             return StreamingResponse(content_generator(), media_type='text/plain')