Update main.py

main.py

@@ -2,7 +2,6 @@ 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
@@ -21,101 +20,86 @@ if not REPLICATE_API_TOKEN:
 # --- FastAPI App Initialization ---
 app = FastAPI(
     title="Replicate to OpenAI Compatibility Layer",
-    version="2.0.0 (Native Streaming & Context Fixed)",
+    version="2.1.0 (Model Input Fixed)",
 )
 
-# --- Pydantic Models for OpenAI Compatibility ---
-
+# --- Pydantic Models ---
 class ModelCard(BaseModel):
-    id: str
-    object: str = "model"
-    created: int = Field(default_factory=lambda: int(time.time()))
-    owned_by: str = "replicate"
+    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] = []
+    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
+    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
-    tools: Optional[List[Tool]] = None
-    tool_choice: Optional[Union[str, Dict]] = None
-
-# --- Replicate Model Mapping ---
+    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
+
+# --- 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.
+    Prepares the input payload for Replicate, handling model-specific formats.
     """
-    # Convert Pydantic message objects to a list of dictionaries
-    messages_for_replicate = [msg.dict() for msg in request.messages]
-
-    payload = {
-        "messages": messages_for_replicate
-    }
+    payload = {}
+    
+    # *** THIS IS THE CRITICAL FIX ***
+    # Claude models on Replicate require a single 'prompt' string.
+    # We must convert the 'messages' array into a formatted string.
+    if "claude" in request.model:
+        prompt_parts = []
+        system_prompt = None
+        image_url = None
+
+        for msg in request.messages:
+            if msg.role == "system":
+                system_prompt = str(msg.content)
+            elif msg.role == "user":
+                if isinstance(msg.content, list): # Vision case
+                    for item in msg.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: # Text-only case
+                    prompt_parts.append(f"User: {msg.content}")
+            elif msg.role == "assistant":
+                prompt_parts.append(f"Assistant: {msg.content}")
+
+        payload["prompt"] = "\n".join(prompt_parts)
+        if system_prompt:
+            payload["system_prompt"] = system_prompt
+        if image_url:
+            payload["image"] = image_url
+            
+    # Other models like Llama-3 accept the 'messages' array directly.
+    else:
+        payload["messages"] = [msg.dict() for msg in request.messages]
 
-    # Add other compatible parameters
+    # Add common 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.
-    """
+    """Connects to Replicate's native SSE stream for 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()
@@ -126,10 +110,13 @@ async def stream_replicate_native_sse(model_id: str, payload: dict):
                 yield json.dumps({"error": {"message": error_detail}})
                 return
         except httpx.HTTPStatusError as e:
-            yield json.dumps({"error": {"message": e.response.text}})
+            try:
+                error_body = e.response.json()
+                yield json.dumps({"error": {"message": json.dumps(error_body)}})
+            except json.JSONDecodeError:
+                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()
@@ -146,11 +133,10 @@ async def stream_replicate_native_sse(model_id: str, payload: dict):
                             }
                             yield json.dumps(chunk)
                         elif current_event == "done":
-                            break # Exit loop when done event is received
+                            break
         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"}]
@@ -158,9 +144,7 @@ async def stream_replicate_native_sse(model_id: str, payload: dict):
     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()])
@@ -186,14 +170,11 @@ async def create_chat_completion(request: OpenAIChatCompletionRequest):
             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)