This commit is contained in:
Nuno Campos
2025-01-22 16:02:09 -08:00
parent 3f2557c9c9
commit c697c2aa04
2 changed files with 86 additions and 14 deletions
+43 -7
View File
@@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph
from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.any import AnyDict, AnyInt
from tests.drain import drain_topics_async
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -198,7 +198,13 @@ async def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": None,
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -265,7 +271,13 @@ async def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -362,7 +374,13 @@ async def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -469,7 +487,13 @@ async def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": None,
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -531,7 +555,13 @@ async def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -649,7 +679,13 @@ async def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -15,7 +15,7 @@ from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.any import AnyDict, AnyInt
from tests.drain import drain_topics
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -197,7 +197,13 @@ def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": None,
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -264,7 +270,13 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_previous": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -361,7 +373,13 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_previous": None,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -467,7 +485,13 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_previous": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": None,
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -529,7 +553,13 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_previous": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -647,7 +677,13 @@ def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]