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