Update main.py

main.py

@@ -9,156 +9,154 @@ 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 environment variables
 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="4.0.0 (Stable & Correct)",
-)
+# FastAPI Init
+app = FastAPI(title="Replicate to OpenAI Compatibility Layer", version="4.0.0 (Docs Compliant)")
 
 # --- 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]]]
-
 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
 
-# --- Model Mapping (Simplified for direct endpoint usage) ---
+# --- Supported Models ---
+# Maps OpenAI-friendly names to Replicate model paths
 SUPPORTED_MODELS = {
-    "llama3-8b-instruct": {
-        "id": "meta/meta-llama-3-8b-instruct",
-        "input_type": "messages"
-    },
-    "claude-4.5-haiku": {
-        "id": "anthropic/claude-4.5-haiku",
-        "input_type": "prompt"
-    }
+    "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, model_details: dict) -> Dict[str, Any]:
-    """Prepares the 'input' dictionary for Replicate, handling model-specific formats."""
-    input_payload = {}
+# --- Core Logic ---
+def prepare_replicate_input(request: OpenAIChatCompletionRequest, replicate_model_id: str) -> Dict[str, Any]:
+    """Formats the input specifically for the requested Replicate model."""
+    payload = {}
     
-    if model_details["input_type"] == "prompt":
+    # Claude on Replicate strictly requires a 'prompt' string, not 'messages' array.
+    if "anthropic/claude" in replicate_model_id:
         prompt_parts = []
         system_prompt = None
         for msg in request.messages:
-            if msg.role == "system": system_prompt = str(msg.content)
-            elif msg.role == "user": prompt_parts.append(f"User: {msg.content}")
-            elif msg.role == "assistant": prompt_parts.append(f"Assistant: {msg.content}")
+            if msg.role == "system":
+                 # Extract system prompt if present
+                system_prompt = str(msg.content)
+            elif msg.role == "user":
+                # Handle both simple string content and list content (for potential future vision support)
+                content = msg.content
+                if isinstance(content, list):
+                     text_parts = [item.get("text", "") for item in content if item.get("type") == "text"]
+                     content = " ".join(text_parts)
+                prompt_parts.append(f"User: {content}")
+            elif msg.role == "assistant":
+                prompt_parts.append(f"Assistant: {msg.content}")
+        
+        # Standard Claude prompting convention
         prompt_parts.append("Assistant:")
-        input_payload["prompt"] = "\n".join(prompt_parts)
-        if system_prompt: input_payload["system_prompt"] = system_prompt
-    else: # "messages"
-        input_payload["messages"] = [msg.dict() for msg in request.messages]
-
-    if request.max_tokens is not None: input_payload["max_new_tokens"] = request.max_tokens
-    if request.temperature is not None: input_payload["temperature"] = request.temperature
-    if request.top_p is not None: input_payload["top_p"] = request.top_p
-    return input_payload
+        payload["prompt"] = "\n\n".join(prompt_parts)
+        if system_prompt:
+             payload["system_prompt"] = system_prompt
+             
+    # Llama 3 and others often support the 'messages' array natively.
+    else:
+         # Convert Pydantic models to pure dicts
+         payload["prompt"] = [msg.dict() for msg in request.messages]
+
+    # Map common OpenAI parameters to Replicate equivalents
+    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
 
-async def stream_replicate_native_sse(model_id: str, input_payload: dict):
-    """Connects to Replicate's native SSE stream using the model-specific endpoint."""
-    url = f"https://api.replicate.com/v1/models/{model_id}/predictions"
+async def stream_replicate_sse(replicate_model_id: str, input_payload: dict):
+    """Handles the full streaming lifecycle using standard Replicate endpoints."""
+    # 1. Start Prediction specifically at the named model endpoint
+    url = f"https://api.replicate.com/v1/models/{replicate_model_id}/predictions"
     headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json"}
     
-    # The request body is now simple and correct
-    request_body = {"input": input_payload, "stream": True}
-    
-    async with httpx.AsyncClient(timeout=300) as client:
-        prediction = None
+    async with httpx.AsyncClient(timeout=60.0) as client:
         try:
-            response = await client.post(url, headers=headers, json=request_body)
+            # Explicitly request stream=True in the body, though often implicit
+            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")
 
             if not stream_url:
-                error_detail = prediction.get("detail", "Failed to get stream URL.")
-                yield json.dumps({"error": {"message": error_detail}})
-                return
+                 yield json.dumps({"error": {"message": "Model did not return a stream URL."}})
+                 return
+
         except httpx.HTTPStatusError as e:
-            try: yield json.dumps({"error": {"message": json.dumps(e.response.json())}})
-            except: yield json.dumps({"error": {"message": e.response.text}})
-            return
-        
-        try:
-            async with client.stream("GET", stream_url, headers={"Accept": "text/event-stream"}) as sse:
-                sse.raise_for_status()
-                current_event = ""
-                async for line in sse.aiter_lines():
-                    if line.startswith("event:"):
-                        current_event = line[len("event:"):].strip()
-                    elif line.startswith("data:"):
-                        data = line[len("data:"):].strip()
-                        if current_event == "output":
+             yield json.dumps({"error": {"message": e.response.text, "type": "upstream_error"}})
+             return
+
+        # 2. Connect to the provided Stream URL
+        async with client.stream("GET", stream_url, headers={"Accept": "text/event-stream"}, timeout=None) as sse:
+            current_event = None
+            async for line in sse.aiter_lines():
+                if line.startswith("event:"):
+                    current_event = line[len("event:"):].strip()
+                elif line.startswith("data:"):
+                    data = line[len("data:"):].strip()
+                    
+                    if current_event == "output":
+                        # CRITICAL: Wrap in try/except to ignore empty keep-alive lines that crash standard parsers
+                        try:
+                            # Replicate sometimes sends raw strings, sometimes JSON. 
+                            # For chat models, it's usually a raw string token.
+                            # We try to load as JSON first, if it fails, use raw data.
                             try:
                                 content = json.loads(data)
+                            except json.JSONDecodeError:
+                                content = data
+                            
+                            if content: # Ensure we don't send empty chunks
                                 chunk = {
-                                    "id": prediction["id"], "object": "chat.completion.chunk", "created": int(time.time()), "model": model_id,
+                                    "id": prediction_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": replicate_model_id,
                                     "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}]
                                 }
                                 yield json.dumps(chunk)
-                            except json.JSONDecodeError:
-                                # Silently ignore malformed or empty data lines
-                                pass
-                        elif current_event == "done":
-                            break
-        except Exception as e:
-            yield json.dumps({"error": {"message": f"Streaming error: {str(e)}"}})
-
-    done_chunk = {
-        "id": prediction["id"] if prediction else "unknown", "object": "chat.completion.chunk", "created": int(time.time()), "model": model_id,
-        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
-    }
-    yield json.dumps(done_chunk)
+                        except Exception:
+                            pass # Safely ignore malformed lines
+                            
+                    elif current_event == "done":
+                        break
+
+    # 3. Send final [DONE] event
+    yield json.dumps({"id": prediction_id, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]})
     yield "[DONE]"
 
