OraLlamaVS Connection Pool Support + Filtering (#19412)

Author: fileames

docs/docs/examples/vector_stores/orallamavs.ipynb

@@ -305,15 +305,6 @@
     "    - vector_stores (list): A list of OracleVS instances.\n",
     "    \"\"\"\n",
     "    for i, vs in enumerate(vector_stores, start=1):\n",
-    "        # Adding texts\n",
-    "        try:\n",
-    "            vs.add_texts(text_nodes, metadata)\n",
-    "            print(f\"\\n\\n\\nAdd texts complete for vector store {i}\\n\\n\\n\")\n",
-    "        except Exception as ex:\n",
-    "            print(\n",
-    "                f\"\\n\\n\\nExpected error on duplicate add for vector store {i}\\n\\n\\n\"\n",
-    "            )\n",
-    "\n",
     "        # Deleting texts using the value of 'id'\n",
     "        vs.delete(\"test-1\")\n",
     "        print(f\"\\n\\n\\nDelete texts complete for vector store {i}\\n\\n\\n\")\n",

llama-index-integrations/vector_stores/llama-index-vector-stores-oracledb/README.md

@@ -123,34 +123,18 @@ orallamavs.create_index(
 
 print("Index created.")
 
-
 # Perform Semantic Search
-query = "What is Oracle AI Vector Store?"
-filter = {"document_id": ["1"]}
-
+embedding = embedder._get_text_embedding("What is Oracle AI Vector Store?")
+query = VectorStoreQuery(query_embedding=embedding, similarity_top_k=1)
 # Similarity search without a filter
-print(vectorstore.similarity_search(query, 1))
-
-# Similarity search with a filter
-print(vectorstore.similarity_search(query, 1, filter=filter))
+print(vectorstore.query(query))
 
-# Similarity search with relevance score
-print(vectorstore.similarity_search_with_score(query, 1))
-
-# Similarity search with relevance score with filter
-print(vectorstore.similarity_search_with_score(query, 1, filter=filter))
-
-# Max marginal relevance search
-print(
-    vectorstore.max_marginal_relevance_search(
-        query, 1, fetch_k=20, lambda_mult=0.5
-    )
+filters = MetadataFilters(
+    filters=[ExactMatchFilter(key="document_id", value="1")]
 )
-
-# Max marginal relevance search with filter
-print(
-    vectorstore.max_marginal_relevance_search(
-        query, 1, fetch_k=20, lambda_mult=0.5, filter=filter
-    )
+query = VectorStoreQuery(
+    query_embedding=embedding, filters=filters, similarity_top_k=1
 )
+# Similarity search with a filter
+print(vectorstore.query(query))
 ```

llama-index-integrations/vector_stores/llama-index-vector-stores-oracledb/llama_index/vector_stores/oracledb/base.py

@@ -36,6 +36,9 @@
     BasePydanticVectorStore,
     VectorStoreQuery,
     VectorStoreQueryResult,
+    FilterOperator,
+    MetadataFilters,
+    MetadataFilter,
 )
 
 if TYPE_CHECKING:
@@ -86,6 +89,31 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
     return cast(T, wrapper)
 
 
+def _get_connection(client: Any) -> Connection | None:
+    # Dynamically import oracledb and the required classes
+    try:
+        import oracledb
+    except ImportError as e:
+        raise ImportError(
+            "Unable to import oracledb, please install with `pip install -U oracledb`."
+        ) from e
+
+    # check if ConnectionPool exists
+    connection_pool_class = getattr(oracledb, "ConnectionPool", None)
+
+    if isinstance(client, oracledb.Connection):
+        return client
+    elif connection_pool_class and isinstance(client, connection_pool_class):
+        return client.acquire()
+    else:
+        valid_types = "oracledb.Connection"
+        if connection_pool_class:
+            valid_types += " or oracledb.ConnectionPool"
+        raise TypeError(
+            f"Expected client of type {valid_types}, got {type(client).__name__}"
+        )
+
+
 def _escape_str(value: str) -> str:
     BS = "\\"
     must_escape = (BS, "'")
@@ -103,7 +131,7 @@ def _escape_str(value: str) -> str:
     },
     "node_info": {
         "type": "JSON",
-        "extract_func": lambda x: json.dumps(x.node_info),
+        "extract_func": lambda x: json.dumps(x.get_node_info()),
     },
     "metadata": {
         "type": "JSON",
@@ -195,10 +223,11 @@ def _create_table(connection: Connection, table_name: str) -> None:
 
 @_handle_exceptions
 def create_index(
-    connection: Connection,
+    client: Any,
     vector_store: OraLlamaVS,
     params: Optional[dict[str, Any]] = None,
 ) -> None:
+    connection = _get_connection(client)
     if params:
         if params["idx_type"] == "HNSW":
             _create_hnsw_index(
@@ -350,7 +379,8 @@ def _create_ivf_index(
 
 
 @_handle_exceptions
-def drop_table_purge(connection: Connection, table_name: str) -> None:
+def drop_table_purge(client: Any, table_name: str) -> None:
+    connection = _get_connection(client)
     if _table_exists(connection, table_name):
         cursor = connection.cursor()
         with cursor:
@@ -427,9 +457,10 @@ def __init__(
                 batch_size=batch_size,
                 params=params,
             )
+            connection = _get_connection(_client)
             # Assign _client to PrivateAttr after the Pydantic initialization
             object.__setattr__(self, "_client", _client)
-            _create_table(_client, table_name)
+            _create_table(connection, table_name)
 
         except oracledb.DatabaseError as db_err:
             logger.exception(f"Database error occurred while create table: {db_err}")
@@ -456,26 +487,82 @@ def client(self) -> Any:
     def class_name(cls) -> str:
         return "OraLlamaVS"
 
+    def _convert_oper_to_sql(
+        self,
+        oper: FilterOperator,
+        metadata_column: str,
+        filter_key: str,
+        value_bind: str,
+    ) -> str:
+        if oper == FilterOperator.IS_EMPTY:
+            return f"NOT JSON_EXISTS({metadata_column}, '$.{filter_key}') OR JSON_EQUAL(JSON_QUERY({metadata_column}, '$.{filter_key}'), '[]') OR JSON_EQUAL(JSON_QUERY({metadata_column}, '$.{filter_key}'), 'null')"
+        elif oper == FilterOperator.CONTAINS:
+            return f"JSON_EXISTS({metadata_column}, '$.{filter_key}[*]?(@ == $val)' PASSING {value_bind} AS \"val\")"
+        else:
+            oper_map = {
+                FilterOperator.EQ: "{0} = {1}",  # default operator (string, int, float)
+                FilterOperator.GT: "{0} > {1}",  # greater than (int, float)
+                FilterOperator.LT: "{0} < {1}",  # less than (int, float)
+                FilterOperator.NE: "{0} != {1}",  # not equal to (string, int, float)
+                FilterOperator.GTE: "{0} >= {1}",  # greater than or equal to (int, float)
+                FilterOperator.LTE: "{0} <= {1}",  # less than or equal to (int, float)
+                FilterOperator.IN: "{0} IN ({1})",  # In array (string or number)
+                FilterOperator.NIN: "{0} NOT IN ({1})",  # Not in array (string or number)
+                FilterOperator.TEXT_MATCH: "{0} LIKE '%' || {1} || '%'",  # full text match (allows you to search for a specific substring, token or phrase within the text field)
+            }
+
+            if oper not in oper_map:
+                raise ValueError(
+                    f"FilterOperation {oper} cannot be used with this vector store."
+                )
+
+            operation_f = oper_map.get(oper)
+
+            return operation_f.format(
+                f"JSON_VALUE({metadata_column}, '$.{filter_key}')", value_bind
+            )
+
+    def _get_filter_string(
+        self, filter: MetadataFilters | MetadataFilter, bind_variables: list
+    ) -> str:
+        if isinstance(filter, MetadataFilter):
+            if not re.match(r"^[a-zA-Z0-9_]+$", filter.key):
+                raise ValueError(f"Invalid metadata key format: {filter.key}")
+
+            value_bind = f""
+            if filter.operator == FilterOperator.IS_EMPTY:
+                # No values needed
+                pass
+            elif isinstance(filter.value, List):
+                # Needs multiple binds for a list https://python-oracledb.readthedocs.io/en/latest/user_guide/bind.html#binding-multiple-values-to-a-sql-where-in-clause
+                value_binds = []
+                for val in filter.value:
+                    value_binds.append(f":value{len(bind_variables)}")
+                    bind_variables.append(val)
+                value_bind = ",".join(value_binds)
+            else:
+                value_bind = f":value{len(bind_variables)}"
+                bind_variables.append(filter.value)
+
+            return self._convert_oper_to_sql(
+                filter.operator, self.metadata_column, filter.key, value_bind
+            )
+
+        # Combine all sub filters
+        filter_strings = [
+            self._get_filter_string(f_, bind_variables) for f_ in filter.filters
+        ]
+
+        return f" {filter.condition.value.upper()} ".join(filter_strings)
+
     def _append_meta_filter_condition(
-        self, where_str: Optional[str], exact_match_filter: list
+        self, where_str: Optional[str], filters: Optional[MetadataFilters]
     ) -> Tuple[str, list]:
         bind_variables = []
-        filter_conditions = []
-
-        # Validate metadata keys (only allow alphanumeric and underscores)
-        for filter_item in exact_match_filter:
-            # Validate the key - only allow safe characters for JSON path
-            if not re.match(r"^[a-zA-Z0-9_]+$", filter_item.key):
-                raise ValueError(f"Invalid metadata key format: {filter_item.key}")
-            # Use JSON_VALUE with parameterized values
-            filter_conditions.append(
-                f"JSON_VALUE({self.metadata_column}, '$.{filter_item.key}') = :value{len(bind_variables)}"
-            )
-            bind_variables.append(filter_item.value)
 
-        # Convert filter conditions to a single string
-        filter_str = " AND ".join(filter_conditions)
+        filter_str = self._get_filter_string(filters, bind_variables)
 
+        # Convert filter conditions to a single string
         if where_str is None:
             where_str = filter_str
         else:
@@ -534,22 +621,25 @@ def add(self, nodes: list[BaseNode], **kwargs: Any) -> list[str]:
         if not nodes:
             return []
 
+        connection = _get_connection(self._client)
+
         for result_batch in iter_batch(nodes, self.batch_size):
             dml, bind_values = self._build_insert(values=result_batch)
 
-            with self._client.cursor() as cursor:
+            with connection.cursor() as cursor:
                 # Use executemany to insert the batch
                 cursor.executemany(dml, bind_values)
-                self._client.commit()
+                connection.commit()
 
         return [node.node_id for node in nodes]
 
     @_handle_exceptions
     def delete(self, ref_doc_id: str, **kwargs: Any) -> None:
-        with self._client.cursor() as cursor:
+        connection = _get_connection(self._client)
+        with connection.cursor() as cursor:
             ddl = f"DELETE FROM {self.table_name} WHERE doc_id = :ref_doc_id"
             cursor.execute(ddl, [ref_doc_id])
-            self._client.commit()
+            connection.commit()
 
     @_handle_exceptions
     def _get_clob_value(self, result: Any) -> str:
@@ -595,7 +685,7 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul
         bind_vars = []
         if query.filters is not None:
             where_str, bind_vars = self._append_meta_filter_condition(
-                where_str, query.filters.filters
+                where_str, query.filters
             )
 
         # build query sql
@@ -625,7 +715,9 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul
         params = {"embedding": embedding}
         for i, value in enumerate(bind_vars):
             params[f"value{i}"] = value
-        with self._client.cursor() as cursor:
+
+        connection = _get_connection(self._client)
+        with connection.cursor() as cursor:
             cursor.execute(query_sql, **params)
             results = cursor.fetchall()
 

llama-index-integrations/vector_stores/llama-index-vector-stores-oracledb/pyproject.toml

@@ -26,7 +26,7 @@ dev = [
 
 [project]
 name = "llama-index-vector-stores-oracledb"
-version = "0.3.1"
+version = "0.3.2"
 description = "llama-index vector_stores oracle database integration"
 authors = [{name = "Your Name", email = "[email protected]"}]
 requires-python = ">=3.9,<3.13"