refac

Co-Authored-By: Classic298 <27028174+Classic298@users.noreply.github.com>

backend/open_webui/retrieval/vector/dbs/milvus.py

@@ -3,7 +3,8 @@ NOTE: This vector database integration is community-supported and maintained on
 """
 
 import logging
-from typing import Optional
+import re
+from typing import Any, Optional
 
 from open_webui.config import (
     MILVUS_DB,
@@ -23,7 +24,7 @@ from open_webui.retrieval.vector.main import (
     VectorDBBase,
     VectorItem,
 )
-from open_webui.retrieval.vector.utils import process_metadata
+from open_webui.retrieval.vector.utils import iter_filter_conditions, process_metadata
 from open_webui.utils.json_codec import JSONCodec
 from pymilvus import DataType
 from pymilvus import MilvusClient as Client
@@ -35,6 +36,36 @@ log = logging.getLogger(__name__)
 # field). Clamp long chunks before insert so one oversized chunk can't fail the
 # whole batch and leave the file with zero embeddings.
 MILVUS_TEXT_MAX_LENGTH = 65535
+_SAFE_METADATA_KEY_RE = re.compile(r'^[A-Za-z_][A-Za-z0-9_]{0,63}$')
+
+
+def _escape_milvus_string(value: str) -> str:
+    if not isinstance(value, str):
+        raise TypeError(f'Expected str, got {type(value).__name__}')
+    return value.replace('\\', '\\\\').replace("'", "\\'")
+
+
+def _milvus_literal(value: Any) -> str:
+    if isinstance(value, str):
+        return f"'{_escape_milvus_string(value)}'"
+    if isinstance(value, bool):
+        return str(value).lower()
+    if isinstance(value, (int, float)):
+        return str(value)
+    raise TypeError(f'Unsupported Milvus filter value type: {type(value).__name__}')
+
+
+def _metadata_exprs(filter: Optional[dict]) -> list[str]:
+    exprs = []
+    for key, op, value in iter_filter_conditions(filter):
+        if not isinstance(key, str) or not _SAFE_METADATA_KEY_RE.fullmatch(key):
+            raise ValueError(f'Invalid Milvus metadata filter key: {key!r}')
+        if op == '$in':
+            items = [f"metadata['{key}'] == {_milvus_literal(item)}" for item in value]
+            exprs.append(f'({" or ".join(items)})' if items else 'false')
+        else:
+            exprs.append(f"metadata['{key}'] == {_milvus_literal(value)}")
+    return exprs
 
 
 class MilvusClient(VectorDBBase):
@@ -193,6 +224,9 @@ class MilvusClient(VectorDBBase):
     ) -> Optional[SearchResult]:
         # Search for the nearest neighbor items based on the vectors and return 'limit' number of results.
         collection_name = collection_name.replace('-', '_')
+        kwargs = {}
+        if filter:
+            kwargs['filter'] = ' and '.join(_metadata_exprs(filter))
         # For some index types like IVF_FLAT, search params like nprobe can be set.
         # Example: search_params = {"nprobe": 10} if using IVF_FLAT
         # For simplicity, not adding configurable search_params here, but could be extended.
@@ -201,6 +235,7 @@ class MilvusClient(VectorDBBase):
             data=vectors,
             limit=limit,
             output_fields=['data', 'metadata'],
+            **kwargs,
             # search_params=search_params # Potentially add later if needed
         )
         return self._result_to_search_result(result)
@@ -364,7 +399,9 @@ class MilvusClient(VectorDBBase):
                 ids=ids,
             )
         elif filter:
-            filter_string = ' && '.join([f'metadata["{key}"] == {JSONCodec.dumps(value)}' for key, value in filter.items()])
+            filter_string = ' && '.join(
+                [f'metadata["{key}"] == {JSONCodec.dumps(value)}' for key, value in filter.items()]
+            )
             log.info(
                 'Deleting items by filter from %s_%s. Filter: %s',
                 self.collection_prefix,