Update main.py

main.py

@@ -46,6 +46,11 @@ class ChatRequest(BaseModel):
     stream: Optional[bool] = False
 
 
+class EmbeddingRequest(BaseModel):
+    input: str
+    model: str = "text-embedding-004"
+
+
 async def verify_authorization(authorization: str = Header(None)):
     if not authorization:
         logger.error("Missing Authorization header")
@@ -152,6 +157,24 @@ async def chat_completion(request: ChatRequest, authorization: str = Header(None
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@app.post("/v1/embeddings")
+@app.post("/hf/v1/embeddings")
+async def embedding(request: EmbeddingRequest, authorization: str = Header(None)):
+    await verify_authorization(authorization)
+    async with key_lock:
+        api_key = next(key_cycle)
+        logger.info(f"Using API key: {api_key}")
+
+    try:
+        client = openai.OpenAI(api_key=api_key, base_url=config.settings.BASE_URL)
+        response = client.embeddings.create(input=request.input, model=request.model)
+        logger.info("Embedding successful")
+        return response
+    except Exception as e:
+        logger.error(f"Error in embedding: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @app.get("/health")
 @app.get("/")
 async def health_check():