Update main.py

main.py

@@ -1,34 +1,14 @@
-import os
-import uuid
-import logging
-import json
+from __future__ import annotations
+
 import re
 import random
 import string
-from datetime import datetime
-from typing import Any, Dict, List, Optional
-
-import httpx
-from fastapi import FastAPI, HTTPException, Depends
+import uuid
+from aiohttp import ClientSession
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from starlette.middleware.cors import CORSMiddleware
-from starlette.responses import StreamingResponse
-
-# Setup logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# FastAPI app setup
-app = FastAPI()
-
-# CORS middleware setup
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+from typing import List, Dict, Any, Optional
+from datetime import datetime
 
 # Mock implementations for ImageResponse and to_data_uri
 class ImageResponse:
@@ -40,7 +20,6 @@ def to_data_uri(image: Any) -> str:
     # Placeholder for actual image encoding
     return "data:image/png;base64,..."  # Replace with actual base64 data
 
-# Define models and providers
 class AsyncGeneratorProvider:
     pass
 
@@ -54,7 +33,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     supports_stream = True
     supports_system_message = True
     supports_message_history = True
-
+    
     default_model = 'blackbox'
     models = [
         'blackbox',
@@ -79,13 +58,13 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
         'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
     }
-
+    
     userSelectedModel = {
         "gpt-4o": "gpt-4o",
         "gemini-pro": "gemini-pro",
         'claude-sonnet-3.5': "claude-sonnet-3.5",
     }
-
+    
     model_aliases = {
         "gemini-flash": "gemini-1.5-flash",
         "flux": "ImageGenerationLV45LJp",
@@ -136,7 +115,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             if not messages[0]['content'].startswith(prefix):
                 messages[0]['content'] = f"{prefix} {messages[0]['content']}"
         
-        async with httpx.AsyncClient(headers=headers) as session:
+        async with ClientSession(headers=headers) as session:
             if image is not None:
                 messages[-1]["data"] = {
                     "fileText": image_name,
@@ -187,14 +166,16 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                     else:
                         raise Exception("Image URL not found in the response")
                 else:
-                    async for chunk in response.aiter_bytes():
+                    async for chunk in response.content.iter_any():
                         if chunk:
                             decoded_chunk = chunk.decode()
                             decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
                             if decoded_chunk.strip():
                                 yield decoded_chunk
 
-# Message and chat request models
+# FastAPI app setup
+app = FastAPI()
+
 class Message(BaseModel):
     role: str
     content: str
@@ -202,68 +183,44 @@ class Message(BaseModel):
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
-    stream: Optional[bool] = False
-
-# Verify app secret (placeholder)
-async def verify_app_secret(app_secret: str):
-    if app_secret != os.getenv("APP_SECRET"):
-        raise HTTPException(status_code=403, detail="Forbidden")
-
-@app.post("/v1/chat/completions")
-async def chat_completions(
-    request: ChatRequest, 
-    app_secret: str = Depends(lambda: os.getenv("APP_SECRET"))
-):
-    # Validate model
-    if request.model not in Blackbox.models:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(Blackbox.models)}",
-        )
-
-    # Generate a UUID for the conversation
-    conversation_id = str(uuid.uuid4()).replace("-", "")
-
-    # Define headers
-    headers = {
-        "Authorization": f"Bearer {app_secret}",
-        "Content-Type": "application/json",
-        "uniqueid": conversation_id,
+    stream: Optional[bool] = False  # Add this for streaming
+
+def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion.chunk",
+        "created": int(datetime.now().timestamp()),
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "delta": {"content": content, "role": "assistant"},
+                "finish_reason": finish_reason,
+            }
+        ],
+        "usage": None,
     }
 
-    json_data = {
-        "attachments": [],
-        "conversationId": conversation_id,
-        "prompt": "\n".join(
-            [f"{msg.role}: {msg.content}" for msg in request.messages]
-        ),
-    }
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatRequest):
+    messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
 
-    async def generate():
-        async with httpx.AsyncClient() as client:
-            try:
-                async with client.stream('POST', Blackbox.api_endpoint, headers=headers, json=json_data) as response:
-                    response.raise_for_status()
-                    async for line in response.aiter_lines():
-                        if line and line != "[DONE]":
-                            content = json.loads(line)["data"]
-                            yield f"data: {json.dumps(content)}\n\n"
-                    yield "data: [DONE]\n\n"
-            except httpx.HTTPStatusError as e:
-                logger.error(f"HTTP error occurred: {e}")
-                raise HTTPException(status_code=e.response.status_code, detail=str(e))
-            except httpx.RequestError as e:
-                logger.error(f"An error occurred while requesting: {e}")
-                raise HTTPException(status_code=500, detail=str(e))
+    async_generator = Blackbox.create_async_generator(
+        model=request.model,
+        messages=messages
+    )
 
     if request.stream:
+        async def generate():
+            async for chunk in async_generator:
+                yield f"data: {json.dumps(create_response(chunk, request.model))}\n\n"
+            yield "data: [DONE]\n\n"
+
         return StreamingResponse(generate(), media_type="text/event-stream")
     else:
-        full_response = ""
-        async for chunk in generate():
-            if chunk.startswith("data: ") and not chunk[6:].startswith("[DONE]"):
-                data = json.loads(chunk[6:])
-                full_response += data.get("choices", [{}])[0].get("delta", {}).get("content", "")
+        response_content = ""
+        async for chunk in async_generator:
+            response_content += chunk if isinstance(chunk, str) else chunk.content  # Concatenate response
 
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
@@ -272,9 +229,12 @@ async def chat_completions(
             "model": request.model,
             "choices": [
                 {
-                    "index": 0,
-                    "message": {"role": "assistant", "content": full_response},
+                    "message": {
+                        "role": "assistant",
+                        "content": response_content
+                    },
                     "finish_reason": "stop",
+                    "index": 0
                 }
             ],
             "usage": None,