Enhancement: initialize DB2VS with connection_args. (#87)

* Enhancement: initialize DB2VS with connection_args.

* Fix format and issues.

* Remove output.

* Improve the code.

* Resolve comments.

* Fixed some problems.

* Change exception type.

Author: yiweiHeOSS

libs/langchain-db2/langchain_db2/db2vs.py

@@ -187,19 +187,39 @@ class DB2VS(VectorStore):
 
     def __init__(
         self,
-        client: Connection,
         embedding_function: Union[
             Callable[[str], List[float]],
             Embeddings,
         ],
         table_name: str,
+        client: Optional[Connection] = None,
         distance_strategy: DistanceStrategy = DistanceStrategy.EUCLIDEAN_DISTANCE,
         query: Optional[str] = "What is a Db2 database",
         params: Optional[Dict[str, Any]] = None,
+        connection_args: Optional[Dict[str, Any]] = None,
     ):
-        try:
+        if client is None:
+            if connection_args is not None:
+                database = connection_args.get("database")
+                host = connection_args.get("host")
+                port = connection_args.get("port")
+                username = connection_args.get("username")
+                password = connection_args.get("password")
+
+                conn_str = f"DATABASE={database};hostname={host};port={port};"
+                f"uid={username};pwd={password};"
+
+                if "security" in connection_args:
+                    security = connection_args.get("security")
+                    conn_str += f"security={security};"
+
+                self.client = ibm_db_dbi.connect(conn_str, "", "")
+            else:
+                raise ValueError("No valid connection or connection_args is passed")
+        else:
             """Initialize with ibm_db_dbi client."""
             self.client = client
+        try:
             """Initialize with necessary components."""
             if not isinstance(embedding_function, Embeddings):
                 logger.warning(

libs/langchain-db2/tests/integration_tests/test_db2vs.py

@@ -245,7 +245,7 @@ def test_add_texts_test() -> None:
         {"id": "101", "link": "Document Example Test 2"},
     ]
     model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    vs_obj = DB2VS(connection, model, "TB1", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB1", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     vs_obj.add_texts(texts, metadata)
     drop_table(connection, "TB1")
 
@@ -256,7 +256,7 @@ def test_add_texts_test() -> None:
         {"link": "Document Example Test 2"},
     ]
     model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    vs_obj = DB2VS(connection, model, "TB2", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB2", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     vs_obj.add_texts(texts, metadataNoID)
     drop_table(connection, "TB2")
 
@@ -268,14 +268,14 @@ def test_add_texts_test() -> None:
         {"link": "Document Example Test 2"},
         {"link": "Document Example Test 3"},
     ]
-    vs_obj = DB2VS(connection, model, "TB2", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB2", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     vs_obj.add_texts(texts1, metadataPartialID)
     drop_table(connection, "TB2")
 
     # 3. Add record but neither metadata nor ids are there
     # Expectation: Successful, new ID will be generated
     model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    vs_obj = DB2VS(connection, model, "TB3", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB3", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     texts2 = ["Sam", "John"]
     vs_obj.add_texts(texts2)
     drop_table(connection, "TB3")
@@ -291,17 +291,17 @@ def test_add_texts_test() -> None:
     # Successful
     # Successful
 
-    vs_obj = DB2VS(connection, model, "TB4", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB4", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     ids4 = ["114", "124"]
     vs_obj.add_texts(texts2, ids=ids4)
     drop_table(connection, "TB4")
 
-    vs_obj = DB2VS(connection, model, "TB5", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB5", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     ids5 = ["", "134"]
     vs_obj.add_texts(texts2, ids=ids5)
     drop_table(connection, "TB5")
 
-    vs_obj = DB2VS(connection, model, "TB6", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB6", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     ids6 = [
         """Good afternoon
     my friends""",
@@ -310,15 +310,15 @@ def test_add_texts_test() -> None:
     vs_obj.add_texts(texts2, ids=ids6)
     drop_table(connection, "TB6")
 
-    vs_obj = DB2VS(connection, model, "TB7", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB7", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     ids7 = ['"Good afternoon"', '"India"']
     vs_obj.add_texts(texts2, ids=ids7)
     drop_table(connection, "TB7")
 
     # 5. Add record with ids option but the id are duplicated
     # Expectations: SQL0803N having duplicate values for the index key
     try:
-        vs_obj = DB2VS(connection, model, "TB8", DistanceStrategy.EUCLIDEAN_DISTANCE)
+        vs_obj = DB2VS(model, "TB8", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
         ids8 = ["118", "118"]
         vs_obj.add_texts(texts2, ids=ids8)
         drop_table(connection, "TB8")
@@ -327,7 +327,7 @@ def test_add_texts_test() -> None:
 
     # 6. Add records with both ids and metadatas
     # Expectation: Successful, the ID will be generated based on ids
-    vs_obj = DB2VS(connection, model, "TB9", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB9", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     texts3 = ["Sam 6", "John 6"]
     ids9 = ["1", "2"]
     metadata = [
@@ -340,7 +340,7 @@ def test_add_texts_test() -> None:
     # This one may run slow before using executemany() <<<<<<<<<<
     # 7. Add 10000 records
     # Expectation:Successful
-    vs_obj = DB2VS(connection, model, "TB10", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB10", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     texts4 = ["Sam{0}".format(i) for i in range(1, 10000)]
     ids10 = ["Hello{0}".format(i) for i in range(1, 10000)]
     vs_obj.add_texts(texts4, ids=ids10)
@@ -352,7 +352,7 @@ def add(val: str) -> None:
         model = HuggingFaceEmbeddings(
             model_name="sentence-transformers/all-mpnet-base-v2"
         )
-        vs_obj = DB2VS(connection, model, "TB11", DistanceStrategy.EUCLIDEAN_DISTANCE)
+        vs_obj = DB2VS(model, "TB11", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
         texts5 = [val]
         ids11 = texts5
         vs_obj.add_texts(texts5, ids=ids11)
@@ -371,7 +371,7 @@ def add1(val: str) -> None:
         model = HuggingFaceEmbeddings(
             model_name="sentence-transformers/all-mpnet-base-v2"
         )
-        vs_obj = DB2VS(connection, model, "TB12", DistanceStrategy.EUCLIDEAN_DISTANCE)
+        vs_obj = DB2VS(model, "TB12", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
         texts = [val]
         ids12 = texts
         vs_obj.add_texts(texts, ids=ids12)
@@ -406,7 +406,7 @@ def test_embed_documents_test() -> None:
     # 1. Embed String Example-'Sam'
     # Expectation: Successful.
     model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    vs_obj = DB2VS(connection, model, "TB7", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB7", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     vs_obj._embed_documents(
         [
             "Sam",
@@ -438,7 +438,7 @@ def test_embed_query_test() -> None:
     # 1. Embed String
     # Expectation: Successful.
     model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    vs_obj = DB2VS(connection, model, "TB8", DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_obj = DB2VS(model, "TB8", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
     vs_obj._embed_query("Sam")
 
     # 2. Embed Empty string
@@ -466,9 +466,9 @@ def test_perform_search_test() -> None:
     model1 = HuggingFaceEmbeddings(
         model_name="sentence-transformers/paraphrase-mpnet-base-v2"
     )
-    vs_1 = DB2VS(connection, model1, "TB10", DistanceStrategy.EUCLIDEAN_DISTANCE)
-    vs_2 = DB2VS(connection, model1, "TB11", DistanceStrategy.DOT_PRODUCT)
-    vs_3 = DB2VS(connection, model1, "TB12", DistanceStrategy.COSINE)
+    vs_1 = DB2VS(model1, "TB10", connection, DistanceStrategy.EUCLIDEAN_DISTANCE)
+    vs_2 = DB2VS(model1, "TB11", connection, DistanceStrategy.DOT_PRODUCT)
+    vs_3 = DB2VS(model1, "TB12", connection, DistanceStrategy.COSINE)
 
     # vector store lists:
     vs_list = [vs_1, vs_2, vs_3]
@@ -535,7 +535,7 @@ def test_get_pks() -> None:
 
     table_name = f"Unique_table_{int(time.time())}"
 
-    db2vs = DB2VS(client=connection, embedding_function=model, table_name=table_name)
+    db2vs = DB2VS(embedding_function=model, table_name=table_name, client=connection)
     pks = db2vs.get_pks()
 
     assert isinstance(pks, list)

libs/langchain-db2/tests/unit_tests/test_db2vs.py

@@ -10,7 +10,7 @@ def test_init() -> None:
     client = MagicMock()
     embedding = DeterministicFakeEmbedding(size=100)
     table_name = "foo"
-    db2vs = DB2VS(client, embedding, table_name)
+    db2vs = DB2VS(embedding, table_name, client)
     assert db2vs is not None
     assert isinstance(db2vs, DB2VS)
     assert len(client.mock_calls) == 3