Update main.py

main.py

@@ -13,7 +13,7 @@ from datetime import datetime
 
 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
 
 # Configure logging
@@ -39,6 +39,13 @@ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
+class ImageResponse:
+    def __init__(self, images: Union[str, List[str]], alt: str = "Generated Image"):
+        if isinstance(images, str):
+            images = [images]
+        self.images = images
+        self.alt = alt
+
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
@@ -184,20 +191,34 @@ class Blackbox:
 
     @staticmethod
     def clean_response(text: str) -> str:
-        pattern = r'^\$\@\$v=undefined-rv1\$\@\$'
+        pattern = r'^\$\@\$v=undefined-rv1\$\@'
         cleaned_text = re.sub(pattern, '', text)
         return cleaned_text
 
     @classmethod
-    async def generate_response(
+    async def create_async_generator(
         cls,
         model: str,
         messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
         websearch: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Union[str, Dict[str, Any]], None]:
+    ) -> AsyncGenerator[Union[str, ImageResponse], 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, ImageResponse]: 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()
@@ -302,9 +323,10 @@ class Blackbox:
                         match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
                         if match:
                             image_url = match.group(1)
-                            yield {"type": "image", "url": image_url}
+                            image_response = ImageResponse(images=image_url, alt="Generated Image")
+                            yield image_response
                         else:
-                            yield {"type": "text", "content": cleaned_response}
+                            yield cleaned_response
                     else:
                         if websearch:
                             match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
@@ -325,12 +347,12 @@ class Blackbox:
                             else:
                                 final_response = cleaned_response
                         else:
-                            if '$~~~$' in cleaned_response:
-                                final_response = cleaned_response.split('$~~~$')[0].strip()
+                            if '$~~~' in cleaned_response:
+                                final_response = cleaned_response.split('$~~~')[0].strip()
                             else:
                                 final_response = cleaned_response
 
-                        yield {"type": "text", "content": final_response}
+                        yield final_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -339,10 +361,11 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                yield {"type": "error", "content": error_text}
+                yield error_text
             except Exception as e:
-                yield {"type": "error", "content": f"Unexpected error during /api/chat request: {str(e)}"}
+                yield f"Unexpected error during /api/chat request: {str(e)}"
 
+            # Not clear what to do with this second request; keeping it for compatibility
             chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
 
             try:
@@ -362,9 +385,9 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                yield {"type": "error", "content": error_text}
+                yield error_text
             except Exception as e:
-                yield {"type": "error", "content": f"Unexpected error during /chat/{chat_id} request: {str(e)}"}
+                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -455,6 +478,7 @@ class Message(BaseModel):
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
+    stream: Optional[bool] = False  # Added for streaming support
     temperature: Optional[float] = 1.0
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
@@ -463,8 +487,17 @@ class ChatRequest(BaseModel):
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
-    stream: Optional[bool] = False
-    websearch: Optional[bool] = False
+
+# Helper function to collect responses from async generator
+async def collect_response_content(generator: AsyncGenerator[Union[str, ImageResponse], None]) -> str:
+    response_content = ''
+    async for chunk in generator:
+        if isinstance(chunk, str):
+            response_content += chunk
+        elif isinstance(chunk, ImageResponse):
+            # Handle image response if needed
+            response_content += f"[Image: {chunk.alt}] {', '.join(chunk.images)}\n"
+    return response_content
 
 @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)):
@@ -481,40 +514,36 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             raise HTTPException(status_code=400, detail="Requested model is not available.")
 
         # Process the request with actual message content, but don't log it
-        response_generator = Blackbox.generate_response(
+        generator = Blackbox.create_async_generator(
             model=request.model,
             messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
             temperature=request.temperature,
-            max_tokens=request.max_tokens,
-            websearch=request.websearch
+            max_tokens=request.max_tokens
         )
 
         if request.stream:
             async def stream_response():
-                async for chunk in response_generator:
-                    if chunk["type"] == "text":
-                        yield f"data: {json.dumps({'choices': [{'delta': {'content': chunk['content']}}]})}\n\n"
-                    elif chunk["type"] == "image":
-                        yield f"data: {json.dumps({'choices': [{'delta': {'image': chunk['url']}}]})}\n\n"
-                    elif chunk["type"] == "error":
-                        yield f"data: {json.dumps({'error': chunk['content']})}\n\n"
-                yield "data: [DONE]\n\n"
-            
+                async for chunk in generator:
+                    if isinstance(chunk, str):
+                        data = json.dumps({
+                            "choices": [{
+                                "delta": {"content": chunk},
+                                "index": 0,
+                                "finish_reason": None
+                            }],
+                            "model": request.model,
+                            "id": f"chatcmpl-{uuid.uuid4()}",
+                            "object": "chat.completion.chunk",
+                            "created": int(datetime.now().timestamp()),
+                        })
+                        yield f"data: {data}\n\n"
+            logger.info(f"Streaming response enabled for API key: {api_key} | IP: {client_ip}")
+
             return StreamingResponse(stream_response(), media_type="text/event-stream")
         else:
-            full_response = ""
-            image_url = None
-            async for chunk in response_generator:
-                if chunk["type"] == "text":
-                    full_response += chunk["content"]
-                elif chunk["type"] == "image":
-                    image_url = chunk["url"]
-                elif chunk["type"] == "error":
-                    raise HTTPException(status_code=500, detail=chunk["content"])
-
+            response_content = await collect_response_content(generator)
             logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
-            
-            response = {
+            return {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
                 "created": int(datetime.now().timestamp()),
@@ -524,23 +553,17 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         "index": 0,
                         "message": {
                             "role": "assistant",
-                            "content": full_response
+                            "content": response_content
                         },
                         "finish_reason": "stop"
                     }
                 ],
                 "usage": {
                     "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
-                    "completion_tokens": len(full_response.split()),
-                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(full_response.split())
+                    "completion_tokens": len(response_content.split()),
+                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
                 },
             }
-
-            if image_url:
-                response["choices"][0]["message"]["image"] = image_url
-
-            return response
-
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))
@@ -584,4 +607,4 @@ async def http_exception_handler(request: Request, exc: HTTPException):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)