Update main.py

main.py

@@ -3,12 +3,12 @@ import os
 import httpx
 import json
 import time
+import asyncio
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
+from fastapi.responses import StreamingResponse, 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()
@@ -17,7 +17,7 @@ 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.1.0 (Enhanced Token Tracking)")
+app = FastAPI(title="Replicate to OpenAI Compatibility Layer", version="9.2.0 (Full OpenAI Compatibility)")
 
 # --- Pydantic Models ---
 class ModelCard(BaseModel):
@@ -25,9 +25,34 @@ class ModelCard(BaseModel):
 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
+    role: Literal["system", "user", "assistant", "tool"]; content: Union[str, List[Dict[str, Any]]]; name: Optional[str] = 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 = {
@@ -38,21 +63,36 @@ SUPPORTED_MODELS = {
 }
 
 # --- Core Logic ---
-def prepare_replicate_input(request: OpenAIChatCompletionRequest) -> Dict[str, Any]:
-    """
-    Formats the input for Replicate's API, flattening the message history into a
-    single 'prompt' string and handling images separately.
-    """
-    payload = {}
+def format_messages_for_replicate(messages: List[ChatMessage], functions: Optional[List[FunctionDefinition]] = None) -> Dict[str, Any]:
+    """Convert OpenAI messages to Replicate-compatible format with function calling support."""
     prompt_parts = []
     system_prompt = None
     image_input = None
     
-    for msg in request.messages:
+    # Add functions to system prompt if provided
+    if functions:
+        functions_text = "\n\nAvailable functions:\n"
+        for func in functions:
+            functions_text += f"- {func.name}: {func.description or 'No 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":
-            prompt_parts.append(f"Assistant: {msg.content}")
+            # Handle tool calls in assistant messages
+            if hasattr(msg, 'tool_calls') and 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":
+            # Handle tool responses
+            prompt_parts.append(f"Tool Response: {msg.content}")
         elif msg.role == "user":
             user_text_content = ""
             if isinstance(msg.content, list):
@@ -67,37 +107,46 @@ def prepare_replicate_input(request: OpenAIChatCompletionRequest) -> Dict[str, A
             prompt_parts.append(f"User: {user_text_content}")
 
     prompt_parts.append("Assistant:")
-    payload["prompt"] = "\n\n".join(prompt_parts)
-    
-    if system_prompt:
-        payload["system_prompt"] = system_prompt
-    if image_input:
-        payload["image"] = image_input
-         
-    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
+    return {
+        "prompt": "\n\n".join(prompt_parts),
+        "system_prompt": system_prompt,
+        "image": image_input
+    }
+
+def parse_function_call(content: str) -> Optional[Dict[str, Any]]:
+    """Parse function call from model response."""
+    try:
+        # Look for JSON-like function call patterns
+        if "function_call" in content or ("name" in content and "arguments" in content):
+            # Extract JSON part
+            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_sse(replicate_model_id: str, input_payload: dict):
-    """Handles the full streaming lifecycle with enhanced token tracking and timing."""
+async def stream_replicate_response(replicate_model_id: str, input_payload: dict, request_id: str):
+    """Stream response with full OpenAI compatibility including tool calls."""
     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  # Rough estimation
+    prompt_tokens = len(input_payload.get("prompt", "")) // 4
     completion_tokens = 0
     
-    async with httpx.AsyncClient(timeout=60.0) as client:
+    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")
-            prediction_id = prediction.get("id", "stream-unknown")
             if not stream_url:
-                yield json.dumps({'error': {'message': 'Model did not return a 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
@@ -105,116 +154,99 @@ async def stream_replicate_sse(replicate_model_id: str, input_payload: dict):
                 error_json = e.response.json()
                 error_details = error_json.get("detail", error_details)
             except json.JSONDecodeError: pass
-            yield json.dumps({'error': {'message': f'Upstream Error: {error_details}', 'type': 'replicate_error'}})
+            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 = ""
+                
                 async for line in sse.aiter_lines():
-                    if not line: # Skip empty lines
-                        continue
+                    if not line: continue
+                    
                     if line.startswith("event:"):
                         current_event = line[len("event:"):].strip()
-                    elif line.startswith("data:"):
-                        # Remove "data:" prefix and optional space
-                        raw_data = line[5:]  # Remove "data:"
-                        if raw_data.startswith(" "):
-                            data_content = raw_data[1:]  # Remove the first space only
-                        else:
-                            data_content = raw_data
+                    elif line.startswith("data:") and current_event == "output":
+                        raw_data = line[5:].strip()
+                        if not raw_data: continue
                         
-                        if current_event == "output":
-                            if not data_content:
-                                continue
-                                
-                            content_token = ""
-                            try:
-                                # Handle JSON-encoded strings properly
-                                content_token = json.loads(data_content)
-                            except (json.JSONDecodeError, TypeError):
-                                # Handle plain text tokens
-                                content_token = data_content
-                            
-                            completion_tokens += 1
-                            chunk = {
-                                "choices": [{
-                                    "delta": {"content": content_token},
-                                    "finish_reason": None,
-                                    "index": 0,
-                                    "logprobs": None,
-                                    "native_finish_reason": None
-                                }],
-                                "created": int(time.time()),
-                                "id": f"gen-{int(time.time())}-{prediction_id[-12:]}",
-                                "model": replicate_model_id,
-                                "object": "chat.completion.chunk",
-                                "provider": "Anthropic" if "anthropic" in replicate_model_id else "Replicate"
-                            }
-                            yield json.dumps(chunk)
-                            
-                        elif current_event == "done":
-                            # Calculate timing
-                            end_time = time.time()
-                            inference_time = end_time - start_time
-                            
-                            # Send usage chunk before done
-                            usage_chunk = {
-                                "choices": [{
-                                    "delta": {},
-                                    "finish_reason": None,
-                                    "index": 0,
-                                    "logprobs": None,
-                                    "native_finish_reason": None
-                                }],
-                                "created": int(time.time()),
-                                "id": f"gen-{int(time.time())}-{prediction_id[-12:]}",
-                                "model": replicate_model_id,
-                                "object": "chat.completion.chunk",
-                                "provider": "Anthropic" if "anthropic" in replicate_model_id else "Replicate",
-                                "usage": {
-                                    "cache_discount": 0,
-                                    "completion_tokens": completion_tokens,
-                                    "completion_tokens_details": {"image_tokens": 0, "reasoning_tokens": 0},
-                                    "cost": 0,
-                                    "cost_details": {
-                                        "upstream_inference_completions_cost": 0,
-                                        "upstream_inference_cost": None,
-                                        "upstream_inference_prompt_cost": 0
-                                    },
-                                    "input_tokens": prompt_tokens,
-                                    "is_byok": False,
-                                    "prompt_tokens": prompt_tokens,
-                                    "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0},
-                                    "total_tokens": prompt_tokens + completion_tokens,
-                                    "inference_time": round(inference_time, 3)
-                                }
-                            }
-                            yield json.dumps(usage_chunk)
+                        content_token = ""
+                        try:
+                            content_token = json.loads(raw_data)
+                        except (json.JSONDecodeError, TypeError):
+                            content_token = raw_data
+                        
+                        accumulated_content += content_token
+                        completion_tokens += 1
+                        
+                        # Check for function calls in accumulated content
+                        function_call = parse_function_call(accumulated_content)
+                        if function_call:
+                            # Send tool call chunk
+                            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:
+                            # Send regular content chunk
+                            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"
                             
-                            # Send final chunk with stop reason
-                            final_chunk = {
-                                "choices": [{
-                                    "delta": {},
-                                    "finish_reason": "stop",
-                                    "index": 0,
-                                    "logprobs": None,
-                                    "native_finish_reason": "end_turn"
-                                }],
-                                "created": int(time.time()),
-                                "id": f"gen-{int(time.time())}-{prediction_id[-12:]}",
-                                "model": replicate_model_id,
-                                "object": "chat.completion.chunk",
-                                "provider": "Anthropic" if "anthropic" in replicate_model_id else "Replicate"
-                            }
-                            yield json.dumps(final_chunk)
-                            break
+                    elif current_event == "done":
+                        # Send final usage chunk
+                        end_time = time.time()
+                        inference_time = end_time - start_time
+                        
+                        usage = Usage(
+                            prompt_tokens=prompt_tokens,
+                            completion_tokens=completion_tokens,
+                            total_tokens=prompt_tokens + completion_tokens,
+                            inference_time=round(inference_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 json.dumps({'error': {'message': 'Stream timed out.', 'type': 'timeout_error'}})
+            yield f"data: {json.dumps({'error': {'message': 'Stream timed out.', 'type': 'timeout_error'}})}\n\n"
             return
 
-    # Send [DONE] event
-    yield "[DONE]"
+    yield "data: [DONE]\n\n"
 
 # --- Endpoints ---
 @app.get("/v1/models")
@@ -222,23 +254,39 @@ 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: OpenAIChatCompletionRequest):
+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_input = prepare_replicate_input(request)
-
+    # Format messages for Replicate
+    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 EventSourceResponse(stream_replicate_sse(SUPPORTED_MODELS[request.model], replicate_input), media_type="text/event-stream")
+        return StreamingResponse(
+            stream_replicate_response(SUPPORTED_MODELS[request.model], replicate_input, request_id),
+            media_type="text/event-stream"
+        )
 
-    # Non-streaming fallback with usage data
+    # Non-streaming response
     url = f"https://api.replicate.com/v1/models/{SUPPORTED_MODELS[request.model]}/predictions"
-    headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json", "Prefer": "wait=120"}
+    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=130.0)
+            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", []))
@@ -246,18 +294,57 @@ async def create_chat_completion(request: OpenAIChatCompletionRequest):
             # Calculate timing and tokens
             end_time = time.time()
             inference_time = end_time - start_time
-            prompt_tokens = len(input_payload.get("prompt", "")) // 4  # Rough estimation
-            completion_tokens = len(output) // 4  # Rough estimation
+            prompt_tokens = len(replicate_input.get("prompt", "")) // 4
+            completion_tokens = len(output) // 4
+            
+            # Parse function call if present
+            tool_calls = None
+            function_call = parse_function_call(output)
+            if function_call:
+                tool_call = ToolCall(
+                    id=f"call_{int(time.time())}",
+                    function=FunctionCall(
+                        name=function_call["name"],
+                        arguments=function_call["arguments"]
+                    )
+                )
+                tool_calls = [tool_call]
             
-            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": {
-                    "prompt_tokens": prompt_tokens,
-                    "completion_tokens": completion_tokens,
-                    "total_tokens": prompt_tokens + completion_tokens,
-                    "inference_time": round(inference_time, 3)
-                }
-            }
+            return ChatCompletion(
+                id=request_id,
+                created=int(time.time()),
+                model=request.model,
+                choices=[Choice(
+                    index=0,
+                    message=ChatMessage(
+                        role="assistant",
+                        content=output if not function_call else None,
+                        tool_calls=tool_calls
+                    ),
+                    finish_reason="tool_calls" if function_call else "stop"
+                )],
+                usage=Usage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=prompt_tokens + completion_tokens,
+                    inference_time=round(inference_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.0"}
+
+# Performance optimization middleware
+@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.0"
+    return response