Update app.py

app.py

@@ -1,10 +1,11 @@
+
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 import numpy as np
 from sentence_transformers import SentenceTransformer
 import asyncio
 import logging
-from typing import List
+from typing import List, Optional, Union
 from pydantic import BaseModel
 import time
 
@@ -14,17 +15,21 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI(title="Embedding API")
 
+MODEL_NAME = "all-MiniLM-L12-v2"
+
 # Load model on startup
 try:
-    model = SentenceTransformer("all-MiniLM-L12-v2")
+    model = SentenceTransformer(MODEL_NAME)
     logger.info("Model loaded successfully")
 except Exception as e:
     logger.error(f"Failed to load model: {e}")
     model = None
 
 # Request batching queue
-class EmbedRequest(BaseModel):
-    texts: List[str]
+class EmbeddingRequest(BaseModel):
+    input: Union[str, List[str]]
+    model: str = MODEL_NAME
+    encoding_format: Optional[str] = None
 
 class BatchItem:
     def __init__(self, texts: List[str]):
@@ -104,38 +109,47 @@ async def startup():
     logger.info("Batch processor started")
 
 
-@app.get("/health")
-async def health():
-    """Health check endpoint"""
+@app.get("/v1/models")
+async def list_models():
+    """OpenAI-compatible models endpoint"""
     if model is None:
-        return JSONResponse({"status": "unhealthy", "reason": "model not loaded"}, status_code=503)
-    return {"status": "healthy"}
-
-
-@app.post("/embed")
-async def embed(request: EmbedRequest):
-    """
-    Embed texts. Automatically batches with concurrent requests.
+        raise HTTPException(status_code=503, detail="Model not ready")
     
-    Example:
-    ```
-    POST /embed
-    {
-        "texts": ["hello world", "foo bar"]
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": MODEL_NAME,
+                "object": "model",
+                "created": 0,
+                "owned_by": "huggingface",
+                "permission": [],
+                "root": MODEL_NAME,
+                "parent": None
+            }
+        ]
     }
-    ```
-    """
+
+
+@app.post("/v1/embeddings")
+async def create_embeddings(request: EmbeddingRequest):
     if model is None:
         raise HTTPException(status_code=503, detail="Model not ready")
     
-    if not request.texts:
-        raise HTTPException(status_code=400, detail="texts cannot be empty")
+    # Normalize input to list
+    if isinstance(request.input, str):
+        texts = [request.input]
+    else:
+        texts = request.input
+    
+    if not texts:
+        raise HTTPException(status_code=400, detail="input cannot be empty")
     
-    if len(request.texts) > 512:
+    if len(texts) > 512:
         raise HTTPException(status_code=400, detail="Cannot embed more than 512 texts at once")
     
-    # Create request item
-    batch_item = BatchItem(request.texts)
+    # Create batch item
+    batch_item = BatchItem(texts)
     
     async with batch_lock:
         pending_batch.append(batch_item)
@@ -144,29 +158,52 @@ async def embed(request: EmbedRequest):
         if len(pending_batch) >= MAX_BATCH_SIZE:
             batch_event.set()
     
-    # Signal processor that there's work (in case it's idle)
+    # Signal processor that there's work
     batch_event.set()
     
     # Wait for result with timeout
     try:
-        embeddings = await asyncio.wait_for(batch_item.future, timeout=30.0)
-        return {
-            "embeddings": embeddings,
-            "model": "all-MiniLM-L12-v2",
-            "count": len(request.texts)
-        }
+        embeddings_list = await asyncio.wait_for(batch_item.future, timeout=30.0)
     except asyncio.TimeoutError:
         raise HTTPException(status_code=504, detail="Embedding request timed out")
+    
+    # Format response as OpenAI-compatible
+    data = []
+    for idx, embedding in enumerate(embeddings_list):
+        data.append({
+            "object": "embedding",
+            "embedding": embedding,
+            "index": idx
+        })
+    
+    return {
+        "object": "list",
+        "data": data,
+        "model": MODEL_NAME,
+        "usage": {
+            "prompt_tokens": sum(len(text.split()) for text in texts),
+            "total_tokens": sum(len(text.split()) for text in texts)
+        }
+    }
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint"""
+    if model is None:
+        return JSONResponse({"status": "unhealthy", "reason": "model not loaded"}, status_code=503)
+    return {"status": "healthy"}
 
 
 @app.get("/")
 async def root():
     """API info"""
     return {
-        "name": "Embedding API",
-        "model": "all-MiniLM-L12-v2",
+        "name": "OpenAI-Compatible Embedding API",
+        "model": MODEL_NAME,
         "endpoints": {
-            "POST /embed": "Embed texts",
+            "GET /v1/models": "List available models",
+            "POST /v1/embeddings": "Create embeddings (OpenAI-compatible)",
             "GET /health": "Health check"
         }
     }
@@ -174,4 +211,4 @@ async def root():
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)