Update apis/chat_api.py

apis/chat_api.py

@@ -335,6 +335,22 @@ class ChatAPIApp:
         options: Optional[dict] = None
 
     def get_embeddings(self, request: EmbeddingRequest, api_key: str = Depends(extract_api_key)):
+        try:
+            model = request.model
+            model_kwargs = request.options
+            embeddings = self.embeddings[model].encode(request.prompt, convert_to_tensor=True)#, **model_kwargs)
+            return {
+                "object":"list", 
+                "data":[
+                    "object": "embedding", "index": 0, "embedding": embeddings.tolist()
+                ], 
+                "model": model, 
+                "usage":{},
+            }
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+
+    def get_embeddings_ollama(self, request: EmbeddingRequest, api_key: str = Depends(extract_api_key)):
         try:
             model = request.model
             model_kwargs = request.options
@@ -383,12 +399,19 @@ class ChatAPIApp:
                 summary="Ollama Chat completions in conversation session",
                 include_in_schema=include_in_schema,
             )(self.chat_completions_ollama)
-            
-            self.app.post(
-                prefix + "/embeddings",
-                summary="Get Embeddings with prompt",
-                include_in_schema=include_in_schema,
-            )(self.get_embeddings)
+
+            if prefix in ["/api"]:
+                self.app.post(
+                    prefix + "/embeddings",
+                    summary="Ollama Get Embeddings with prompt",
+                    include_in_schema=True,
+                )(self.get_embeddings_ollama)
+            else:
+                self.app.post(
+                    prefix + "/embeddings",
+                    summary="Get Embeddings with prompt",
+                    include_in_schema=include_in_schema,
+                )(self.get_embeddings)
 
         self.app.get(
             "/api/tags",