Update app.py

app.py

@@ -1,56 +1,68 @@
-from fastapi import FastAPI, HTTPException, Depends, Header, Request
+from fastapi import FastAPI, HTTPException, Depends, Header
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from typing import List, Optional, Literal
 import json
 import g4f
 from g4f.Provider import OpenaiAccount, RetryProvider
-from g4f.models import ModelUtils
 
 app = FastAPI()
 
-# Complete list of available models from G4F
-models = [
-    # OpenAI models
-    "gpt-4", "gpt-4-turbo", "gpt-4o", "gpt-3.5-turbo",
+# Organized model list (updated with all G4F available models)
+MODELS = {
+    # OpenAI
+    "openai": [
+        "gpt-4o", "gpt-4o-mini", "gpt-4", "gpt-4-turbo", 
+        "gpt-3.5-turbo", "gpt-3.5-turbo-16k"
+    ],
     
-    # Anthropic models
-    "claude-3-opus", "claude-3-sonnet", "claude-3-haiku", "claude-2.1",
+    # Anthropic
+    "anthropic": [
+        "claude-3-opus", "claude-3-sonnet", "claude-3-haiku",
+        "claude-3.5", "claude-3.7-sonnet", "claude-2.1"
+    ],
     
-    # Google models
-    "gemini-pro", "gemini-1.5-pro", "gemini-1.5-flash",
+    # Google
+    "google": [
+        "gemini-pro", "gemini-1.5-pro", "gemini-1.5-flash",
+        "gemini-2.5-pro-exp-03-25"
+    ],
     
-    # Meta models
-    "llama-2-70b", "llama-2-13b", "llama-2-7b", "llama-3-70b", "llama-3-8b",
+    # Meta
+    "meta": [
+        "llama-3-70b", "llama-3-8b", "llama-3.1-405b",
+        "llama-2-70b", "llama-2-13b", "llama-2-7b"
+    ],
     
-    # Other providers
-    "mistral-7b", "mixtral-8x7b", "command-r-plus", "cohere-command-r",
-    "deepseek-chat", "deepseek-coder", "code-llama-34b", "code-llama-70b",
+    # XAI (Grok)
+    "xai": [
+        "grok-1", "grok-1.5", "grok-2", "grok-3"
+    ],
     
-    # Specialized models
-    "grok-1", "grok-1.5", "grok-2", "o1", "o3-mini", "flux", "flux-pro"
-]
-
-# Configure G4F backend
-class CustomBackend(g4f.Provider.BackendApi):
-    working = True
-    ssl = False
-    url = "https://ahe.hopto.org"
-    headers = {"Authorization": "Basic Z2dnOmc0Zl8="}
+    # Other
+    "other": [
+        "o1", "o3-mini", "mistral-7b", "mixtral-8x7b",
+        "command-r-plus", "deepseek-chat", "code-llama-70b"
+    ],
     
-    @classmethod
-    async def create_async_generator(
-        cls,
-        model: str,
-        messages: g4f.typing.Messages,
-        **kwargs
-    ) -> g4f.typing.AsyncResult:
-        if model in OpenaiAccount.get_models():
-            kwargs["provider"] = OpenaiAccount
-        async for chunk in super().create_async_generator(model, messages, **kwargs):
-            yield chunk
+    # Image Models
+    "image": [
+        "dall-e-3", "stable-diffusion-xl",
+        "flux", "flux-pro", "playground-v2.5"
+    ]
+}
+
+# Flattened list for API endpoint
+ALL_MODELS = [
+    *MODELS["openai"],
+    *MODELS["anthropic"],
+    *MODELS["google"],
+    *MODELS["meta"],
+    *MODELS["xai"],
+    *MODELS["other"]
+]
 
-# Pydantic models
+# Pydantic Models
 class Message(BaseModel):
     role: Literal["system", "user", "assistant"]
     content: str
@@ -61,53 +73,36 @@ class ChatRequest(BaseModel):
     temperature: Optional[float] = 0.7
     max_tokens: Optional[int] = None
     top_p: Optional[float] = 0.9
-    streaming: bool = True
-
-class ChatResponse(BaseModel):
-    role: str = "assistant"
-    content: str
-    model: Optional[str] = None
+    stream: Optional[bool] = True
 
 class ModelListResponse(BaseModel):
