Update main.py

main.py

@@ -9,6 +9,7 @@ import asyncio
 import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, Union, AsyncGenerator
+from datetime import datetime
 
 from aiohttp import ClientSession, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
@@ -38,14 +39,6 @@ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
-# Define ImageResponse for handling image outputs
-class ImageResponseModel(BaseModel):
-    images: str
-    alt: str
-
-# Define Messages type for better type hinting
-Messages = List[Dict[str, Any]]
-
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
@@ -200,12 +193,10 @@ class Blackbox:
         cls,
         model: str,
         messages: List[Dict[str, str]],
-        temperature: Optional[float] = 1.0,
-        max_tokens: Optional[int] = 1024,
-        websearch: bool = False,
         proxy: Optional[str] = None,
+        websearch: bool = False,
         **kwargs
-    ) -> Union[str, ImageResponseModel]:
+    ) -> AsyncGenerator[Union[str, Dict[str, Any]], None]:
         model = cls.get_model(model)
         chat_id = cls.generate_random_string()
         next_action = cls.generate_next_action()
@@ -269,9 +260,9 @@ class Blackbox:
             "trendingAgentMode": trending_agent_mode,
             "isMicMode": False,
             "userSystemPrompt": None,
-            "maxTokens": max_tokens,
+            "maxTokens": 1024,
             "playgroundTopP": 0.9,
-            "playgroundTemperature": temperature,
+            "playgroundTemperature": 0.5,
             "isChromeExt": False,
             "githubToken": None,
             "clickedAnswer2": False,
@@ -283,6 +274,18 @@ 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:
             try:
                 async with session.post(
@@ -299,10 +302,9 @@ class Blackbox:
                         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
+                            yield {"type": "image", "url": image_url}
                         else:
-                            return cleaned_response
+                            yield {"type": "text", "content": cleaned_response}
                     else:
                         if websearch:
                             match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
@@ -328,7 +330,30 @@ class Blackbox:
                             else:
                                 final_response = cleaned_response
 
-                        return final_response
+                        yield {"type": "text", "content": 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 {"type": "error", "content": error_text}
+            except Exception as e:
+                yield {"type": "error", "content": 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:
@@ -337,9 +362,9 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                return error_text
+                yield {"type": "error", "content": error_text}
             except Exception as e:
-                return f"Unexpected error during /api/chat 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):
@@ -433,12 +458,13 @@ class ChatRequest(BaseModel):
     temperature: Optional[float] = 1.0
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
-    max_tokens: Optional[int] = 1024
+    max_tokens: Optional[int] = None
     presence_penalty: Optional[float] = 0.0
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
-    websearch: Optional[bool] = False  # Added websearch parameter
+    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)):
@@ -446,7 +472,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} | Websearch: {request.websearch}")
+    logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
 
     try:
         # Validate that the requested model is available
@@ -454,8 +480,8 @@ 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.")
 
-        # Generate response using the updated Blackbox class
-        response_content = await Blackbox.generate_response(
+        # Process the request with actual message content, but don't log it
+        response_generator = Blackbox.generate_response(
             model=request.model,
             messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
             temperature=request.temperature,
@@ -463,34 +489,32 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             websearch=request.websearch
         )
 
-        # Handle image responses
-        if isinstance(response_content, ImageResponseModel):
-            logger.info(f"Generated image 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.images,
-                            "alt": response_content.alt
-                        },
-                        "finish_reason": "stop"
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
-                    "completion_tokens": len(response_content.images.split()),
-                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.images.split())
-                },
-            }
+        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"
+            
+            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"])
+
             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()),
@@ -500,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))
@@ -523,14 +553,14 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 
 # Endpoint: GET /v1/models
 @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
-async def get_models(req: Request, api_key: str = Depends(get_api_key)):
+async def get_models(req: Request):
     client_ip = req.client.host
     logger.info(f"Fetching available models from IP: {client_ip}")
     return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
 
 # Endpoint: GET /v1/health
 @app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
-async def health_check(req: Request, api_key: str = Depends(get_api_key)):
+async def health_check(req: Request):
     client_ip = req.client.host
     logger.info(f"Health check requested from IP: {client_ip}")
     return {"status": "ok"}
@@ -554,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)