Add query endpoint for embedding and refactor embedding models

app.py

@@ -85,6 +85,70 @@ def create_app() -> FastAPI:
 
 app = create_app()
 
+@app.post("/query", response_model=Union[EmbedResponse, SparseEmbedResponse])
+async def create_query(request: EmbedRequest):
+    """Create a single dense or sparse query embedding for the given text.
+
+    The request must include `model_id`. For sparse models (config type
+    "sparse-embeddings") the endpoint returns a `SparseEmbedResponse`,
+    otherwise a dense `EmbedResponse` is returned.
+
+    Args:
+        request: `EmbedRequest` pydantic model with text, prompt and model_id.
+    Returns:
+        Union[EmbedResponse, SparseEmbedResponse]: The embedding response.
+    Raises:
+        HTTPException: on validation or internal errors with appropriate
+            HTTP status codes.
+    """
+    
+    if not request.model_id:
+        raise HTTPException(status_code=400, detail="model_id is required")
+
+    try:
+        assert model_manager is not None
+        model = model_manager.get_model(request.model_id)
+        start_time = time.time()
+
+        config = model_manager.model_configs[request.model_id]
+
+        if config.type == "sparse-embeddings":
+            # Sparse embedding
+            sparse_result = model.query_embed(text=[request.text], prompt=request.prompt)
+            processing_time = time.time() - start_time
+
+            if isinstance(sparse_result, dict) and "indices" in sparse_result:
+                sparse_embedding = SparseEmbedding(
+                    text=request.text,
+                    indices=sparse_result["indices"],
+                    values=sparse_result["values"],
+                )
+            else:
+                raise ValueError(f"Unexpected sparse result format: {sparse_result}")
+
+            return SparseEmbedResponse(
+                sparse_embedding=sparse_embedding,
+                model_id=request.model_id,
+                processing_time=processing_time,
+            )
+
+        # Dense embedding
+        embedding = model.query_embed(text=[request.text], prompt=request.prompt)[0]
+        processing_time = time.time() - start_time
+
+        return EmbedResponse(
+            embedding=embedding,
+            dimension=len(embedding),
+            model_id=request.model_id,
+            processing_time=processing_time,
+        )
+
+    except AssertionError:
+        logger.exception("Model manager is not initialized")
+        raise HTTPException(status_code=500, detail="Server not ready")
+    except Exception:
+        logger.exception("Error creating query embedding")
+        raise HTTPException(status_code=500, detail="Failed to create query embedding")
 
 @app.post("/embed", response_model=Union[EmbedResponse, SparseEmbedResponse])
 async def create_embedding(request: EmbedRequest):
@@ -117,7 +181,7 @@ async def create_embedding(request: EmbedRequest):
 
         if config.type == "sparse-embeddings":
             # Sparse embedding
-            sparse_result = model.embed(request.text, prompt=request.prompt)
+            sparse_result = model.embed_documents(text=[request.text], prompt=request.prompt)
             processing_time = time.time() - start_time
 
             if isinstance(sparse_result, dict) and "indices" in sparse_result:
@@ -136,7 +200,7 @@ async def create_embedding(request: EmbedRequest):
             )
 
         # Dense embedding
