feat/support passing in metadata to write to pinecone payload (#156)

* support passing in metadata to write to pinecone payload

* add test

* lint

* lint

* tidy

* fix: force help text to be str

* fix help text

* pass in str of tuple into field constructor

* skip octoai test due to service not availaible

---------

Co-authored-by: Yao You <[email protected]>

Author: rbiseck3

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.0.24
+
+### Enhancements
+
+* **Support dynamic metadata mapping in Pinecone uploader**
+
 ## 0.0.23
 
 ### Fixes

test/processes/connectors/test_pinecone.py

@@ -0,0 +1,44 @@
+import pytest
+
+from unstructured_ingest.v2.processes.connectors.pinecone import (
+    PineconeUploadStager,
+    PineconeUploadStagerConfig,
+)
+
+
+@pytest.fixture
+def test_element_dict():
+    return {
+        "embeddings": [0, 1],
+        "text": "test dict",
+        "metadata": {
+            "text_as_html": "text as html",
+            "foo": "foo",
+        },
+    }
+
+
+@pytest.mark.parametrize(
+    ("metadata_fields", "expected_to_exist", "not_expected_to_exist"),
+    [
+        (None, ["text_as_html"], ["foo"]),
+        (("foo",), ["foo"], ["text_as_html"]),
+    ],
+)
+def test_conform_dict(
+    monkeypatch, test_element_dict, metadata_fields, expected_to_exist, not_expected_to_exist
+):
+    if metadata_fields is not None:
+        stager = PineconeUploadStager(
+            upload_stager_config=PineconeUploadStagerConfig(metadata_fields=metadata_fields)
+        )
+    else:
+        stager = PineconeUploadStager()
+    results = stager.conform_dict(test_element_dict.copy())
+    results.pop("id")
+    assert test_element_dict["embeddings"] == results.pop("values")
+
+    assert all(
+        results["metadata"][key] == test_element_dict["metadata"][key] for key in expected_to_exist
+    )
+    assert all(key not in results["metadata"] for key in not_expected_to_exist)

test_e2e/test-src.sh

@@ -63,7 +63,8 @@ all_tests=(
   'hubspot.sh'
   'local-embed.sh'
   'local-embed-bedrock.sh'
-  'local-embed-octoai.sh'
+  # NOTE (yao): octoai url is giving 404
+  # 'local-embed-octoai.sh'
   'local-embed-vertexai.sh'
   'local-embed-voyageai.sh'
   'local-embed-mixedbreadai.sh'

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.0.23"  # pragma: no cover
+__version__ = "0.0.24"  # pragma: no cover

unstructured_ingest/v2/cli/utils/model_conversion.py

@@ -155,14 +155,14 @@ def _get_type_from_field(field: FieldInfo) -> click.ParamType:
 
 def get_option_from_field(option_name: str, field_info: FieldInfo) -> Option:
     param_decls = [option_name]
-    help = field_info.description or ""
+    help_text = field_info.description or ""
     if examples := field_info.examples:
-        help += f" [Examples: {', '.join(examples)}]"
+        help_text += f" [Examples: {', '.join(examples)}]"
     option_kwargs = {
         "type": _get_type_from_field(field_info),
         "default": get_default_value_from_field(field_info),
         "required": field_info.is_required(),
-        "help": help,
+        "help": str(help_text),
         "is_flag": is_boolean_flag(field_info),
         "show_default": field_info.default is not PydanticUndefined,
     }

unstructured_ingest/v2/processes/connectors/pinecone.py

@@ -58,20 +58,6 @@ def get_index(self, **index_kwargs) -> "PineconeIndex":
         return index
 
 
-class PineconeUploadStagerConfig(UploadStagerConfig):
-    pass
-
-
-class PineconeUploaderConfig(UploaderConfig):
-    batch_size: Optional[int] = Field(
-        default=None,
-        description="Optional number of records per batch. Will otherwise limit by size.",
-    )
-    pool_threads: Optional[int] = Field(
-        default=1, description="Optional limit on number of threads to use for upload"
-    )
-
-
 ALLOWED_FIELDS = (
     "element_id",
     "text",
@@ -86,31 +72,56 @@ class PineconeUploaderConfig(UploaderConfig):
     "is_continuation",
     "link_urls",
     "link_texts",
+    "text_as_html",
 )
 
 
+class PineconeUploadStagerConfig(UploadStagerConfig):
+    metadata_fields: list[str] = Field(
+        default=str(ALLOWED_FIELDS),
+        description=(
+            "which metadata from the source element to map to the payload metadata being sent to "
+            "Pinecone."
+        ),
+    )
+
+
+class PineconeUploaderConfig(UploaderConfig):
+    batch_size: Optional[int] = Field(
+        default=None,
+        description="Optional number of records per batch. Will otherwise limit by size.",
+    )
+    pool_threads: Optional[int] = Field(
+        default=1, description="Optional limit on number of threads to use for upload"
+    )
+
+
 @dataclass
 class PineconeUploadStager(UploadStager):
     upload_stager_config: PineconeUploadStagerConfig = field(
         default_factory=lambda: PineconeUploadStagerConfig()
     )
 
-    @staticmethod
-    def conform_dict(element_dict: dict) -> dict:
+    def conform_dict(self, element_dict: dict) -> dict:
         embeddings = element_dict.pop("embeddings", None)
         metadata: dict[str, Any] = element_dict.pop("metadata", {})
         data_source = metadata.pop("data_source", {})
         coordinates = metadata.pop("coordinates", {})
-
-        element_dict.update(metadata)
-        element_dict.update(data_source)
-        element_dict.update(coordinates)
+        pinecone_metadata = {}
+        for possible_meta in [element_dict, metadata, data_source, coordinates]:
+            pinecone_metadata.update(
+                {
+                    k: v
+                    for k, v in possible_meta.items()
+                    if k in self.upload_stager_config.metadata_fields
+                }
+            )
 
         return {
             "id": str(uuid.uuid4()),
             "values": embeddings,
             "metadata": flatten_dict(
-                {k: v for k, v in element_dict.items() if k in ALLOWED_FIELDS},
+                pinecone_metadata,
                 separator="-",
                 flatten_lists=True,
                 remove_none=True,