mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 22:52:29 +02:00
530 lines
19 KiB
Python
530 lines
19 KiB
Python
# type: ignore
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from conftest import DEFAULT_URI # type: ignore
|
|
|
|
from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
|
|
from langgraph.store.postgres import AsyncPostgresStore
|
|
|
|
|
|
class MockAsyncCursor:
|
|
def __init__(self, fetch_result: Any) -> None:
|
|
self.fetch_result = fetch_result
|
|
self.execute = AsyncMock()
|
|
self.fetchall = AsyncMock(return_value=self.fetch_result)
|
|
|
|
|
|
class MockAsyncConnection:
|
|
def __init__(self) -> None:
|
|
self.cursor = MagicMock()
|
|
self.pipeline = MagicMock(
|
|
return_value=AsyncMock(__aenter__=AsyncMock(), __aexit__=AsyncMock())
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_connection() -> MockAsyncConnection:
|
|
return MockAsyncConnection()
|
|
|
|
|
|
@pytest.fixture
|
|
async def store(mock_connection: MockAsyncConnection) -> AsyncPostgresStore:
|
|
return AsyncPostgresStore(mock_connection)
|
|
|
|
|
|
async def test_abatch_order(store: AsyncPostgresStore) -> None:
|
|
mock_connection = store.conn
|
|
mock_get_cursor = MockAsyncCursor(
|
|
[
|
|
{
|
|
"key": "key1",
|
|
"value": '{"data": "value1"}',
|
|
"created_at": datetime.now(),
|
|
"updated_at": datetime.now(),
|
|
"prefix": "test.foo",
|
|
},
|
|
{
|
|
"key": "key2",
|
|
"value": '{"data": "value2"}',
|
|
"created_at": datetime.now(),
|
|
"updated_at": datetime.now(),
|
|
"prefix": "test.bar",
|
|
},
|
|
]
|
|
)
|
|
mock_search_cursor = MockAsyncCursor(
|
|
[
|
|
{
|
|
"key": "key1",
|
|
"value": '{"data": "value1"}',
|
|
"created_at": datetime.now(),
|
|
"updated_at": datetime.now(),
|
|
"prefix": "test.foo",
|
|
},
|
|
]
|
|
)
|
|
mock_list_namespaces_cursor = MockAsyncCursor(
|
|
[
|
|
{"truncated_prefix": b"\x01test"},
|
|
]
|
|
)
|
|
|
|
failures = []
|
|
|
|
def cursor_side_effect(binary: bool = False) -> Any:
|
|
cursor = MagicMock()
|
|
|
|
async def execute_side_effect(query: str, *params: Any) -> None:
|
|
# My super sophisticated database.
|
|
if "SELECT prefix, key," in query:
|
|
cursor.fetchall = mock_search_cursor.fetchall
|
|
elif "SELECT DISTINCT ON (truncated_prefix)" in query:
|
|
cursor.fetchall = mock_list_namespaces_cursor.fetchall
|
|
elif "WHERE prefix = %s AND key" in query:
|
|
cursor.fetchall = mock_get_cursor.fetchall
|
|
elif "INSERT INTO " in query:
|
|
pass
|
|
else:
|
|
e = ValueError(f"Unmatched query: {query}")
|
|
failures.append(e)
|
|
raise e
|
|
|
|
cursor.execute = AsyncMock(side_effect=execute_side_effect)
|
|
return cursor
|
|
|
|
mock_connection.cursor.side_effect = cursor_side_effect # type: ignore
|
|
|
|
ops = [
|
|
GetOp(namespace=("test",), key="key1"),
|
|
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
|
|
SearchOp(
|
|
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
|
|
),
|
|
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
|
|
GetOp(namespace=("test",), key="key3"),
|
|
]
|
|
results = await store.abatch(ops)
|
|
assert not failures
|
|
assert len(results) == 5
|
|
assert isinstance(results[0], Item)
|
|
assert isinstance(results[0].value, dict)
|
|
assert results[0].value == {"data": "value1"}
|
|
assert results[0].key == "key1"
|
|
assert results[1] is None
|
|
assert isinstance(results[2], list)
|
|
assert len(results[2]) == 1
|
|
assert isinstance(results[3], list)
|
|
assert results[3] == [("test",)]
|
|
assert results[4] is None
|
|
|
|
ops_reordered = [
|
|
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
|
|
GetOp(namespace=("test",), key="key2"),
|
|
ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
|
|
PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
|
|
GetOp(namespace=("test",), key="key1"),
|
|
]
|
|
|
|
results_reordered = await store.abatch(ops_reordered)
|
|
assert not failures
|
|
assert len(results_reordered) == 5
|
|
assert isinstance(results_reordered[0], list)
|
|
assert len(results_reordered[0]) == 1
|
|
assert isinstance(results_reordered[1], Item)
|
|
assert results_reordered[1].value == {"data": "value2"}
|
|
assert results_reordered[1].key == "key2"
|
|
assert isinstance(results_reordered[2], list)
|
|
assert results_reordered[2] == [("test",)]
|
|
assert results_reordered[3] is None
|
|
assert isinstance(results_reordered[4], Item)
|
|
assert results_reordered[4].value == {"data": "value1"}
|
|
assert results_reordered[4].key == "key1"
|
|
|
|
|
|
async def test_batch_get_ops(store: AsyncPostgresStore) -> None:
|
|
mock_connection = store.conn
|
|
mock_cursor = MockAsyncCursor(
|
|
[
|
|
{
|
|
"key": "key1",
|
|
"value": '{"data": "value1"}',
|
|
"created_at": datetime.now(),
|
|
"updated_at": datetime.now(),
|
|
"prefix": "test.foo",
|
|
},
|
|
{
|
|
"key": "key2",
|
|
"value": '{"data": "value2"}',
|
|
"created_at": datetime.now(),
|
|
"updated_at": datetime.now(),
|
|
"prefix": "test.bar",
|
|
},
|
|
]
|
|
)
|
|
mock_connection.cursor.return_value = mock_cursor
|
|
|
|
ops = [
|
|
GetOp(namespace=("test",), key="key1"),
|
|
GetOp(namespace=("test",), key="key2"),
|
|
GetOp(namespace=("test",), key="key3"),
|
|
]
|
|
|
|
results = await store.abatch(ops)
|
|
|
|
assert len(results) == 3
|
|
assert results[0] is not None
|
|
assert results[1] is not None
|
|
assert results[2] is None
|
|
assert results[0].key == "key1"
|
|
assert results[1].key == "key2"
|
|
|
|
|
|
async def test_batch_put_ops(store: AsyncPostgresStore) -> None:
|
|
mock_connection = store.conn
|
|
mock_cursor = MockAsyncCursor([])
|
|
mock_connection.cursor.return_value = mock_cursor
|
|
|
|
ops = [
|
|
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
|
|
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
|
|
PutOp(namespace=("test",), key="key3", value=None),
|
|
]
|
|
|
|
results = await store.abatch(ops)
|
|
|
|
assert len(results) == 3
|
|
assert all(result is None for result in results)
|
|
assert mock_cursor.execute.call_count == 2
|
|
|
|
|
|
async def test_batch_search_ops(store: AsyncPostgresStore) -> None:
|
|
mock_connection = store.conn
|
|
mock_cursor = MockAsyncCursor(
|
|
[
|
|
{
|
|
"key": "key1",
|
|
"value": '{"data": "value1"}',
|
|
"created_at": datetime.now(),
|
|
"updated_at": datetime.now(),
|
|
"prefix": "test.foo",
|
|
},
|
|
{
|
|
"key": "key2",
|
|
"value": '{"data": "value2"}',
|
|
"created_at": datetime.now(),
|
|
"updated_at": datetime.now(),
|
|
"prefix": "test.bar",
|
|
},
|
|
]
|
|
)
|
|
mock_connection.cursor.return_value = mock_cursor
|
|
|
|
ops = [
|
|
SearchOp(
|
|
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
|
|
),
|
|
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
|
|
]
|
|
|
|
results = await store.abatch(ops)
|
|
|
|
assert len(results) == 2
|
|
assert len(results[0]) == 2
|
|
assert len(results[1]) == 2
|
|
|
|
|
|
async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
|
|
mock_connection = store.conn
|
|
mock_cursor = MockAsyncCursor(
|
|
[
|
|
{"truncated_prefix": b"\x01test.namespace1"},
|
|
{"truncated_prefix": b"\x01test.namespace2"},
|
|
]
|
|
)
|
|
mock_connection.cursor.return_value = mock_cursor
|
|
|
|
ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
|
|
|
|
results = await store.abatch(ops)
|
|
|
|
assert len(results) == 1
|
|
assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
|
|
|
|
|
|
# The following use the actual DB connection
|
|
|
|
|
|
class TestAsyncPostgresStore:
|
|
@pytest.fixture(autouse=True)
|
|
async def setup(self) -> None:
|
|
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
|
|
await store.setup()
|
|
|
|
async def test_basic_store_ops(self) -> None:
|
|
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
|
|
namespace = ("test", "documents")
|
|
item_id = "doc1"
|
|
item_value = {"title": "Test Document", "content": "Hello, World!"}
|
|
|
|
await store.aput(namespace, item_id, item_value)
|
|
item = await store.aget(namespace, item_id)
|
|
|
|
assert item
|
|
assert item.namespace == namespace
|
|
assert item.key == item_id
|
|
assert item.value == item_value
|
|
|
|
updated_value = {
|
|
"title": "Updated Test Document",
|
|
"content": "Hello, LangGraph!",
|
|
}
|
|
await store.aput(namespace, item_id, updated_value)
|
|
updated_item = await store.aget(namespace, item_id)
|
|
|
|
assert updated_item.value == updated_value
|
|
assert updated_item.updated_at > item.updated_at
|
|
different_namespace = ("test", "other_documents")
|
|
item_in_different_namespace = await store.aget(different_namespace, item_id)
|
|
assert item_in_different_namespace is None
|
|
|
|
new_item_id = "doc2"
|
|
new_item_value = {"title": "Another Document", "content": "Greetings!"}
|
|
await store.aput(namespace, new_item_id, new_item_value)
|
|
|
|
search_results = await store.asearch(["test"], limit=10)
|
|
items = search_results
|
|
assert len(items) == 2
|
|
assert any(item.key == item_id for item in items)
|
|
assert any(item.key == new_item_id for item in items)
|
|
|
|
namespaces = await store.alist_namespaces(prefix=["test"])
|
|
assert ("test", "documents") in namespaces
|
|
|
|
await store.adelete(namespace, item_id)
|
|
await store.adelete(namespace, new_item_id)
|
|
deleted_item = await store.aget(namespace, item_id)
|
|
assert deleted_item is None
|
|
|
|
deleted_item = await store.aget(namespace, new_item_id)
|
|
assert deleted_item is None
|
|
|
|
empty_search_results = await store.asearch(["test"], limit=10)
|
|
assert len(empty_search_results) == 0
|
|
|
|
async def test_list_namespaces(self) -> None:
|
|
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
|
|
test_pref = str(uuid.uuid4())
|
|
test_namespaces = [
|
|
(test_pref, "test", "documents", "public", test_pref),
|
|
(test_pref, "test", "documents", "private", test_pref),
|
|
(test_pref, "test", "images", "public", test_pref),
|
|
(test_pref, "test", "images", "private", test_pref),
|
|
(test_pref, "prod", "documents", "public", test_pref),
|
|
(
|
|
test_pref,
|
|
"prod",
|
|
"documents",
|
|
"some",
|
|
"nesting",
|
|
"public",
|
|
test_pref,
|
|
),
|
|
(test_pref, "prod", "documents", "private", test_pref),
|
|
]
|
|
|
|
for namespace in test_namespaces:
|
|
await store.aput(namespace, "dummy", {"content": "dummy"})
|
|
|
|
prefix_result = await store.alist_namespaces(prefix=[test_pref, "test"])
|
|
assert len(prefix_result) == 4
|
|
assert all([ns[1] == "test" for ns in prefix_result])
|
|
|
|
specific_prefix_result = await store.alist_namespaces(
|
|
prefix=[test_pref, "test", "documents"]
|
|
)
|
|
assert len(specific_prefix_result) == 2
|
|
assert all(
|
|
[ns[1:3] == ("test", "documents") for ns in specific_prefix_result]
|
|
)
|
|
|
|
suffix_result = await store.alist_namespaces(suffix=["public", test_pref])
|
|
assert len(suffix_result) == 4
|
|
assert all(ns[-2] == "public" for ns in suffix_result)
|
|
|
|
prefix_suffix_result = await store.alist_namespaces(
|
|
prefix=[test_pref, "test"], suffix=["public", test_pref]
|
|
)
|
|
assert len(prefix_suffix_result) == 2
|
|
assert all(
|
|
ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
|
|
)
|
|
|
|
wildcard_prefix_result = await store.alist_namespaces(
|
|
prefix=[test_pref, "*", "documents"]
|
|
)
|
|
assert len(wildcard_prefix_result) == 5
|
|
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
|
|
|
|
wildcard_suffix_result = await store.alist_namespaces(
|
|
suffix=["*", "public", test_pref]
|
|
)
|
|
assert len(wildcard_suffix_result) == 4
|
|
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
|
|
wildcard_single = await store.alist_namespaces(
|
|
suffix=["some", "*", "public", test_pref]
|
|
)
|
|
assert len(wildcard_single) == 1
|
|
assert wildcard_single[0] == (
|
|
test_pref,
|
|
"prod",
|
|
"documents",
|
|
"some",
|
|
"nesting",
|
|
"public",
|
|
test_pref,
|
|
)
|
|
|
|
max_depth_result = await store.alist_namespaces(max_depth=3)
|
|
assert all([len(ns) <= 3 for ns in max_depth_result])
|
|
max_depth_result = await store.alist_namespaces(
|
|
max_depth=4, prefix=[test_pref, "*", "documents"]
|
|
)
|
|
assert (
|
|
len(set(tuple(res) for res in max_depth_result))
|
|
== len(max_depth_result)
|
|
== 5
|
|
)
|
|
|
|
limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
|
|
assert len(limit_result) == 3
|
|
|
|
offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
|
|
assert len(offset_result) == len(test_namespaces) - 3
|
|
|
|
empty_prefix_result = await store.alist_namespaces(prefix=[test_pref])
|
|
assert len(empty_prefix_result) == len(test_namespaces)
|
|
assert set(tuple(ns) for ns in empty_prefix_result) == set(
|
|
tuple(ns) for ns in test_namespaces
|
|
)
|
|
|
|
for namespace in test_namespaces:
|
|
await store.adelete(namespace, "dummy")
|
|
|
|
async def test_search(self):
|
|
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
|
|
test_namespaces = [
|
|
("test_search", "documents", "user1"),
|
|
("test_search", "documents", "user2"),
|
|
("test_search", "reports", "department1"),
|
|
("test_search", "reports", "department2"),
|
|
]
|
|
test_items = [
|
|
{"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
|
|
{"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
|
|
{"title": "Report A", "author": "John Doe", "tags": ["final"]},
|
|
{"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
|
|
]
|
|
empty = await store.asearch(
|
|
(
|
|
"scoped",
|
|
"assistant_id",
|
|
"shared",
|
|
"6c5356f6-63ab-4158-868d-cd9fd14c736e",
|
|
),
|
|
limit=10,
|
|
offset=0,
|
|
)
|
|
assert len(empty) == 0
|
|
|
|
for namespace, item in zip(test_namespaces, test_items):
|
|
await store.aput(namespace, f"item_{namespace[-1]}", item)
|
|
|
|
docs_result = await store.asearch(["test_search", "documents"])
|
|
assert len(docs_result) == 2
|
|
assert all([item.namespace[1] == "documents" for item in docs_result]), [
|
|
item.namespace for item in docs_result
|
|
]
|
|
|
|
reports_result = await store.asearch(["test_search", "reports"])
|
|
assert len(reports_result) == 2
|
|
assert all(item.namespace[1] == "reports" for item in reports_result)
|
|
|
|
limited_result = await store.asearch(["test_search"], limit=2)
|
|
assert len(limited_result) == 2
|
|
offset_result = await store.asearch(["test_search"])
|
|
assert len(offset_result) == 4
|
|
|
|
offset_result = await store.asearch(["test_search"], offset=2)
|
|
assert len(offset_result) == 2
|
|
assert all(item not in limited_result for item in offset_result)
|
|
|
|
john_doe_result = await store.asearch(
|
|
["test_search"], filter={"author": "John Doe"}
|
|
)
|
|
assert len(john_doe_result) == 2
|
|
assert all(item.value["author"] == "John Doe" for item in john_doe_result)
|
|
|
|
draft_result = await store.asearch(
|
|
["test_search"], filter={"tags": ["draft"]}
|
|
)
|
|
assert len(draft_result) == 2
|
|
assert all("draft" in item.value["tags"] for item in draft_result)
|
|
|
|
page1 = await store.asearch(["test_search"], limit=2, offset=0)
|
|
page2 = await store.asearch(["test_search"], limit=2, offset=2)
|
|
all_items = page1 + page2
|
|
assert len(all_items) == 4
|
|
assert len(set(item.key for item in all_items)) == 4
|
|
empty = await store.asearch(
|
|
(
|
|
"scoped",
|
|
"assistant_id",
|
|
"shared",
|
|
"again",
|
|
"maybe",
|
|
"some-long",
|
|
"6be5cb0e-2eb4-42e6-bb6b-fba3c269db25",
|
|
),
|
|
limit=10,
|
|
offset=0,
|
|
)
|
|
assert len(empty) == 0
|
|
|
|
# Test with a namespace beginning with a number (like a UUID)
|
|
uuid_namespace = (str(uuid.uuid4()), "documents")
|
|
uuid_item_id = "uuid_doc"
|
|
uuid_item_value = {
|
|
"title": "UUID Document",
|
|
"content": "This document has a UUID namespace.",
|
|
}
|
|
|
|
# Insert the item with the UUID namespace
|
|
await store.aput(uuid_namespace, uuid_item_id, uuid_item_value)
|
|
|
|
# Retrieve the item to verify it was stored correctly
|
|
retrieved_item = await store.aget(uuid_namespace, uuid_item_id)
|
|
assert retrieved_item is not None
|
|
assert retrieved_item.namespace == uuid_namespace
|
|
assert retrieved_item.key == uuid_item_id
|
|
assert retrieved_item.value == uuid_item_value
|
|
|
|
# Search for the item using the UUID namespace
|
|
search_result = await store.asearch([uuid_namespace[0]])
|
|
assert len(search_result) == 1
|
|
assert search_result[0].key == uuid_item_id
|
|
assert search_result[0].value == uuid_item_value
|
|
|
|
# Clean up: delete the item with the UUID namespace
|
|
await store.adelete(uuid_namespace, uuid_item_id)
|
|
|
|
# Verify the item was deleted
|
|
deleted_item = await store.aget(uuid_namespace, uuid_item_id)
|
|
assert deleted_item is None
|
|
|
|
for namespace in test_namespaces:
|
|
await store.adelete(namespace, f"item_{namespace[-1]}")
|