fix: fix Pinecone types + add py.typed (#1993)

Author: anakin87

.github/workflows/pinecone.yml

@@ -52,11 +52,9 @@ jobs:
       - name: Install Hatch
         run: pip install --upgrade hatch
 
-    # TODO: Once this integration is properly typed, use hatch run test:types
-    # https://github.com/deepset-ai/haystack-core-integrations/issues/1771
       - name: Lint
         if: matrix.python-version == '3.9' && runner.os == 'Linux'
-        run: hatch run fmt-check && hatch run lint:typing
+        run: hatch run fmt-check && hatch run test:types
 
       - name: Generate docs
         if: matrix.python-version == '3.9' && runner.os == 'Linux'

integrations/pinecone/pyproject.toml

@@ -74,12 +74,14 @@ all = 'pytest {args:tests}'
 # with pytest-xdist (https://pytest-xdist.readthedocs.io/en/stable/distribution.html)
 cov-retry = 'all -n auto --maxprocesses=2 --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
 
-[tool.hatch.envs.lint]
-installer = "uv"
-detached = true
-dependencies = ["pip", "mypy>=1.0.0", "ruff>=0.0.243", "numpy"]
-[tool.hatch.envs.lint.scripts]
-typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+types = """mypy -p haystack_integrations.components.retrievers.pinecone \
+-p haystack_integrations.document_stores.pinecone {args}"""
+
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
 [tool.hatch.metadata]
 allow-direct-references = true
@@ -165,8 +167,4 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-markers = ["unit: unit tests", "integration: integration tests"]
-
-[[tool.mypy.overrides]]
-module = ["pinecone.*", "haystack.*", "haystack_integrations.*", "pytest.*", "pytest_asyncio.*"]
-ignore_missing_imports = true
+markers = ["integration: integration tests"]

integrations/pinecone/src/haystack_integrations/components/retrievers/pinecone/embedding_retriever.py

@@ -116,7 +116,7 @@ def run(
         query_embedding: List[float],
         filters: Optional[Dict[str, Any]] = None,
         top_k: Optional[int] = None,
-    ):
+    ) -> Dict[str, List[Document]]:
         """
         Retrieve documents from the `PineconeDocumentStore`, based on their dense embeddings.
 
@@ -145,7 +145,7 @@ async def run_async(
         query_embedding: List[float],
         filters: Optional[Dict[str, Any]] = None,
         top_k: Optional[int] = None,
-    ):
+    ) -> Dict[str, List[Document]]:
         """
         Asynchronously retrieve documents from the `PineconeDocumentStore`, based on their dense embeddings.
 

integrations/pinecone/src/haystack_integrations/components/retrievers/py.typed

@@ -2,14 +2,15 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 from copy import copy
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional, Union
 
 from haystack import default_from_dict, default_to_dict, logging
 from haystack.dataclasses import Document
 from haystack.document_stores.types import DuplicatePolicy
 from haystack.utils import Secret, deserialize_secrets_inplace
 
 from pinecone import Pinecone, PineconeAsyncio, PodSpec, ServerlessSpec
+from pinecone.db_data import _Index, _IndexAsyncio
 
 from .filters import _normalize_filters, _validate_filters
 
