main.py

import os
import httpx
import json
import time
import asyncio
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional, Union, Literal
from dotenv import load_dotenv

# 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.2.2 (Spacing 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]]]
    name: Optional[str] = None
    tool_calls: Optional[List[Any]] = None

class FunctionDefinition(BaseModel):
    name: str
    description: Optional[str] = None
    parameters: Optional[Dict[str, Any]] = None

class ToolDefinition(BaseModel):
    type: Literal["function"]
    function: FunctionDefinition

class FunctionCall(BaseModel):
    name: str
    arguments: str

class ToolCall(BaseModel):
    id: str
    type: Literal["function"] = "function"
    function: FunctionCall

class ChatCompletionRequest(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
    stop: Optional[Union[str, List[str]]] = None
    tools: Optional[List[ToolDefinition]] = None
    tool_choice: Optional[Union[str, Dict[str, Any]]] = None
    functions: Optional[List[FunctionDefinition]] = None
    function_call: Optional[Union[str, Dict[str, str]]] = None

class Choice(BaseModel):
    index: int
    message: ChatMessage
    finish_reason: Optional[str] = None

class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    inference_time: Optional[float] = None

class ChatCompletion(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: List[Choice]
    usage: Usage

class DeltaMessage(BaseModel):
    role: Optional[str] = None
    content: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None

class ChoiceDelta(BaseModel):
    index: int
    delta: DeltaMessage
    finish_reason: Optional[str] = None

class ChatCompletionChunk(BaseModel):
    id: str
    object: str = "chat.completion.chunk"
    created: int
    model: str
    choices: List[ChoiceDelta]
    usage: Optional[Usage] = None

# --- 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 format_messages_for_replicate(messages: List[ChatMessage], functions: Optional[List[FunctionDefinition]] = None) -> Dict[str, Any]:
    prompt_parts = []
    system_prompt = None
    image_input = None
    
    if functions:
        functions_text = "You have access to the following tools. Use them if required to answer the user's question.\n\n"
        for func in functions:
            functions_text += f"- Function: {func.name}\n"
            if func.description: functions_text += f"  Description: {func.description}\n"
            if func.parameters: functions_text += f"  Parameters: {json.dumps(func.parameters)}\n"
        prompt_parts.append(functions_text)

    for msg in messages:
        if msg.role == "system":
            system_prompt = str(msg.content)
        elif msg.role == "assistant":
            if msg.tool_calls:
                tool_calls_text = "\nTool calls:\n"
                for tool_call in msg.tool_calls:
                    tool_calls_text += f"- {tool_call.function.name}({tool_call.function.arguments})\n"
                prompt_parts.append(f"Assistant: {tool_calls_text}")
            else:
                prompt_parts.append(f"Assistant: {msg.content}")
        elif msg.role == "tool":
            prompt_parts.append(f"Tool Response: {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:") # Let the model generate the space after this
    return {
        "prompt": "\n\n".join(prompt_parts),
        "system_prompt": system_prompt,
        "image": image_input
    }

def parse_function_call(content: str) -> Optional[Dict[str, Any]]:
    try:
        if "function_call" in content or ("name" in content and "arguments" in content):
            start = content.find("{")
            end = content.rfind("}") + 1
            if start != -1 and end > start:
                json_str = content[start:end]
                parsed = json.loads(json_str)
                if "name" in parsed and "arguments" in parsed:
                    return parsed
    except (json.JSONDecodeError, Exception):
        pass
    return None

async def stream_replicate_response(replicate_model_id: str, input_payload: dict, request_id: str):
    url = f"https://api.replicate.com/v1/models/{replicate_model_id}/predictions"
    headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json"}
    
    start_time = time.time()
    prompt_tokens = len(input_payload.get("prompt", "")) // 4
    completion_tokens = 0
    
    async with httpx.AsyncClient(timeout=300.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")
            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_details = e.response.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
                accumulated_content = ""
                first_token = True
                
                async for line in sse.aiter_lines():
                    if not line: continue
                    
                    if line.startswith("event:"):
                        current_event = line[len("event:"):].strip()
                    elif line.startswith("data:") and current_event == "output":
                        raw_data = line[5:].strip()
                        if not raw_data: continue

                        try:
                            content_token = json.loads(raw_data)
                        except (json.JSONDecodeError, TypeError):
                            content_token = raw_data
                        
                        # ### MAJOR FIX HERE ###
                        # This logic robustly handles the leading space by only stripping
                        # the very first non-empty token of the entire stream.
                        if first_token:
                            content_token = content_token.lstrip()
                            # Only flip the flag if we've actually processed a token with content.
                            if content_token:
                                first_token = False
                        
                        accumulated_content += content_token
                        completion_tokens += 1
                        
                        function_call = parse_function_call(accumulated_content)
                        if function_call:
                            tool_call = ToolCall(id=f"call_{int(time.time())}", function=FunctionCall(name=function_call["name"], arguments=function_call["arguments"]))
                            chunk = ChatCompletionChunk(id=request_id, created=int(time.time()), model=replicate_model_id, choices=[ChoiceDelta(index=0, delta=DeltaMessage(tool_calls=[tool_call]), finish_reason=None)])
                            yield f"data: {chunk.json()}\n\n"
                        else:
                            # Only yield a chunk if there is content to send.
                            if content_token:
                                chunk = ChatCompletionChunk(id=request_id, created=int(time.time()), model=replicate_model_id, choices=[ChoiceDelta(index=0, delta=DeltaMessage(content=content_token), finish_reason=None)])
                                yield f"data: {chunk.json()}\n\n"
                            
                    elif current_event == "done":
                        end_time = time.time()
                        usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, inference_time=round(end_time - start_time, 3))
                        usage_chunk = ChatCompletionChunk(id=request_id, created=int(time.time()), model=replicate_model_id, choices=[ChoiceDelta(index=0, delta=DeltaMessage(), finish_reason="stop")], usage=usage)
                        yield f"data: {usage_chunk.json()}\n\n"
                        break
                        
        except httpx.ReadTimeout:
            yield f"data: {json.dumps({'error': {'message': 'Stream timed out.', 'type': 'timeout_error'}})}\n\n"
            return

    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: ChatCompletionRequest):
    if request.model not in SUPPORTED_MODELS:
        raise HTTPException(status_code=404, detail=f"Model not found. Available models: {list(SUPPORTED_MODELS.keys())}")
    
    replicate_model_id = SUPPORTED_MODELS[request.model]
    formatted = format_messages_for_replicate(request.messages, request.functions)
    replicate_input = {
        "prompt": formatted["prompt"],
        "max_new_tokens": request.max_tokens or 512,
        "temperature": request.temperature or 0.7,
        "top_p": request.top_p or 1.0
    }
    if formatted["system_prompt"]: replicate_input["system_prompt"] = formatted["system_prompt"]
    if formatted["image"]: replicate_input["image"] = formatted["image"]
    
    request_id = f"chatcmpl-{int(time.time())}"
    
    if request.stream:
        return StreamingResponse(
            stream_replicate_response(replicate_model_id, replicate_input, request_id),
            media_type="text/event-stream"
        )

    # Non-streaming response
    url = f"https://api.replicate.com/v1/models/{replicate_model_id}/predictions"
    headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json"}
    start_time = time.time()
    
    async with httpx.AsyncClient() as client:
        try:
            resp = await client.post(url, headers=headers, json={"input": replicate_input}, timeout=300.0)
            resp.raise_for_status()
            pred = resp.json()
            output = "".join(pred.get("output", []))
            
            output = output.strip() # Clean up any leading/trailing whitespace
            
            end_time = time.time()
            prompt_tokens = len(replicate_input.get("prompt", "")) // 4
            completion_tokens = len(output) // 4
            
            tool_calls = None
            finish_reason = "stop"
            message_content = output
            
            function_call = parse_function_call(output)
            if function_call:
                tool_calls = [ToolCall(id=f"call_{int(time.time())}", function=FunctionCall(name=function_call["name"], arguments=function_call["arguments"]))]
                finish_reason = "tool_calls"
                message_content = None # OpenAI standard: content is null when tool_calls are present
            
            return ChatCompletion(
                id=request_id,
                created=int(time.time()),
                model=request.model,
                choices=[Choice(
                    index=0,
                    message=ChatMessage(role="assistant", content=message_content, tool_calls=tool_calls),
                    finish_reason=finish_reason
                )],
                usage=Usage(
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=prompt_tokens + completion_tokens,
                    inference_time=round(end_time - start_time, 3)
                )
            )
        except httpx.HTTPStatusError as e:
            raise HTTPException(status_code=e.response.status_code, detail=f"Error from Replicate API: {e.response.text}")
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}")

@app.get("/")
async def root():
    return {"message": "Replicate to OpenAI Compatibility Layer API", "version": "9.2.2"}

@app.middleware("http")
async def add_performance_headers(request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(round(process_time, 3))
    response.headers["X-API-Version"] = "9.2.2"
    return response