Update main.py

main.py

@@ -5,7 +5,7 @@ import json
 import time
 import asyncio
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 from typing import List, Dict, Any, Optional, Union, Literal
 from dotenv import 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.2.0 (Full OpenAI Compatibility)")
+app = FastAPI(title="Replicate to OpenAI Compatibility Layer", version="9.2.1 (Spacing Fixed)")
 
 # --- Pydantic Models ---
 class ModelCard(BaseModel):
@@ -106,6 +106,7 @@ def format_messages_for_replicate(messages: List[ChatMessage], functions: Option
                 user_text_content = str(msg.content)
             prompt_parts.append(f"User: {user_text_content}")
 
+    # Fix: Don't add trailing space, let model decide spacing
     prompt_parts.append("Assistant:")
     return {
         "prompt": "\n\n".join(prompt_parts),
@@ -161,6 +162,7 @@ async def stream_replicate_response(replicate_model_id: str, input_payload: dict
             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
@@ -177,6 +179,11 @@ async def stream_replicate_response(replicate_model_id: str, input_payload: dict
                         except (json.JSONDecodeError, TypeError):
                             content_token = raw_data
                         
+                        # Fix: Handle spacing properly - don't prepend space to first token
+                        if first_token:
+                            content_token = content_token.lstrip()
+                            first_token = False
+                        
                         accumulated_content += content_token
                         completion_tokens += 1
                         
@@ -291,6 +298,9 @@ async def create_chat_completion(request: ChatCompletionRequest):
             pred = resp.json()
             output = "".join(pred.get("output", []))
             
+            # Fix: Clean up leading/trailing whitespace
+            output = output.strip()
+            
             # Calculate timing and tokens
             end_time = time.time()
             inference_time = end_time - start_time
@@ -337,7 +347,7 @@ async def create_chat_completion(request: ChatCompletionRequest):
 
 @app.get("/")
 async def root():
-    return {"message": "Replicate to OpenAI Compatibility Layer API", "version": "9.2.0"}
+    return {"message": "Replicate to OpenAI Compatibility Layer API", "version": "9.2.1"}
 
 # Performance optimization middleware
 @app.middleware("http")
@@ -346,5 +356,5 @@ async def add_performance_headers(request, call_next):
     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
+    response.headers["X-API-Version"] = "9.2.1"
+    return response