@@ -70,8 +71,8 @@ def __init__(
         self.dimension = dimension
         self.index_name = index
 
-        self._index = None
-        self._async_index = None
+        self._index: Optional[_Index] = None
+        self._async_index: Optional[_IndexAsyncio] = None
         self._dummy_vector = [-10.0] * self.dimension
 
     def _initialize_index(self):
@@ -137,7 +138,7 @@ async def _initialize_async_index(self):
         await async_client.close()
 
     @staticmethod
-    def _convert_dict_spec_to_pinecone_object(spec: Dict[str, Any]):
+    def _convert_dict_spec_to_pinecone_object(spec: Dict[str, Any]) -> Union[ServerlessSpec, PodSpec]:
         """Convert the spec dictionary to a Pinecone spec object"""
 
         if "serverless" in spec:
@@ -447,12 +448,12 @@ def _convert_query_result_to_documents(self, query_result: Dict[str, Any]) -> Li
         return documents
 
     @staticmethod
-    def _discard_invalid_meta(document: Document):
+    def _discard_invalid_meta(document: Document) -> None:
         """
         Remove metadata fields with unsupported types from the document.
         """
 
-        def valid_type(value: Any):
+        def valid_type(value: Any) -> bool:
             return isinstance(value, METADATA_SUPPORTED_TYPES) or (
                 isinstance(value, list) and all(isinstance(i, str) for i in value)
             )
@@ -475,8 +476,6 @@ def valid_type(value: Any):
 
             document.meta = new_meta
 
-        return document
-
     def _convert_documents_to_pinecone_format(self, documents: List[Document]) -> List[Dict[str, Any]]:
         documents_for_pinecone = []
         for document in documents:
@@ -491,7 +490,7 @@ def _convert_documents_to_pinecone_format(self, documents: List[Document]) -> Li
             if document.meta:
                 self._discard_invalid_meta(document)
 
-            doc_for_pinecone = {"id": document.id, "values": embedding, "metadata": dict(document.meta)}
+            doc_for_pinecone: Dict[str, Any] = {"id": document.id, "values": embedding, "metadata": dict(document.meta)}
 
             # we save content as metadata
             if document.content is not None:
@@ -506,10 +505,10 @@ def _convert_documents_to_pinecone_format(self, documents: List[Document]) -> Li
                 )
             if hasattr(document, "sparse_embedding") and document.sparse_embedding is not None:
                 logger.warning(
-                    "Document %s has the `sparse_embedding` field set,"
+                    "Document {document_id} has the `sparse_embedding` field set,"
                     "but storing sparse embeddings in Pinecone is not currently supported."
                     "The `sparse_embedding` field will be ignored.",
-                    document.id,
+                    document_id=document.id,
                 )
 
             documents_for_pinecone.append(doc_for_pinecone)

integrations/pinecone/src/haystack_integrations/document_stores/pinecone/document_store.py

@@ -158,13 +158,13 @@ def test_discard_invalid_meta_invalid():
             ],
         },
     )
-    pinecone_doc = PineconeDocumentStore._discard_invalid_meta(invalid_metadata_doc)
+    PineconeDocumentStore._discard_invalid_meta(invalid_metadata_doc)
 
-    assert pinecone_doc.meta["source_id"] == "62049ba1d1e1d5ebb1f6230b0b00c5356b8706c56e0b9c36b1dfc86084cd75f0"
-    assert pinecone_doc.meta["page_number"] == 1
-    assert pinecone_doc.meta["split_id"] == 0
-    assert pinecone_doc.meta["split_idx_start"] == 0
-    assert "_split_overlap" not in pinecone_doc.meta
+    assert invalid_metadata_doc.meta["source_id"] == "62049ba1d1e1d5ebb1f6230b0b00c5356b8706c56e0b9c36b1dfc86084cd75f0"
+    assert invalid_metadata_doc.meta["page_number"] == 1
+    assert invalid_metadata_doc.meta["split_id"] == 0
+    assert invalid_metadata_doc.meta["split_idx_start"] == 0
+    assert "_split_overlap" not in invalid_metadata_doc.meta
 
 
 def test_discard_invalid_meta_valid():
@@ -175,10 +175,10 @@ def test_discard_invalid_meta_valid():
             "page_number": 1,
         },
     )
-    pinecone_doc = PineconeDocumentStore._discard_invalid_meta(valid_metadata_doc)
+    PineconeDocumentStore._discard_invalid_meta(valid_metadata_doc)
 
-    assert pinecone_doc.meta["source_id"] == "62049ba1d1e1d5ebb1f6230b0b00c5356b8706c56e0b9c36b1dfc86084cd75f0"
-    assert pinecone_doc.meta["page_number"] == 1
+    assert valid_metadata_doc.meta["source_id"] == "62049ba1d1e1d5ebb1f6230b0b00c5356b8706c56e0b9c36b1dfc86084cd75f0"
+    assert valid_metadata_doc.meta["page_number"] == 1
 
 
 def test_convert_meta_to_int():