community[minor]: Add async methods to CassandraLoader (#20609)

Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
This commit is contained in:
Christophe Bornet
2024-04-18 21:45:20 +02:00
committed by GitHub
parent 8c29b7bf35
commit d2d01370bc
3 changed files with 54 additions and 8 deletions

View File

@@ -3,6 +3,7 @@ from __future__ import annotations
from typing import (
TYPE_CHECKING,
Any,
AsyncIterator,
Callable,
Iterator,
Optional,
@@ -13,6 +14,7 @@ from typing import (
from langchain_core.documents import Document
from langchain_community.document_loaders.base import BaseLoader
from langchain_community.utilities.cassandra import wrapped_response_future
_NOT_SET = object()
@@ -112,3 +114,15 @@ class CassandraLoader(BaseLoader):
yield Document(
page_content=self.page_content_mapper(row), metadata=metadata
)
async def alazy_load(self) -> AsyncIterator[Document]:
for row in await wrapped_response_future(
self.session.execute_async,
self.query,
**self.query_kwargs,
):
metadata = self.metadata.copy()
metadata.update(self.metadata_mapper(row))
yield Document(
page_content=self.page_content_mapper(row), metadata=metadata
)