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="4.1.0 (Context Fixed)")

# --- 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"
    # You can add more models here
}

# --- Core Logic ---
def prepare_replicate_input(request: OpenAIChatCompletionRequest) -> Dict[str, Any]:
    """
    Formats the input for Replicate API, preserving the conversational context.
    """
    payload = {}
    
    # --- CONTEXT FIX START ---
    # Modern chat models on Replicate (like Llama 3 and Claude 4.5) expect
    # the 'messages' array directly, just like OpenAI.
    # We no longer need to flatten the conversation into a single prompt string.
    
    # Extract system prompt if it exists, as some models take it as a separate parameter.
    messages_for_payload = []
    system_prompt = None
    for msg in request.messages:
        if msg.role == "system":
            # Claude and some other models prefer a dedicated system_prompt field.
            system_prompt = str(msg.content)
        else:
            # Handle user/assistant roles. Convert Pydantic model to a standard dict.
            messages_for_payload.append(msg.dict())

    # The main input for conversation is the 'messages' array.
    payload["messages"] = messages_for_payload
    
    # Add system_prompt to the payload if it was found.
    if system_prompt:
         payload["system_prompt"] = system_prompt
         
    # --- CONTEXT FIX END ---

    # Map common OpenAI parameters to Replicate equivalents
    # Note: Replicate's parameter for max tokens is often 'max_new_tokens'
    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
    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:
            # Request a streaming prediction
            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 json.dumps({"error": {"message": "Model did not return a stream URL."}})
                 return

        except httpx.HTTPStatusError as e:
             error_details = e.response.text
             try:
                 # Try to parse the error for a cleaner message
                 error_json = e.response.json()
                 error_details = error_json.get("detail", error_details)
             except json.JSONDecodeError:
                 pass # Use raw text if not JSON
             yield json.dumps({"error": {"message": f"Upstream Error: {error_details}", "type": "replicate_error"}})
             return

        # 2. Connect to the provided Stream URL and process Server-Sent Events (SSE)
        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 line.startswith("event:"):
                        current_event = line[len("event:"):].strip()
                    elif line.startswith("data:"):
                        data = line[len("data:"):].strip()
                        
                        if current_event == "output":
                            # The 'output' event for chat models sends one token at a time as a plain string.
                            # We don't need to parse it as JSON.
                            if data: # 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": data}, "finish_reason": None}]
                                }
                                yield json.dumps(chunk)
                                
                        elif current_event == "done":
                            # The 'done' event signals the end of the stream.
                            break
        except httpx.ReadTimeout:
            # Handle cases where the stream times out
            yield json.dumps({"error": {"message": "Stream timed out.", "type": "timeout_error"}})
            return


    # 3. Send the final termination chunk in OpenAI format
    final_chunk = {
        "id": prediction_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": replicate_model_id,
        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
    }
    yield json.dumps(final_chunk)
    # Some clients (like curl) expect a final "[DONE]" message to close the connection.
    yield "[DONE]"

# --- Endpoints ---
@app.get("/v1/models")
async def list_models():
    """Lists the currently supported 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):
    """Handles chat completion requests, streaming or non-streaming."""
    if request.model not in SUPPORTED_MODELS:
        raise HTTPException(status_code=404, detail=f"Model not found. Available models: {list(SUPPORTED_MODELS.keys())}")
    
    replicate_id = SUPPORTED_MODELS[request.model]
    replicate_input = prepare_replicate_input(request)

    if request.stream:
        # Return a streaming response
        return EventSourceResponse(stream_replicate_sse(replicate_id, replicate_input), media_type="text/event-stream")

    # 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=120"} # Increased wait time
    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()
            # The output of chat models is typically a list of strings (tokens)
            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": { # Placeholder usage object
                    "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}")