Spaces:
Paused
Paused
File size: 5,085 Bytes
791d9c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """API routes - OpenAI compatible endpoints"""
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from typing import List
import base64
import re
import json
from ..core.auth import verify_api_key_header
from ..core.models import ChatCompletionRequest
from ..services.generation_handler import GenerationHandler, MODEL_CONFIG
router = APIRouter()
# Dependency injection will be set up in main.py
generation_handler: GenerationHandler = None
def set_generation_handler(handler: GenerationHandler):
"""Set generation handler instance"""
global generation_handler
generation_handler = handler
@router.get("/v1/models")
async def list_models(api_key: str = Depends(verify_api_key_header)):
"""List available models"""
models = []
for model_id, config in MODEL_CONFIG.items():
description = f"{config['type'].capitalize()} generation"
if config['type'] == 'image':
description += f" - {config['model_name']}"
else:
description += f" - {config['model_key']}"
models.append({
"id": model_id,
"object": "model",
"owned_by": "flow2api",
"description": description
})
return {
"object": "list",
"data": models
}
@router.post("/v1/chat/completions")
async def create_chat_completion(
request: ChatCompletionRequest,
api_key: str = Depends(verify_api_key_header)
):
"""Create chat completion (unified endpoint for image and video generation)"""
try:
# Extract prompt from messages
if not request.messages:
raise HTTPException(status_code=400, detail="Messages cannot be empty")
last_message = request.messages[-1]
content = last_message.content
# Handle both string and array format (OpenAI multimodal)
prompt = ""
images: List[bytes] = []
if isinstance(content, str):
# Simple text format
prompt = content
elif isinstance(content, list):
# Multimodal format
for item in content:
if item.get("type") == "text":
prompt = item.get("text", "")
elif item.get("type") == "image_url":
# Extract base64 image
image_url = item.get("image_url", {}).get("url", "")
if image_url.startswith("data:image"):
# Parse base64
match = re.search(r"base64,(.+)", image_url)
if match:
image_base64 = match.group(1)
image_bytes = base64.b64decode(image_base64)
images.append(image_bytes)
# Fallback to deprecated image parameter
if request.image and not images:
if request.image.startswith("data:image"):
match = re.search(r"base64,(.+)", request.image)
if match:
image_base64 = match.group(1)
image_bytes = base64.b64decode(image_base64)
images.append(image_bytes)
if not prompt:
raise HTTPException(status_code=400, detail="Prompt cannot be empty")
# Call generation handler
if request.stream:
# Streaming response
async def generate():
async for chunk in generation_handler.handle_generation(
model=request.model,
prompt=prompt,
images=images if images else None,
stream=True
):
yield chunk
# Send [DONE] signal
yield "data: [DONE]\n\n"
return StreamingResponse(
generate(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no"
}
)
else:
# Non-streaming response
result = None
async for chunk in generation_handler.handle_generation(
model=request.model,
prompt=prompt,
images=images if images else None,
stream=False
):
result = chunk
if result:
# Parse the result JSON string
try:
result_json = json.loads(result)
return JSONResponse(content=result_json)
except json.JSONDecodeError:
# If not JSON, return as-is
return JSONResponse(content={"result": result})
else:
raise HTTPException(status_code=500, detail="Generation failed: No response from handler")
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
|