Address review feedback: lazy import, dimension probe, batch pipeline, tag escaping, typed exceptions

Signed-off-by: Daria Korenieva <daric2612@gmail.com>

Author: daric93

application/vectorstore/valkey.py

@@ -6,8 +6,11 @@
 import uuid
 from typing import Any, Dict, List, Optional
 
+_GLIDE_AVAILABLE = False
 try:
     from glide_sync import (
+        Batch,
+        ConnectionError as GlideConnectionError,
         DataType,
         DistanceMetricType,
         Field,
@@ -17,31 +20,33 @@
         GlideClient,
         GlideClientConfiguration,
         NodeAddress,
+        RequestError,
         ReturnField,
         ServerCredentials,
         TagField,
         TextField,
+        TimeoutError as GlideTimeoutError,
         VectorAlgorithm,
         VectorField,
         VectorFieldAttributesFlat,
         VectorFieldAttributesHnsw,
         VectorType,
         ft,
     )
+
+    _GLIDE_AVAILABLE = True
 except ImportError:
-    raise ImportError(
-        "Could not import valkey-glide-sync. "
-        "Please install with `pip install valkey-glide-sync`."
-    )
+    pass
 
-from application.core.settings import settings
-from application.vectorstore.base import BaseVectorStore
-from application.vectorstore.document_class import Document
+from application.core.settings import settings  # noqa: E402
+from application.vectorstore.base import BaseVectorStore  # noqa: E402
+from application.vectorstore.document_class import Document  # noqa: E402
 
 logger = logging.getLogger(__name__)
 
 # Characters that must be escaped in Valkey tag field query values.
-_TAG_SPECIAL_CHARS = set(r".,<>{}[]\"':;!@#$%^&*()-+=~ /|")
+# Includes '?' which is a single-character wildcard in valkey-search TAG queries.
+_TAG_SPECIAL_CHARS = set(r".,<>{}[]\"':;!@#$%^&*()-+=~ /|?")
 
 # Batch size for DELETE operations in delete_index.
 _DELETE_BATCH_SIZE = 100
