WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-29 21:15:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

Lint and fix

Browse Source
This commit is contained in:
William Fu-Hinthorn
2024-12-06 10:05:29 -08:00
parent 4406828c39
commit 054ca2f78d
5 changed files with 215 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/tests/test_async_store.py
+82 -91
View File
@@ -4,6 +4,7 @@ import itertools
import sys
import uuid
from collections.abc import AsyncIterator
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from typing import Any, Optional
@@ -64,104 +65,94 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
await conn.execute(f"DROP DATABASE {database}")
async def test_large_batches(store: AsyncPostgresStore) -> None:
N = 100
def test_large_batches(store: AsyncPostgresStore) -> None:
N = 1000
M = 10
coros = []
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor(max_workers=10) as executor:
for m in range(M):
ops = []
for i in range(N):
for i in range(N):
coros.append(
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
)
)
coros.append(
executor.submit(
store.get,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
)
)
coros.append(
executor.submit(
store.list_namespaces,
prefix=None,
max_depth=m + 1,
)
)
coros.append(
executor.submit(
store.search,
("test",),
)
)
coros.append(
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
)
)
coros.append(
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
None,
)
)
# ops.extend(
# [
# PutOp(
# namespace=("test", "foo", "bar", "baz", str(m % 2)),
# key=f"key{i}", # {m}",
# value=None,
# ),
# GetOp(namespace=("test",), key=f"key{i}{m}"),
# ListNamespacesOp(
# match_conditions=None, max_depth=i + 1, limit=m, offset=0
# ),
# SearchOp(
# namespace_prefix=("test",),
# filter=None,
# limit=10,
# offset=0,
# ),
# ]
#
# ops.extend(
# [
# # PutOp(
# # namespace=("test", "foo", "bar", "baz", str(m % 2)),
# # key=f"key{i}", # {m}",
# # value={"data": f"value{i}{m}"},
# # ),
# # GetOp(namespace=("test",), key=f"key{i}{m}"),
# # ListNamespacesOp(
# # match_conditions=None, max_depth=i + 1, limit=m, offset=0
# # ),
# # SearchOp(
# # namespace_prefix=("test",),
# # filter={"data": f"value{i}{m}"},
# # limit=10,
# # offset=0,
# # ),
# ]
# )
_ = [
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
),
executor.submit(
store.get,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
),
executor.submit(
store.list_namespaces,
prefix=None,
max_depth=m + 1,
),
executor.submit(
store.search,
("test",),
),
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
),
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
None,
),
]
# coros.extend(ops)
# executor.map(store.batch, [[op] for op in coros])
# await asyncio.gather(*coros) # *[store.abatch(ops) for ops in coros])
async def test_large_batches_async(store: AsyncPostgresStore) -> None:
N = 1000
M = 10
coros = []
for m in range(M):
for i in range(N):
coros.append(
store.aput(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
)
)
coros.append(
store.aget(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
)
)
coros.append(
store.alist_namespaces(
prefix=None,
max_depth=m + 1,
)
)
coros.append(
store.asearch(
("test",),
)
)
coros.append(
store.aput(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
)
)
coros.append(
store.adelete(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
)
)
await asyncio.gather(*coros)
async def test_abatch_order(store: AsyncPostgresStore) -> None:
libs/checkpoint-postgres/tests/test_store.py
+93 -5
View File
@@ -1,5 +1,7 @@
# type: ignore
import asyncio
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from typing import Any, Optional
from uuid import uuid4
@@ -17,11 +19,7 @@ from langgraph.store.base import (
SearchOp,
)
from langgraph.store.postgres import PostgresStore
from tests.conftest import (
DEFAULT_URI,
VECTOR_TYPES,
CharacterEmbeddings,
)
from tests.conftest import DEFAULT_URI, VECTOR_TYPES, CharacterEmbeddings
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -59,6 +57,96 @@ def store(request) -> PostgresStore:
conn.execute(f"DROP DATABASE {database}")
def test_large_batches(store: PostgresStore) -> None:
N = 1000
M = 10
with ThreadPoolExecutor(max_workers=10) as executor:
for m in range(M):
for i in range(N):
_ = [
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
),
executor.submit(
store.get,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
),
executor.submit(
store.list_namespaces,
prefix=None,
max_depth=m + 1,
),
executor.submit(
store.search,
("test",),
),
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
),
executor.submit(
store.put,
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
None,
),
]
async def test_large_batches_async(store: PostgresStore) -> None:
N = 1000
M = 10
coros = []
for m in range(M):
for i in range(N):
coros.append(
store.aput(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
)
)
coros.append(
store.aget(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
)
)
coros.append(
store.alist_namespaces(
prefix=None,
max_depth=m + 1,
)
)
coros.append(
store.asearch(
("test",),
)
)
coros.append(
store.aput(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
value={"foo": "bar" + str(i)},
)
)
coros.append(
store.adelete(
("test", "foo", "bar", "baz", str(m % 2)),
f"key{i}",
)
)
await asyncio.gather(*coros)
def test_batch_order(store: PostgresStore) -> None:
# Setup test data
store.put(("test", "foo"), "key1", {"data": "value1"})