-        embedding = model.embed([request.text], request.prompt)[0]
+        embedding = model.embed_documents(text=[request.text], prompt=request.prompt)[0]
         processing_time = time.time() - start_time
 
         return EmbedResponse(

core/config.py

@@ -0,0 +1,9 @@
+from typing import Any, Dict
+
+
+class ModelConfig:
+    def __init__(self, model_id: str, config: Dict[str, Any]):
+        self.id = model_id
+        self.name = config["name"]
+        self.type = config["type"] 
+        self.repository = config["repository"]

core/embedding.py

@@ -1,19 +1,9 @@
-from loguru import logger
-from typing import Dict, List, Optional, Any
+from typing import List, Optional
 from sentence_transformers import SentenceTransformer
-from sentence_transformers import SparseEncoder
+from loguru import logger
 
+from .config import ModelConfig
 
-class ModelConfig:
-    def __init__(self, model_id: str, config: Dict[str, Any]):
-        self.id = model_id
-        self.name = config["name"]
-        self.type = config["type"]  # "embedding" or "sparse"
-        self.dimension = int(config["dimension"])
-        self.max_tokens = int(config["max_tokens"])
-        self.description = config["description"]
-        self.language = config["language"]
-        self.repository = config["repository"]
 
 class EmbeddingModel:
     """
@@ -43,115 +33,44 @@ class EmbeddingModel:
             logger.error(f"Failed to load embedding model {self.config.id}: {e}")
             raise
 
-    def embed(self, texts: List[str], prompt: Optional[str] = None) -> List[List[float]]:
+    def query_embed(self, text: List[str], prompt: Optional[str] = None) -> List[float]:
         """
-        method to generate embeddings for a list of texts.
+        method to generate embedding for a single text.
         
         Args:
-            texts: List of input texts
+            text: Input text
             prompt: Optional prompt for instruction-based models
         
-        Returns:            
-        List of embedding vectors
+        Returns:
+            Embedding vector
         """
         if not self._loaded:
             self.load()
             
         try:
-            embeddings = self.model.encode(texts, prompt=prompt)
-            return [embedding.tolist() for embedding in embeddings]
+            embedding = self.model.encode_query(text, prompt=prompt)
+            return embedding[0].tolist()
         except Exception as e:
             logger.error(f"Embedding generation failed: {e}")
             raise
 
-class SparseEmbeddingModel:
-    """
-    Sparse embedding model wrapper.
-    
-    Attributes:
-        config: ModelConfig instance
-        model: SparseEncoder instance
-        _loaded: Flag indicating if the model is loaded
-    """
-    
-    def __init__(self, config: ModelConfig):
-        self.config = config
-        self.model: Optional[SparseEncoder] = None
-        self._loaded = False
-    
-    def load(self) -> None:
-        """Load the sparse embedding model."""
-        if self._loaded:
-            return
-            
-        logger.info(f"Loading sparse model: {self.config.name}")
-        try:
-            self.model = SparseEncoder(self.config.name)
-            self._loaded = True
-            logger.success(f"Loaded sparse model: {self.config.id}")
-        except Exception as e:
-            logger.error(f"Failed to load sparse model {self.config.id}: {e}")
-            raise
-
-    def embed(self, text: str, prompt: Optional[str] = None) -> Dict[Any, Any]:
-        """
-        Generate a sparse embedding for a single text.
-        
-        Args:
-            text: Input text
-            prompt: Optional prompt for instruction-based models
-        
-        Returns:
-            Sparse embedding as a dictionary with 'indices' and 'values' keys.
-        """
-        
-        try:
-            tensor = self.model.encode([text])
-            
-            values = tensor[0].coalesce().values().tolist()
-            indices = tensor[0].coalesce().indices()[0].tolist()
-            
-            return {
-                "indices": indices,
-                "values": values
-            }
-            
-
-        except Exception as e:
-            logger.error(f"Embedding error: {e}")
-            raise
-
-    def embed_batch(self, texts: List[str], prompt: Optional[str] = None) -> List[Dict[str, Any]]:
+    def embed_documents(self, texts: List[str], prompt: Optional[str] = None) -> List[List[float]]:
         """
-        Generate sparse embeddings for a batch of texts.
+        method to generate embeddings for a list of texts.
         
         Args:
             texts: List of input texts
             prompt: Optional prompt for instruction-based models
         
-        Returns:
-            List of sparse embeddings as dictionaries with 'text' and 'sparse_embedding' keys.
+        Returns:            
+        List of embedding vectors
         """
         if not self._loaded:
             self.load()
             
         try:
-            tensors = self.model.encode(texts)
-            results = []
-            
-            for i, tensor in enumerate(tensors):
-                values = tensor.coalesce().values().tolist()
-                indices = tensor.coalesce().indices()[0].tolist()
-                
-                results.append({
-                    "text": texts[i],
-                    "sparse_embedding": {
-                        "indices": indices,
-                        "values": values
-                    }
-                })
-            
-            return results
+            embeddings = self.model.encode_document(texts, prompt=prompt)
+            return [embedding.tolist() for embedding in embeddings]
         except Exception as e:
-            logger.error(f"Sparse embedding generation failed: {e}")
+            logger.error(f"Embedding generation failed: {e}")
             raise

core/model_manager.py

@@ -1,9 +1,11 @@
 import yaml
 from pathlib import Path
 from loguru import logger
-from typing import Dict, List, Any, Optional, Union
+from typing import Dict, List, Any, Union
 from threading import Lock
-from .embedding import ModelConfig, EmbeddingModel, SparseEmbeddingModel
+from .embedding import EmbeddingModel
+from .sparse import SparseEmbeddingModel
+from .config import ModelConfig
 
 class ModelManager:
     """
@@ -20,7 +22,6 @@ class ModelManager:
     def __init__(self, config_path: str = "config.yaml"):
         self.models: Dict[str, Union[EmbeddingModel, SparseEmbeddingModel]] = {}
         self.model_configs: Dict[str, ModelConfig] = {}
-        self.default_model_id: Optional[str] = None
         self._lock = Lock()  # For thread safety
         self._preload_complete = False
         
@@ -40,9 +41,6 @@ class ModelManager:
             for model_id, model_cfg in config["models"].items():
                 self.model_configs[model_id] = ModelConfig(model_id, model_cfg)
                 
-            if "default" in config and "model" in config["default"]:
-                self.default_model_id = config["default"]["model"]
-                
             logger.info(f"Loaded {len(self.model_configs)} model configurations")
             
         except Exception as e:
@@ -140,10 +138,6 @@ class ModelManager:
             "id": config.id,
             "name": config.name,
             "type": config.type,
-            "dimension": config.dimension,
-            "max_tokens": config.max_tokens,
-            "description": config.description,
-            "language": config.language,
             "loaded": is_loaded,
             "repository": config.repository,
         }
