Update main.py

main.py

@@ -39,13 +39,6 @@ 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"
@@ -191,34 +184,20 @@ 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 create_async_generator(
+    async def generate_response(
         cls,
         model: str,
         messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
         websearch: bool = False,
         **kwargs
-    ) -> 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.
-        """
+    ) -> AsyncGenerator[Union[str, Dict[str, Any]], None]:
         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()
@@ -323,10 +302,9 @@ class Blackbox:
                         match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
                         if match:
                             image_url = match.group(1)
-                            image_response = ImageResponse(images=image_url, alt="Generated Image")
-                            yield image_response
+                            yield {"type": "image", "url": image_url}
                         else:
-                            yield cleaned_response
+                            yield {"type": "text", "content": cleaned_response}
                     else:
                         if websearch:
                             match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
@@ -347,12 +325,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 final_response
+                        yield {"type": "text", "content": final_response}
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -361,11 +339,10 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                yield error_text
+                yield {"type": "error", "content": error_text}
             except Exception as e:
-                yield f"Unexpected error during /api/chat request: {str(e)}"
+                yield {"type": "error", "content": 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:
@@ -385,9 +362,9 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                yield error_text
+                yield {"type": "error", "content": error_text}
             except Exception as e:
-                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
+                yield {"type": "error", "content": f"Unexpected error during /chat/{chat_id} request: {str(e)}"}
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -478,7 +455,6 @@ 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
@@ -487,17 +463,8 @@ class ChatRequest(BaseModel):
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
-
-# 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
+    stream: Optional[bool] = False
+    websearch: Optional[bool] = False
 
 @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)):
@@ -514,36 +481,40 @@ 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
-        generator = Blackbox.create_async_generator(
+        response_generator = 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
+            max_tokens=request.max_tokens,
+            websearch=request.websearch
         )
 
         if request.stream:
             async def stream_response():
-                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}")
-
+                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"
+            
             return StreamingResponse(stream_response(), media_type="text/event-stream")
         else:
-            response_content = await collect_response_content(generator)
+            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"])
+
             logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
-            return {
+            
+            response = {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
                 "created": int(datetime.now().timestamp()),
@@ -553,17 +524,23 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         "index": 0,
                         "message": {
                             "role": "assistant",
-                            "content": response_content
+                            "content": full_response
                         },
                         "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())
+                    "completion_tokens": len(full_response.split()),
+                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(full_response.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))
@@ -607,4 +584,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)