github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-09-02 19:47:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

24
libs/community/tests/integration_tests/document_loaders/test_cassandra.py
View File

@@ -55,9 +55,9 @@ def keyspace() -> Iterator[str]:
session.execute(f"DROP TABLE IF EXISTS {keyspace}.{CASSANDRA_TABLE}")
def test_loader_table(keyspace: str) -> None:
async def test_loader_table(keyspace: str) -> None:
loader = CassandraLoader(table=CASSANDRA_TABLE)
assert loader.load() == [
expected = [
Document(
page_content="Row(row_id='id1', body_blob='text1')",
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
@@ -67,24 +67,28 @@ def test_loader_table(keyspace: str) -> None:
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
),
]
assert loader.load() == expected
assert await loader.aload() == expected
def test_loader_query(keyspace: str) -> None:
async def test_loader_query(keyspace: str) -> None:
loader = CassandraLoader(
query=f"SELECT body_blob FROM {keyspace}.{CASSANDRA_TABLE}"
)
assert loader.load() == [
expected = [
Document(page_content="Row(body_blob='text1')"),
Document(page_content="Row(body_blob='text2')"),
]
assert loader.load() == expected
assert await loader.aload() == expected
def test_loader_page_content_mapper(keyspace: str) -> None:
async def test_loader_page_content_mapper(keyspace: str) -> None:
def mapper(row: Any) -> str:
return str(row.body_blob)
loader = CassandraLoader(table=CASSANDRA_TABLE, page_content_mapper=mapper)
assert loader.load() == [
expected = [
Document(
page_content="text1",
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
@@ -94,14 +98,16 @@ def test_loader_page_content_mapper(keyspace: str) -> None:
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
),
]
assert loader.load() == expected
assert await loader.aload() == expected
def test_loader_metadata_mapper(keyspace: str) -> None:
async def test_loader_metadata_mapper(keyspace: str) -> None:
def mapper(row: Any) -> dict:
return {"id": row.row_id}
loader = CassandraLoader(table=CASSANDRA_TABLE, metadata_mapper=mapper)
assert loader.load() == [
expected = [
Document(
page_content="Row(row_id='id1', body_blob='text1')",
metadata={
@@ -119,3 +125,5 @@ def test_loader_metadata_mapper(keyspace: str) -> None:
},
),
]
assert loader.load() == expected
assert await loader.aload() == expected