Skip to content

Commit d1bd0cd

Browse files
committed
Implement aget_nodes in PGVectorStore
1 parent 4e924af commit d1bd0cd

File tree

2 files changed

+65
-1
lines changed
  • llama-index-integrations/vector_stores/llama-index-vector-stores-postgres

2 files changed

+65
-1
lines changed

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

+64
Original file line numberDiff line numberDiff line change
@@ -953,6 +953,18 @@ def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
953953
session.execute(stmt)
954954
session.commit()
955955

956+
async def adelete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
957+
from sqlalchemy import delete
958+
959+
self._initialize()
960+
async with self._async_session() as session, session.begin():
961+
stmt = delete(self._table_class).where(
962+
self._table_class.metadata_["doc_id"].astext == ref_doc_id
963+
)
964+
965+
await session.execute(stmt)
966+
await session.commit()
967+
956968
def delete_nodes(
957969
self,
958970
node_ids: Optional[List[str]] = None,
@@ -1087,6 +1099,58 @@ def get_nodes(
10871099

10881100
return nodes
10891101

1102+
async def aget_nodes(
1103+
self,
1104+
node_ids: Optional[List[str]] = None,
1105+
filters: Optional[MetadataFilters] = None,
1106+
) -> List[BaseNode]:
1107+
"""Get nodes asynchronously from vector store."""
1108+
assert (
1109+
node_ids is not None or filters is not None
1110+
), "Either node_ids or filters must be provided"
1111+
1112+
self._initialize()
1113+
from sqlalchemy import select
1114+
1115+
stmt = select(
1116+
self._table_class.node_id,
1117+
self._table_class.text,
1118+
self._table_class.metadata_,
1119+
self._table_class.embedding,
1120+
)
1121+
1122+
if node_ids:
1123+
stmt = stmt.where(self._table_class.node_id.in_(node_ids))
1124+
1125+
if filters:
1126+
filter_clause = self._recursively_apply_filters(filters)
1127+
stmt = stmt.where(filter_clause)
1128+
1129+
nodes: List[BaseNode] = []
1130+
1131+
async with self._async_session() as session, session.begin():
1132+
res = (await session.execute(stmt)).fetchall()
1133+
for item in res:
1134+
node_id = item.node_id
1135+
text = item.text
1136+
metadata = item.metadata_
1137+
embedding = item.embedding
1138+
1139+
try:
1140+
node = metadata_dict_to_node(metadata)
1141+
node.set_content(str(text))
1142+
node.embedding = embedding
1143+
except Exception:
1144+
node = TextNode(
1145+
id_=node_id,
1146+
text=text,
1147+
metadata=metadata,
1148+
embedding=embedding,
1149+
)
1150+
nodes.append(node)
1151+
1152+
return nodes
1153+
10901154

10911155
def _dedup_results(results: List[DBEmbeddingRow]) -> List[DBEmbeddingRow]:
10921156
seen_ids = set()

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
2727
license = "MIT"
2828
name = "llama-index-vector-stores-postgres"
2929
readme = "README.md"
30-
version = "0.5.0"
30+
version = "0.5.1"
3131

3232
[tool.poetry.dependencies]
3333
python = ">=3.9,<4.0"

0 commit comments

Comments
 (0)