main.py

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="9.0.0 (Definitive Streaming Fix)")

# --- 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 ---
SUPPORTED_MODELS = {
    "llama3-8b-instruct": "meta/meta-llama-3-8b-instruct",
    "claude-4.5-haiku": "anthropic/claude-4.5-haiku",
    "claude-4.5-sonnet": "anthropic/claude-4.5-sonnet",
    "llava-13b": "yorickvp/llava-13b:e272157381e2a3bf12df3a8edd1f38d1dbd736bbb7437277c8b34175f8fce358"
}

# --- Core Logic ---
def prepare_replicate_input(request: OpenAIChatCompletionRequest) -> Dict[str, Any]:
    """
    Formats the input for Replicate's API, flattening the message history into a
    single 'prompt' string and handling images separately.
    """
    payload = {}
    prompt_parts = []
    system_prompt = None
    image_input = None
    
    for msg in request.messages:
        if msg.role == "system":
            system_prompt = str(msg.content)
        elif msg.role == "assistant":
            prompt_parts.append(f"Assistant: {msg.content}")
        elif msg.role == "user":
            user_text_content = ""
            if isinstance(msg.content, list):
                for item in msg.content:
                    if item.get("type") == "text":
                        user_text_content += item.get("text", "")
                    elif item.get("type") == "image_url":
                        image_url_data = item.get("image_url", {})
                        image_input = image_url_data.get("url")
            else:
                user_text_content = str(msg.content)
            prompt_parts.append(f"User: {user_text_content}")

    prompt_parts.append("Assistant:")
    payload["prompt"] = "\n\n".join(prompt_parts)
    
    if system_prompt:
        payload["system_prompt"] = system_prompt
    if image_input:
        payload["image"] = image_input
         
    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 with correct whitespace preservation."""
    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:
            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", "stream-unknown")
            if not stream_url:
                yield f"data: {json.dumps({'error': {'message': 'Model did not return a stream URL.'}})}\n\n"
                return
        except httpx.HTTPStatusError as e:
            error_details = e.response.text
            try:
                error_json = e.response.json()
                error_details = error_json.get("detail", error_details)
            except json.JSONDecodeError: pass
            yield f"data: {json.dumps({'error': {'message': f'Upstream Error: {error_details}', 'type': 'replicate_error'}})}\n\n"
            return

        try:
            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 not line: # Skip empty lines
                        continue
                    if line.startswith("event:"):
                        current_event = line[len("event:"):].strip()
                    elif line.startswith("data:"):
                        # FIXED: Preserve all whitespace including leading/trailing spaces
                        raw_data = line[5:]  # Remove "data:" prefix
                        
                        # Handle empty data lines (preserve them)
                        if not raw_data:
                            continue
                        
                        # Remove only the optional single space after data: if present
                        # This is per SSE spec and preserves actual content spaces
                        if raw_data.startswith(" "):
                            data_content = raw_data[1:]  # Remove the first space only
                        else:
                            data_content = raw_data
                        
                        if current_event == "output":
                            if not data_content:
                                continue
                                
                            content_token = ""
                            try:
                                # Handle JSON-encoded strings properly (including spaces)
                                content_token = json.loads(data_content)
                            except (json.JSONDecodeError, TypeError):
                                # Handle plain text tokens (preserve as-is)
                                content_token = data_content
                            
                            # Create chunk with exact format you specified
                            chunk = {
                                "choices": [{
                                    "delta": {"content": content_token},
                                    "finish_reason": None,
                                    "index": 0,
                                    "logprobs": None,
                                    "native_finish_reason": None
                                }],
                                "created": int(time.time()),
                                "id": f"gen-{int(time.time())}-{prediction_id[-12:]}",  # Format like your example
                                "model": replicate_model_id,
                                "object": "chat.completion.chunk",
                                "provider": "Anthropic" if "anthropic" in replicate_model_id else "Replicate"
                            }
                            yield f"data: {json.dumps(chunk)}\n\n"
                            
                        elif current_event == "done":
                            # Send usage chunk before done
                            usage_chunk = {
                                "choices": [{
                                    "delta": {},
                                    "finish_reason": None,
                                    "index": 0,
                                    "logprobs": None,
                                    "native_finish_reason": None
                                }],
                                "created": int(time.time()),
                                "id": f"gen-{int(time.time())}-{prediction_id[-12:]}",
                                "model": replicate_model_id,
                                "object": "chat.completion.chunk",
                                "provider": "Anthropic" if "anthropic" in replicate_model_id else "Replicate",
                                "usage": {
                                    "cache_discount": 0,
                                    "completion_tokens": 0,
                                    "completion_tokens_details": {"image_tokens": 0, "reasoning_tokens": 0},
                                    "cost": 0,
                                    "cost_details": {
                                        "upstream_inference_completions_cost": 0,
                                        "upstream_inference_cost": None,
                                        "upstream_inference_prompt_cost": 0
                                    },
                                    "input_tokens": 0,
                                    "is_byok": False,
                                    "prompt_tokens": 0,
                                    "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0},
                                    "total_tokens": 0
                                }
                            }
                            yield f"data: {json.dumps(usage_chunk)}\n\n"
                            
                            # Send final chunk with stop reason
                            final_chunk = {
                                "choices": [{
                                    "delta": {},
                                    "finish_reason": "stop",
                                    "index": 0,
                                    "logprobs": None,
                                    "native_finish_reason": "end_turn"
                                }],
                                "created": int(time.time()),
                                "id": f"gen-{int(time.time())}-{prediction_id[-12:]}",
                                "model": replicate_model_id,
                                "object": "chat.completion.chunk",
                                "provider": "Anthropic" if "anthropic" in replicate_model_id else "Replicate"
                            }
                            yield f"data: {json.dumps(final_chunk)}\n\n"
                            break
        except httpx.ReadTimeout:
            yield f"data: {json.dumps({'error': {'message': 'Stream timed out.', 'type': 'timeout_error'}})}\n\n"
            return

    # Send [DONE] event
    yield "data: [DONE]\n\n"

# --- 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(status_code=404, detail=f"Model not found. Available models: {list(SUPPORTED_MODELS.keys())}")
    
    replicate_input = prepare_replicate_input(request)

    if request.stream:
        return EventSourceResponse(stream_replicate_sse(SUPPORTED_MODELS[request.model], replicate_input), media_type="text/event-stream")

    # Non-streaming fallback
    url = f"https://api.replicate.com/v1/models/{SUPPORTED_MODELS[request.model]}/predictions"
    headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json", "Prefer": "wait=120"}
    async with httpx.AsyncClient() as client:
        try:
            resp = await client.post(url, headers=headers, json={"input": replicate_input}, timeout=130.0)
            resp.raise_for_status()
            pred = resp.json()
            output = "".join(pred.get("output", []))
            return {
                "id": pred.get("id"), "object": "chat.completion", "created": int(time.time()), "model": request.model,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": output}, "finish_reason": "stop"}],
                "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
            }
        except httpx.HTTPStatusError as e:
            raise HTTPException(status_code=e.response.status_code, detail=f"Error from Replicate API: {e.response.text}")