WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add test for update_state when using custom state class

Browse Source
This commit is contained in:
Nuno Campos
2024-08-05 08:54:16 -07:00
parent 4f2f5f7cbc
commit e8c2fefb29
3 changed files with 281 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/tests/__snapshots__/test_pregel_async.ambr
+121
View File
@@ -36,6 +36,127 @@
+---------+
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2.1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'anyOf': list([
dict({
'type': 'string',
}),
dict({
'type': 'null',
}),
]),
'default': None,
'title': 'Answer',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2.2
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'anyOf': list([
dict({
'type': 'string',
}),
dict({
'type': 'null',
}),
]),
'default': None,
'title': 'Answer',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch
'''
+-----------+
libs/langgraph/tests/test_pregel.py
+21
View File
@@ -6894,6 +6894,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
return {"answer": ",".join(data.docs)}
def decider(data: State) -> str:
print("decider", data)
assert isinstance(data, State)
return "retriever_two"
@@ -6960,6 +6961,16 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert app_w_interrupt.update_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
}
}
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
snapshot: SnapshotAssertion,
@@ -7072,6 +7083,16 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert app_w_interrupt.update_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
}
}
def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
def sorted_add(
libs/langgraph/tests/test_pregel_async.py
+139
View File
@@ -5606,6 +5606,145 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert await app_w_interrupt.aupdate_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
}
}
async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
snapshot: SnapshotAssertion,
) -> None:
from pydantic import BaseModel, ValidationError
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
if isinstance(y[0], tuple):
for rem, _ in y:
x.remove(rem)
y = [t[1] for t in y]
return sorted(operator.add(x, y))
class InnerObject(BaseModel):
yo: int
class State(BaseModel):
query: str
inner: InnerObject
answer: Optional[str] = None
docs: Annotated[list[str], sorted_add]
class StateUpdate(BaseModel):
query: Optional[str] = None
answer: Optional[str] = None
docs: Optional[list[str]] = None
async def rewrite_query(data: State) -> State:
return {"query": f"query: {data.query}"}
async def analyzer_one(data: State) -> State:
return StateUpdate(query=f"analyzed: {data.query}")
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
await asyncio.sleep(0.1)
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data.docs)}
async def decider(data: State) -> str:
assert isinstance(data, State)
return "retriever_two"
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("analyzer_one", analyzer_one)
workflow.add_node("retriever_one", retriever_one)
workflow.add_node("retriever_two", retriever_two)
workflow.add_node("qa", qa)
workflow.set_entry_point("rewrite_query")
workflow.add_edge("rewrite_query", "analyzer_one")
workflow.add_edge("analyzer_one", "retriever_one")
workflow.add_conditional_edges(
"rewrite_query", decider, {"retriever_two": "retriever_two"}
)
workflow.add_edge(["retriever_one", "retriever_two"], "qa")
workflow.set_finish_point("qa")
app = workflow.compile()
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.get_input_schema().schema() == snapshot
assert app.get_output_schema().schema() == snapshot
with pytest.raises(ValidationError):
await app.ainvoke({"query": {}})
assert await app.ainvoke(
{"query": "what is weather in sf", "inner": {"yo": 1}}
) == {
"query": "analyzed: query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
"inner": {"yo": 1},
}
assert [
c
async for c in app.astream(
{"query": "what is weather in sf", "inner": {"yo": 1}}
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"query": "what is weather in sf", "inner": {"yo": 1}}, config
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert await app_w_interrupt.aupdate_state(
config, {"docs": ["doc5"]}, as_node="rewrite_query"
) == {
"configurable": {
"thread_id": "1",
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
}
}
async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
def sorted_add(