mirror of
https://github.com/hwchase17/langchain.git
synced 2025-05-19 22:11:41 +00:00
a0c2281540
```python
"""python scripts/update_mypy_ruff.py"""
import glob
import tomllib
from pathlib import Path
import toml
import subprocess
import re
ROOT_DIR = Path(__file__).parents[1]
def main():
for path in glob.glob(str(ROOT_DIR / "libs/**/pyproject.toml"), recursive=True):
print(path)
with open(path, "rb") as f:
pyproject = tomllib.load(f)
try:
pyproject["tool"]["poetry"]["group"]["typing"]["dependencies"]["mypy"] = (
"^1.10"
)
pyproject["tool"]["poetry"]["group"]["lint"]["dependencies"]["ruff"] = (
"^0.5"
)
except KeyError:
continue
with open(path, "w") as f:
toml.dump(pyproject, f)
cwd = "/".join(path.split("/")[:-1])
completed = subprocess.run(
"poetry lock --no-update; poetry install --with typing; poetry run mypy . --no-color",
cwd=cwd,
shell=True,
capture_output=True,
text=True,
)
logs = completed.stdout.split("\n")
to_ignore = {}
for l in logs:
if re.match("^(.*)\:(\d+)\: error:.*\[(.*)\]", l):
path, line_no, error_type = re.match(
"^(.*)\:(\d+)\: error:.*\[(.*)\]", l
).groups()
if (path, line_no) in to_ignore:
to_ignore[(path, line_no)].append(error_type)
else:
to_ignore[(path, line_no)] = [error_type]
print(len(to_ignore))
for (error_path, line_no), error_types in to_ignore.items():
all_errors = ", ".join(error_types)
full_path = f"{cwd}/{error_path}"
try:
with open(full_path, "r") as f:
file_lines = f.readlines()
except FileNotFoundError:
continue
file_lines[int(line_no) - 1] = (
file_lines[int(line_no) - 1][:-1] + f" # type: ignore[{all_errors}]\n"
)
with open(full_path, "w") as f:
f.write("".join(file_lines))
subprocess.run(
"poetry run ruff format .; poetry run ruff --select I --fix .",
cwd=cwd,
shell=True,
capture_output=True,
text=True,
)
if __name__ == "__main__":
main()
```
131 lines
3.8 KiB
Python
131 lines
3.8 KiB
Python
"""
|
|
Test of Cassandra document loader class `CassandraLoader`
|
|
"""
|
|
|
|
import os
|
|
from typing import Any, Iterator
|
|
|
|
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() -> Iterator[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}")
|
|
|
|
|
|
async def test_loader_table(keyspace: str) -> None:
|
|
loader = CassandraLoader(table=CASSANDRA_TABLE)
|
|
expected = [
|
|
Document(
|
|
page_content="Row(row_id='id1', body_blob='text1')",
|
|
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
|
|
),
|
|
Document(
|
|
page_content="Row(row_id='id2', body_blob='text2')",
|
|
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
|
|
),
|
|
]
|
|
assert loader.load() == expected
|
|
assert await loader.aload() == expected
|
|
|
|
|
|
async def test_loader_query(keyspace: str) -> None:
|
|
loader = CassandraLoader(
|
|
query=f"SELECT body_blob FROM {keyspace}.{CASSANDRA_TABLE}"
|
|
)
|
|
expected = [
|
|
Document(page_content="Row(body_blob='text1')"),
|
|
Document(page_content="Row(body_blob='text2')"),
|
|
]
|
|
assert loader.load() == expected
|
|
assert await loader.aload() == expected
|
|
|
|
|
|
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)
|
|
expected = [
|
|
Document(
|
|
page_content="text1",
|
|
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
|
|
),
|
|
Document(
|
|
page_content="text2",
|
|
metadata={"table": CASSANDRA_TABLE, "keyspace": keyspace},
|
|
),
|
|
]
|
|
assert loader.load() == expected
|
|
assert await loader.aload() == expected
|
|
|
|
|
|
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)
|
|
expected = [
|
|
Document(
|
|
page_content="Row(row_id='id1', body_blob='text1')",
|
|
metadata={
|
|
"table": CASSANDRA_TABLE,
|
|
"keyspace": keyspace,
|
|
"id": "id1",
|
|
},
|
|
),
|
|
Document(
|
|
page_content="Row(row_id='id2', body_blob='text2')",
|
|
metadata={
|
|
"table": CASSANDRA_TABLE,
|
|
"keyspace": keyspace,
|
|
"id": "id2",
|
|
},
|
|
),
|
|
]
|
|
assert loader.load() == expected
|
|
assert await loader.aload() == expected
|