Update app.py

app.py

@@ -1,11 +1,10 @@
-from fastapi import FastAPI, HTTPException, Depends, Header
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 from typing import List, Optional, Literal
 import json
 import g4f
 from g4f.Provider import Blackbox, RetryProvider
-from g4f.models import ModelUtils
 
 app = FastAPI()
 
@@ -13,18 +12,33 @@ app = FastAPI()
 g4f.Provider.Blackbox.url = "https://www.blackbox.ai/api/chat"
 g4f.Provider.Blackbox.working = True
 
-# Available Models
+# All available models from Blackbox provider
 TEXT_MODELS = [
-    "blackboxai", "blackboxai-pro", "gpt-4o-mini", "deepseek-chat",
-    "deepseek-v3", "deepseek-r1", "gpt-4o", "o1", "o3-mini",
-    "claude-3.7-sonnet", "llama-3.3-70b", "mixtral-small-24b", "qwq-32b"
+    # Blackbox models
+    "blackbox", "blackbox-pro", "blackbox-70b", "blackbox-180b",
+    
+    # OpenAI compatible
+    "gpt-4", "gpt-4-turbo", "gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo",
+    
+    # Anthropic
+    "claude-3-opus", "claude-3-sonnet", "claude-3-haiku", "claude-3.5", "claude-3.7-sonnet",
+    
+    # Meta
+    "llama-3-70b", "llama-3-8b", "llama-3.3-70b", "llama-2-70b",
+    
+    # DeepSeek
+    "deepseek-chat", "deepseek-v3", "deepseek-r1", "deepseek-coder",
+    
+    # Other
+    "o1", "o3-mini", "mixtral-8x7b", "mixtral-small-24b", "qwq-32b",
+    "command-r-plus", "code-llama-70b", "gemini-pro", "gemini-1.5-flash"
 ]
 
 IMAGE_MODELS = [
-    "flux", "flux-pro", "dall-e-3", "stable-diffusion-xl"
+    "flux", "flux-pro", "dall-e-3", "stable-diffusion-xl", "playground-v2.5",
+    "kandinsky-3", "deepfloyd-if", "sdxl-turbo"
 ]
 
-# Pydantic Models
 class Message(BaseModel):
     role: Literal["system", "user", "assistant"]
     content: str
@@ -40,10 +54,11 @@ class ImageRequest(BaseModel):
     model: str
     prompt: str
     size: Optional[str] = "1024x1024"
+    quality: Optional[Literal["standard", "hd"]] = "standard"
 
 @app.get("/v1/models")
 async def get_models():
-    """Return available models"""
+    """Return all available models"""
     return {
         "text_models": TEXT_MODELS,
         "image_models": IMAGE_MODELS
@@ -51,7 +66,7 @@ async def get_models():
 
 @app.post("/v1/chat/completions")
 async def chat_completion(request: ChatRequest):
-    """Handle Blackbox and other text generation requests"""
+    """Handle text generation with Blackbox and other models"""
     if request.model not in TEXT_MODELS:
         raise HTTPException(
             status_code=400,
@@ -63,21 +78,12 @@ async def chat_completion(request: ChatRequest):
     try:
         if request.stream:
             async def stream_generator():
-                # Special handling for Blackbox models
-                if request.model in ["blackboxai", "blackboxai-pro"]:
-                    provider = Blackbox
-                else:
-                    provider = RetryProvider([
-                        g4f.Provider.Blackbox,
-                        g4f.Provider.DeepSeek,
-                        g4f.Provider.OpenaiChat
-                    ])
-                
                 response = await g4f.ChatCompletion.create_async(
                     model=request.model,
                     messages=messages,
-                    provider=provider,
+                    provider=RetryProvider([Blackbox]),
                     temperature=request.temperature,
+                    max_tokens=request.max_tokens,
                     stream=True
                 )
                 
@@ -85,20 +91,20 @@ async def chat_completion(request: ChatRequest):
                     if isinstance(chunk, str):
                         yield f"data: {json.dumps({'content': chunk})}\n\n"
                     elif hasattr(chunk, 'choices'):
-                        yield f"data: {json.dumps({'content': chunk.choices[0].delta.content})}\n\n"
+                        content = chunk.choices[0].delta.get('content', '')
+                        yield f"data: {json.dumps({'content': content})}\n\n"
                 yield "data: [DONE]\n\n"
 
             return StreamingResponse(stream_generator(), media_type="text/event-stream")
         
         else:
-            # Non-streaming response
             response = await g4f.ChatCompletion.create_async(
                 model=request.model,
                 messages=messages,
-                provider=Blackbox if request.model in ["blackboxai", "blackboxai-pro"] else None,
-                temperature=request.temperature
+                provider=RetryProvider([Blackbox]),
+                temperature=request.temperature,
+                max_tokens=request.max_tokens
             )
-            
             return {"content": str(response)}
 
     except Exception as e:
@@ -106,7 +112,7 @@ async def chat_completion(request: ChatRequest):
 
 @app.post("/v1/images/generations")
 async def generate_image(request: ImageRequest):
-    """Handle Flux and other image generation"""
+    """Handle image generation with Flux and other models"""
     if request.model not in IMAGE_MODELS:
         raise HTTPException(
             status_code=400,
@@ -118,10 +124,12 @@ async def generate_image(request: ImageRequest):
             image_data = g4f.ImageGeneration.create(
                 prompt=request.prompt,
                 model=request.model,
-                provider=g4f.Provider.Blackbox
+                provider=Blackbox,
+                size=request.size
             )
             return JSONResponse({
-                "url": f"data:image/png;base64,{image_data.decode('utf-8')}"
+                "url": f"data:image/png;base64,{image_data.decode('utf-8')}",
+                "model": request.model
             })
         else:
             # Implementation for other image providers
@@ -134,4 +142,4 @@ async def generate_image(request: ImageRequest):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=7860)