github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-09-04 04:28:58 +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

73
libs/community/langchain_community/document_loaders/astradb.py
View File

@@ -2,8 +2,6 @@ from __future__ import annotations
import json
import logging
import threading
from queue import Queue
from typing import (
TYPE_CHECKING,
Any,
@@ -16,7 +14,6 @@ from typing import (
)
from langchain_core.documents import Document
from langchain_core.runnables import run_in_executor
from langchain_community.document_loaders.base import BaseLoader
from langchain_community.utilities.astradb import _AstraDBEnvironment
@@ -33,6 +30,7 @@ class AstraDBLoader(BaseLoader):
def __init__(
self,
collection_name: str,
*,
token: Optional[str] = None,
api_endpoint: Optional[str] = None,
astra_db_client: Optional[AstraDB] = None,
@@ -65,38 +63,27 @@ class AstraDBLoader(BaseLoader):
return list(self.lazy_load())
def lazy_load(self) -> Iterator[Document]:
queue = Queue(self.nb_prefetched) # type: ignore
t = threading.Thread(target=self.fetch_results, args=(queue,))
t.start()
while True:
doc = queue.get()
if doc is None:
break
yield doc
t.join()
for doc in self.collection.paginated_find(
filter=self.filter,
options=self.find_options,
projection=self.projection,
sort=None,
prefetched=self.nb_prefetched,
):
yield Document(
page_content=self.extraction_function(doc),
metadata={
"namespace": self.collection.astra_db.namespace,
"api_endpoint": self.collection.astra_db.base_url,
"collection": self.collection_name,
},
)
async def aload(self) -> List[Document]:
"""Load data into Document objects."""
return [doc async for doc in self.alazy_load()]
async def alazy_load(self) -> AsyncIterator[Document]:
if not self.astra_env.async_astra_db:
iterator = run_in_executor(
None,
self.collection.paginated_find,
filter=self.filter,
options=self.find_options,
projection=self.projection,
sort=None,
prefetched=True,
)
done = object()
while True:
item = await run_in_executor(None, lambda it: next(it, done), iterator)
if item is done:
break
yield item # type: ignore[misc]
return
async_collection = await self.astra_env.async_astra_db.collection(
self.collection_name
)
@@ -105,7 +92,7 @@ class AstraDBLoader(BaseLoader):
options=self.find_options,
projection=self.projection,
sort=None,
prefetched=True,
prefetched=self.nb_prefetched,
):
yield Document(
page_content=self.extraction_function(doc),
@@ -115,29 +102,3 @@ class AstraDBLoader(BaseLoader):
"collection": self.collection_name,
},
)
def fetch_results(self, queue: Queue): # type: ignore[no-untyped-def]
self.fetch_page_result(queue)
while self.find_options.get("pageState"):
self.fetch_page_result(queue)
queue.put(None)
def fetch_page_result(self, queue: Queue): # type: ignore[no-untyped-def]
res = self.collection.find(
filter=self.filter,
options=self.find_options,
projection=self.projection,
sort=None,
)
self.find_options["pageState"] = res["data"].get("nextPageState")
for doc in res["data"]["documents"]:
queue.put(
Document(
page_content=self.extraction_function(doc),
metadata={
"namespace": self.collection.astra_db.namespace,
"api_endpoint": self.collection.astra_db.base_url,
"collection": self.collection.collection_name,
},
)
)