github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-09-09 06:53:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[patch]: Use astrapy built-in pagination prefetch in AstraDBLoader (#17569)

Browse Source
This commit is contained in:
Christophe Bornet
2024-02-15 15:52:56 +01:00
committed by GitHub
parent 387cacb881
commit 789cd5198d
2 changed files with 25 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
libs/community/tests/integration_tests/document_loaders/test_astradb.py
View File

@@ -15,7 +15,7 @@ from __future__ import annotations
import json
import os
import uuid
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, AsyncIterator, Iterator
import pytest
@@ -37,12 +37,12 @@ def _has_env_vars() -> bool:
@pytest.fixture
def astra_db_collection() -> AstraDBCollection:
def astra_db_collection() -> Iterator[AstraDBCollection]:
from astrapy.db import AstraDB
astra_db = AstraDB(
token=ASTRA_DB_APPLICATION_TOKEN,
api_endpoint=ASTRA_DB_API_ENDPOINT,
token=ASTRA_DB_APPLICATION_TOKEN or "",
api_endpoint=ASTRA_DB_API_ENDPOINT or "",
namespace=ASTRA_DB_KEYSPACE,
)
collection_name = f"lc_test_loader_{str(uuid.uuid4()).split('-')[0]}"
@@ -58,12 +58,12 @@ def astra_db_collection() -> AstraDBCollection:
@pytest.fixture
async def async_astra_db_collection() -> AsyncAstraDBCollection:
async def async_astra_db_collection() -> AsyncIterator[AsyncAstraDBCollection]:
from astrapy.db import AsyncAstraDB
astra_db = AsyncAstraDB(
token=ASTRA_DB_APPLICATION_TOKEN,
api_endpoint=ASTRA_DB_API_ENDPOINT,
token=ASTRA_DB_APPLICATION_TOKEN or "",
api_endpoint=ASTRA_DB_API_ENDPOINT or "",
namespace=ASTRA_DB_KEYSPACE,
)
collection_name = f"lc_test_loader_{str(uuid.uuid4()).split('-')[0]}"
@@ -167,5 +167,5 @@ class TestAstraDB:
find_options={"limit": 30},
extraction_function=lambda x: x["foo"],
)
doc = await anext(loader.alazy_load()) # type: ignore[name-defined]
doc = await loader.alazy_load().__anext__()
assert doc.page_content == "bar"