Address Traian and Pouyan's feedback on redundant None-checks and defaults in EmbeddingsCacheConfig

Author: tgasser-nv

nemoguardrails/embeddings/basic.py

@@ -15,7 +15,7 @@
 
 import asyncio
 import logging
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, cast
 
 from annoy import AnnoyIndex  # type: ignore
 
@@ -73,8 +73,14 @@ def __init__(
         self._model: Optional[EmbeddingModel] = None
         self._items: List[IndexItem] = []
         self._embeddings: List[List[float]] = []
-        self.embedding_model: Optional[str] = embedding_model
-        self.embedding_engine: Optional[str] = embedding_engine
+        self.embedding_model: str = (
+            embedding_model
+            if embedding_model
+            else "sentence-transformers/all-MiniLM-L6-v2"
+        )
+        self.embedding_engine: str = (
+            embedding_engine if embedding_engine else "SentenceTransformers"
+        )
         self.embedding_params = embedding_params or {}
         self._embedding_size = 0
         self.search_threshold = search_threshold or float("inf")
@@ -124,9 +130,8 @@ def embeddings(self):
 
     def _init_model(self):
         """Initialize the model used for computing the embeddings."""
-        # Provide defaults if not specified
-        model = self.embedding_model or "sentence-transformers/all-MiniLM-L6-v2"
-        engine = self.embedding_engine or "SentenceTransformers"
+        model = self.embedding_model
+        engine = self.embedding_engine
 
         self._model = init_embedding_model(
             embedding_model=model,
@@ -152,9 +157,9 @@ async def _get_embeddings(self, texts: List[str]) -> List[List[float]]:
         if self._model is None:
             self._init_model()
 
-        if not self._model:
-            raise Exception("Couldn't initialize embedding model")
-        embeddings = await self._model.encode_async(texts)
+        # self._model can't be None here, or self._init_model() would throw a ValueError
+        model: EmbeddingModel = cast(EmbeddingModel, self._model)
+        embeddings = await model.encode_async(texts)
         return embeddings
 
     async def add_item(self, item: IndexItem):
@@ -218,7 +223,6 @@ async def _run_batch(self):
         if not self._current_batch_finished_event:
             raise Exception("self._current_batch_finished_event not initialized")
 
-        assert self._current_batch_finished_event is not None
         batch_event: asyncio.Event = self._current_batch_finished_event
         self._current_batch_finished_event = None
 

nemoguardrails/embeddings/cache.py

@@ -244,9 +244,9 @@ def from_config(cls, config: EmbeddingsCacheConfig):
 
     def get_config(self):
         return EmbeddingsCacheConfig(
-            key_generator=self._key_generator.name if self._key_generator else "sha256",
-            store=self._cache_store.name if self._cache_store else "filesystem",
-            store_config=self._store_config,
+            key_generator=self._key_generator.name if self._key_generator else None,
+            store=self._cache_store.name if self._cache_store else None,
+            store_config=self._store_config if self._store_config else None,
         )
 
     @singledispatchmethod

nemoguardrails/rails/llm/config.py

@@ -358,15 +358,15 @@ class EmbeddingsCacheConfig(BaseModel):
         default=False,
         description="Whether caching of the embeddings should be enabled or not.",
     )
-    key_generator: str = Field(
+    key_generator: Optional[str] = Field(
         default="sha256",
         description="The method to use for generating the cache keys.",
     )
-    store: str = Field(
+    store: Optional[str] = Field(
         default="filesystem",
         description="What type of store to use for the cached embeddings.",
     )
-    store_config: Dict[str, Any] = Field(
+    store_config: Optional[Dict[str, Any]] = Field(
         default_factory=dict,
         description="Any additional configuration options required for the store. "
         "For example, path for `filesystem` or `host`/`port`/`db` for redis.",