@@ -210,7 +204,6 @@ High-performance API for generating text embeddings using multiple model archite
                 loaded_models.append({
                     "id": model_id,
                     "type": self.model_configs[model_id].type,
-                    "dimension": model.config.dimension,
                     "name": model.config.name
                 })
         

core/sparse.py

@@ -0,0 +1,125 @@
+from typing import Any, Dict, List, Optional
+from sentence_transformers import SparseEncoder
+from loguru import logger
+
+from .config import ModelConfig
+
+
+class SparseEmbeddingModel:
+    """
+    Sparse embedding model wrapper.
+    
+    Attributes:
+        config: ModelConfig instance
+        model: SparseEncoder instance
+        _loaded: Flag indicating if the model is loaded
+    """
+    
+    def __init__(self, config: ModelConfig):
+        self.config = config
+        self.model: Optional[SparseEncoder] = None
+        self._loaded = False
+    
+    def load(self) -> None:
+        """Load the sparse embedding model."""
+        if self._loaded:
+            return
+            
+        logger.info(f"Loading sparse model: {self.config.name}")
+        try:
+            self.model = SparseEncoder(self.config.name)
+            self._loaded = True
+            logger.success(f"Loaded sparse model: {self.config.id}")
+        except Exception as e:
+            logger.error(f"Failed to load sparse model {self.config.id}: {e}")
+            raise
+
+    def query_embed(self, text: List[str], prompt: Optional[str] = None) -> Dict[Any, Any]:
+        """
+        Generate a sparse embedding for a single text.
+        
+        Args:
+            text: Input text
+            prompt: Optional prompt for instruction-based models
+        Returns:
+            Sparse embedding as a dictionary with 'indices' and 'values' keys.
+        """
+        if not self._loaded:
+            self.load()
+            
+        try:
+            tensor = self.model.encode_query(text)
+            
+            values = tensor[0].coalesce().values().tolist()
+            indices = tensor[0].coalesce().indices()[0].tolist()
+            
+            return {
+                "indices": indices,
+                "values": values
+            }
+        except Exception as e:
+            logger.error(f"Embedding error: {e}")
+            raise
+        
+    def embed_documents(self, text: List[str], prompt: Optional[str] = None) -> Dict[Any, Any]:
+        """
+        Generate a sparse embedding for a single text.
+        
+        Args:
+            text: Input text
+            prompt: Optional prompt for instruction-based models
+        
+        Returns:
+            Sparse embedding as a dictionary with 'indices' and 'values' keys.
+        """
+        
+        try:
+            tensor = self.model.encode(text)
+            
+            values = tensor[0].coalesce().values().tolist()
+            indices = tensor[0].coalesce().indices()[0].tolist()
+            
+            return {
+                "indices": indices,
+                "values": values
+            }
+            
+
+        except Exception as e:
+            logger.error(f"Embedding error: {e}")
+            raise
+
+    def embed_batch(self, texts: List[str], prompt: Optional[str] = None) -> List[Dict[str, Any]]:
+        """
+        Generate sparse embeddings for a batch of texts.
+        
+        Args:
+            texts: List of input texts
+            prompt: Optional prompt for instruction-based models
+        
+        Returns:
+            List of sparse embeddings as dictionaries with 'text' and 'sparse_embedding' keys.
+        """
+        if not self._loaded:
+            self.load()
+            
+        try:
+            tensors = self.model.encode(texts)
+            results = []
+            
+            for i, tensor in enumerate(tensors):
+                values = tensor.coalesce().values().tolist()
+                indices = tensor.coalesce().indices()[0].tolist()
+                
+                results.append({
+                    "text": texts[i],
+                    "sparse_embedding": {
+                        "indices": indices,
+                        "values": values
+                    }
+                })
+            
+            return results
+        except Exception as e:
+            logger.error(f"Sparse embedding generation failed: {e}")
+            raise