main.py

import os
import httpx
import json
import time
import asyncio
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 from .env file
load_dotenv()

# --- Configuration ---
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
if not REPLICATE_API_TOKEN:
    raise ValueError("REPLICATE_API_TOKEN environment variable not set.")

# --- FastAPI App Initialization ---
app = FastAPI(
    title="Replicate to OpenAI Compatibility Layer",
    version="2.0.0 (Native Streaming & Context Fixed)",
)

# --- Pydantic Models for OpenAI Compatibility ---

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 ToolFunction(BaseModel):
    name: str
    description: str
    parameters: Dict[str, Any]

class Tool(BaseModel):
    type: Literal["function"]
    function: ToolFunction

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
    tools: Optional[List[Tool]] = None
    tool_choice: Optional[Union[str, Dict]] = None

# --- Replicate Model Mapping ---
SUPPORTED_MODELS = {
    "llama3-8b-instruct": "meta/meta-llama-3-8b-instruct",
    "claude-4.5-haiku": "anthropic/claude-4.5-haiku"
}


# --- Helper Functions ---

def prepare_replicate_input(request: OpenAIChatCompletionRequest) -> Dict[str, Any]:
    """
    Prepares the input payload for Replicate's chat models.
    This now correctly passes the messages array for context.
    """
    # Convert Pydantic message objects to a list of dictionaries
    messages_for_replicate = [msg.dict() for msg in request.messages]

    payload = {
        "messages": messages_for_replicate
    }

    # Add other compatible parameters
    if request.max_tokens is not None:
        payload["max_new_tokens"] = request.max_tokens
    if request.temperature is not None:
        payload["temperature"] = request.temperature
    if request.top_p is not None:
        payload["top_p"] = request.top_p

    # Vision support: Find image URL in the last user message if present
    last_user_message = next((m for m in reversed(request.messages) if m.role == 'user'), None)
    if last_user_message and isinstance(last_user_message.content, list):
        for item in last_user_message.content:
            if item.get("type") == "image_url":
                payload["image"] = item.get("image_url", {}).get("url")
                # Reformat messages to be a simple prompt string for vision models if needed,
                # as some might not support the `messages` format with images.
                # For Claude Haiku, a prompt string is more reliable with images.
                if "claude" in request.model:
                    text_prompts = [item.get('text', '') for item in last_user_message.content if item.get('type') == 'text']
                    payload["prompt"] = " ".join(text_prompts)
                    del payload["messages"]
                break
                
    return payload

async def stream_replicate_native_sse(model_id: str, payload: dict):
    """
    Connects to Replicate's native SSE stream for true token-by-token streaming.
    """
    url = f"https://api.replicate.com/v1/models/{model_id}/predictions"
    headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json"}
    
    async with httpx.AsyncClient(timeout=300) as client:
        # 1. Create the prediction to get the stream URL
        try:
            # Add stream=True to the outer payload for Replicate
            response = await client.post(url, headers=headers, json={"input": payload, "stream": True})
            response.raise_for_status()
            prediction = response.json()
            stream_url = prediction.get("urls", {}).get("stream")

            if not stream_url:
                error_detail = prediction.get("detail", "Failed to get stream URL.")
                yield json.dumps({"error": {"message": error_detail}})
                return
        except httpx.HTTPStatusError as e:
            yield json.dumps({"error": {"message": e.response.text}})
            return
        
        # 2. Connect to the SSE stream and yield OpenAI-compatible chunks
        try:
            async with client.stream("GET", stream_url, headers={"Accept": "text/event-stream"}) as sse:
                sse.raise_for_status()
                current_event = ""
                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":
                            chunk = {
                                "id": prediction["id"], "object": "chat.completion.chunk", "created": int(time.time()), "model": model_id,
                                "choices": [{"index": 0, "delta": {"content": json.loads(data)}, "finish_reason": None}]
                            }
                            yield json.dumps(chunk)
                        elif current_event == "done":
                            break # Exit loop when done event is received
        except Exception as e:
            yield json.dumps({"error": {"message": f"Streaming error: {str(e)}"}})

    # 3. Send the final DONE chunk
    done_chunk = {
        "id": prediction["id"], "object": "chat.completion.chunk", "created": int(time.time()), "model": model_id,
        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
    }
    yield json.dumps(done_chunk)
    yield "[DONE]"


# --- API Endpoints ---

@app.get("/v1/models", response_model=ModelList)
async def list_models():
    return ModelList(data=[ModelCard(id=model_name) for model_name in SUPPORTED_MODELS.keys()])

@app.post("/v1/chat/completions")
async def create_chat_completion(request: OpenAIChatCompletionRequest):
    model_key = request.model
    if model_key not in SUPPORTED_MODELS:
        raise HTTPException(status_code=404, detail=f"Model not found. Supported models: {list(SUPPORTED_MODELS.keys())}")
    
    replicate_model_id = SUPPORTED_MODELS[model_key]
    replicate_input = prepare_replicate_input(request)

    if request.stream:
        return EventSourceResponse(stream_replicate_native_sse(replicate_model_id, replicate_input))
    
    # Synchronous request
    url = f"https://api.replicate.com/v1/models/{replicate_model_id}/predictions"
    headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json", "Prefer": "wait=120"}

    async with httpx.AsyncClient(timeout=150) as client:
        try:
            response = await client.post(url, headers=headers, json={"input": replicate_input})
            response.raise_for_status()
            prediction = response.json()
            
            output = "".join(prediction.get("output", []))
            
            return JSONResponse(content={
                "id": prediction["id"], "object": "chat.completion", "created": int(time.time()), "model": model_key,
                "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=e.response.text)