-# --- API Endpoints ---
-@app.get("/v1/models", response_model=ModelList)
+# --- Endpoints ---
+@app.get("/v1/models")
 async def list_models():
-    return ModelList(data=[ModelCard(id=model_name) for model_name in SUPPORTED_MODELS.keys()])
+    return ModelList(data=[ModelCard(id=k) for k 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())}")
+    if request.model not in SUPPORTED_MODELS:
+        raise HTTPException(404, f"Model not found. Available: {list(SUPPORTED_MODELS.keys())}")
     
-    model_details = SUPPORTED_MODELS[model_key]
-    replicate_input = prepare_replicate_input(request, model_details)
+    replicate_id = SUPPORTED_MODELS[request.model]
+    replicate_input = prepare_replicate_input(request, replicate_id)
 
     if request.stream:
-        return EventSourceResponse(stream_replicate_native_sse(model_details["id"], replicate_input))
-    
-    # Synchronous Request
-    url = f"https://api.replicate.com/v1/models/{model_details['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", []))
-            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)
+        return EventSourceResponse(stream_replicate_sse(replicate_id, replicate_input))
+
+    # Non-streaming fallback
+    url = f"https://api.replicate.com/v1/models/{replicate_id}/predictions"
+    headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}", "Content-Type": "application/json", "Prefer": "wait=60"}
+    async with httpx.AsyncClient() as client:
+        resp = await client.post(url, headers=headers, json={"input": replicate_input})
+        if resp.is_error: raise HTTPException(resp.status_code, resp.text)
+        pred = resp.json()
+        output = "".join(pred.get("output", []))
+        return {"id": pred["id"], "choices": [{"message": {"role": "assistant", "content": output}, "finish_reason": "stop"}]}