feat: add embedding_params to BasicEmbeddingsIndex

- Added `embedding_params` attribute to `BasicEmbeddingsIndex` class.
- Updated the constructor to accept `embedding_params`.
- Modified `_init_model` method to pass `embedding_params` to `init_embedding_model`.
- Updated `init_embedding_model` function to handle `embedding_params`.
- Adjusted `NIMEmbeddingModel` and `OpenAIEmbeddingModel` to accept additional parameters.
- Updated `LLMRails` to handle default embedding parameters.

improve style

Author: Pouyanpi

nemoguardrails/embeddings/basic.py

@@ -47,6 +47,7 @@ class BasicEmbeddingsIndex(EmbeddingsIndex):
 
     embedding_model: str
     embedding_engine: str
+    embedding_params: Dict[str, Any]
     index: AnnoyIndex
     embedding_size: int
     cache_config: EmbeddingsCacheConfig
@@ -60,6 +61,7 @@ def __init__(
         self,
         embedding_model=None,
         embedding_engine=None,
+        embedding_params=None,
         index=None,
         cache_config: Union[EmbeddingsCacheConfig, Dict[str, Any]] = None,
         search_threshold: float = None,
@@ -83,6 +85,7 @@ def __init__(
         self._embeddings = []
         self.embedding_model = embedding_model
         self.embedding_engine = embedding_engine
+        self.embedding_params = embedding_params or {}
         self._embedding_size = 0
         self.search_threshold = search_threshold or float("inf")
         if isinstance(cache_config, Dict):
@@ -132,7 +135,9 @@ def embeddings_index(self, index):
     def _init_model(self):
         """Initialize the model used for computing the embeddings."""
         self._model = init_embedding_model(
-            embedding_model=self.embedding_model, embedding_engine=self.embedding_engine
+            embedding_model=self.embedding_model,
+            embedding_engine=self.embedding_engine,
+            embedding_params=self.embedding_params,
         )
 
     @cache_embeddings

nemoguardrails/embeddings/providers/__init__.py

@@ -70,12 +70,15 @@ def register_embedding_provider(
 register_embedding_provider(nim.NVIDIAAIEndpointsEmbeddingModel)
 
 
-def init_embedding_model(embedding_model: str, embedding_engine: str) -> EmbeddingModel:
+def init_embedding_model(
+    embedding_model: str, embedding_engine: str, embedding_params: dict = {}
+) -> EmbeddingModel:
     """Initialize the embedding model.
 
     Args:
         embedding_model (str): The path or name of the embedding model.
         embedding_engine (str): The name of the embedding engine.
+        embedding_params (dict): Additional parameters for the embedding model.
 
     Returns:
         EmbeddingModel: An instance of the initialized embedding model.
@@ -84,10 +87,16 @@ def init_embedding_model(embedding_model: str, embedding_engine: str) -> Embeddi
         ValueError: If the embedding engine is invalid.
     """
 
-    model_key = f"{embedding_engine}-{embedding_model}"
+    embedding_params_str = (
+        "_".join([f"{key}={value}" for key, value in embedding_params.items()])
+        or "default"
+    )
+
+    model_key = f"{embedding_engine}-{embedding_model}-{embedding_params_str}"
 
     if model_key not in _embedding_model_cache:
-        model = EmbeddingProviderRegistry().get(embedding_engine)(embedding_model)
+        provider_class = EmbeddingProviderRegistry().get(embedding_engine)
+        model = provider_class(embedding_model=embedding_model, **embedding_params)
         _embedding_model_cache[model_key] = model
 
     return _embedding_model_cache[model_key]

nemoguardrails/embeddings/providers/nim.py

@@ -33,12 +33,12 @@ class NIMEmbeddingModel(EmbeddingModel):
 
     engine_name = "nim"
 
-    def __init__(self, embedding_model: str):
+    def __init__(self, embedding_model: str, **kwargs):
         try:
             from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
 
             self.model = embedding_model
-            self.document_embedder = NVIDIAEmbeddings(model=embedding_model)
+            self.document_embedder = NVIDIAEmbeddings(model=embedding_model, **kwargs)
 
         except ImportError:
             raise ImportError(

nemoguardrails/embeddings/providers/openai.py

@@ -43,6 +43,7 @@ class OpenAIEmbeddingModel(EmbeddingModel):
     def __init__(
         self,
         embedding_model: str,
+        **kwargs,
     ):
         try:
             import openai
@@ -59,7 +60,7 @@ def __init__(
             )
 
         self.model = embedding_model
-        self.client = OpenAI()
+        self.client = OpenAI(**kwargs)
 
         self.embedding_size_dict = {
             "text-embedding-ada-002": 1536,

nemoguardrails/rails/llm/llmrails.py

@@ -103,6 +103,7 @@ def __init__(
         # The default embeddings model is using FastEmbed
         self.default_embedding_model = "all-MiniLM-L6-v2"
         self.default_embedding_engine = "FastEmbed"
+        self.default_embedding_params = {}
 
         # We keep a cache of the events history associated with a sequence of user messages.
         # TODO: when we update the interface to allow to return a "state object", this
@@ -212,6 +213,7 @@ def __init__(
             if model.type == "embeddings":
                 self.default_embedding_model = model.model
                 self.default_embedding_engine = model.engine
+                self.default_embedding_params = model.parameters or {}
                 break
 
         # InteractionLogAdapters used for tracing
@@ -429,6 +431,9 @@ def _get_embeddings_search_provider_instance(
                 embedding_engine=esp_config.parameters.get(
                     "embedding_engine", self.default_embedding_engine
                 ),
+                embedding_params=esp_config.parameters.get(
+                    "embedding_parameters", self.default_embedding_params
+                ),
                 cache_config=esp_config.cache,
                 # We make sure we also pass additional relevant params.
                 **{