main.py

import os
import httpx
import json
import time
from fastapi import FastAPI, Request, HTTPException, Header
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="1.0.0",
)

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

# /v1/models endpoint
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] = []

# /v1/chat/completions endpoint
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 ---
# We hardcode the models we want to expose.
SUPPORTED_MODELS = {
    "llama3-8b-instruct": "meta/meta-llama-3-8b-instruct",
    "claude-4.5-haiku": "anthropic/claude-4.5-haiku"
}


# --- Helper Functions ---

def format_tools_for_prompt(tools: List[Tool]) -> str:
    """Converts OpenAI tools to a string for the system prompt."""
    if not tools:
        return ""
    
    prompt = "You have access to the following tools. To use a tool, respond with a JSON object in the following format:\n"
    prompt += '{"type": "tool_call", "name": "tool_name", "arguments": {"arg_name": "value"}}\n\n'
    prompt += "Available tools:\n"
    for tool in tools:
        prompt += json.dumps(tool.function.dict(), indent=2) + "\n"
    return prompt

def prepare_replicate_input(request: OpenAIChatCompletionRequest) -> Dict[str, Any]:
    """Prepares the input payload for the Replicate API."""
    input_data = {}
    prompt_parts = []
    system_prompt = ""

    # Handle messages, separating system, user, assistant and vision content
    image_url = None
    for message in request.messages:
        if message.role == "system":
            system_prompt += message.content + "\n"
        elif message.role == "user":
            if isinstance(message.content, list): # Vision support
                for item in message.content:
                    if item.get("type") == "text":
                        prompt_parts.append(f"User: {item.get('text', '')}")
                    elif item.get("type") == "image_url":
                        image_url = item.get("image_url", {}).get("url")
            else:
                prompt_parts.append(f"User: {message.content}")
        elif message.role == "assistant":
            prompt_parts.append(f"Assistant: {message.content}")

    # Add tool instructions to system prompt
    if request.tools:
        tool_prompt = format_tools_for_prompt(request.tools)
        system_prompt += "\n" + tool_prompt

    input_data["prompt"] = "\n".join(prompt_parts)
    if system_prompt:
        input_data["system_prompt"] = system_prompt
    if image_url:
        input_data["image"] = image_url

    # Map other parameters
    if request.temperature is not None:
        input_data["temperature"] = request.temperature
    if request.top_p is not None:
        input_data["top_p"] = request.top_p
    if request.max_tokens is not None:
        # Replicate uses `max_new_tokens` or `max_tokens` depending on model
        input_data["max_new_tokens"] = request.max_tokens

    return input_data


async def stream_replicate_response(model_id: str, payload: dict):
    """Generator for streaming Replicate responses."""
    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 and get the stream URL
        payload["stream"] = True
        try:
            response = await client.post(url, headers=headers, json={"input": payload})
            response.raise_for_status()
            prediction = response.json()
            stream_url = prediction.get("urls", {}).get("stream")

            if not stream_url:
                yield f"data: {json.dumps({'error': 'Failed to get stream URL'})}\n\n"
                return
        except httpx.HTTPStatusError as e:
            yield f"data: {json.dumps({'error': str(e.response.text)})}\n\n"
            return

        # 2. Connect to the SSE stream
        try:
            async with client.stream("GET", stream_url, headers={"Accept": "text/event-stream"}) as sse:
                async for line in sse.aiter_lines():
                    if line.startswith("data:"):
                        event_data = line[len("data:"):].strip()
                        try:
                            data = json.loads(event_data)
                            # Format as OpenAI chunk
                            chunk = {
                                "id": prediction["id"],
                                "object": "chat.completion.chunk",
                                "created": int(time.time()),
                                "model": model_id,
                                "choices": [{
                                    "index": 0,
                                    "delta": {"content": data},
                                    "finish_reason": None
                                }]
                            }
                            yield f"data: {json.dumps(chunk)}\n\n"
                        except json.JSONDecodeError:
                            continue # Skip non-json lines
        except Exception as e:
            yield f"data: {json.dumps({'error': f'Streaming error: {str(e)}'})}\n\n"

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


# --- API Endpoints ---

@app.get("/v1/models", response_model=ModelList)
async def list_models():
    """Lists the available models that this compatibility layer supports."""
    model_cards = [
        ModelCard(id=model_name) for model_name in SUPPORTED_MODELS.keys()
    ]
    return ModelList(data=model_cards)

@app.post("/v1/chat/completions")
async def create_chat_completion(request: OpenAIChatCompletionRequest):
    """Creates a chat completion, either streaming or synchronous."""
    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_response(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" # Wait up to 120 seconds for a response
    }

    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 = prediction.get("output", "")
            if isinstance(output, list):
                output = "".join(output)

            # Check for tool call
            try:
                # A simple check if the output is a JSON for a tool call
                tool_call_data = json.loads(output)
                if tool_call_data.get("type") == "tool_call":
                    message_content = None
                    tool_calls = [{
                        "id": f"call_{int(time.time())}",
                        "type": "function",
                        "function": {
                            "name": tool_call_data["name"],
                            "arguments": json.dumps(tool_call_data["arguments"])
                        }
                    }]
                else:
                    message_content = output
                    tool_calls = None
            except (json.JSONDecodeError, TypeError):
                message_content = output
                tool_calls = None

            # Format response in OpenAI format
            completion_response = {
                "id": prediction["id"],
                "object": "chat.completion",
                "created": int(time.time()),
                "model": model_key,
                "choices": [{
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": message_content,
                        "tool_calls": tool_calls,
                    },
                    "finish_reason": "stop" # Or map from Replicate if available
                }],
                "usage": { # Note: Replicate doesn't provide token usage in the same way
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                }
            }
            return JSONResponse(content=completion_response)

        except httpx.HTTPStatusError as e:
            raise HTTPException(status_code=e.response.status_code, detail=e.response.text)