chore: get method of pg vectorstore exact same as chroma

gRedHeadphone · gRedHeadphone · commit e82e26a3877e · 2025-11-12T07:46:55.000Z
diff --git a/examples/pg_vectorstore.ipynb b/examples/pg_vectorstore.ipynb
@@ -258,22 +258,22 @@
     "\n",
     "docs = [\n",
     "    Document(\n",
-    "        id=uuid.uuid4(),\n",
+    "        id=str(uuid.uuid4()),\n",
     "        page_content=\"there are cats in the pond\",\n",
     "        metadata={\"likes\": 1, \"location\": \"pond\", \"topic\": \"animals\"},\n",
     "    ),\n",
     "    Document(\n",
-    "        id=uuid.uuid4(),\n",
+    "        id=str(uuid.uuid4()),\n",
     "        page_content=\"ducks are also found in the pond\",\n",
     "        metadata={\"likes\": 30, \"location\": \"pond\", \"topic\": \"animals\"},\n",
     "    ),\n",
     "    Document(\n",
-    "        id=uuid.uuid4(),\n",
+    "        id=str(uuid.uuid4()),\n",
     "        page_content=\"fresh apples are available at the market\",\n",
     "        metadata={\"likes\": 20, \"location\": \"market\", \"topic\": \"food\"},\n",
     "    ),\n",
     "    Document(\n",
-    "        id=uuid.uuid4(),\n",
+    "        id=str(uuid.uuid4()),\n",
     "        page_content=\"the market also sells fresh oranges\",\n",
     "        metadata={\"likes\": 5, \"location\": \"market\", \"topic\": \"food\"},\n",
     "    ),\n",
@@ -287,9 +287,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Get documents\n",
+    "## Get collection\n",
     "\n",
-    "Get documents from the vectorstore using filters and parameters."
+    "Get collection from the vectorstore using filters and parameters."
    ]
   },
   {
@@ -298,13 +298,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "documents_with_apple = await vectorstore.aget({\"content\": {\"$ilike\": \"%apple%\"}})\n",
-    "first_three_documents = await vectorstore.aget(limit=3)\n",
-    "rest_of_documents = await vectorstore.aget(limit=5, offset=3)\n",
+    "documents_with_apple = await vectorstore.aget(where_document={\"$ilike\": \"%apple%\"}, include=\"documents\")\n",
+    "paginated_ids = await vectorstore.aget(limit=3, offset=3)\n",
     "\n",
-    "print([doc.page_content for doc in documents_with_apple])\n",
-    "print([doc.page_content for doc in first_three_documents])\n",
-    "print([doc.page_content for doc in rest_of_documents])"
+    "print(documents_with_apple[\"documents\"])\n",
+    "print(paginated_ids[\"ids\"])"
    ]
   },
   {
diff --git a/examples/pg_vectorstore_how_to.ipynb b/examples/pg_vectorstore_how_to.ipynb
@@ -331,9 +331,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Get documents\n",
+    "### Get collection\n",
     "\n",
-    "Get documents from the vectorstore using filters and parameters."
+    "Get collection from the vectorstore using filters and parameters."
    ]
   },
   {
@@ -342,13 +342,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "documents_with_apple = await store.aget({\"content\": {\"$ilike\": \"%apple%\"}})\n",
-    "first_three_documents = await store.aget(limit=3)\n",
-    "rest_of_documents = await store.aget(limit=5, offset=3)\n",
+    "documents_with_apple = await store.aget(where_document={\"$ilike\": \"%apple%\"}, include=\"documents\")\n",
+    "paginated_ids = await store.aget(limit=3, offset=3)\n",
     "\n",
-    "print([doc.page_content for doc in documents_with_apple])\n",
-    "print([doc.page_content for doc in first_three_documents])\n",
-    "print([doc.page_content for doc in rest_of_documents])"
+    "print(documents_with_apple[\"documents\"])\n",
+    "print(paginated_ids[\"ids\"])"
    ]
   },
   {
diff --git a/langchain_postgres/v2/async_vectorstore.py b/langchain_postgres/v2/async_vectorstore.py
@@ -678,18 +678,11 @@ async def __query_collection_with_filter(
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         filter: Optional[dict] = None,
+        columns: Optional[list[str]] = None,
         **kwargs: Any,
     ) -> Sequence[RowMapping]:
         """Asynchronously query the database collection using filters and parameters and return matching rows."""
 
-        columns = [
-            self.id_column,
-            self.content_column,
-            self.embedding_column,
-        ] + self.metadata_columns
-        if self.metadata_json_column:
-            columns.append(self.metadata_json_column)
-
         column_names = ", ".join(f'"{col}"' for col in columns)
 
         safe_filter = None
@@ -1037,35 +1030,68 @@ async def is_valid_index(
 
     async def aget(
         self,
-        filter: Optional[dict] = None,
+        ids: Optional[Sequence[str]] = None,
+        where: Optional[dict] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
+        where_document: Optional[dict] = None,
+        include: Optional[list[str]] = None,
         **kwargs: Any,
-    ) -> list[Document]:
+    ) -> dict[str, Any]:
         """Retrieve documents from the collection using filters and parameters."""
+        filter = {}
+        if ids:
+            filter.update({self.id_column: {"$in": ids}})
+        if where:
+            filter.update(where)
+        if where_document:
+            filter.update({self.content_column: where_document})
+
+        if include is None:
+            include = ["metadatas", "documents"]
+
+        fields_mapping = {
+            "embeddings": [self.embedding_column],
+            "metadatas": self.metadata_columns + [self.metadata_json_column]
+            if self.metadata_json_column
+            else self.metadata_columns,
+            "documents": [self.content_column],
+        }
+
+        included_fields = ["ids"]
+        columns = [self.id_column]
+
+        for field, cols in fields_mapping.items():
+            if field in include:
+                included_fields.append(field)
+                columns.extend(cols)
 
         results = await self.__query_collection_with_filter(
-            limit=limit, offset=offset, filter=filter, **kwargs
+            limit=limit, offset=offset, filter=filter, columns=columns, **kwargs
         )
 
-        documents = []
+        final_results = {field: [] for field in included_fields}
+
         for row in results:
-            metadata = (
-                row[self.metadata_json_column]
-                if self.metadata_json_column and row[self.metadata_json_column]
-                else {}
-            )
-            for col in self.metadata_columns:
-                metadata[col] = row[col]
-            documents.append(
-                Document(
-                    page_content=row[self.content_column],
-                    metadata=metadata,
-                    id=str(row[self.id_column]),
-                ),
-            )
+            final_results["ids"].append(str(row[self.id_column]))
 
-        return documents
+            if "metadatas" in final_results:
+                metadata = (
+                    row.get(self.metadata_json_column) or {}
+                    if self.metadata_json_column
+                    else {}
+                )
+                for col in self.metadata_columns:
+                    metadata[col] = row[col]
+                final_results["metadatas"].append(metadata)
+
+            if "documents" in final_results:
+                final_results["documents"].append(row[self.content_column])
+
+            if "embeddings" in final_results:
+                final_results["embeddings"].append(row[self.embedding_column])
+
+        return final_results
 
     async def aget_by_ids(self, ids: Sequence[str]) -> list[Document]:
         """Get documents by ids."""
@@ -1323,11 +1349,14 @@ def _create_filter_clause(self, filters: Any) -> tuple[str, dict]:
 
     def get(
         self,
-        filter: Optional[dict] = None,
+        ids: Optional[Sequence[str]] = None,
+        where: Optional[dict] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
+        where_document: Optional[dict] = None,
+        include: Optional[list[str]] = None,
         **kwargs: Any,
-    ) -> list[Document]:
+    ) -> dict[str, Any]:
         raise NotImplementedError(
             "Sync methods are not implemented for AsyncPGVectorStore. Use PGVectorStore interface instead."
         )
