Update main.py

main.py

@@ -12,7 +12,7 @@ from typing import List, Dict, Any, Optional, Union, AsyncGenerator
 
 from aiohttp import ClientSession, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 from datetime import datetime
 
@@ -317,6 +317,199 @@ class Blackbox:
             except Exception as e:
                 return f"Unexpected error during /api/chat request: {str(e)}"
 
+    @classmethod
+    async def create_async_generator(
+        cls,
+        model: str,
+        messages: List[Dict[str, str]],
+        proxy: Optional[str] = None,
+        websearch: bool = False,
+        **kwargs
+    ) -> AsyncGenerator[Union[str, ImageResponseModel], None]:
+        """
+        Creates an asynchronous generator for streaming responses from Blackbox AI.
+
+        Parameters:
+            model (str): Model to use for generating responses.
+            messages (List[Dict[str, str]]): Message history.
+            proxy (Optional[str]): Proxy URL, if needed.
+            websearch (bool): Enables or disables web search mode.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Union[str, ImageResponseModel]: Segments of the generated response or ImageResponse objects.
+        """
+        model = cls.get_model(model)
+
+        chat_id = cls.generate_random_string()
+        next_action = cls.generate_next_action()
+        next_router_state_tree = cls.generate_next_router_state_tree()
+
+        agent_mode = cls.agentMode.get(model, {})
+        trending_agent_mode = cls.trendingAgentMode.get(model, {})
+
+        prefix = cls.model_prefixes.get(model, "")
+        
+        formatted_prompt = ""
+        for message in messages:
+            role = message.get('role', '').capitalize()
+            content = message.get('content', '')
+            if role and content:
+                formatted_prompt += f"{role}: {content}\n"
+        
+        if prefix:
+            formatted_prompt = f"{prefix} {formatted_prompt}".strip()
+
+        referer_path = cls.model_referers.get(model, f"/?model={model}")
+        referer_url = f"{cls.url}{referer_path}"
+
+        common_headers = {
+            'accept': '*/*',
+            'accept-language': 'en-US,en;q=0.9',
+            'cache-control': 'no-cache',
+            'origin': cls.url,
+            'pragma': 'no-cache',
+            'priority': 'u=1, i',
+            'sec-ch-ua': '"Chromium";v="129", "Not=A?Brand";v="8"',
+            'sec-ch-ua-mobile': '?0',
+            'sec-ch-ua-platform': '"Linux"',
+            'sec-fetch-dest': 'empty',
+            'sec-fetch-mode': 'cors',
+            'sec-fetch-site': 'same-origin',
+            'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) '
+                          'AppleWebKit/537.36 (KHTML, like Gecko) '
+                          'Chrome/129.0.0.0 Safari/537.36'
+        }
+
+        headers_api_chat = {
+            'Content-Type': 'application/json',
+            'Referer': referer_url
+        }
+        headers_api_chat_combined = {**common_headers, **headers_api_chat}
+
+        payload_api_chat = {
+            "messages": [
+                {
+                    "id": chat_id,
+                    "content": formatted_prompt,
+                    "role": "user"
+                }
+            ],
+            "id": chat_id,
+            "previewToken": None,
+            "userId": None,
+            "codeModelMode": True,
+            "agentMode": agent_mode,
+            "trendingAgentMode": trending_agent_mode,
+            "isMicMode": False,
+            "userSystemPrompt": None,
+            "maxTokens": 1024,
+            "playgroundTopP": 0.9,
+            "playgroundTemperature": 0.5,
+            "isChromeExt": False,
+            "githubToken": None,
+            "clickedAnswer2": False,
+            "clickedAnswer3": False,
+            "clickedForceWebSearch": False,
+            "visitFromDelta": False,
+            "mobileClient": False,
+            "webSearchMode": websearch,
+            "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:
+            try:
+                async with session.post(
+                    cls.api_endpoint,
+                    headers=headers_api_chat_combined,
+                    json=payload_api_chat,
+                    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)
+                            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."
+                            else:
+                                final_response = cleaned_response
+                        else:
+                            if '$~~~$' in cleaned_response:
+                                final_response = cleaned_response.split('$~~~$')[0].strip()
+                            else:
+                                final_response = cleaned_response
+
+                        yield final_response
+            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 /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):
@@ -414,6 +607,7 @@ class ChatRequest(BaseModel):
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
+    stream: Optional[bool] = False  # Added stream parameter
 
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
 async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
@@ -421,7 +615,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
     # Redact user messages only for logging purposes
     redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
 
-    logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
+    logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages} | Stream: {request.stream}")
 
     try:
         # Validate that the requested model is available
@@ -429,36 +623,35 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
             raise HTTPException(status_code=400, detail="Requested model is not available.")
 
-        # Process the request with actual message content, but don't log it
-        response_content = await Blackbox.generate_response(
-            model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-            temperature=request.temperature,
-            max_tokens=request.max_tokens
-        )
-
-        logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
-        return {
-            "id": f"chatcmpl-{uuid.uuid4()}",
-            "object": "chat.completion",
-            "created": int(datetime.now().timestamp()),
-            "model": request.model,
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": response_content
-                    },
-                    "finish_reason": "stop"
-                }
-            ],
-            "usage": {
-                "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
-                "completion_tokens": len(response_content.split()),
-                "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
-            },
-        }
+        if request.stream:
+            # Streaming response
+            async def content_generator():
+                async for chunk in Blackbox.create_async_generator(
+                    model=request.model,
+                    messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
+                    proxy=None,  # Add proxy if needed
+                    websearch=False  # Modify if websearch is needed
+                ):
+                    if isinstance(chunk, ImageResponseModel):
+                        yield f"Image URL: {chunk.images}\n"
+                    else:
+                        yield f"{chunk}\n"
+
+            logger.info(f"Initiating streaming response for API key: {api_key} | IP: {client_ip}")
+            return StreamingResponse(content_generator(), media_type='text/plain')
+        else:
+            # Non-streaming response
+            response_content = await Blackbox.generate_response(
+                model=request.model,
+                messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
+                temperature=request.temperature,
+                max_tokens=request.max_tokens
+            )
+
+            logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
+            return {
+                "content": response_content
+            }
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))