Update apis/chat_api.py

apis/chat_api.py

@@ -25,6 +25,7 @@ from networks.huggingface_streamer import HuggingfaceStreamer
 from networks.huggingchat_streamer import HuggingchatStreamer
 from networks.openai_streamer import OpenaiStreamer
 
+from sentence_transformers import SentenceTransformer
 
 class ChatAPIApp:
     def __init__(self):
@@ -35,6 +36,10 @@ class ChatAPIApp:
             version=CONFIG["version"],
         )
         self.setup_routes()
+        self.embeddings = {
+            "mxbai-embed-large":SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1"),
+            "nomic-embed-text": SentenceTransformer("nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True)
+        }
 
     def get_available_models(self):
         return {"object": "list", "data": AVAILABLE_MODELS_DICTS}
@@ -136,6 +141,22 @@ class ChatAPIApp:
         except Exception as e:
             raise HTTPException(status_code=500, detail=str(e))
 
+
+    class EmbeddingRequest(BaseModel):
+        model: str
+        prompt: str
+        options: Optional[dict] = None
+
+    def get_embeddings(self, item: EmbeddingRequest, api_key: str = Depends(extract_api_key)):
+        try:
+            model = request.model
+            model_kwargs = request.options
+            embeddings = models[model].encode(request.prompt, convert_to_tensor=True)#, **model_kwargs)
+            return {"embedding": embeddings.tolist()}
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+        
+
     def get_readme(self):
         readme_path = Path(__file__).parents[1] / "README.md"
         with open(readme_path, "r", encoding="utf-8") as rf:
@@ -163,6 +184,12 @@ class ChatAPIApp:
                 summary="Chat completions in conversation session",
                 include_in_schema=include_in_schema,
             )(self.chat_completions)
+
+        self.app.post(
+            "/api/embeddings",
+            summary="Get Embeddings with prompt",
+            include_in_schema=True,
+        )(self.get_embeddings)
         self.app.get(
             "/readme",
             summary="README of HF LLM API",