diff --git a/langchain_postgres/v2/vectorstores.py b/langchain_postgres/v2/vectorstores.py
@@ -877,26 +877,48 @@ def get_by_ids(self, ids: Sequence[str]) -> list[Document]:
 
     async def aget(
         self,
-        filter: Optional[dict] = None,
+        ids: Optional[Sequence[str]] = None,
+        where: Optional[dict] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
+        where_document: Optional[dict] = None,
+        include: Optional[list[str]] = None,
         **kwargs: Any,
-    ) -> list[Document]:
+    ) -> dict[str, Any]:
         """Retrieve documents from the collection using filters and parameters."""
         return await self._engine._run_as_async(
-            self.__vs.aget(filter=filter, limit=limit, offset=offset, **kwargs)
+            self.__vs.aget(
+                ids=ids,
+                where=where,
+                limit=limit,
+                offset=offset,
+                where_document=where_document,
+                include=include,
+                **kwargs,
+            )
         )
 
     def get(
         self,
-        filter: Optional[dict] = None,
+        ids: Optional[Sequence[str]] = None,
+        where: Optional[dict] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
+        where_document: Optional[dict] = None,
+        include: Optional[list[str]] = None,
         **kwargs: Any,
-    ) -> list[Document]:
+    ) -> dict[str, Any]:
         """Retrieve documents from the collection using filters and parameters."""
         return self._engine._run_as_sync(
-            self.__vs.aget(filter=filter, limit=limit, offset=offset, **kwargs)
+            self.__vs.aget(
+                ids=ids,
+                where=where,
+                limit=limit,
+                offset=offset,
+                where_document=where_document,
+                include=include,
+                **kwargs,
+            )
         )
 
     def get_table_name(self) -> str:
diff --git a/tests/unit_tests/v2/test_async_pg_vectorstore_search.py b/tests/unit_tests/v2/test_async_pg_vectorstore_search.py
@@ -370,6 +370,24 @@ async def test_vectorstore_with_metadata_filters(
         )
         assert [doc.metadata["code"] for doc in docs] == expected_ids, test_filter
 
