Compare commits

...

1 Commits

6 changed files with 41 additions and 25 deletions
@@ -29,14 +29,11 @@ def _write_config_file(data: dict[str, Any]) -> Path:
def _run_async(coro):
try:
loop = asyncio.get_event_loop()
if loop.is_closed():
raise RuntimeError
if loop.is_running():
return asyncio.run(coro)
except Exception: # pylint: disable=broad-except
loop = asyncio.get_running_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
if loop.is_running():
return asyncio.run(coro)
return loop.run_until_complete(coro)
+24 -11
View File
@@ -25,6 +25,18 @@ _DEF_CHECKPOINT_DIR = Path("/tmp/langgraph_checkpoints")
_TEST_LOOP = asyncio.new_event_loop()
asyncio.set_event_loop(_TEST_LOOP)
_TEST_EVENT_LOOP = None
def _get_test_loop():
"""Safe event loop accessor for synchronous test step functions."""
global _TEST_EVENT_LOOP
try:
return asyncio.get_running_loop()
except RuntimeError:
if _TEST_EVENT_LOOP is None:
_TEST_EVENT_LOOP = asyncio.new_event_loop()
asyncio.set_event_loop(_TEST_EVENT_LOOP)
return _TEST_EVENT_LOOP
def _standard_topological_levels(nodes: dict[str, object]) -> dict[int, set[str]]:
middle_nodes = {name for name in nodes if name not in {"start", "end"}}
@@ -153,14 +165,11 @@ def _ensure_topological_levels(graph):
def _run_async(coro):
try:
loop = asyncio.get_event_loop()
if loop.is_closed():
raise RuntimeError
if loop.is_running():
return asyncio.run(coro)
except Exception: # pylint: disable=broad-except
loop = asyncio.get_running_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
if loop.is_running():
return asyncio.run(coro)
return loop.run_until_complete(coro)
@@ -435,7 +444,11 @@ def step_graph_route_idle_time(context):
@given("I have a graph state that has been idle too long")
def step_graph_state_idle(context):
context.graph_state = GraphState()
current_time = asyncio.get_event_loop().time()
try:
current_time = asyncio.get_running_loop().time()
except RuntimeError:
import time as _t2
current_time = _t2.monotonic()
context.graph_state.metadata = {"last_updated": current_time - 10.0}
context.graph_state.messages = []
@@ -673,7 +686,7 @@ def step_attempt_downgrade(context):
)
return
try:
context.downgrade_result = asyncio.get_event_loop().run_until_complete(
context.downgrade_result = _get_test_loop().run_until_complete(
context.route_bridge.downgrade_graph_to_stream(
context.route_config, context.langgraph
)
@@ -709,7 +722,7 @@ def step_route_upgrade_downgrade(context):
@when("I perform an upgrade and then downgrade cycle")
def step_cycle_upgrade_downgrade(context):
context.langgraph = asyncio.get_event_loop().run_until_complete(
context.langgraph = _get_test_loop().run_until_complete(
context.route_bridge.upgrade_stream_to_graph(
context.route_config, context.stream_message
)
@@ -731,7 +744,7 @@ def step_cycle_upgrade_downgrade(context):
name=name, type=getattr(node, "type", NodeType.FUNCTION)
)
context.stream_config = asyncio.get_event_loop().run_until_complete(
context.stream_config = _get_test_loop().run_until_complete(
context.route_bridge.downgrade_graph_to_stream(
RouteConfig.from_graph_config(
GraphConfig(
+1 -3
View File
@@ -54,9 +54,7 @@ class Agent(ABC):
def process_message_sync(
self, message: Any, context: dict[str, Any] | None = None
) -> Any:
return asyncio.get_event_loop().run_until_complete(
self.process_message(message, context or {})
)
return asyncio.run(self.process_message(message, context or {}))
@abstractmethod
async def process_message(
+4 -1
View File
@@ -224,7 +224,10 @@ class RxPyLangGraphBridge:
)
def create_future_task(msg: StreamMessage) -> asyncio.Future[StreamMessage]:
loop = asyncio.get_event_loop()
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
task: asyncio.Future[StreamMessage] = asyncio.ensure_future(
execute_graph(msg), loop=loop
)
+2 -2
View File
@@ -106,7 +106,7 @@ class Node: # pylint: disable=too-many-instance-attributes
async def execute(self, state: GraphState) -> dict[str, Any]: # pylint: disable=too-many-branches
self.execution_count += 1
loop = asyncio.get_event_loop()
loop = asyncio.get_running_loop()
start_time = loop.time()
try:
state_updates = {"current_node": self.name}
@@ -271,7 +271,7 @@ class Node: # pylint: disable=too-many-instance-attributes
if asyncio.iscoroutinefunction(fn):
result = await fn(state)
else:
loop = asyncio.get_event_loop()
loop = asyncio.get_running_loop()
result = await loop.run_in_executor(None, lambda: fn(state))
if asyncio.iscoroutine(result):
result = await result
+6 -1
View File
@@ -84,9 +84,14 @@ class RouteBridge:
if "idle_time" in conditions:
idle_threshold = conditions["idle_time"]
last_updated = state.metadata.get("last_updated")
try:
now_time = asyncio.get_running_loop().time()
except RuntimeError:
import time as _t_mod
now_time = _t_mod.monotonic()
if (
last_updated
and (asyncio.get_event_loop().time() - last_updated) > idle_threshold
and (now_time - last_updated) > idle_threshold
):
return True
if (