-    models: List[str]
-
-# API Key Verification
-async def verify_api_key(x_api_key: str = Header(...)):
-    if x_api_key != "fb207532285886a5568298b4b4e61124":
-        raise HTTPException(status_code=403, detail="Invalid API key")
+    openai: List[str]
+    anthropic: List[str]
+    google: List[str]
+    meta: List[str]
+    xai: List[str]
+    other: List[str]
+    image: List[str]
 
-@app.get("/v1/models", response_model=ModelListResponse, tags=["Models"])
+# API Endpoints
+@app.get("/v1/models", response_model=ModelListResponse)
 async def get_models():
-    """Get list of all available models"""
-    return ModelListResponse(models=models)
+    """Get all available models categorized by provider"""
+    return ModelListResponse(**MODELS)
 
-@app.post("/v1/chat/completions", response_model=Optional[ChatResponse], tags=["Chat"])
-async def chat_completion(
-    request: ChatRequest,
-    api_key: str = Depends(verify_api_key)
-):
-    """
-    Handle chat completion requests with streaming support.
-    
-    Args:
-        request: ChatRequest containing model, messages and parameters
-        api_key: Verified API key
-        
-    Returns:
-        Either a StreamingResponse or direct ChatResponse
-    """
-    # Validate model
-    if request.model not in models:
+@app.post("/v1/chat/completions")
+async def chat_completion(request: ChatRequest):
+    """Handle chat completion requests"""
+    if request.model not in ALL_MODELS:
         raise HTTPException(
             status_code=400,
-            detail=f"Invalid model. Available models: {', '.join(models)}"
+            detail=f"Invalid model. Available: {ALL_MODELS}"
         )
-
-    # Prepare messages
+    
     messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
     
     try:
-        if request.streaming:
+        if request.stream:
             async def stream_generator():
                 response = await g4f.ChatCompletion.create_async(
                     model=request.model,
@@ -115,45 +110,53 @@ async def chat_completion(
                     temperature=request.temperature,
                     top_p=request.top_p,
                     max_tokens=request.max_tokens,
-                    provider=RetryProvider([CustomBackend])
+                    provider=RetryProvider([g4f.Provider.BackendApi])
                 )
                 
                 async for chunk in response:
-                    if isinstance(chunk, dict):
-                        yield f"data: {json.dumps(chunk)}\n\n"
-                    else:
-                        yield f"data: {json.dumps({'content': str(chunk)})}\n\n"
+                    yield f"data: {json.dumps({'content': str(chunk)})}\n\n"
                 yield "data: [DONE]\n\n"
-
-            return StreamingResponse(
-                stream_generator(),
-                media_type="text/event-stream"
-            )
+            
+            return StreamingResponse(stream_generator(), media_type="text/event-stream")
         else:
             response = await g4f.ChatCompletion.create_async(
                 model=request.model,
                 messages=messages,
                 temperature=request.temperature,
                 top_p=request.top_p,
-                max_tokens=request.max_tokens,
-                provider=RetryProvider([CustomBackend])
+                provider=RetryProvider([g4f.Provider.BackendApi])
             )
-            
-            if isinstance(response, str):
-                return ChatResponse(content=response, model=request.model)
-            elif isinstance(response, dict):
-                return ChatResponse(
-                    content=response.get("choices", [{}])[0].get("message", {}).get("content", ""),
-                    model=request.model
-                )
-            return ChatResponse(content=str(response), model=request.model
-
+            return {"content": str(response)}
+    
     except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Image Generation Endpoint
+@app.post("/v1/images/generations")
+async def generate_image(prompt: str, model: str = "dall-e-3"):
+    if model not in MODELS["image"]:
         raise HTTPException(
-            status_code=500,
-            detail=f"Error processing request: {str(e)}"
+            status_code=400,
+            detail=f"Invalid image model. Available: {MODELS['image']}"
         )
+    
+    try:
+        if model in ["flux", "flux-pro"]:
+            image_data = g4f.ImageGeneration.create(
+                prompt=prompt,
+                model=model,
+                provider=g4f.Provider.BackendApi
+            )
+            return {"url": f"data:image/png;base64,{image_data.decode('utf-8')}"}
+        else:
+            # Implementation for other image providers
+            raise HTTPException(
+                status_code=501,
+                detail=f"{model} implementation pending"
+            )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)