Implement aget_nodes and adelete in PGVectorStore (#18515)

Author: GinkREAL

llama-index-integrations/vector_stores/llama-index-vector-stores-postgres/llama_index/vector_stores/postgres/base.py

@@ -100,6 +100,11 @@ class HybridAbstractData(base):  # type: ignore
             model.text_search_tsv,  # type: ignore
             postgresql_using="gin",
         )
+        Index(
+            f"{indexname}_1",
+            model.metadata_["ref_doc_id"].astext,  # type: ignore
+            postgresql_using="btree",
+        )
     else:
 
         class AbstractData(base):  # type: ignore
@@ -116,6 +121,12 @@ class AbstractData(base):  # type: ignore
             {"__tablename__": tablename, "__table_args__": {"schema": schema_name}},
         )
 
+        Index(
+            f"{indexname}_1",
+            model.metadata_["ref_doc_id"].astext,  # type: ignore
+            postgresql_using="btree",
+        )
+
     return model
 
 
@@ -953,6 +964,18 @@ def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
             session.execute(stmt)
             session.commit()
 
+    async def adelete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
+        from sqlalchemy import delete
+
+        self._initialize()
+        async with self._async_session() as session, session.begin():
+            stmt = delete(self._table_class).where(
+                self._table_class.metadata_["doc_id"].astext == ref_doc_id
+            )
+
+            await session.execute(stmt)
+            await session.commit()
+
     def delete_nodes(
         self,
         node_ids: Optional[List[str]] = None,
@@ -1087,6 +1110,58 @@ def get_nodes(
 
         return nodes
 
+    async def aget_nodes(
+        self,
+        node_ids: Optional[List[str]] = None,
+        filters: Optional[MetadataFilters] = None,
+    ) -> List[BaseNode]:
+        """Get nodes asynchronously from vector store."""
+        assert (
+            node_ids is not None or filters is not None
+        ), "Either node_ids or filters must be provided"
+
+        self._initialize()
+        from sqlalchemy import select
+
+        stmt = select(
+            self._table_class.node_id,
+            self._table_class.text,
+            self._table_class.metadata_,
+            self._table_class.embedding,
+        )
+
+        if node_ids:
+            stmt = stmt.where(self._table_class.node_id.in_(node_ids))
+
+        if filters:
+            filter_clause = self._recursively_apply_filters(filters)
+            stmt = stmt.where(filter_clause)
+
+        nodes: List[BaseNode] = []
+
+        async with self._async_session() as session, session.begin():
+            res = (await session.execute(stmt)).fetchall()
+            for item in res:
+                node_id = item.node_id
+                text = item.text
+                metadata = item.metadata_
+                embedding = item.embedding
+
+                try:
+                    node = metadata_dict_to_node(metadata)
+                    node.set_content(str(text))
+                    node.embedding = embedding
+                except Exception:
+                    node = TextNode(
+                        id_=node_id,
+                        text=text,
+                        metadata=metadata,
+                        embedding=embedding,
+                    )
+                nodes.append(node)
+
+            return nodes
+
 
 def _dedup_results(results: List[DBEmbeddingRow]) -> List[DBEmbeddingRow]:
     seen_ids = set()

llama-index-integrations/vector_stores/llama-index-vector-stores-postgres/pyproject.toml

@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-vector-stores-postgres"
 readme = "README.md"
-version = "0.5.0"
+version = "0.5.1"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"