@@ -70,7 +75,15 @@ def __init__(
             source_id: Identifier for the document source, used to
                 namespace and filter documents.
             embeddings_key: Key name or API key for the embeddings provider.
+
+        Raises:
+            ImportError: If valkey-glide-sync is not installed.
         """
+        if not _GLIDE_AVAILABLE:
+            raise ImportError(
+                "Could not import valkey-glide-sync. "
+                "Please install with `pip install valkey-glide-sync`."
+            )
         super().__init__()
         self._source_id = str(source_id).replace("application/indexes/", "").rstrip("/")
         self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key)
@@ -127,7 +140,10 @@ def _ensure_index_exists(self):
         Uses VALKEY_DISTANCE_METRIC, VALKEY_VECTOR_TYPE, and VALKEY_VECTOR_ALGORITHM
         from settings. Falls back to cosine/float32/hnsw if values are unrecognized.
         """
-        embedding_dim = getattr(self._embedding, "dimension", 768)
+        embedding_dim = getattr(self._embedding, "dimension", None)
+        if embedding_dim is None:
+            probe = self._embedding.embed_query("dimension probe")
+            embedding_dim = len(probe)
 
         distance_metric = self._resolve_distance_metric(settings.VALKEY_DISTANCE_METRIC)
         vector_type = self._resolve_vector_type(settings.VALKEY_VECTOR_TYPE)
@@ -310,7 +326,7 @@ def search(self, question: str, k: int = 2, *args, **kwargs) -> List[Document]:
 
             return self._parse_search_results(results)
 
-        except Exception as e:
+        except (RequestError, GlideConnectionError, GlideTimeoutError) as e:
             logger.error(f"Error searching Valkey: {e}", exc_info=True)
             return []
 
@@ -406,6 +422,8 @@ def add_texts(
         metadatas = metadatas or [{}] * len(texts)
         doc_ids: List[str] = []
 
+        # Use non-atomic Batch (pipeline) to reduce network round trips.
+        batch = Batch(False)
         for text, embedding, metadata in zip(texts, embeddings, metadatas):
             doc_id = str(uuid.uuid4())
             key = self._doc_key(doc_id)
@@ -418,15 +436,17 @@ def add_texts(
                 "embedding": vector_bytes,
             }
 
-            try:
-                self._client.hset(key, fields)
-                doc_ids.append(doc_id)
-            except Exception as e:
-                logger.error(
-                    f"Error adding document to Valkey (wrote {len(doc_ids)}/{len(texts)} "
-                    f"before failure): {e}"
-                )
-                raise
+            batch.hset(key, fields)
+            doc_ids.append(doc_id)
+
+        try:
+            self._client.exec(batch, raise_on_error=True)
+        except (RequestError, GlideConnectionError, GlideTimeoutError) as e:
+            logger.error(
+                f"Error adding documents to Valkey via pipeline "
+                f"({len(doc_ids)} documents): {e}"
+            )
+            raise
 
         return doc_ids
 
@@ -488,7 +508,7 @@ def delete_index(self, *args, **kwargs):
                 batch = keys[i : i + _DELETE_BATCH_SIZE]
                 self._client.delete(batch)
 
-        except Exception as e:
+        except (RequestError, GlideConnectionError, GlideTimeoutError) as e:
             logger.error(f"Error deleting index from Valkey: {e}", exc_info=True)
 
     def save_local(self, *args, **kwargs):
@@ -550,7 +570,7 @@ def get_chunks(self) -> List[Dict[str, Any]]:
 
             return chunks
 
-        except Exception as e:
+        except (RequestError, GlideConnectionError, GlideTimeoutError) as e:
             logger.error(f"Error getting chunks from Valkey: {e}", exc_info=True)
             return []
 
@@ -586,7 +606,7 @@ def add_chunk(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> str
         try:
             self._client.hset(key, fields)
             return doc_id
-        except Exception as e:
+        except (RequestError, GlideConnectionError, GlideTimeoutError) as e:
             logger.error(f"Error adding chunk to Valkey: {e}")
             raise
 
@@ -603,6 +623,6 @@ def delete_chunk(self, chunk_id: str) -> bool:
             key = self._doc_key(chunk_id)
             result = self._client.delete([key])
             return result > 0
-        except Exception as e:
+        except (RequestError, GlideConnectionError, GlideTimeoutError) as e:
             logger.error(f"Error deleting chunk from Valkey: {e}", exc_info=True)
             return False

tests/vectorstore/test_valkey.py

@@ -5,6 +5,8 @@
 
 import pytest
 
+from glide_sync import RequestError
+
 
 def _make_store(source_id="test-source", embeddings_key="key"):
     """Helper to create a ValkeyStore with all external deps mocked."""
@@ -136,7 +138,7 @@ def test_search_returns_documents(self):
     def test_search_returns_empty_on_error(self):
         store, mock_client, _ = _make_store()
 
-        with patch("application.vectorstore.valkey.ft.search", side_effect=Exception("connection lost")):
+        with patch("application.vectorstore.valkey.ft.search", side_effect=RequestError("connection lost")):
             results = store.search("query")
             assert results == []
 
@@ -178,12 +180,12 @@ class TestValkeyStoreAddTexts:
     def test_add_texts_returns_ids(self):
         store, mock_client, mock_emb = _make_store()
         mock_emb.embed_documents.return_value = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
-        mock_client.hset = Mock(return_value=3)
+        mock_client.exec = Mock(return_value=["OK", "OK"])
 
         ids = store.add_texts(["text1", "text2"], [{"a": 1}, {"b": 2}])
 
         assert len(ids) == 2
-        assert mock_client.hset.call_count == 2
+        mock_client.exec.assert_called_once()
 
     def test_add_texts_empty_returns_empty(self):
         store, _, _ = _make_store()
@@ -192,35 +194,43 @@ def test_add_texts_empty_returns_empty(self):
     def test_add_texts_default_metadatas(self):
         store, mock_client, mock_emb = _make_store()
         mock_emb.embed_documents.return_value = [[0.1, 0.2, 0.3]]
-        mock_client.hset = Mock(return_value=3)
+        mock_client.exec = Mock(return_value=["OK"])
 
         ids = store.add_texts(["text1"])
         assert len(ids) == 1
 
     def test_add_texts_raises_on_error(self):
         store, mock_client, mock_emb = _make_store()
         mock_emb.embed_documents.return_value = [[0.1, 0.2, 0.3]]
-        mock_client.hset = Mock(side_effect=Exception("write failed"))
+        mock_client.exec = Mock(side_effect=RequestError("write failed"))
 
-        with pytest.raises(Exception, match="write failed"):
+        with pytest.raises(RequestError, match="write failed"):
             store.add_texts(["text1"])
 
     def test_add_texts_stores_correct_fields(self):
         store, mock_client, mock_emb = _make_store(source_id="src1")
         mock_emb.embed_documents.return_value = [[0.1, 0.2, 0.3]]
-        mock_client.hset = Mock(return_value=3)
+        mock_client.exec = Mock(return_value=["OK"])
 
-        store.add_texts(["hello"], [{"key": "val"}])
+        with patch("application.vectorstore.valkey.Batch") as MockBatch:
+            mock_batch_instance = MagicMock()
+            MockBatch.return_value = mock_batch_instance
 
-        call_args = mock_client.hset.call_args
-        key = call_args[0][0]
-        fields = call_args[0][1]
+            store.add_texts(["hello"], [{"key": "val"}])
+
+            # Verify the batch was created as non-atomic (pipeline)
+            MockBatch.assert_called_once_with(False)
+
+            # Verify hset was called on the batch with correct fields
+            call_args = mock_batch_instance.hset.call_args
+            key = call_args[0][0]
+            fields = call_args[0][1]
 
-        assert key.startswith("doc:")
-        assert fields["content"] == "hello"
-        assert fields["source_id"] == "src1"
-        assert json.loads(fields["metadata"]) == {"key": "val"}
-        assert isinstance(fields["embedding"], bytes)
+            assert key.startswith("doc:")
+            assert fields["content"] == "hello"
+            assert fields["source_id"] == "src1"
+            assert json.loads(fields["metadata"]) == {"key": "val"}
+            assert isinstance(fields["embedding"], bytes)
 
 
 @pytest.mark.unit
@@ -277,7 +287,7 @@ def test_delete_index_paginates_large_sets(self):
     def test_delete_index_handles_error(self):
         store, mock_client, _ = _make_store()
 
-        with patch("application.vectorstore.valkey.ft.search", side_effect=Exception("fail")):
+        with patch("application.vectorstore.valkey.ft.search", side_effect=RequestError("fail")):
             # Should not raise
             store.delete_index()
 
@@ -310,7 +320,7 @@ def test_get_chunks(self):
     def test_get_chunks_returns_empty_on_error(self):
         store, mock_client, _ = _make_store()
 
-        with patch("application.vectorstore.valkey.ft.search", side_effect=Exception("fail")):
+        with patch("application.vectorstore.valkey.ft.search", side_effect=RequestError("fail")):
             assert store.get_chunks() == []
 
     def test_get_chunks_uses_return_fields(self):
@@ -378,7 +388,7 @@ def test_delete_chunk_not_found(self):
 
     def test_delete_chunk_returns_false_on_error(self):
         store, mock_client, _ = _make_store()
-        mock_client.delete = Mock(side_effect=Exception("fail"))
+        mock_client.delete = Mock(side_effect=RequestError("fail"))
 
         result = store.delete_chunk("uuid-123")
         assert result is False