Functions to simplify Pinecone 3x client API parameters with Domino (#…

…139)

* Added functions to create Pinecone3x API call parameters

* Added more test cases

* Updated the comment

* fixed the style error

* fixed the build error

* updated data-services submodule

domino_data/vectordb.py

@@ -1,3 +1,5 @@
+from typing import Dict
+
 import os
 
 _import_error_message = (
@@ -14,6 +16,9 @@
     else:
         raise
 
+HEADER_DOMINO_DATASOURCE = "X-Domino-Datasource"
+HEADER_PINECONE_INDEX = "X-Domino-Pinecone-Index"
+
 
 class DominoPineconeConfiguration(OpenApiConfiguration):
     def __init__(
@@ -51,8 +56,43 @@ def __init__(
         )
 
         self.proxy = os.getenv("DOMINO_DATA_API_GATEWAY", "http://127.0.0.1:8766")
-        self.proxy_headers = {"X-Domino-Datasource": datasource}
+        self.proxy_headers = {HEADER_DOMINO_DATASOURCE: datasource}
 
     def get_host_from_settings(self, index, variables=None, servers=None):
         url = super().get_host_from_settings(index, variables, servers)
         return url.replace("https://", "http://")
+
+
+def domino_pinecone3x_init_params(datasource_name: str) -> Dict[str, str]:
+    """Wrap the parameters to initialize a Pinecone 3.x client
+
+    Args:
+        datasource_name: the name of the Pinecone data source
+
+    Returns:
+        A dictionary of parameters to initialize the Pinecone client
+    """
+    return {
+        "api_key": "domino",
+        "host": os.getenv("DOMINO_DATA_API_GATEWAY", "http://127.0.0.1:8766"),
+        "additional_headers": {HEADER_DOMINO_DATASOURCE: datasource_name},
+    }
+
+
+def domino_pinecone3x_index_params(datasource_name: str, index_name: str) -> Dict[str, str]:
+    """Wrap the parameters to target an index in the Pinecone 3.x client
+
+    Args:
+        datasource_name: the name of the Pinecone data source
+        index_name: the name of the index
+
+    Returns:
+        A dictionary of parameters to target the Pinecone index for vector operations
+    """
+    return {
+        "host": os.getenv("DOMINO_DATA_API_GATEWAY", "http://127.0.0.1:8766"),
+        "additional_headers": {
+            HEADER_DOMINO_DATASOURCE: datasource_name,
+            HEADER_PINECONE_INDEX: index_name,
+        },
+    }

tests/test_vectordb.py

@@ -1,7 +1,13 @@
 """Test Domino customized Pinecone OpenAPI Configuration"""
 import os
 
-from domino_data.vectordb import DominoPineconeConfiguration
+from domino_data.vectordb import (
+    HEADER_DOMINO_DATASOURCE,
+    HEADER_PINECONE_INDEX,
+    DominoPineconeConfiguration,
+    domino_pinecone3x_index_params,
+    domino_pinecone3x_init_params,
+)
 
 
 def test_get_host_from_setting():
@@ -34,3 +40,22 @@ def test_get_host_from_setting():
 
     host = domino_conf.get_host_from_settings(None, variables={})
     assert host == f"http://unknown-unknown.svc.unknown.pinecone.io"
+
+
+def test_get_domino_pinecone3x_init_params():
+    test_data_source = "test_pine_cone"
+    os.environ["DOMINO_DATA_API_GATEWAY"] = "http://127.0.0.1:8766"
+    init_params = domino_pinecone3x_init_params(test_data_source)
+    assert init_params["host"] == "http://127.0.0.1:8766"
+    assert init_params["api_key"] == "domino"
+    assert init_params["additional_headers"][HEADER_DOMINO_DATASOURCE] == test_data_source
+
+
+def test_get_domino_pinecone3x_index_params():
+    test_data_source = "test_pine_cone"
+    test_index_name = "semantic-search-fast"
+    os.environ["DOMINO_DATA_API_GATEWAY"] = "http://127.0.0.1:8766"
+    init_params = domino_pinecone3x_index_params(test_data_source, test_index_name)
+    assert init_params["host"] == "http://127.0.0.1:8766"
+    assert init_params["additional_headers"][HEADER_DOMINO_DATASOURCE] == test_data_source
+    assert init_params["additional_headers"][HEADER_PINECONE_INDEX] == test_index_name