Update main.py

main.py

@@ -1,6 +1,10 @@
 import os
 import uuid
 import logging
+import json
+import re
+import random
+import string
 from datetime import datetime
 from typing import Any, Dict, List, Optional
 
@@ -8,7 +12,23 @@ import httpx
 from fastapi import FastAPI, HTTPException, Depends
 from pydantic import BaseModel
 from starlette.middleware.cors import CORSMiddleware
-from starlette.responses import StreamingResponse, Response
+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=["*"],
+)
 
 # Mock implementations for ImageResponse and to_data_uri
 class ImageResponse:
@@ -20,6 +40,7 @@ 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
 
@@ -33,7 +54,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     supports_stream = True
     supports_system_message = True
     supports_message_history = True
-    
+
     default_model = 'blackbox'
     models = [
         'blackbox',
@@ -58,13 +79,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",
@@ -115,7 +136,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             if not messages[0]['content'].startswith(prefix):
                 messages[0]['content'] = f"{prefix} {messages[0]['content']}"
         
-        async with ClientSession(headers=headers) as session:
+        async with httpx.AsyncClient(headers=headers) as session:
             if image is not None:
                 messages[-1]["data"] = {
                     "fileText": image_name,
@@ -166,16 +187,14 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                     else:
                         raise Exception("Image URL not found in the response")
                 else:
-                    async for chunk in response.content.iter_any():
+                    async for chunk in response.aiter_bytes():
                         if chunk:
                             decoded_chunk = chunk.decode()
                             decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
                             if decoded_chunk.strip():
                                 yield decoded_chunk
 
-# FastAPI app setup
-app = FastAPI()
-
+# Message and chat request models
 class Message(BaseModel):
     role: str
     content: str
@@ -183,19 +202,22 @@ class Message(BaseModel):
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
+    stream: Optional[bool] = False
 
-from fastapi.responses import Response
+# 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(verify_app_secret)
-):
+async def chat_completions(request: ChatRequest, app_secret: str = Depends(verify_app_secret)):
     logger.info(f"Received chat completion request for model: {request.model}")
 
-    if request.model not in [model['id'] for model in ALLOWED_MODELS]:
+    # 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(model['id'] for model in ALLOWED_MODELS)}",
+            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(Blackbox.models)}",
         )
 
     # Generate a UUID for the conversation
@@ -205,10 +227,7 @@ async def chat_completions(
         "attachments": [],
         "conversationId": conversation_id,
         "prompt": "\n".join(
-            [
-                f"{'User' if msg.role == 'user' else 'Assistant'}: {msg.content}"
-                for msg in request.messages
-            ]
+            [f"{msg.role}: {msg.content}" for msg in request.messages]
         ),
     }
 
@@ -217,13 +236,12 @@ async def chat_completions(
     async def generate():
         async with httpx.AsyncClient() as client:
             try:
-                async with client.stream('POST', f'{BASE_URL}/api/chat/gpt4o/chat', headers=headers, json=json_data, timeout=120.0) as response:
+                async with client.stream('POST', f'{Blackbox.api_endpoint}', headers=headers, json=json_data, timeout=120.0) 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(create_chat_completion_data(content['message'], request.model))}\n\n"
-                    yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
+                            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}")
@@ -233,16 +251,13 @@ async def chat_completions(
                 raise HTTPException(status_code=500, detail=str(e))
 
     if request.stream:
-        logger.info("Streaming response")
         return StreamingResponse(generate(), media_type="text/event-stream")
     else:
-        logger.info("Non-streaming response")
         full_response = ""
         async for chunk in generate():
             if chunk.startswith("data: ") and not chunk[6:].startswith("[DONE]"):
                 data = json.loads(chunk[6:])
-                if data["choices"][0]["delta"].get("content"):
-                    full_response += data["choices"][0]["delta"]["content"]
+                full_response += data.get("choices", [{}])[0].get("delta", {}).get("content", "")
 
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",