import os import httpx import json import time from fastapi import FastAPI, HTTPException from fastapi.responses import JSONResponse from pydantic import BaseModel, Field from typing import List, Dict, Any, Optional, Union, Literal from dotenv import load_dotenv from sse_starlette.sse import EventSourceResponse # Load environment variables load_dotenv() REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN") if not REPLICATE_API_TOKEN: raise ValueError("REPLICATE_API_TOKEN environment variable not set.") # FastAPI Init app = FastAPI(title="Replicate to OpenAI Compatibility Layer", version="4.0.0 (Docs Compliant)") # --- Pydantic Models --- class ModelCard(BaseModel): id: str; object: str = "model"; created: int = Field(default_factory=lambda: int(time.time())); owned_by: str = "replicate" class ModelList(BaseModel): object: str = "list"; data: List[ModelCard] = [] class ChatMessage(BaseModel): role: Literal["system", "user", "assistant", "tool"]; content: Union[str, List[Dict[str, Any]]] class OpenAIChatCompletionRequest(BaseModel): model: str; messages: List[ChatMessage]; temperature: Optional[float] = 0.7; top_p: Optional[float] = 1.0; max_tokens: Optional[int] = None; stream: Optional[bool] = False # --- Supported Models --- # Maps OpenAI-friendly names to Replicate model paths SUPPORTED_MODELS = { "llama3-8b-instruct": "meta/meta-llama-3-8b-instruct", "claude-4.5-haiku": "anthropic/claude-4.5-haiku" } # --- Core Logic --- def prepare_replicate_input(request: OpenAIChatCompletionRequest, replicate_model_id: str) -> Dict[str, Any]: """Formats the input specifically for the requested Replicate model.""" payload = {} # Claude on Replicate strictly requires a 'prompt' string, not 'messages' array. if "anthropic/claude" in replicate_model_id: prompt_parts = [] system_prompt = None for msg in request.messages: if msg.role == "system": # Extract system prompt if present system_prompt = str(msg.content) elif msg.role == "user": # Handle both simple string content and list content (for potential future vision support) content = msg.content if isinstance(content, list): text_parts = [item.get("text", "") for item in content if item.get("type") == "text"] content = " ".join(text_parts) prompt_parts.append(f"User: {content}") elif msg.role == "assistant": prompt_parts.append(f"Assistant: {msg.content}") # Standard Claude prompting convention prompt_parts.append("Assistant:") payload["prompt"] = "\n\n".join(prompt_parts) if system_prompt: payload["system_prompt"] = system_prompt # Llama 3 and others often support the 'messages' array natively. else: # Convert Pydantic models to pure dicts payload["prompt"] = [msg.dict() for msg in request.messages] # Map common OpenAI parameters to Replicate equivalents if request.max_tokens: payload["max_new_tokens"] = request.max_tokens if request.temperature: payload["temperature"] = request.temperature if request.top_p: payload["top_p"] = request.top_p return payload async def stream_replicate_sse(replicate_model_id: str, input_payload: dict): """Handles the full streaming lifecycle using standard Replicate endpoints.""" # 1. Start Prediction specifically at the named model endpoint url = f"https://api.replicate.com/v1/models/{replicate_model_id}/predictions" headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json"} async with httpx.AsyncClient(timeout=60.0) as client: try: # Explicitly request stream=True in the body, though often implicit response = await client.post(url, headers=headers, json={"input": input_payload, "stream": True}) response.raise_for_status() prediction = response.json() stream_url = prediction.get("urls", {}).get("stream") prediction_id = prediction.get("id") if not stream_url: yield json.dumps({"error": {"message": "Model did not return a stream URL."}}) return except httpx.HTTPStatusError as e: yield json.dumps({"error": {"message": e.response.text, "type": "upstream_error"}}) return # 2. Connect to the provided Stream URL async with client.stream("GET", stream_url, headers={"Accept": "text/event-stream"}, timeout=None) as sse: current_event = None async for line in sse.aiter_lines(): if line.startswith("event:"): current_event = line[len("event:"):].strip() elif line.startswith("data:"): data = line[len("data:"):].strip() if current_event == "output": # CRITICAL: Wrap in try/except to ignore empty keep-alive lines that crash standard parsers try: # Replicate sometimes sends raw strings, sometimes JSON. # For chat models, it's usually a raw string token. # We try to load as JSON first, if it fails, use raw data. try: content = json.loads(data) except json.JSONDecodeError: content = data if content: # Ensure we don't send empty chunks chunk = { "id": prediction_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": replicate_model_id, "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}] } yield json.dumps(chunk) except Exception: pass # Safely ignore malformed lines elif current_event == "done": break # 3. Send final [DONE] event yield json.dumps({"id": prediction_id, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}) yield "[DONE]" # --- Endpoints --- @app.get("/v1/models") async def list_models(): return ModelList(data=[ModelCard(id=k) for k in SUPPORTED_MODELS.keys()]) @app.post("/v1/chat/completions") async def create_chat_completion(request: OpenAIChatCompletionRequest): if request.model not in SUPPORTED_MODELS: raise HTTPException(404, f"Model not found. Available: {list(SUPPORTED_MODELS.keys())}") replicate_id = SUPPORTED_MODELS[request.model] replicate_input = prepare_replicate_input(request, replicate_id) if request.stream: return EventSourceResponse(stream_replicate_sse(replicate_id, replicate_input)) # Non-streaming fallback url = f"https://api.replicate.com/v1/models/{replicate_id}/predictions" headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json", "Prefer": "wait=60"} async with httpx.AsyncClient() as client: resp = await client.post(url, headers=headers, json={"input": replicate_input}) if resp.is_error: raise HTTPException(resp.status_code, resp.text) pred = resp.json() output = "".join(pred.get("output", [])) return {"id": pred["id"], "choices": [{"message": {"role": "assistant", "content": output}, "finish_reason": "stop"}]}