Add OracleEmbeddings optimization

Author: fileames

libs/oracledb/langchain_oracledb/vectorstores/oraclevs.py

@@ -48,6 +48,8 @@
 from langchain_core.embeddings import Embeddings
 from langchain_core.vectorstores import VectorStore
 
+from ..embeddings import OracleEmbeddings
+
 logger = logging.getLogger(__name__)
 log_level = os.getenv("LOG_LEVEL", "ERROR").upper()
 logging.basicConfig(
@@ -862,9 +864,9 @@ def _get_similarity_search_query(
     k: int,
     db_filter: Optional[FilterGroup | FilterCondition] = None,
     return_embeddings: bool = False,
-) -> str:
+) -> Tuple[str, list[str]]:
     where_clause = ""
-    bind_variables = []
+    bind_variables: list[str] = []
     if db_filter:
         where_clause = _generate_where_clause(db_filter, bind_variables)
 
@@ -1142,33 +1144,61 @@ def add_texts(
         texts = list(texts)
         processed_ids = get_processed_ids(texts, metadatas, ids)
 
-        embeddings = self._embed_documents(texts)
         if not metadatas:
             metadatas = [{} for _ in texts]
 
-        docs: List[Tuple[Any, Any, Any, Any]] = [
-            (
-                id_,
-                array.array("f", embedding),
-                metadata,
-                text,
-            )
-            for id_, embedding, metadata, text in zip(
-                processed_ids, embeddings, metadatas, texts
-            )
-        ]
+        docs: Any
+        if not isinstance(self.embeddings, OracleEmbeddings):
+            embeddings = self._embed_documents(texts)
+
+            docs = [
+                (
+                    id_,
+                    array.array("f", embedding),
+                    metadata,
+                    text,
+                )
+                for id_, embedding, metadata, text in zip(
+                    processed_ids, embeddings, metadatas, texts
+                )
+            ]
+        else:
+            docs = list(zip(processed_ids, metadatas, texts))
 
         connection = _get_connection(self.client)
         if connection is None:
             raise ValueError("Failed to acquire a connection.")
         with connection.cursor() as cursor:
-            cursor.setinputsizes(None, None, oracledb.DB_TYPE_JSON, None)
-            cursor.executemany(
-                f"INSERT INTO {self.table_name} (id, embedding, metadata, "
-                f"text) VALUES (:1, :2, :3, :4)",
-                docs,
-            )
-            connection.commit()
+            if not isinstance(self.embeddings, OracleEmbeddings):
+                cursor.setinputsizes(None, None, oracledb.DB_TYPE_JSON, None)
+                cursor.executemany(
+                    f"INSERT INTO {self.table_name} (id, embedding, metadata, "
+                    f"text) VALUES (:1, :2, :3, :4)",
+                    docs,
+                )
+                connection.commit()
+            else:
+                if self.embeddings.proxy:
+                    cursor.execute(
+                        "begin utl_http.set_proxy(:proxy); end;",
+                        proxy=self.embeddings.proxy,
+                    )
+
+                cursor.setinputsizes(None, oracledb.DB_TYPE_JSON, None)
+                cursor.executemany(
+                    f"INSERT INTO {self.table_name} (id, metadata, "
+                    f"text) VALUES (:1, :2, :3)",
+                    docs,
+                )
+
+                cursor.setinputsizes(oracledb.DB_TYPE_JSON)
+                update_sql = (
+                    f"UPDATE {self.table_name} "
+                    "SET embedding = dbms_vector_chain.utl_to_embedding(text, json(:1))"
+                )
+                cursor.execute(update_sql, [self.embeddings.params])
+                connection.commit()
+
         return processed_ids
 
     @_ahandle_exceptions
@@ -1192,33 +1222,60 @@ async def aadd_texts(
         texts = list(texts)
         processed_ids = get_processed_ids(texts, metadatas, ids)
 
-        embeddings = await self._aembed_documents(texts)
         if not metadatas:
             metadatas = [{} for _ in texts]
 
-        docs: List[Tuple[Any, Any, Any, Any]] = [
-            (
-                id_,
-                array.array("f", embedding),
-                metadata,
-                text,
-            )
-            for id_, embedding, metadata, text in zip(
-                processed_ids, embeddings, metadatas, texts
-            )
-        ]
+        docs: Any
+        if not isinstance(self.embeddings, OracleEmbeddings):
+            embeddings = await self._aembed_documents(texts)
+
+            docs = [
+                (
+                    id_,
+                    array.array("f", embedding),
+                    metadata,
+                    text,
+                )
+                for id_, embedding, metadata, text in zip(
+                    processed_ids, embeddings, metadatas, texts
+                )
+            ]
+        else:
+            docs = list(zip(processed_ids, metadatas, texts))
 
         async def context(connection: Any) -> None:
             if connection is None:
                 raise ValueError("Failed to acquire a connection.")
             with connection.cursor() as cursor:
-                cursor.setinputsizes(None, None, oracledb.DB_TYPE_JSON, None)
-                await cursor.executemany(
-                    f"INSERT INTO {self.table_name} (id, embedding, metadata, "
-                    f"text) VALUES (:1, :2, :3, :4)",
-                    docs,
-                )
-                await connection.commit()
+                if not isinstance(self.embeddings, OracleEmbeddings):
+                    cursor.setinputsizes(None, None, oracledb.DB_TYPE_JSON, None)
+                    await cursor.executemany(
+                        f"INSERT INTO {self.table_name} (id, embedding, metadata, "
+                        f"text) VALUES (:1, :2, :3, :4)",
+                        docs,
+                    )
+                    await connection.commit()
+                else:
+                    if self.embeddings.proxy:
+                        await cursor.execute(
+                            "begin utl_http.set_proxy(:proxy); end;",
+                            proxy=self.embeddings.proxy,
+                        )
+
+                    cursor.setinputsizes(None, oracledb.DB_TYPE_JSON, None)
+                    await cursor.executemany(
+                        f"INSERT INTO {self.table_name} (id, metadata, "
+                        f"text) VALUES (:1, :2, :3)",
+                        docs,
+                    )
+
+                    cursor.setinputsizes(oracledb.DB_TYPE_JSON)
+                    update_sql = (
+                        f"UPDATE {self.table_name} "
+                        "SET embedding = dbms_vector_chain.utl_to_embedding(text, json(:1))"  # noqa: E501
+                    )
+                    await cursor.execute(update_sql, [self.embeddings.params])
+                    await connection.commit()
 
         await _handle_context(self.client, context)
 
@@ -1628,7 +1685,6 @@ async def amax_marginal_relevance_search_with_score_by_vector(
     def max_marginal_relevance_search_by_vector(
         self,
         embedding: List[float],
-        *,
         k: int = 4,
         fetch_k: int = 20,
         lambda_mult: float = 0.5,

libs/oracledb/tests/integration_tests/vectorstores/test_oraclevs.py

@@ -18,6 +18,7 @@
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores.utils import DistanceStrategy
 
+from langchain_oracledb.embeddings import OracleEmbeddings
 from langchain_oracledb.vectorstores.oraclevs import (
     FilterCondition,
     FilterGroup,
@@ -2479,6 +2480,80 @@ async def test_index_table_case_async(caplog: pytest.LogCaptureFixture) -> None:
     await adrop_table_purge(connection, "TB2")
 
 
+##################################
+##### test_oracle_embeddings  ####
+##################################
+
+
+def test_oracle_embeddings() -> None:
+    try:
+        connection = oracledb.connect(user=username, password=password, dsn=dsn)
+    except Exception:
+        sys.exit(1)
+
+    drop_table_purge(connection, "TB1")
+
+    texts = ["Database Document", "Code Document"]
+    metadata = [
+        {"id": "100", "link": "Document Example Test 1"},
+        {"id": "101", "link": "Document Example Test 2"},
+    ]
+    embedder_params = {"provider": "database", "model": "allminilm"}
+    proxy = ""
+
+    # instance
+    model = OracleEmbeddings(conn=connection, params=embedder_params, proxy=proxy)
+
+    vs_obj = OracleVS(connection, model, "TB1", DistanceStrategy.EUCLIDEAN_DISTANCE)
+
+    vs_obj.add_texts(texts, metadata)
+    res = vs_obj.similarity_search("database", 1)
+
+    assert "Database" in res[0].page_content
+
+    drop_table_purge(connection, "TB1")
+
+    connection.close()
+
+
+@pytest.mark.asyncio
+async def test_oracle_embeddings_async(caplog: pytest.LogCaptureFixture) -> None:
+    try:
+        connection = await oracledb.connect_async(
+            user=username, password=password, dsn=dsn
+        )
+
+        connection_sync = oracledb.connect(user=username, password=password, dsn=dsn)
+    except Exception:
+        sys.exit(1)
+
+    await adrop_table_purge(connection, "TB1")
+
+    texts = ["Database Document", "Code Document"]
+    metadata = [
+        {"id": "100", "link": "Document Example Test 1"},
+        {"id": "101", "link": "Document Example Test 2"},
+    ]
+    embedder_params = {"provider": "database", "model": "allminilm"}
+    proxy = ""
+
+    # instance
+    model = OracleEmbeddings(conn=connection_sync, params=embedder_params, proxy=proxy)
+
+    vs_obj = await OracleVS.acreate(
+        connection, model, "TB1", DistanceStrategy.EUCLIDEAN_DISTANCE
+    )
+
+    await vs_obj.aadd_texts(texts, metadata)
+    res = await vs_obj.asimilarity_search("database", 1)
+
+    assert "Database" in res[0].page_content
+
+    await adrop_table_purge(connection, "TB1")
+
+    await connection.close()
+
+
 ##################################
 ##### test_quote_identifier  #####
 ##################################