Update main.py

main.py

@@ -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()
@@ -213,7 +234,7 @@ class Blackbox:
             content = message.get('content', '')
             if role and content:
                 formatted_prompt += f"{role}: {content}\n"
-        
+
         if prefix:
             formatted_prompt = f"{prefix} {formatted_prompt}".strip()
 
@@ -295,42 +316,14 @@ 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)
-                            yield {"type": "image", "url": image_url}
-                        else:
-                            yield {"type": "text", "content": 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 {"type": "text", "content": final_response}
+                    # We update this part to stream the response incrementally
+                    # Instead of waiting for the full response, we read the response as it arrives
+                    async for chunk in response_api_chat.content.iter_chunked(1024):
+                        if not chunk:
+                            continue
+                        text = chunk.decode('utf-8', errors='ignore')
+                        cleaned_chunk = cls.clean_response(text)
+                        yield cleaned_chunk
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -339,10 +332,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 +356,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):
@@ -393,7 +387,7 @@ async def rate_limiter_per_ip(request: Request):
     current_time = time.time()
 
     # Initialize or update the count and timestamp
-    if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
+    if current_time - rate_limit_store[client_ip]["timest(由Dream 提供Free https://opus.gptuu.com)amp"] > RATE_LIMIT_WINDOW:
         rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
     else:
         if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
@@ -455,6 +449,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 +458,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 +485,43 @@ 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"
+                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()),
+                        })
+                        # Ensure that each chunk is sent immediately
+                        yield f"data: {data}\n\n"
+                    elif isinstance(chunk, ImageResponse):
+                        # Handle image responses here if needed
+                        pass  # For now, we skip image handling in streaming
+                # Send the termination message
                 yield "data: [DONE]\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 +531,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 +585,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)