Add api for embedding

genai_stack/constant.py

@@ -3,6 +3,7 @@
 STACK = "/stack"
 COMPONENT = "/component"
 SESSION = "/session"
+EMBEDDING = "/embedding"
 RETRIEVER = "/retriever"
 VECTORDB = "/vectordb"
 ETL = "/etl"

genai_stack/genai_server/models/embedding_models.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel
+
+
+class EmbeddingRequestModel(BaseModel):
+    text: str
+    session_id: int
+
+
+class EmbeddingResponseModel(BaseModel):
+    embedding: list

genai_stack/genai_server/routers/embedding_routes.py

@@ -0,0 +1,15 @@
+from fastapi import APIRouter
+
+from genai_stack.constant import API, EMBEDDING
+from genai_stack.genai_server.settings.settings import settings
+from genai_stack.genai_server.models.embedding_models import EmbeddingResponseModel, EmbeddingRequestModel
+from genai_stack.genai_server.services.embedding_service import EmbeddingService
+
+service = EmbeddingService(store=settings.STORE)
+
+router = APIRouter(prefix=API + EMBEDDING, tags=["embedding"])
+
+
+@router.post("/embed-text")
+def retrieve(data: EmbeddingRequestModel) -> EmbeddingResponseModel:
+    return service.embed_text(data)

genai_stack/genai_server/services/embedding_service.py

@@ -0,0 +1,22 @@
+from fastapi import HTTPException
+from sqlalchemy.orm import Session
+
+from genai_stack.genai_platform.services import BaseService
+from genai_stack.genai_server.routers.embedding_routes import EmbeddingRequestModel, EmbeddingResponseModel
+from genai_stack.genai_server.utils import get_current_stack
+from genai_stack.genai_server.settings.config import stack_config
+from genai_stack.genai_store.schemas import StackSessionSchema
+
+
+class EmbeddingService(BaseService):
+
+    def embed_text(self, data: EmbeddingRequestModel) -> EmbeddingResponseModel:
+        with Session(self.engine) as session:
+                stack_session = session.get(StackSessionSchema, data.session_id)
+                if stack_session is None:
+                    raise HTTPException(status_code=404, detail=f"Session {data.session_id} not found")
+                stack = get_current_stack(config=stack_config, session=stack_session)
+                stack.embedding.load()
+                return EmbeddingResponseModel(
+                    embedding=stack.embedding.embed_text(data.text)
+                )

tests/api/test_genai_server/test_embedding.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+
+"""Tests for `genai_server`."""
+import unittest
+import requests
+
+
+class TestEmbeddingServerAPI(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.base_url = "http://127.0.0.1:5000/api/embedding"
+
+    def test_embed_test(self):
+        response = requests.post(
+            url=self.base_url + "/embed-text",
+            json={
+                "text": "New Delhi is the capital of India.",
+                "session_id": 1
+            }
+        )
+        assert response.status_code == 200
+        assert response.json()
+        data = response.json()
+        assert "embedding" in data
+        assert type(data["embedding"]) == list

tests/test_embedding.py

@@ -16,4 +16,5 @@ def test_huggingface_embedding(self):
             "encode_kwargs": {"normalize_embeddings": False},
         }
         embedding = LangchainEmbedding.from_kwargs(name="HuggingFaceEmbeddings", fields=config)
+        embedding.load()
         embedding.embed_text("something")