main.py

import os
import httpx
import json
import time
import asyncio
from fastapi import FastAPI, Request, 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.")

POLLING_INTERVAL_SECONDS = 1  # How often to poll for updates

# --- FastAPI App Initialization ---
app = FastAPI(
    title="Replicate to OpenAI Compatibility Layer",
    version="1.2.0 (Streaming 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 format_tools_for_prompt(tools: List[Tool]) -> str:
    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]:
    input_data = {}
    prompt_parts = []
    system_prompt = ""
    image_url = None

    for message in request.messages:
        if message.role == "system":
            system_prompt += str(message.content) + "\n"
        elif message.role == "user":
            content = message.content
            if isinstance(content, list):
                for item in 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: {str(content)}")
        elif message.role == "assistant":
            prompt_parts.append(f"Assistant: {str(message.content)}")

    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

    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:
        input_data["max_new_tokens"] = request.max_tokens

    return input_data

async def stream_replicate_with_polling(model_id: str, payload: dict):
    """
    Creates a prediction and polls the 'get' URL to stream back results.
    Yields raw JSON strings for EventSourceResponse to handle.
    """
    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:
        prediction = None
        try:
            response = await client.post(url, headers=headers, json={"input": payload})
            response.raise_for_status()
            prediction = response.json()
            get_url = prediction.get("urls", {}).get("get")

            if not get_url:
                error_detail = prediction.get("detail", "Failed to start prediction.")
                error_chunk = {"error": {"message": error_detail, "type": "api_error", "code": 500}}
                yield json.dumps(error_chunk)
                return
        except httpx.HTTPStatusError as e:
            error_chunk = {"error": {"message": e.response.text, "type": "api_error", "code": e.response.status_code}}
            yield json.dumps(error_chunk)
            return
        
        previous_output = ""
        status = ""
        while status not in ["succeeded", "failed", "canceled"]:
            await asyncio.sleep(POLLING_INTERVAL_SECONDS)
            try:
                poll_response = await client.get(get_url, headers=headers)
                poll_response.raise_for_status()
                prediction_update = poll_response.json()
                status = prediction_update["status"]

                if status == "failed":
                    error_detail = prediction_update.get("error", "Prediction failed.")
                    chunk = {"choices": [{"delta": {"content": f"\n\n[ERROR: {error_detail}]"}, "finish_reason": "error"}]}
                    yield json.dumps(chunk)
                    break

                if "output" in prediction_update and prediction_update["output"] is not None:
                    current_output = "".join(prediction_update["output"])
                    new_chunk_text = current_output[len(previous_output):]
                    
                    if new_chunk_text:
                        chunk = {
                            "id": prediction["id"], "object": "chat.completion.chunk", "created": int(time.time()), "model": model_id,
                            "choices": [{"index": 0, "delta": {"content": new_chunk_text}, "finish_reason": None}]
                        }
                        yield json.dumps(chunk) # *** FIX: Yield raw JSON string
                        previous_output = current_output
            
            except Exception as e:
                error_chunk = {"error": {"message": f"Polling error: {str(e)}", "type": "internal_error", "code": 500}}
                yield json.dumps(error_chunk)
                break
    
    # Send the final 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" if status == "succeeded" else "error"}]
    }
    yield json.dumps(done_chunk) # *** FIX: Yield raw JSON string
    yield "[DONE]" # *** FIX: Yield the special [DONE] marker


# --- 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_with_polling(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", []))
            
            try:
                tool_call_data = json.loads(output)
                if tool_call_data.get("type") == "tool_call":
                    message_content, tool_calls = None, [{"id": f"call_{int(time.time())}", "type": "function", "function": {"name": tool_call_data["name"], "arguments": json.dumps(tool_call_data["arguments"])}}]
                else:
                    message_content, tool_calls = output, None
            except (json.JSONDecodeError, TypeError):
                message_content, tool_calls = output, None

            return JSONResponse(content={
                "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"}],
                "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)