+    async def test_async_vectorstore_get_ids(
+        self,
+        vs_custom_filter: AsyncPGVectorStore
+    ) -> None:
+        """Test end to end construction and filter."""
+
+        res = await vs_custom_filter.aget(ids=ids[:2])
+        assert set(res["ids"]) == set(ids[:2])
+
+    async def test_async_vectorstore_get_docs(
+        self,
+        vs_custom_filter: AsyncPGVectorStore
+    ) -> None:
+        """Test end to end construction and filter."""
+
+        res = await vs_custom_filter.aget(where_document={"$in": texts[:2]})
+        assert set(res["documents"]) == set(texts[:2])
+
     @pytest.mark.parametrize("test_filter, expected_ids", FILTERING_TEST_CASES)
     async def test_vectorstore_get(
         self,
@@ -378,8 +396,8 @@ async def test_vectorstore_get(
         expected_ids: list[str],
     ) -> None:
         """Test end to end construction and filter."""
-        docs = await vs_custom_filter.aget(test_filter)
-        assert set([doc.metadata["code"] for doc in docs]) == set(expected_ids), (
+        res = await vs_custom_filter.aget(where=test_filter)
+        assert set([r["code"] for r in res["metadatas"]]) == set(expected_ids), (
             test_filter
         )
 
@@ -389,14 +407,14 @@ async def test_vectorstore_get_limit_offset(
     ) -> None:
         """Test limit and offset parameters of get method"""
 
-        all_docs = await vs_custom_filter.aget()
-        docs_from_combining = (
-            (await vs_custom_filter.aget(limit=1))
-            + (await vs_custom_filter.aget(limit=1, offset=1))
-            + (await vs_custom_filter.aget(offset=2))
+        all_ids = (await vs_custom_filter.aget())["ids"]
+        ids_from_combining = (
+            (await vs_custom_filter.aget(limit=1))["ids"]
+            + (await vs_custom_filter.aget(limit=1, offset=1))["ids"]
+            + (await vs_custom_filter.aget(offset=2))["ids"]
         )
 
-        assert all_docs == docs_from_combining
+        assert all_ids == ids_from_combining
 
     async def test_asimilarity_hybrid_search(self, vs: AsyncPGVectorStore) -> None:
         results = await vs.asimilarity_search(
diff --git a/tests/unit_tests/v2/test_pg_vectorstore_search.py b/tests/unit_tests/v2/test_pg_vectorstore_search.py
@@ -429,6 +429,24 @@ def test_sync_vectorstore_with_metadata_filters(
         docs = vs_custom_filter_sync.similarity_search("meow", k=5, filter=test_filter)
         assert [doc.metadata["code"] for doc in docs] == expected_ids, test_filter
 
+    def test_sync_vectorstore_get_ids(
+        self,
+        vs_custom_filter_sync: PGVectorStore
+    ) -> None:
+        """Test end to end construction and filter."""
+
+        res = vs_custom_filter_sync.get(ids=ids[:2])
+        assert set(res["ids"]) == set(ids[:2])
+
+    def test_sync_vectorstore_get_docs(
+        self,
+        vs_custom_filter_sync: PGVectorStore
+    ) -> None:
+        """Test end to end construction and filter."""
+
+        res = vs_custom_filter_sync.get(where_document={"$in": texts[:2]})
+        assert set(res["documents"]) == set(texts[:2])
+
     @pytest.mark.parametrize("test_filter, expected_ids", FILTERING_TEST_CASES)
     def test_sync_vectorstore_get(
         self,
@@ -438,8 +456,8 @@ def test_sync_vectorstore_get(
     ) -> None:
         """Test end to end construction and filter."""
 
-        docs = vs_custom_filter_sync.get(filter=test_filter)
-        assert set([doc.metadata["code"] for doc in docs]) == set(expected_ids), (
+        res = vs_custom_filter_sync.get(where=test_filter)
+        assert set([r["code"] for r in res["metadatas"]]) == set(expected_ids), (
             test_filter
         )
 
@@ -449,14 +467,14 @@ def test_sync_vectorstore_get_limit_offset(
     ) -> None:
         """Test limit and offset parameters of get method"""
 
-        all_docs = vs_custom_filter_sync.get()
-        docs_from_combining = (
-            vs_custom_filter_sync.get(limit=1)
-            + vs_custom_filter_sync.get(limit=1, offset=1)
-            + vs_custom_filter_sync.get(offset=2)
+        all_ids = vs_custom_filter_sync.get()["ids"]
+        ids_from_combining = (
+            vs_custom_filter_sync.get(limit=1)["ids"]
+            + vs_custom_filter_sync.get(limit=1, offset=1)["ids"]
+            + vs_custom_filter_sync.get(offset=2)["ids"]
         )
 
-        assert all_docs == docs_from_combining
+        assert all_ids == ids_from_combining
 
     @pytest.mark.parametrize("test_filter", NEGATIVE_TEST_CASES)
     def test_metadata_filter_negative_tests(
