Update main.py

main.py

@@ -9,11 +9,10 @@ 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
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
 # Configure logging
@@ -39,13 +38,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,7 +183,7 @@ 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
 
@@ -203,19 +195,9 @@ class Blackbox:
         proxy: Optional[str] = None,
         websearch: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponse], None]:
+    ) -> AsyncGenerator[Union[str, Dict[str, Any]], 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)
 
@@ -234,7 +216,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,18 +277,6 @@ 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(
@@ -316,14 +286,42 @@ class Blackbox:
                     proxy=proxy
                 ) as response_api_chat:
                     response_api_chat.raise_for_status()
-                    # 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
+                    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, "alt": "Generated Image"}
+                        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:
@@ -336,79 +334,6 @@ class Blackbox:
             except Exception as 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:
-                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):
-        self.model = model
-        self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
-        super().__init__(self.message)
-
-async def cleanup_rate_limit_stores():
-    """
-    Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
-    """
-    while True:
-        current_time = time.time()
-        ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
-        for ip in ips_to_delete:
-            del rate_limit_store[ip]
-            logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
-        await asyncio.sleep(CLEANUP_INTERVAL)
-
-async def rate_limiter_per_ip(request: Request):
-    """
-    Rate limiter that enforces a limit based on the client's IP address.
-    """
-    client_ip = request.client.host
-    current_time = time.time()
-
-    # Initialize or update the count and timestamp
-    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:
-            logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
-            raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
-        rate_limit_store[client_ip]["count"] += 1
-
-async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
-    """
-    Dependency to extract and validate the API key from the Authorization header.
-    """
-    client_ip = request.client.host
-    if authorization is None or not authorization.startswith('Bearer '):
-        logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
-        raise HTTPException(status_code=401, detail='Invalid authorization header format')
-    api_key = authorization[7:]
-    if api_key not in API_KEYS:
-        logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
-        raise HTTPException(status_code=401, detail='Invalid API key')
-    return api_key
-
 # FastAPI app setup
 app = FastAPI()
 
@@ -449,7 +374,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
@@ -459,17 +383,6 @@ class ChatRequest(BaseModel):
     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
-
 @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)):
     client_ip = req.client.host
@@ -485,63 +398,35 @@ 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_content = await 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
         )
 
-        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()),
-                        })
-                        # 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:
-            response_content = await collect_response_content(generator)
-            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())
-                },
-            }
+        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())
+            },
+        }
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))
@@ -582,7 +467,3 @@ async def http_exception_handler(request: Request, exc: HTTPException):
             }
         },
     )
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)