community[minor]: Add Cassandra document loader (#16215)

- **Description:** document loader for Apache Cassandra
  - **Twitter handle:** cbornet_
This commit is contained in:
Christophe Bornet
2024-01-19 03:49:02 +01:00
committed by GitHub
parent fc84083ce5
commit 3ccbe11363
3 changed files with 251 additions and 0 deletions

View File

@@ -0,0 +1,123 @@
import json
from typing import (
TYPE_CHECKING,
Any,
Callable,
Iterator,
List,
Optional,
Sequence,
Union,
)
from langchain_core.documents import Document
from langchain_community.document_loaders.base import BaseLoader
def default_page_content_mapper(row: Any) -> str:
if hasattr(row, "_asdict"):
return json.dumps(row._asdict())
return json.dumps(row)
_NOT_SET = object()
if TYPE_CHECKING:
from cassandra.cluster import Session
from cassandra.pool import Host
from cassandra.query import Statement
class CassandraLoader(BaseLoader):
def __init__(
self,
table: Optional[str] = None,
session: Optional["Session"] = None,
keyspace: Optional[str] = None,
query: Optional[Union[str, "Statement"]] = None,
page_content_mapper: Callable[[Any], str] = default_page_content_mapper,
metadata_mapper: Callable[[Any], dict] = lambda _: {},
*,
query_parameters: Union[dict, Sequence] = None,
query_timeout: Optional[float] = _NOT_SET,
query_trace: bool = False,
query_custom_payload: dict = None,
query_execution_profile: Any = _NOT_SET,
query_paging_state: Any = None,
query_host: "Host" = None,
query_execute_as: str = None,
) -> None:
"""
Document Loader for Apache Cassandra.
Args:
table: The table to load the data from.
(do not use together with the query parameter)
session: The cassandra driver session.
If not provided, the cassio resolved session will be used.
keyspace: The keyspace of the table.
If not provided, the cassio resolved keyspace will be used.
query: The query used to load the data.
(do not use together with the table parameter)
page_content_mapper: a function to convert a row to string page content.
query_parameters: The query parameters used when calling session.execute .
query_timeout: The query timeout used when calling session.execute .
query_custom_payload: The query custom_payload used when calling
session.execute .
query_execution_profile: The query execution_profile used when calling
session.execute .
query_host: The query host used when calling session.execute .
query_execute_as: The query execute_as used when calling session.execute .
"""
if query and table:
raise ValueError("Cannot specify both query and table.")
if not query and not table:
raise ValueError("Must specify query or table.")
if not session or (table and not keyspace):
try:
from cassio.config import check_resolve_keyspace, check_resolve_session
except (ImportError, ModuleNotFoundError):
raise ImportError(
"Could not import a recent cassio package."
"Please install it with `pip install --upgrade cassio`."
)
if table:
_keyspace = keyspace or check_resolve_keyspace(keyspace)
self.query = f"SELECT * FROM {_keyspace}.{table};"
self.metadata = {"table": table, "keyspace": _keyspace}
else:
self.query = query
self.metadata = {}
self.session = session or check_resolve_session(session)
self.page_content_mapper = page_content_mapper
self.metadata_mapper = metadata_mapper
self.query_kwargs = {
"parameters": query_parameters,
"trace": query_trace,
"custom_payload": query_custom_payload,
"paging_state": query_paging_state,
"host": query_host,
"execute_as": query_execute_as,
}
if query_timeout is not _NOT_SET:
self.query_kwargs["timeout"] = query_timeout
if query_execution_profile is not _NOT_SET:
self.query_kwargs["execution_profile"] = query_execution_profile
def load(self) -> List[Document]:
return list(self.lazy_load())
def lazy_load(self) -> Iterator[Document]:
for row in self.session.execute(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
)

View File

@@ -14,6 +14,13 @@ ASTRA_DB_APPLICATION_TOKEN=AstraCS:your_astra_db_application_token
# ASTRA_DB_KEYSPACE=your_astra_db_namespace
# cassandra
CASSANDRA_CONTACT_POINTS=127.0.0.1
# CASSANDRA_USERNAME=your_cassandra_username
# CASSANDRA_PASSWORD=your_cassandra_password
# CASSANDRA_KEYSPACE=your_cassandra_keyspace
# pinecone
# your api key from left menu "API Keys" in https://app.pinecone.io
PINECONE_API_KEY=your_pinecone_api_key_here

View File

@@ -0,0 +1,121 @@
"""
Test of Cassandra document loader class `CassandraLoader`
"""
import os
from typing import Any
import pytest
from langchain_core.documents import Document
from langchain_community.document_loaders.cassandra import CassandraLoader
CASSANDRA_DEFAULT_KEYSPACE = "docloader_test_keyspace"
CASSANDRA_TABLE = "docloader_test_table"
@pytest.fixture(autouse=True, scope="session")
def keyspace() -> str:
import cassio
from cassandra.cluster import Cluster
from cassio.config import check_resolve_session, resolve_keyspace
from cassio.table.tables import PlainCassandraTable
if any(
env_var in os.environ
for env_var in [
"CASSANDRA_CONTACT_POINTS",
"ASTRA_DB_APPLICATION_TOKEN",
"ASTRA_DB_INIT_STRING",
]
):
cassio.init(auto=True)
session = check_resolve_session()
else:
cluster = Cluster()
session = cluster.connect()
keyspace = resolve_keyspace() or CASSANDRA_DEFAULT_KEYSPACE
cassio.init(session=session, keyspace=keyspace)
session.execute(
(
f"CREATE KEYSPACE IF NOT EXISTS {keyspace} "
f"WITH replication = {{'class': 'SimpleStrategy', 'replication_factor': 1}}"
)
)
# We use a cassio table by convenience to seed the DB
table = PlainCassandraTable(
table=CASSANDRA_TABLE, keyspace=keyspace, session=session
)
table.put(row_id="id1", body_blob="text1")
table.put(row_id="id2", body_blob="text2")
yield keyspace
session.execute(f"DROP TABLE IF EXISTS {keyspace}.{CASSANDRA_TABLE}")
def test_loader_table(keyspace: str) -> None:
loader = CassandraLoader(table=CASSANDRA_TABLE)
assert loader.load() == [
Document(
page_content='{"row_id": "id1", "body_blob": "text1"}',
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
),
Document(
page_content='{"row_id": "id2", "body_blob": "text2"}',
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
),
]
def test_loader_query(keyspace: str) -> None:
loader = CassandraLoader(
query=f"SELECT body_blob FROM {keyspace}.{CASSANDRA_TABLE}"
)
assert loader.load() == [
Document(page_content='{"body_blob": "text1"}'),
Document(page_content='{"body_blob": "text2"}'),
]
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() == [
Document(
page_content="text1",
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
),
Document(
page_content="text2",
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
),
]
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() == [
Document(
page_content='{"row_id": "id1", "body_blob": "text1"}',
metadata={
"table": CASSANDRA_TABLE,
"keyspace": keyspace,
"id": "id1",
},
),
Document(
page_content='{"row_id": "id2", "body_blob": "text2"}',
metadata={
"table": CASSANDRA_TABLE,
"keyspace": keyspace,
"id": "id2",
},
),
]