diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..f36d689 --- /dev/null +++ b/.env.example @@ -0,0 +1,18 @@ +# Copy to `.env` (which is git-ignored) and fill in the keys you have. +# The integration tests under tests_integration/ read these via python-dotenv. +# +# cp .env.example .env +# +# Only the providers whose keys are set will run; the rest fail/skip. + +# Required by the OpenAI-based tests (incl. concurrent_invocation_example.py). +OPENAI_API_KEY= + +# Other LLM providers used by tests_integration/. +ANTHROPIC_API_KEY= +GEMINI_API_KEY= +DEEPSEEK_API_KEY= +OPENROUTER_API_KEY= + +# Search tool used by some function-call tests. +TAVILY_API_KEY= diff --git a/grafi/workflows/impl/event_driven_workflow.py b/grafi/workflows/impl/event_driven_workflow.py index 8070bac..a3627c3 100644 --- a/grafi/workflows/impl/event_driven_workflow.py +++ b/grafi/workflows/impl/event_driven_workflow.py @@ -1,26 +1,18 @@ -import asyncio -from collections import deque from typing import Any from typing import AsyncGenerator from typing import Dict from typing import List -from typing import Set -from loguru import logger from openinference.semconv.trace import OpenInferenceSpanKindValues from pydantic import PrivateAttr from grafi.common.containers.container import container from grafi.common.decorators.record_decorators import record_workflow_invoke -from grafi.common.events.event import Event from grafi.common.events.topic_events.consume_from_topic_event import ( ConsumeFromTopicEvent, ) from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent -from grafi.common.events.topic_events.topic_event import TopicEvent -from grafi.common.exceptions import NodeExecutionError from grafi.common.exceptions import WorkflowError -from grafi.common.models.invoke_context import InvokeContext from grafi.nodes.node import Node from grafi.nodes.node_base import NodeBase from grafi.tools.function_calls.function_call_tool import FunctionCallTool @@ -28,14 +20,9 @@ from grafi.topics.expressions.topic_expression import extract_topics from grafi.topics.topic_base import TopicBase from grafi.topics.topic_factory import TopicFactory -from grafi.topics.topic_impl.in_workflow_input_topic import InWorkflowInputTopic from grafi.topics.topic_impl.in_workflow_output_topic import InWorkflowOutputTopic from grafi.topics.topic_types import TopicType -from grafi.workflows.impl.async_node_tracker import AsyncNodeTracker -from grafi.workflows.impl.async_output_queue import AsyncOutputQueue -from grafi.workflows.impl.utils import get_async_output_events -from grafi.workflows.impl.utils import get_node_input -from grafi.workflows.impl.utils import publish_events +from grafi.workflows.impl.workflow_run import WorkflowRun from grafi.workflows.workflow import Workflow from grafi.workflows.workflow import WorkflowBuilder @@ -44,8 +31,12 @@ class EventDrivenWorkflow(Workflow): """ An event-driven workflow that invokes a directed graph of Nodes in response to topic publish events. - This workflow can handle streaming events via `StreamTopicEvent` and relay them to a custom - `stream_event_handler`. + This class is the immutable *definition*: it owns the topology (nodes, topic + configuration, and the topic→node subscription map) and links function-tool + specs onto LLM nodes. It owns no runtime state. Each ``invoke()`` creates a + :class:`~grafi.workflows.impl.workflow_run.WorkflowRun` that holds all mutable + per-invocation state (topic queues, tracker, ready-queue, stop flag), so + concurrent invocations of one instance are isolated from each other. """ name: str = "EventDrivenWorkflow" @@ -54,20 +45,16 @@ class EventDrivenWorkflow(Workflow): # OpenInference semantic attribute oi_span_type: OpenInferenceSpanKindValues = OpenInferenceSpanKindValues.AGENT - # Topics known to this workflow (e.g., "agent_input", "agent_stream_output") + # Topics known to this workflow (e.g., "agent_input", "agent_stream_output"). + # This is the definition's topic *template*; each run gets its own copies with + # fresh queues. _topics: Dict[str, TopicBase] = PrivateAttr(default_factory=dict) # Mapping of topic_name -> list of node_names that subscribe to that topic _topic_nodes: Dict[str, List[str]] = PrivateAttr(default_factory=dict) - # Event graph for this workflow - # Queue of nodes that are ready to invoke (in response to published events) - _invoke_queue: deque[NodeBase] = PrivateAttr(default_factory=deque) - - _tracker: AsyncNodeTracker = PrivateAttr(default_factory=AsyncNodeTracker) - - # Optional callback that handles output events - # Including agent output event, stream event and hil event + # In-flight runs, keyed by id(run), so stop() can reach them. + _active_runs: Dict[int, WorkflowRun] = PrivateAttr(default_factory=dict) def model_post_init(self, _context: Any) -> None: self._add_topics() @@ -76,10 +63,13 @@ def model_post_init(self, _context: Any) -> None: def stop(self) -> None: """ Stop the workflow execution. - Overrides base class to also trigger force stop on the tracker. + + Sets the definition-level stop flag and forwards the stop to every + in-flight run so their trackers and node tasks wind down. """ super().stop() - self._tracker.force_stop_sync() + for run in list(self._active_runs.values()): + run.stop() @classmethod def builder(cls) -> WorkflowBuilder: @@ -145,8 +135,6 @@ def _handle_function_calling_nodes(self) -> None: # Map each topic -> the nodes that publish to it published_topics_to_nodes: Dict[str, List[NodeBase]] = {} - published_topics_to_nodes = {} - for node in self.nodes.values(): if isinstance(node.tool, LLM): # If the node is an LLM node, we need to check its published topics @@ -183,640 +171,24 @@ def _handle_function_calling_nodes(self) -> None: function_node.tool.get_function_specs() ) - # Workflow invoke methods - - async def _get_output_events(self) -> List[ConsumeFromTopicEvent]: - consumed_events: List[ConsumeFromTopicEvent] = [] - - output_topics = [ - topic - for topic in self._topics.values() - if topic.type == TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE - or topic.type == TopicType.AGENT_OUTPUT_TOPIC_TYPE - ] - - for output_topic in output_topics: - if await output_topic.can_consume(self.name): - events = await output_topic.consume(self.name) - for event in events: - consumed_events.append( - ConsumeFromTopicEvent( - name=event.name, - type=event.type, - consumer_name=self.name, - consumer_type=self.type, - invoke_context=event.invoke_context, - offset=event.offset, - data=event.data, - ) - ) - - return consumed_events - - async def _commit_events( - self, - consumer_name: str, - topic_events: List[ConsumeFromTopicEvent], - track_commit: bool = True, - ) -> None: - if not topic_events: - return - topic_max_offset: Dict[str, int] = {} - - for topic_event in topic_events: - topic_max_offset[topic_event.name] = max( - topic_max_offset.get(topic_event.name, 0), topic_event.offset - ) - - for topic, offset in topic_max_offset.items(): - await self._topics[topic].commit(consumer_name, offset) - - # Notify tracker that messages have been committed - # (skip if already tracked elsewhere, e.g., by output listener) - if track_commit: - logger.debug( - f"Committing {len(topic_events)} events for {consumer_name}, track_commit={track_commit}" - ) - await self._tracker.on_messages_committed( - len(topic_events), source=f"commit:{consumer_name}" - ) - - def _topic_consumers(self, topic_name: str) -> List[str]: - """Consumers that will commit each event published to ``topic_name``. - - These are the subscribing nodes, plus the workflow itself for output - topics (the output listener / ``_get_output_events`` drains them under - ``self.name``). This is the single source of truth for fan-out: one event - published here yields one delivery -- and one eventual commit -- per - consumer. Deduplicated, since a node consumes a topic once regardless of - how many of its subscription expressions reference it. - """ - consumers = list(dict.fromkeys(self._topic_nodes.get(topic_name, []))) - topic = self._topics.get(topic_name) - if topic is not None and topic.type in ( - TopicType.AGENT_OUTPUT_TOPIC_TYPE, - TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE, - ): - consumers.append(self.name) - return consumers - - async def _count_pending_consumable(self) -> int: - """Count deliveries still awaiting consumption across all topics. - - For each topic this sums the unconsumed events per consumer. The total - equals the number of commit operations the resumed run must perform to - drain the restored state, which is exactly what the quiescence tracker - needs seeded on recovery. - """ - total = 0 - for topic_name, topic in self._topics.items(): - for consumer_name in self._topic_consumers(topic_name): - total += await topic.unconsumed_count(consumer_name) - return total - - async def _progress_possible(self) -> bool: - """Whether the parallel run can still make progress. - - True while any node is actively processing or any event is still awaiting - consumption. A node can only be invoked when one of its subscribed topics - has an unconsumed event, so a zero pending-consumable count already - implies no node can invoke -- no separate can_invoke() sweep is needed. - - When this is False but the tracker is non-quiescent, the outstanding - deliveries are parked -- e.g. a node with an ``A AND B`` subscription that - consumed ``A`` but whose ``B`` will never arrive -- so the output queue - ends iteration instead of hanging on a commit that will never come. - (Per-consumer delivery accounting prevents premature termination; this - only catches the genuinely stuck case.) - """ - if not await self._tracker.is_idle(): - return True - return await self._count_pending_consumable() > 0 - - async def _add_to_invoke_queue(self, event: TopicEvent) -> None: - topic_name = event.name - - if topic_name not in self._topic_nodes: - return - - topic = self._topics[topic_name] - - # Get all nodes subscribed to this topic - subscribed_nodes = self._topic_nodes[topic_name] - - for node_name in subscribed_nodes: - node = self.nodes[node_name] - # add unprocessed node to the invoke queue - if await topic.can_consume(node_name) and await node.can_invoke(): - self._invoke_queue.append(node) - - async def invoke_sequential( - self, input_data: PublishToTopicEvent - ) -> AsyncGenerator[ConsumeFromTopicEvent, None]: - """ - Invoke the workflow with the given context and input. - Returns results when all nodes complete processing. - """ - invoke_context = input_data.invoke_context - - consumed_events: List[ConsumeFromTopicEvent] = [] - try: - # Process nodes until invoke queue is empty or workflow is stopped - while self._invoke_queue: - # Check if workflow should be stopped - if self._stop_requested: - logger.info("Workflow execution stopped by assistant request") - break - - node = self._invoke_queue.popleft() - - # Given node, collect all the messages can be linked to it - - node_consumed_events: List[ConsumeFromTopicEvent] = ( - await get_node_input(node) - ) - - # Invoke node with collected inputs - if node_consumed_events: - try: - published_events: List[PublishToTopicEvent] = [] - async for result in node.invoke( - invoke_context, node_consumed_events - ): - published_events.extend( - await publish_events( - node, - result, - self._tracker, - self._topic_consumers, - ) - ) - - for event in published_events: - await self._add_to_invoke_queue(event) - - events: List[TopicEvent] = [] - events.extend(node_consumed_events) - events.extend(published_events) - - await container.event_store.record_events(events) - except Exception as e: - raise NodeExecutionError( - node_name=node.name, - message=f"Node execution failed: {e}", - invoke_context=invoke_context, - cause=e, - ) from e - - consumed_events = await self._get_output_events() - - for event in consumed_events: - yield event - finally: - if consumed_events: - await container.event_store.record_events(consumed_events) - - async def invoke_parallel( - self, input_data: PublishToTopicEvent - ) -> AsyncGenerator[ConsumeFromTopicEvent, None]: - invoke_context = input_data.invoke_context - logger.debug( - f"invoke_parallel: tracker_id={id(self._tracker)}, metrics={await self._tracker.get_metrics()}" - ) - - # Start a background task to process all nodes (including streaming generators) - node_processing_task = [ - asyncio.create_task( - self._invoke_node( - invoke_context=invoke_context, - node=node, - ), - name=node.name, - ) - for node in self.nodes.values() - ] - - # Get output topics - output_topics: list[TopicBase] = [ - topic - for topic in self._topics.values() - if topic.type == TopicType.AGENT_OUTPUT_TOPIC_TYPE - or topic.type == TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE - ] - - # Create AsyncOutputQueue with output topics and tracker - output_queue = AsyncOutputQueue( - output_topics, - self.name, - self._tracker, - progress_possible=self._progress_possible, - ) - await output_queue.start_listeners() - - consumed_output_events: List[ConsumeFromTopicEvent] = [] - - # Wait for either new data or completion, with a timeout to check stop flag - try: - async for event in output_queue: - # Check if any node task has failed before yielding events - for i, task in enumerate(node_processing_task): - if task.done() and not task.cancelled(): - try: - result = task.result() - except Exception as task_error: - node_name = ( - list(self.nodes.keys())[i] - if i < len(self.nodes) - else f"node_{i}" - ) - logger.error( - f"Node {node_name} failed during execution: {task_error}" - ) - # Cancel remaining tasks and stop workflow - for t in node_processing_task: - if not t.done(): - t.cancel() - self.stop() - - raise NodeExecutionError( - node_name=node_name, - message=f"Node {node_name} execution failed during workflow: {task_error}", - invoke_context=invoke_context, - cause=task_error, - ) from task_error - - # Now yield the data after committing - consumed_event = ConsumeFromTopicEvent( - name=event.name, - type=event.type, - consumer_name=self.name, - consumer_type=self.type, - invoke_context=event.invoke_context, - offset=event.offset, - data=event.data, - ) - yield consumed_event - - consumed_output_events.append(consumed_event) - finally: - await output_queue.stop_listeners() - - # Commit all consumed output events to topics - # (tracking already done by output listener, so skip tracker update) - await self._commit_events( - consumer_name=self.name, - topic_events=consumed_output_events, - track_commit=False, - ) - - # 4. graceful shutdown all the nodes - self.stop() - - # process events after stopping - if consumed_output_events: - await container.event_store.record_events( - get_async_output_events(consumed_output_events) - ) - - # Wait for all node tasks to complete with proper error handling - for t in node_processing_task: - t.cancel() - node_results = await asyncio.gather( - *node_processing_task, return_exceptions=True - ) - - # Check for exceptions from node tasks and raise NodeExecutionError - for i, result in enumerate(node_results): - if isinstance(result, Exception) and not isinstance( - result, asyncio.CancelledError - ): - node_name = ( - list(self.nodes.keys())[i] - if i < len(self.nodes) - else f"node_{i}" - ) - logger.error(f"Node {node_name} failed with exception: {result}") - raise NodeExecutionError( - node_name=node_name, - message=f"Node {node_name} execution failed: {result}", - invoke_context=invoke_context, - cause=result, - ) from result - - async def _invoke_node(self, invoke_context: InvokeContext, node: NodeBase) -> None: - """Enhanced node invocation with better async patterns and error handling.""" - buffer: Dict[str, List[TopicEvent]] = {} - active_tasks: List[asyncio.Task] = [] - - async def _wait_and_buffer(consumer_name: str, topic: TopicBase) -> None: - """ - Block until *at least one* new record is available on `topic`, - put **all** currently‑available new records into its buffer. - """ - recs = await topic.consume(consumer_name) - - if topic.name not in buffer: - buffer[topic.name] = [] - buffer[topic.name].extend(recs) - - async def _ignore_cancel(task: asyncio.Task) -> None: - try: - await task - except asyncio.CancelledError: - pass - - async def wait_node_invoke(node: NodeBase) -> None: - while not node.can_invoke_with_topics(list(buffer.keys())): - # Check for stop request before creating new tasks - if self._stop_requested: - return - - # for every topic that *doesn't* have data yet, start one waiter - tasks = [ - asyncio.create_task(_wait_and_buffer(node.name, topic)) - for topic in node.subscribed_topics - ] - active_tasks.extend(tasks) - - _, pending = await asyncio.wait( - tasks, return_when=asyncio.FIRST_COMPLETED - ) - - # finished waiters have already filled their buffer inside - # _wait_and_buffer(); we just cancel the rest for this cycle - for t in pending: - t.cancel() - # silence "task was destroyed but it is pending" - asyncio.create_task(_ignore_cancel(t)) - - # Remove completed/cancelled tasks from active_tasks - active_tasks[:] = [t for t in active_tasks if not t.done()] - - def _cancel_all_active_tasks() -> None: - """Cancel all active tasks.""" - for task in active_tasks: - if not task.done(): - task.cancel() - active_tasks.clear() - - try: - while not self._stop_requested: - # Check if node can be invoked - - await wait_node_invoke(node) - - # Check again after wait_node_invoke in case stop was requested - if self._stop_requested: - break - - await self._tracker.enter(node.name) - - try: - consumed_events: List[ConsumeFromTopicEvent] = [] - - for events in buffer.values(): - for event in events: - consumed_event = ConsumeFromTopicEvent( - invoke_context=event.invoke_context, - name=event.name, - type=event.type, - consumer_name=node.name, - consumer_type=node.type, - offset=event.offset, - data=event.data, - ) - consumed_events.append(consumed_event) - - # publish before commit - node_output_events: List[PublishToTopicEvent] = [] - if consumed_events: - async for event in node.invoke(invoke_context, consumed_events): - node_output_events.extend( - await publish_events( - node=node, - publish_event=event, - tracker=self._tracker, - consumers_of=self._topic_consumers, - ) - ) - - await self._commit_events( - consumer_name=node.name, topic_events=consumed_events - ) - await container.event_store.record_events(consumed_events) - await container.event_store.record_events( - get_async_output_events(node_output_events) - ) - - except Exception as node_error: - logger.error(f"Error processing node {node.name}: {node_error}") - # Force stop the tracker so the workflow terminates - await self._tracker.force_stop() - raise NodeExecutionError( - node_name=node.name, - message=f"Async node execution failed: {node_error}", - invoke_context=invoke_context, - cause=node_error, - ) from node_error - finally: - await self._tracker.leave(node.name) - buffer.clear() # Clear buffer for next iteration - - except asyncio.CancelledError: - logger.info(f"Node {node.name} was cancelled") - _cancel_all_active_tasks() - raise - except NodeExecutionError: - _cancel_all_active_tasks() - raise # Re-raise NodeExecutionError as-is - except Exception as e: - logger.error(f"Fatal error in node {node.name} execution: {e}") - _cancel_all_active_tasks() - raise NodeExecutionError( - node_name=node.name, - message=f"Fatal error in node execution: {e}", - invoke_context=invoke_context, - cause=e, - ) from e - finally: - _cancel_all_active_tasks() - buffer.clear() # Clear buffer for next iteration - @record_workflow_invoke async def invoke( self, input_data: PublishToTopicEvent, is_sequential: bool = False ) -> AsyncGenerator[ConsumeFromTopicEvent, None]: """ Run the workflow with streaming output. + + Each call executes on its own :class:`WorkflowRun`, so the definition + instance carries no per-invocation state and concurrent invocations are + isolated. """ - invoke_context = input_data.invoke_context + run = WorkflowRun(self, container.event_store) + self._active_runs[id(run)] = run try: - # Reset stop flag at the beginning of new execution - self.reset_stop_flag() - - await self.init_workflow(input_data, is_sequential) - - if is_sequential: - # If sequential, we just call the sequential method - async for event in self.invoke_sequential(input_data): - yield event - else: - async for event in self.invoke_parallel(input_data): - yield event - - except NodeExecutionError: - raise # Re-raise NodeExecutionError as-is - except Exception as e: - raise WorkflowError( - message=f"Workflow {self.name} async execution failed: {e}", - invoke_context=invoke_context, - cause=e, - ) from e - - async def init_workflow( - self, input_data: PublishToTopicEvent, is_sequential: bool = False - ) -> Any: - # 1 – initial seeding - logger.debug( - f"init_workflow: is_sequential={is_sequential}, tracker_id={id(self._tracker)}" - ) - if not is_sequential: - self._tracker.reset() - - for topic in self._topics.values(): - await topic.reset() - - invoke_context = input_data.invoke_context - - events = [ - event - for event in await container.event_store.get_agent_events( - invoke_context.assistant_request_id - ) - if isinstance(event, TopicEvent) - ] - - if len(events) == 0: - input_topics: List[TopicBase] = [ - topic - for topic in self._topics.values() - if topic.type == TopicType.AGENT_INPUT_TOPIC_TYPE - ] - - events_to_record: List[Event] = [] - # One delivery per consumer of each seeded input topic (not one per - # topic), so a fan-out input feeding several nodes is not declared - # quiescent after only the first node commits. Accumulated here while - # ``event`` is narrowed to PublishToTopicEvent. - seeded_delivery_count = 0 - for input_topic in input_topics: - event = await input_topic.publish_data( - input_data.model_copy( - update={ - "publisher_name": self.name, - "publisher_type": self.type, - }, - deep=True, - ) - ) - if event: - events_to_record.append(event) - seeded_delivery_count += len(self._topic_consumers(event.name)) - if is_sequential: - await self._add_to_invoke_queue(event) - - logger.debug( - f"init_workflow: events_to_record={len(events_to_record)}, input_topics={len(input_topics)}" - ) - if events_to_record: - # Track initial input messages for quiescence detection. - if not is_sequential and seeded_delivery_count: - logger.debug( - f"init_workflow: on_messages_published({seeded_delivery_count})" - ) - await self._tracker.on_messages_published( - seeded_delivery_count, source="init_workflow" - ) - await container.event_store.record_events(events_to_record) - else: - # When there is unfinished workflow, we need to restore the workflow topics - for topic_event in events: - await self._topics[topic_event.name].restore_topic(topic_event) - if is_sequential and isinstance(topic_event, PublishToTopicEvent): - await self._add_to_invoke_queue(topic_event) - - # Parallel execution terminates on tracker quiescence, which is - # (no active nodes) AND (uncommitted_messages == 0). After restoring - # topics the tracker counter is still zero, so without re-seeding it - # the resumed run would be declared quiescent immediately and yield - # nothing. Seed it with the messages that are still pending - # consumption so the workflow drains the restored work. - if not is_sequential: - pending = await self._count_pending_consumable() - if pending: - await self._tracker.on_messages_published( - pending, source="recovery_restore" - ) - - # Process in-workflow topics - in_workflow_output_topic_names: Set[str] = set() - consumed_event_ids = input_data.consumed_event_ids - consumed_events = [ - event for event in events if event.event_id in consumed_event_ids - ] - in_workflow_output_topic_names = set( - [ - event.name - for event in consumed_events - if event.type == TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE - ] - ) - - for in_workflow_output_topic_name in in_workflow_output_topic_names: - in_workflow_output_topic = self._topics.get( - in_workflow_output_topic_name - ) - if in_workflow_output_topic and isinstance( - in_workflow_output_topic, InWorkflowOutputTopic - ): - # if the topic is human request topic, we need to produce a new topic event - for ( - paired_in_workflow_input_topic_name - ) in in_workflow_output_topic.paired_in_workflow_input_topic_names: - paired_in_workflow_input_topic = self._topics.get( - paired_in_workflow_input_topic_name - ) - if paired_in_workflow_input_topic and isinstance( - paired_in_workflow_input_topic, InWorkflowInputTopic - ): - paired_event = ( - await paired_in_workflow_input_topic.publish_data( - input_data.model_copy( - update={ - "publisher_name": self.name, - "publisher_type": self.type, - }, - deep=True, - ) - ) - ) - if paired_event: - # Track one delivery per consumer of the paired - # input topic for quiescence detection. - if not is_sequential: - delivery_count = len( - self._topic_consumers( - paired_in_workflow_input_topic_name - ) - ) - if delivery_count: - await self._tracker.on_messages_published( - delivery_count, - source="restore_paired_input", - ) - if is_sequential: - await self._add_to_invoke_queue(paired_event) - await container.event_store.record_event(paired_event) + async for event in run.run(input_data, is_sequential): + yield event + finally: + self._active_runs.pop(id(run), None) def to_dict(self) -> dict[str, Any]: return { diff --git a/grafi/workflows/impl/parallel_engine.py b/grafi/workflows/impl/parallel_engine.py new file mode 100644 index 0000000..bf2b640 --- /dev/null +++ b/grafi/workflows/impl/parallel_engine.py @@ -0,0 +1,255 @@ +"""Parallel execution engine for a :class:`WorkflowRun`. + +Runs one ``asyncio.Task`` per node and multiplexes the output topics through an +:class:`~grafi.workflows.impl.async_output_queue.AsyncOutputQueue`, terminating +on tracker quiescence (or the no-progress backstop). ``_invoke_node`` is the +per-node loop: wait for a satisfying set of inputs, consume, invoke, publish, +commit, and persist. +""" + +import asyncio +from typing import TYPE_CHECKING +from typing import AsyncGenerator +from typing import Dict +from typing import List + +from loguru import logger + +from grafi.common.events.topic_events.consume_from_topic_event import ( + ConsumeFromTopicEvent, +) +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.events.topic_events.topic_event import TopicEvent +from grafi.common.exceptions import NodeExecutionError +from grafi.common.models.invoke_context import InvokeContext +from grafi.nodes.node_base import NodeBase +from grafi.topics.topic_base import TopicBase +from grafi.topics.topic_types import TopicType +from grafi.workflows.impl.async_output_queue import AsyncOutputQueue +from grafi.workflows.impl.utils import get_async_output_events +from grafi.workflows.impl.utils import publish_events + +if TYPE_CHECKING: + from grafi.workflows.impl.workflow_run import WorkflowRun + + +async def invoke_parallel( + run: "WorkflowRun", input_data: PublishToTopicEvent +) -> AsyncGenerator[ConsumeFromTopicEvent, None]: + invoke_context = input_data.invoke_context + + node_processing_task = [ + asyncio.create_task( + _invoke_node(run, invoke_context=invoke_context, node=node), + name=node.name, + ) + for node in run.nodes.values() + ] + + output_topics: List[TopicBase] = [ + topic + for topic in run.topics.values() + if topic.type == TopicType.AGENT_OUTPUT_TOPIC_TYPE + or topic.type == TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE + ] + + output_queue = AsyncOutputQueue( + output_topics, + run.name, + run.tracker, + progress_possible=run._progress_possible, + ) + await output_queue.start_listeners() + + consumed_output_events: List[ConsumeFromTopicEvent] = [] + try: + async for event in output_queue: + for i, task in enumerate(node_processing_task): + if task.done() and not task.cancelled(): + try: + task.result() + except Exception as task_error: + node_name = ( + list(run.nodes.keys())[i] + if i < len(run.nodes) + else f"node_{i}" + ) + logger.error( + f"Node {node_name} failed during execution: {task_error}" + ) + for t in node_processing_task: + if not t.done(): + t.cancel() + run.stop() + raise NodeExecutionError( + node_name=node_name, + message=f"Node {node_name} execution failed during workflow: {task_error}", + invoke_context=invoke_context, + cause=task_error, + ) from task_error + + consumed_event = ConsumeFromTopicEvent( + name=event.name, + type=event.type, + consumer_name=run.name, + consumer_type=run.type, + invoke_context=event.invoke_context, + offset=event.offset, + data=event.data, + ) + yield consumed_event + consumed_output_events.append(consumed_event) + finally: + await output_queue.stop_listeners() + + await run._commit_events( + consumer_name=run.name, + topic_events=consumed_output_events, + track_commit=False, + ) + + run.stop() + + if consumed_output_events: + await run.event_store.record_events( + get_async_output_events(consumed_output_events) + ) + + for t in node_processing_task: + t.cancel() + node_results = await asyncio.gather( + *node_processing_task, return_exceptions=True + ) + + for i, result in enumerate(node_results): + if isinstance(result, Exception) and not isinstance( + result, asyncio.CancelledError + ): + node_name = ( + list(run.nodes.keys())[i] if i < len(run.nodes) else f"node_{i}" + ) + logger.error(f"Node {node_name} failed with exception: {result}") + raise NodeExecutionError( + node_name=node_name, + message=f"Node {node_name} execution failed: {result}", + invoke_context=invoke_context, + cause=result, + ) from result + + +async def _invoke_node( + run: "WorkflowRun", invoke_context: InvokeContext, node: NodeBase +) -> None: + """Node invocation loop for the parallel engine.""" + buffer: Dict[str, List[TopicEvent]] = {} + active_tasks: List[asyncio.Task] = [] + + async def _wait_and_buffer(consumer_name: str, topic: TopicBase) -> None: + recs = await topic.consume(consumer_name) + if topic.name not in buffer: + buffer[topic.name] = [] + buffer[topic.name].extend(recs) + + async def _ignore_cancel(task: asyncio.Task) -> None: + try: + await task + except asyncio.CancelledError: + pass + + async def wait_node_invoke(node: NodeBase) -> None: + while not node.can_invoke_with_topics(list(buffer.keys())): + if run._stop_requested: + return + # Wait on this run's copy of each subscribed topic. + tasks = [ + asyncio.create_task(_wait_and_buffer(node.name, run.topics[topic.name])) + for topic in node.subscribed_topics + ] + active_tasks.extend(tasks) + _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for t in pending: + t.cancel() + asyncio.create_task(_ignore_cancel(t)) + active_tasks[:] = [t for t in active_tasks if not t.done()] + + def _cancel_all_active_tasks() -> None: + for task in active_tasks: + if not task.done(): + task.cancel() + active_tasks.clear() + + try: + while not run._stop_requested: + await wait_node_invoke(node) + + if run._stop_requested: + break + + await run.tracker.enter(node.name) + try: + consumed_events: List[ConsumeFromTopicEvent] = [] + for events in buffer.values(): + for event in events: + consumed_events.append( + ConsumeFromTopicEvent( + invoke_context=event.invoke_context, + name=event.name, + type=event.type, + consumer_name=node.name, + consumer_type=node.type, + offset=event.offset, + data=event.data, + ) + ) + + node_output_events: List[PublishToTopicEvent] = [] + if consumed_events: + async for event in node.invoke(invoke_context, consumed_events): + node_output_events.extend( + await publish_events( + node=node, + publish_event=event, + tracker=run.tracker, + consumers_of=run._topic_consumers, + topics=run.topics, + ) + ) + + await run._commit_events( + consumer_name=node.name, topic_events=consumed_events + ) + await run.event_store.record_events(consumed_events) + await run.event_store.record_events( + get_async_output_events(node_output_events) + ) + except Exception as node_error: + logger.error(f"Error processing node {node.name}: {node_error}") + await run.tracker.force_stop() + raise NodeExecutionError( + node_name=node.name, + message=f"Async node execution failed: {node_error}", + invoke_context=invoke_context, + cause=node_error, + ) from node_error + finally: + await run.tracker.leave(node.name) + buffer.clear() + except asyncio.CancelledError: + logger.info(f"Node {node.name} was cancelled") + _cancel_all_active_tasks() + raise + except NodeExecutionError: + _cancel_all_active_tasks() + raise + except Exception as e: + logger.error(f"Fatal error in node {node.name} execution: {e}") + _cancel_all_active_tasks() + raise NodeExecutionError( + node_name=node.name, + message=f"Fatal error in node execution: {e}", + invoke_context=invoke_context, + cause=e, + ) from e + finally: + _cancel_all_active_tasks() + buffer.clear() diff --git a/grafi/workflows/impl/run_recovery.py b/grafi/workflows/impl/run_recovery.py new file mode 100644 index 0000000..4b9e0cf --- /dev/null +++ b/grafi/workflows/impl/run_recovery.py @@ -0,0 +1,167 @@ +"""Seeding and event-replay recovery for :meth:`WorkflowRun.init`. + +For a fresh request (no prior events for the ``assistant_request_id``) this seeds +the input topics. Otherwise it restores topic state from the persisted events and +re-seeds the tracker with the work still pending consumption, so a resumed +parallel run drains the restored work instead of declaring immediate quiescence. + +The run starts with empty queues and a fresh tracker, so -- unlike the old +definition-level ``init_workflow`` -- this performs no resets. +""" + +from typing import TYPE_CHECKING +from typing import List +from typing import Set + +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.events.topic_events.topic_event import TopicEvent +from grafi.topics.topic_base import TopicBase +from grafi.topics.topic_impl.in_workflow_input_topic import InWorkflowInputTopic +from grafi.topics.topic_impl.in_workflow_output_topic import InWorkflowOutputTopic +from grafi.topics.topic_types import TopicType + +if TYPE_CHECKING: + from grafi.workflows.impl.workflow_run import WorkflowRun + + +async def init_run( + run: "WorkflowRun", input_data: PublishToTopicEvent, is_sequential: bool +) -> None: + invoke_context = input_data.invoke_context + + events = [ + event + for event in await run.event_store.get_agent_events( + invoke_context.assistant_request_id + ) + if isinstance(event, TopicEvent) + ] + + if len(events) == 0: + await _seed_fresh(run, input_data, is_sequential) + else: + await _restore(run, input_data, events, is_sequential) + + +async def _seed_fresh( + run: "WorkflowRun", input_data: PublishToTopicEvent, is_sequential: bool +) -> None: + """Publish the input to each agent-input topic for a brand-new request.""" + input_topics: List[TopicBase] = [ + topic + for topic in run.topics.values() + if topic.type == TopicType.AGENT_INPUT_TOPIC_TYPE + ] + + events_to_record: List[TopicEvent] = [] + # One delivery per consumer of each seeded input topic (not one per topic), + # so a fan-out input feeding several nodes is not declared quiescent after + # only the first node commits. + seeded_delivery_count = 0 + for input_topic in input_topics: + event = await input_topic.publish_data( + input_data.model_copy( + update={ + "publisher_name": run.name, + "publisher_type": run.type, + }, + deep=True, + ) + ) + if event: + events_to_record.append(event) + seeded_delivery_count += len(run._topic_consumers(event.name)) + if is_sequential: + await run._add_to_invoke_queue(event) + + if events_to_record: + if not is_sequential and seeded_delivery_count: + await run.tracker.on_messages_published( + seeded_delivery_count, source="init_workflow" + ) + await run.event_store.record_events(events_to_record) + + +async def _restore( + run: "WorkflowRun", + input_data: PublishToTopicEvent, + events: List[TopicEvent], + is_sequential: bool, +) -> None: + """Restore topic state from persisted events and resume.""" + for topic_event in events: + await run.topics[topic_event.name].restore_topic(topic_event) + if is_sequential and isinstance(topic_event, PublishToTopicEvent): + await run._add_to_invoke_queue(topic_event) + + # Re-seed the tracker with messages still pending consumption so the resumed + # parallel run drains restored work instead of declaring immediate quiescence. + if not is_sequential: + pending = await run._count_pending_consumable() + if pending: + await run.tracker.on_messages_published(pending, source="recovery_restore") + + await _reissue_paired_inputs(run, input_data, events, is_sequential) + + +async def _reissue_paired_inputs( + run: "WorkflowRun", + input_data: PublishToTopicEvent, + events: List[TopicEvent], + is_sequential: bool, +) -> None: + """Re-issue paired in-workflow inputs for any consumed in-workflow output + (e.g. a human-in-the-loop response resuming the flow).""" + consumed_event_ids = input_data.consumed_event_ids + consumed_events = [ + event for event in events if event.event_id in consumed_event_ids + ] + in_workflow_output_topic_names: Set[str] = set( + event.name + for event in consumed_events + if event.type == TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE + ) + + for in_workflow_output_topic_name in in_workflow_output_topic_names: + in_workflow_output_topic = run.topics.get(in_workflow_output_topic_name) + if not ( + in_workflow_output_topic + and isinstance(in_workflow_output_topic, InWorkflowOutputTopic) + ): + continue + + for ( + paired_in_workflow_input_topic_name + ) in in_workflow_output_topic.paired_in_workflow_input_topic_names: + paired_in_workflow_input_topic = run.topics.get( + paired_in_workflow_input_topic_name + ) + if not ( + paired_in_workflow_input_topic + and isinstance(paired_in_workflow_input_topic, InWorkflowInputTopic) + ): + continue + + paired_event = await paired_in_workflow_input_topic.publish_data( + input_data.model_copy( + update={ + "publisher_name": run.name, + "publisher_type": run.type, + }, + deep=True, + ) + ) + if not paired_event: + continue + + if not is_sequential: + delivery_count = len( + run._topic_consumers(paired_in_workflow_input_topic_name) + ) + if delivery_count: + await run.tracker.on_messages_published( + delivery_count, source="restore_paired_input" + ) + if is_sequential: + await run._add_to_invoke_queue(paired_event) + await run.event_store.record_event(paired_event) diff --git a/grafi/workflows/impl/sequential_engine.py b/grafi/workflows/impl/sequential_engine.py new file mode 100644 index 0000000..930e58b --- /dev/null +++ b/grafi/workflows/impl/sequential_engine.py @@ -0,0 +1,80 @@ +"""Sequential execution engine for a :class:`WorkflowRun`. + +Drains the run's ready-queue one node at a time (concurrency = 1): consume the +node's inputs, invoke it, publish, enqueue any newly-ready nodes, and persist the +consumed + published events. Output is collected once the queue empties. +""" + +from typing import TYPE_CHECKING +from typing import AsyncGenerator +from typing import List + +from loguru import logger + +from grafi.common.events.topic_events.consume_from_topic_event import ( + ConsumeFromTopicEvent, +) +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.events.topic_events.topic_event import TopicEvent +from grafi.common.exceptions import NodeExecutionError +from grafi.workflows.impl.utils import get_node_input +from grafi.workflows.impl.utils import publish_events + +if TYPE_CHECKING: + from grafi.workflows.impl.workflow_run import WorkflowRun + + +async def invoke_sequential( + run: "WorkflowRun", input_data: PublishToTopicEvent +) -> AsyncGenerator[ConsumeFromTopicEvent, None]: + invoke_context = input_data.invoke_context + consumed_events: List[ConsumeFromTopicEvent] = [] + try: + while run.invoke_queue: + if run._stop_requested: + logger.info("Workflow execution stopped by assistant request") + break + + node = run.invoke_queue.popleft() + node_consumed_events: List[ConsumeFromTopicEvent] = await get_node_input( + node, run.topics + ) + + if node_consumed_events: + try: + published_events: List[PublishToTopicEvent] = [] + async for result in node.invoke( + invoke_context, node_consumed_events + ): + published_events.extend( + await publish_events( + node, + result, + run.tracker, + run._topic_consumers, + run.topics, + ) + ) + + for event in published_events: + await run._add_to_invoke_queue(event) + + events: List[TopicEvent] = [] + events.extend(node_consumed_events) + events.extend(published_events) + + await run.event_store.record_events(events) + except Exception as e: + raise NodeExecutionError( + node_name=node.name, + message=f"Node execution failed: {e}", + invoke_context=invoke_context, + cause=e, + ) from e + + consumed_events = await run._get_output_events() + for event in consumed_events: + yield event + finally: + if consumed_events: + await run.event_store.record_events(consumed_events) diff --git a/grafi/workflows/impl/utils.py b/grafi/workflows/impl/utils.py index 7a6a11f..7db2001 100644 --- a/grafi/workflows/impl/utils.py +++ b/grafi/workflows/impl/utils.py @@ -10,6 +10,7 @@ from grafi.common.events.topic_events.topic_event import TopicEvent from grafi.common.models.message import Message from grafi.nodes.node_base import NodeBase +from grafi.topics.topic_base import TopicBase from grafi.workflows.impl.async_node_tracker import AsyncNodeTracker @@ -88,10 +89,16 @@ async def publish_events( publish_event: PublishToTopicEvent, tracker: AsyncNodeTracker, consumers_of: Callable[[str], Sequence[str]], + topics: Dict[str, TopicBase], ) -> List[PublishToTopicEvent]: """ Publish events to all topics the node publishes to. + ``topics`` maps topic name to the *current run's* topic instance (with its + per-run queue); a node's ``publish_to`` only identifies which topics, by name, + it targets. Routing through the run's topics keeps concurrent invocations + isolated. + The tracker counts outstanding *deliveries* -- one per consumer of each published event -- not publications. ``consumers_of(topic_name)`` returns the consumers the workflow topology will route a topic's events to (subscribing @@ -103,12 +110,13 @@ async def publish_events( published_events: List[PublishToTopicEvent] = [] for topic in node.publish_to: + run_topic = topics[topic.name] delivery_count = len(consumers_of(topic.name)) if delivery_count: await tracker.on_messages_published( delivery_count, source=f"node:{node.name}->{topic.name}" ) - event = await topic.publish_data(publish_event) + event = await run_topic.publish_data(publish_event) if event: published_events.append(event) elif delivery_count: @@ -121,17 +129,20 @@ async def publish_events( return published_events -async def get_node_input(node: NodeBase) -> List[ConsumeFromTopicEvent]: +async def get_node_input( + node: NodeBase, topics: Dict[str, TopicBase] +) -> List[ConsumeFromTopicEvent]: """ Get input events for a node from its subscribed topics. - NO CHANGES NEEDED - consumption tracking happens at commit time. + ``topics`` maps topic name to the current run's topic instance; the node's + subscriptions only identify which topics, by name, it reads. Consumption + tracking happens at commit time. """ consumed_events: List[ConsumeFromTopicEvent] = [] - node_subscribed_topics = node._subscribed_topics.values() - - for subscribed_topic in node_subscribed_topics: + for topic_name in node._subscribed_topics: + subscribed_topic = topics[topic_name] if await subscribed_topic.can_consume(node.name): node_consumed_events = await subscribed_topic.consume(node.name) for event in node_consumed_events: diff --git a/grafi/workflows/impl/workflow_run.py b/grafi/workflows/impl/workflow_run.py new file mode 100644 index 0000000..e7c70a0 --- /dev/null +++ b/grafi/workflows/impl/workflow_run.py @@ -0,0 +1,210 @@ +"""One invocation's mutable runtime state. + +A :class:`~grafi.workflows.impl.event_driven_workflow.EventDrivenWorkflow` is a +long-lived *definition*: topology, topic configuration, tools, and LLM clients. +A ``WorkflowRun`` is created per ``invoke()`` call and owns everything that +mutates during execution -- the per-run topic queues, the quiescence tracker, +the ready-queue of invokable nodes, and the stop flag. + +Because each invocation gets its own run, two concurrent invocations of one +workflow instance share only the immutable definition (and the reentrant tools / +clients hanging off it) and never each other's runtime state. The run is +discarded when the invocation finishes. + +The topic *config* (name, type, condition) is shared by reference via a shallow +copy; only the queue instance is fresh, rebuilt as ``type(queue)()`` so a topic's +queue kind is preserved without cloning the node/tool graph. + +This module holds the run's state plus the small helpers the engines share +(topology/quiescence accounting, output collection, commit). The execution +itself lives in focused modules: :mod:`grafi.workflows.impl.run_recovery` +(seeding/recovery), :mod:`grafi.workflows.impl.sequential_engine`, and +:mod:`grafi.workflows.impl.parallel_engine`. +""" + +from collections import deque +from typing import Any +from typing import AsyncGenerator +from typing import Dict +from typing import List + +from grafi.common.events.topic_events.consume_from_topic_event import ( + ConsumeFromTopicEvent, +) +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.events.topic_events.topic_event import TopicEvent +from grafi.common.exceptions import NodeExecutionError +from grafi.common.exceptions import WorkflowError +from grafi.nodes.node_base import NodeBase +from grafi.topics.topic_base import TopicBase +from grafi.topics.topic_types import TopicType +from grafi.workflows.impl.async_node_tracker import AsyncNodeTracker +from grafi.workflows.impl.parallel_engine import invoke_parallel +from grafi.workflows.impl.run_recovery import init_run +from grafi.workflows.impl.sequential_engine import invoke_sequential + + +class WorkflowRun: + """Mutable runtime state for a single workflow invocation.""" + + def __init__(self, definition: Any, event_store: Any) -> None: + # Read-only references into the shared definition. + self.name: str = definition.name + self.type: str = definition.type + self.nodes: Dict[str, NodeBase] = definition.nodes + self.topic_nodes: Dict[str, List[str]] = definition._topic_nodes + self.event_store = event_store + + # Per-run topic instances: same config, fresh queue. + self.topics: Dict[str, TopicBase] = { + name: topic.model_copy(update={"event_queue": type(topic.event_queue)()}) + for name, topic in definition._topics.items() + } + + # Per-run runtime state. + self.tracker: AsyncNodeTracker = AsyncNodeTracker() + self.invoke_queue: deque[NodeBase] = deque() + self._stop_requested: bool = False + + # ------------------------------------------------------------------ # + # Control + # ------------------------------------------------------------------ # + + def stop(self) -> None: + """Stop this run. Sets the stop flag and force-stops the tracker so the + parallel output loop and node tasks wind down.""" + self._stop_requested = True + self.tracker.force_stop_sync() + + # ------------------------------------------------------------------ # + # Topology / quiescence accounting (operate on this run's topics) + # ------------------------------------------------------------------ # + + def _topic_consumers(self, topic_name: str) -> List[str]: + """Consumers that will commit each event published to ``topic_name``: + the subscribing nodes, plus the workflow itself for output topics. The + single source of truth for fan-out -- one published event yields one + delivery (and one eventual commit) per consumer. Deduplicated.""" + consumers = list(dict.fromkeys(self.topic_nodes.get(topic_name, []))) + topic = self.topics.get(topic_name) + if topic is not None and topic.type in ( + TopicType.AGENT_OUTPUT_TOPIC_TYPE, + TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE, + ): + consumers.append(self.name) + return consumers + + async def _count_pending_consumable(self) -> int: + """Deliveries still awaiting consumption across all topics: the number of + commits a resumed run must perform to drain restored state.""" + total = 0 + for topic_name, topic in self.topics.items(): + for consumer_name in self._topic_consumers(topic_name): + total += await topic.unconsumed_count(consumer_name) + return total + + async def _progress_possible(self) -> bool: + """Whether the parallel run can still make progress: a node is active, or + some event is still awaiting consumption. When False while the tracker is + non-quiescent, the outstanding deliveries are parked (e.g. an unsatisfied + ``A AND B`` subscription), so the output queue ends iteration instead of + hanging.""" + if not await self.tracker.is_idle(): + return True + return await self._count_pending_consumable() > 0 + + async def _node_can_invoke(self, node: NodeBase) -> bool: + """``node.can_invoke`` evaluated against *this run's* topics (not the + definition's, whose queues are unused during a run).""" + names_with_data = [ + topic.name + for topic in node.subscribed_topics + if await self.topics[topic.name].can_consume(node.name) + ] + return node.can_invoke_with_topics(names_with_data) + + async def _add_to_invoke_queue(self, event: TopicEvent) -> None: + topic_name = event.name + if topic_name not in self.topic_nodes: + return + topic = self.topics[topic_name] + for node_name in self.topic_nodes[topic_name]: + node = self.nodes[node_name] + if await topic.can_consume(node_name) and await self._node_can_invoke(node): + self.invoke_queue.append(node) + + # ------------------------------------------------------------------ # + # Output collection / commit (shared by both engines) + # ------------------------------------------------------------------ # + + async def _get_output_events(self) -> List[ConsumeFromTopicEvent]: + consumed_events: List[ConsumeFromTopicEvent] = [] + output_topics = [ + topic + for topic in self.topics.values() + if topic.type == TopicType.IN_WORKFLOW_OUTPUT_TOPIC_TYPE + or topic.type == TopicType.AGENT_OUTPUT_TOPIC_TYPE + ] + for output_topic in output_topics: + if await output_topic.can_consume(self.name): + events = await output_topic.consume(self.name) + for event in events: + consumed_events.append( + ConsumeFromTopicEvent( + name=event.name, + type=event.type, + consumer_name=self.name, + consumer_type=self.type, + invoke_context=event.invoke_context, + offset=event.offset, + data=event.data, + ) + ) + return consumed_events + + async def _commit_events( + self, + consumer_name: str, + topic_events: List[ConsumeFromTopicEvent], + track_commit: bool = True, + ) -> None: + if not topic_events: + return + topic_max_offset: Dict[str, int] = {} + for topic_event in topic_events: + topic_max_offset[topic_event.name] = max( + topic_max_offset.get(topic_event.name, 0), topic_event.offset + ) + for topic, offset in topic_max_offset.items(): + await self.topics[topic].commit(consumer_name, offset) + if track_commit: + await self.tracker.on_messages_committed( + len(topic_events), source=f"commit:{consumer_name}" + ) + + # ------------------------------------------------------------------ # + # Lifecycle (thin dispatch to the seeding/recovery + engine modules) + # ------------------------------------------------------------------ # + + async def init(self, input_data: PublishToTopicEvent, is_sequential: bool) -> None: + """Seed input topics for a fresh request, or restore state on recovery.""" + await init_run(self, input_data, is_sequential) + + async def run( + self, input_data: PublishToTopicEvent, is_sequential: bool + ) -> AsyncGenerator[ConsumeFromTopicEvent, None]: + """Seed/recover state, then execute sequentially or in parallel.""" + invoke_context = input_data.invoke_context + try: + await self.init(input_data, is_sequential) + engine = invoke_sequential if is_sequential else invoke_parallel + async for event in engine(self, input_data): + yield event + except NodeExecutionError: + raise + except Exception as e: + raise WorkflowError( + message=f"Workflow {self.name} async execution failed: {e}", + invoke_context=invoke_context, + cause=e, + ) from e diff --git a/tests/workflow/test_event_driven_workflow.py b/tests/workflow/test_event_driven_workflow.py index 05a5795..4831394 100644 --- a/tests/workflow/test_event_driven_workflow.py +++ b/tests/workflow/test_event_driven_workflow.py @@ -6,6 +6,7 @@ import pytest from openinference.semconv.trace import OpenInferenceSpanKindValues +from grafi.common.event_stores.event_store_in_memory import EventStoreInMemory from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent from grafi.common.exceptions import WorkflowError from grafi.common.models.invoke_context import InvokeContext @@ -17,6 +18,7 @@ from grafi.topics.topic_impl.input_topic import InputTopic from grafi.topics.topic_impl.output_topic import OutputTopic from grafi.workflows.impl.event_driven_workflow import EventDrivenWorkflow +from grafi.workflows.impl.workflow_run import WorkflowRun from grafi.workflows.workflow import WorkflowBuilder @@ -284,12 +286,10 @@ async def test_invoke_basic_flow(self, async_workflow): @pytest.mark.asyncio async def test_invoke_with_async_output_queue(self, async_workflow): """Test that invoke uses AsyncOutputQueue.""" - # We can verify that the workflow has the necessary components - assert hasattr(async_workflow, "_tracker") + # The definition owns topology; runtime state (incl. the AsyncOutputQueue) + # is created per invocation in WorkflowRun. assert hasattr(async_workflow, "_topics") - - # The AsyncOutputQueue should be created during invoke execution - # This is more of an integration test ensuring the components work together + assert callable(async_workflow.invoke) class TestEventDrivenWorkflowInitialWorkflow: @@ -310,9 +310,10 @@ def workflow_for_initial_test(self): return EventDrivenWorkflow(nodes={"test_node": node}) def test_init_workflow_method_exists(self, workflow_for_initial_test): - """init_workflow should be available for restoring workflow state.""" - assert hasattr(workflow_for_initial_test, "init_workflow") - assert callable(workflow_for_initial_test.init_workflow) + """A run's init() restores/seeds workflow state for an invocation.""" + run = WorkflowRun(workflow_for_initial_test, EventStoreInMemory()) + assert callable(run.init) + assert callable(run.run) class TestEventDrivenWorkflowToDict: @@ -357,37 +358,32 @@ def workflow_with_tracker(self): return EventDrivenWorkflow(nodes={"test_node": node}) def test_workflow_has_tracker(self, workflow_with_tracker): - """Test that workflow has AsyncNodeTracker.""" - assert hasattr(workflow_with_tracker, "_tracker") + """Each run owns an AsyncNodeTracker; runs have independent trackers.""" from grafi.workflows.impl.async_node_tracker import AsyncNodeTracker - assert isinstance(workflow_with_tracker._tracker, AsyncNodeTracker) + run = WorkflowRun(workflow_with_tracker, EventStoreInMemory()) + assert isinstance(run.tracker, AsyncNodeTracker) + + # Independent trackers per run are the basis of invocation isolation. + run2 = WorkflowRun(workflow_with_tracker, EventStoreInMemory()) + assert run.tracker is not run2.tracker @pytest.mark.asyncio - async def test_tracker_reset_on_init(self, workflow_with_tracker): - """Test that tracker is reset on workflow initialization.""" - # Add some activity to tracker - await workflow_with_tracker._tracker.enter("test_node") - assert not await workflow_with_tracker._tracker.is_idle() + async def test_run_init_leaves_no_active_nodes(self, workflow_with_tracker): + """A fresh run starts idle, and init() seeds work without marking any + node active (nodes only become active once they enter processing).""" + run = WorkflowRun(workflow_with_tracker, EventStoreInMemory()) + assert await run.tracker.is_idle() - # Call init_workflow which should reset tracker invoke_context = InvokeContext( conversation_id="test", invoke_id="test", assistant_request_id="test" ) - with patch( - "grafi.workflows.impl.event_driven_workflow.container" - ) as mock_container: - mock_event_store = Mock() - mock_event_store.get_agent_events = AsyncMock(return_value=[]) - mock_event_store.record_events = AsyncMock() - mock_event_store.record_event = AsyncMock() - mock_container.event_store = mock_event_store - await workflow_with_tracker.init_workflow( - PublishToTopicEvent(invoke_context=invoke_context, data=[]) - ) + await run.init( + PublishToTopicEvent(invoke_context=invoke_context, data=[]), + is_sequential=False, + ) - # Tracker should be reset - assert await workflow_with_tracker._tracker.is_idle() + assert await run.tracker.is_idle() class TestEventDrivenWorkflowStopFlag: diff --git a/tests/workflow/test_fanout_quiescence.py b/tests/workflow/test_fanout_quiescence.py index 9e4bd69..adb4724 100644 --- a/tests/workflow/test_fanout_quiescence.py +++ b/tests/workflow/test_fanout_quiescence.py @@ -23,6 +23,7 @@ from grafi.topics.topic_impl.output_topic import OutputTopic from grafi.topics.topic_types import TopicType from grafi.workflows.impl.event_driven_workflow import EventDrivenWorkflow +from grafi.workflows.impl.workflow_run import WorkflowRun class LabelTool(Tool): @@ -81,11 +82,12 @@ def test_topic_consumers_counts_every_subscriber(): """The topology helper reports both subscribers of the fan-out input topic and the workflow itself as the consumer of the output topic.""" workflow = _build_fanout_workflow() + run = WorkflowRun(workflow, EventStoreInMemory()) - assert set(workflow._topic_consumers("agent_input")) == {"node_a", "node_b"} - assert workflow._topic_consumers("agent_output") == [workflow.name] + assert set(run._topic_consumers("agent_input")) == {"node_a", "node_b"} + assert run._topic_consumers("agent_output") == [workflow.name] # Unknown topic -> no consumers (cannot hang execution). - assert workflow._topic_consumers("does_not_exist") == [] + assert run._topic_consumers("does_not_exist") == [] @pytest.mark.asyncio diff --git a/tests/workflow/test_invocation_isolation.py b/tests/workflow/test_invocation_isolation.py new file mode 100644 index 0000000..3c07b53 --- /dev/null +++ b/tests/workflow/test_invocation_isolation.py @@ -0,0 +1,130 @@ +"""Concurrent-invocation isolation (the 'one assistant, many invokes' goal). + +Characterization test for Defect #3 / runtime Gap 00: two ``invoke()`` calls on +the *same* workflow instance, in flight at the same time, must complete +independently. Each run owns its own runtime state (topic queues, tracker, ready +queue, stop flag); they share only the immutable definition. + +Before per-invocation isolation this FAILS: the two runs share one tracker and +one set of topic queues, and ``init_workflow`` resets them, so the second invoke +drains/resets the first mid-flight (one run yields nothing, or the run hangs). +""" + +import asyncio +from unittest.mock import Mock + +import pytest +from openinference.semconv.trace import OpenInferenceSpanKindValues + +from grafi.common.event_stores.event_store_in_memory import EventStoreInMemory +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.nodes.node import Node +from grafi.tools.tool import Tool +from grafi.topics.expressions.topic_expression import TopicExpr +from grafi.topics.topic_impl.input_topic import InputTopic +from grafi.topics.topic_impl.output_topic import OutputTopic +from grafi.workflows.impl.event_driven_workflow import EventDrivenWorkflow + + +class MarkerTool(Tool): + """Yields a marker derived from the invoke context, after a small await. + + The marker lets a run's output be traced back to the request that produced + it; the ``sleep`` forces the two concurrent runs to interleave so any shared + runtime state corrupts deterministically. + """ + + oi_span_type: OpenInferenceSpanKindValues = OpenInferenceSpanKindValues.TOOL + + async def invoke(self, invoke_context, input_data): + await asyncio.sleep(0.05) + yield [ + Message(role="assistant", content=f"out:{invoke_context.conversation_id}") + ] + + @classmethod + async def from_dict(cls, data): # pragma: no cover - not exercised here + return cls(oi_span_type=OpenInferenceSpanKindValues.TOOL) + + +def _build_workflow() -> EventDrivenWorkflow: + input_topic = InputTopic(name="agent_input") + output_topic = OutputTopic(name="agent_output") + node = Node( + name="marker_node", + type="Node", + tool=MarkerTool(), + subscribed_expressions=[TopicExpr(topic=input_topic)], + publish_to=[output_topic], + ) + return EventDrivenWorkflow.builder().node(node).build() + + +def _input(marker: str) -> PublishToTopicEvent: + return PublishToTopicEvent( + invoke_context=InvokeContext( + conversation_id=marker, + invoke_id=f"invoke-{marker}", + assistant_request_id=f"req-{marker}", + ), + data=[Message(role="user", content=f"hello {marker}")], + ) + + +async def _drain( + workflow: EventDrivenWorkflow, event: PublishToTopicEvent, sequential: bool +): + return [out async for out in workflow.invoke(event, is_sequential=sequential)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sequential", [False, True]) +async def test_concurrent_invokes_on_one_instance_are_isolated(sequential, monkeypatch): + """Two concurrent invokes on one workflow instance complete independently.""" + workflow = _build_workflow() + + # A shared store is fine: events are partitioned by assistant_request_id. + store = EventStoreInMemory() + fake_container = Mock() + fake_container.event_store = store + monkeypatch.setattr( + "grafi.workflows.impl.event_driven_workflow.container", fake_container + ) + + results_a, results_b = await asyncio.wait_for( + asyncio.gather( + _drain(workflow, _input("alpha"), sequential), + _drain(workflow, _input("beta"), sequential), + ), + timeout=10, + ) + + # Each run yields exactly its own output -- no drops, no cross-talk. + assert [m.content for e in results_a for m in e.data] == ["out:alpha"] + assert [m.content for e in results_b for m in e.data] == ["out:beta"] + + +@pytest.mark.asyncio +async def test_stop_does_not_leak_across_concurrent_invokes(monkeypatch): + """Stopping is scoped: one run finishing/stopping must not halt another. + + Here both runs simply complete; the assertion is that running them + concurrently does not strand either one (a shared stop flag / tracker would). + """ + workflow = _build_workflow() + store = EventStoreInMemory() + fake_container = Mock() + fake_container.event_store = store + monkeypatch.setattr( + "grafi.workflows.impl.event_driven_workflow.container", fake_container + ) + + batches = await asyncio.wait_for( + asyncio.gather(*[_drain(workflow, _input(f"r{i}"), False) for i in range(4)]), + timeout=10, + ) + + for i, batch in enumerate(batches): + assert [m.content for e in batch for m in e.data] == [f"out:r{i}"] diff --git a/tests/workflow/test_recovery.py b/tests/workflow/test_recovery.py index 9265947..ef02d9d 100644 --- a/tests/workflow/test_recovery.py +++ b/tests/workflow/test_recovery.py @@ -22,6 +22,7 @@ from grafi.topics.topic_impl.output_topic import OutputTopic from grafi.topics.topic_types import TopicType from grafi.workflows.impl.event_driven_workflow import EventDrivenWorkflow +from grafi.workflows.impl.workflow_run import WorkflowRun class EchoTool(Tool): @@ -73,12 +74,12 @@ async def test_count_pending_consumable_reflects_restored_work(): request_id = "recovery-count" event = _pending_input_event(request_id) - # Restore the topic state as init_workflow would. - for topic in workflow._topics.values(): - await topic.reset() - await workflow._topics["agent_input"].restore_topic(event) + # Per-run state: a fresh run starts with empty queues; restore into the run's + # own topic, as WorkflowRun.init would on recovery. + run = WorkflowRun(workflow, EventStoreInMemory()) + await run.topics["agent_input"].restore_topic(event) - pending = await workflow._count_pending_consumable() + pending = await run._count_pending_consumable() assert pending == 1 diff --git a/tests/workflow/test_utils.py b/tests/workflow/test_utils.py index bb8a456..3208450 100644 --- a/tests/workflow/test_utils.py +++ b/tests/workflow/test_utils.py @@ -254,8 +254,10 @@ def consumers_of(topic_name: str): consumed_event_ids=[event.event_id for event in consumed_events], ) + # The run routes I/O through its own topic instances, keyed by name. + topics = {"topic1": mock_topic1, "topic2": mock_topic2} published_events = await publish_events( - node, publish_to_event, tracker, consumers_of + node, publish_to_event, tracker, consumers_of, topics ) assert len(published_events) == 1 @@ -306,7 +308,8 @@ async def test_get_node_input_async(self): mock_topic1.consume.return_value = [mock_event] - consumed_events = await get_node_input(node) + topics = {"topic1": mock_topic1, "topic2": mock_topic2} + consumed_events = await get_node_input(node, topics) assert len(consumed_events) == 1 assert isinstance(consumed_events[0], ConsumeFromTopicEvent) diff --git a/tests/workflow/test_workflow_run.py b/tests/workflow/test_workflow_run.py new file mode 100644 index 0000000..7d53be7 --- /dev/null +++ b/tests/workflow/test_workflow_run.py @@ -0,0 +1,376 @@ +"""Unit tests for WorkflowRun and its engine/recovery helpers. + +These exercise the per-invocation runtime object directly (no LLM): the +isolation invariants that make concurrent invocation safe, the topology / +commit / progress helpers shared by both engines, error wrapping in ``run``, +stop-forwarding from the definition, and the seeding/recovery in ``init``. +""" + +import uuid +from unittest.mock import Mock + +import pytest +from openinference.semconv.trace import OpenInferenceSpanKindValues + +from grafi.common.event_stores.event_store_in_memory import EventStoreInMemory +from grafi.common.events.topic_events.consume_from_topic_event import ( + ConsumeFromTopicEvent, +) +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.exceptions import NodeExecutionError +from grafi.common.exceptions import WorkflowError +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.nodes.node import Node +from grafi.tools.tool import Tool +from grafi.topics.expressions.subscription_builder import SubscriptionBuilder +from grafi.topics.expressions.topic_expression import TopicExpr +from grafi.topics.queue_impl.in_mem_topic_event_queue import InMemTopicEventQueue +from grafi.topics.topic_impl.input_topic import InputTopic +from grafi.topics.topic_impl.output_topic import OutputTopic +from grafi.topics.topic_impl.topic import Topic +from grafi.topics.topic_types import TopicType +from grafi.workflows.impl.event_driven_workflow import EventDrivenWorkflow +from grafi.workflows.impl.workflow_run import WorkflowRun + + +class EchoTool(Tool): + oi_span_type: OpenInferenceSpanKindValues = OpenInferenceSpanKindValues.TOOL + + async def invoke(self, invoke_context, input_data): + yield [Message(role="assistant", content="ok")] + + @classmethod + async def from_dict(cls, data): # pragma: no cover - not exercised here + return cls(oi_span_type=OpenInferenceSpanKindValues.TOOL) + + +class BoomTool(Tool): + """A tool whose invocation always raises (an async generator that raises).""" + + oi_span_type: OpenInferenceSpanKindValues = OpenInferenceSpanKindValues.TOOL + + async def invoke(self, invoke_context, input_data): + raise RuntimeError("boom") + yield # pragma: no cover - unreachable; makes this an async generator + + @classmethod + async def from_dict(cls, data): # pragma: no cover - not exercised here + return cls(oi_span_type=OpenInferenceSpanKindValues.TOOL) + + +def _ctx() -> InvokeContext: + request_id = uuid.uuid4().hex + return InvokeContext( + conversation_id=request_id, + invoke_id=request_id, + assistant_request_id=request_id, + ) + + +def _pub(content: str = "hi", name: str = "") -> PublishToTopicEvent: + return PublishToTopicEvent( + name=name, + invoke_context=_ctx(), + data=[Message(role="user", content=content)], + ) + + +def _workflow(tool: Tool = None) -> EventDrivenWorkflow: + """input(agent_input) -> node -> output(agent_output).""" + input_topic = InputTopic(name="agent_input") + output_topic = OutputTopic(name="agent_output") + node = Node( + name="node", + type="Node", + tool=tool or EchoTool(), + subscribed_expressions=[TopicExpr(topic=input_topic)], + publish_to=[output_topic], + ) + return EventDrivenWorkflow.builder().node(node).build() + + +def _and_workflow() -> EventDrivenWorkflow: + """A node subscribing to (agent_input AND gate); gate is a plain Topic.""" + input_topic = InputTopic(name="agent_input") + gate = Topic(name="gate") + output_topic = OutputTopic(name="agent_output") + node = Node( + name="stuck", + type="Node", + tool=EchoTool(), + subscribed_expressions=[ + SubscriptionBuilder() + .subscribed_to(input_topic) + .and_() + .subscribed_to(gate) + .build() + ], + publish_to=[output_topic], + ) + return EventDrivenWorkflow.builder().node(node).build() + + +def _patch_container(monkeypatch, store: EventStoreInMemory) -> None: + fake = Mock() + fake.event_store = store + monkeypatch.setattr("grafi.workflows.impl.event_driven_workflow.container", fake) + + +# --------------------------------------------------------------------------- # +# Construction / isolation invariants +# --------------------------------------------------------------------------- # + + +class TestWorkflowRunConstruction: + def test_per_run_topics_are_independent_copies(self): + wf = _workflow() + run1 = WorkflowRun(wf, EventStoreInMemory()) + run2 = WorkflowRun(wf, EventStoreInMemory()) + + # Distinct topic objects, distinct from the definition and each other. + assert run1.topics["agent_input"] is not wf._topics["agent_input"] + assert run1.topics["agent_input"] is not run2.topics["agent_input"] + # And distinct queue instances (the isolated runtime state). + assert ( + run1.topics["agent_input"].event_queue + is not run2.topics["agent_input"].event_queue + ) + assert ( + run1.topics["agent_input"].event_queue + is not wf._topics["agent_input"].event_queue + ) + + def test_queue_kind_is_preserved(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + assert isinstance(run.topics["agent_input"].event_queue, InMemTopicEventQueue) + + def test_runs_have_independent_runtime_state(self): + wf = _workflow() + r1 = WorkflowRun(wf, EventStoreInMemory()) + r2 = WorkflowRun(wf, EventStoreInMemory()) + assert r1.tracker is not r2.tracker + assert r1.invoke_queue is not r2.invoke_queue + r1._stop_requested = True + assert r2._stop_requested is False + + @pytest.mark.asyncio + async def test_publishing_to_one_run_does_not_affect_another(self): + wf = _workflow() + r1 = WorkflowRun(wf, EventStoreInMemory()) + r2 = WorkflowRun(wf, EventStoreInMemory()) + + await r1.topics["agent_input"].publish_data(_pub()) + + assert await r1.topics["agent_input"].can_consume("node") + assert not await r2.topics["agent_input"].can_consume("node") + + +# --------------------------------------------------------------------------- # +# Control +# --------------------------------------------------------------------------- # + + +class TestWorkflowRunStop: + def test_stop_sets_flag_and_force_stops_tracker(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + assert not run._stop_requested + assert not run.tracker._force_stopped + + run.stop() + + assert run._stop_requested + assert run.tracker._force_stopped + + +# --------------------------------------------------------------------------- # +# Topology / quiescence helpers +# --------------------------------------------------------------------------- # + + +class TestNodeCanInvoke: + @pytest.mark.asyncio + async def test_and_subscription_requires_all_topics(self): + wf = _and_workflow() + run = WorkflowRun(wf, EventStoreInMemory()) + node = run.nodes["stuck"] + + assert not await run._node_can_invoke(node) # nothing published + + await run.topics["agent_input"].publish_data(_pub()) + assert not await run._node_can_invoke(node) # only one branch satisfied + + await run.topics["gate"].publish_data(_pub("go")) + assert await run._node_can_invoke(node) # both branches satisfied + + +class TestAddToInvokeQueue: + @pytest.mark.asyncio + async def test_adds_ready_subscriber(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + event = await run.topics["agent_input"].publish_data(_pub()) + await run._add_to_invoke_queue(event) + assert run.nodes["node"] in run.invoke_queue + + @pytest.mark.asyncio + async def test_unknown_topic_is_noop(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + await run._add_to_invoke_queue(_pub(name="does_not_exist")) + assert len(run.invoke_queue) == 0 + + +class TestCommitEvents: + def _consume_event(self) -> ConsumeFromTopicEvent: + return ConsumeFromTopicEvent( + name="agent_input", + type=TopicType.AGENT_INPUT_TOPIC_TYPE, + consumer_name="node", + consumer_type="Node", + offset=0, + invoke_context=_ctx(), + data=[Message(role="user", content="x")], + ) + + @pytest.mark.asyncio + async def test_commit_releases_tracked_deliveries(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + await run.tracker.on_messages_published(1) + assert not await run.tracker.is_quiescent() + + await run._commit_events("node", [self._consume_event()], track_commit=True) + assert await run.tracker.is_quiescent() + + @pytest.mark.asyncio + async def test_commit_skips_tracker_when_disabled(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + await run.tracker.on_messages_published(1) + + await run._commit_events("node", [self._consume_event()], track_commit=False) + assert not await run.tracker.is_quiescent() # still 1 uncommitted + + @pytest.mark.asyncio + async def test_commit_empty_is_noop(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + await run.tracker.on_messages_published(1) + await run._commit_events("node", [], track_commit=True) + assert not await run.tracker.is_quiescent() + + +class TestGetOutputEvents: + @pytest.mark.asyncio + async def test_returns_output_topic_events_under_workflow_consumer(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + await run.topics["agent_output"].publish_data( + PublishToTopicEvent( + invoke_context=_ctx(), + data=[Message(role="assistant", content="done")], + ) + ) + + out = await run._get_output_events() + + assert len(out) == 1 + assert out[0].name == "agent_output" + assert out[0].consumer_name == run.name + + +class TestProgressPossible: + @pytest.mark.asyncio + async def test_false_when_idle_and_nothing_pending(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + assert await run._progress_possible() is False + + @pytest.mark.asyncio + async def test_true_when_event_pending_consumption(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + await run.topics["agent_input"].publish_data(_pub()) + assert await run._progress_possible() is True + + @pytest.mark.asyncio + async def test_true_when_a_node_is_active(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + await run.tracker.enter("node") + assert await run._progress_possible() is True + + +# --------------------------------------------------------------------------- # +# run() error handling +# --------------------------------------------------------------------------- # + + +class TestRunErrorHandling: + @pytest.mark.asyncio + @pytest.mark.parametrize("sequential", [False, True]) + async def test_node_failure_raises_node_execution_error( + self, sequential, monkeypatch + ): + wf = _workflow(tool=BoomTool()) + _patch_container(monkeypatch, EventStoreInMemory()) + + with pytest.raises(NodeExecutionError): + async for _ in wf.invoke(_pub(), is_sequential=sequential): + pass + + @pytest.mark.asyncio + async def test_unexpected_error_is_wrapped_in_workflow_error(self): + class BoomStore: + async def get_agent_events(self, assistant_request_id): + raise RuntimeError("store down") + + run = WorkflowRun(_workflow(), BoomStore()) + with pytest.raises(WorkflowError): + async for _ in run.run(_pub(), is_sequential=False): + pass + + +# --------------------------------------------------------------------------- # +# Definition-level stop forwarding / active-run bookkeeping +# --------------------------------------------------------------------------- # + + +class TestStopForwarding: + def test_stop_forwards_to_active_runs(self): + wf = _workflow() + run = WorkflowRun(wf, EventStoreInMemory()) + wf._active_runs[id(run)] = run + + wf.stop() + + assert wf._stop_requested + assert run._stop_requested + + @pytest.mark.asyncio + async def test_active_runs_cleared_after_invoke(self, monkeypatch): + wf = _workflow() + _patch_container(monkeypatch, EventStoreInMemory()) + + _ = [event async for event in wf.invoke(_pub(), is_sequential=False)] + + assert wf._active_runs == {} + + +# --------------------------------------------------------------------------- # +# Seeding / recovery (init) +# --------------------------------------------------------------------------- # + + +class TestInit: + @pytest.mark.asyncio + async def test_fresh_seeds_input_records_event_and_tracker(self): + store = EventStoreInMemory() + run = WorkflowRun(_workflow(), store) + + await run.init(_pub(), is_sequential=False) + + # Input seeded so the subscriber can consume. + assert await run.topics["agent_input"].can_consume("node") + # Tracker seeded with the delivery (not quiescent yet). + assert not await run.tracker.is_quiescent() + # The seeded publish event was persisted. + assert len(await store.get_events()) >= 1 + + @pytest.mark.asyncio + async def test_fresh_sequential_enqueues_ready_node(self): + run = WorkflowRun(_workflow(), EventStoreInMemory()) + await run.init(_pub(), is_sequential=True) + assert run.nodes["node"] in run.invoke_queue diff --git a/tests_integration/function_call_assistant/concurrent_function_call_invocation_example.py b/tests_integration/function_call_assistant/concurrent_function_call_invocation_example.py new file mode 100644 index 0000000..47db96c --- /dev/null +++ b/tests_integration/function_call_assistant/concurrent_function_call_invocation_example.py @@ -0,0 +1,101 @@ +"""Integration test: concurrent invocations on ONE multi-node function-call assistant. + +The hardest concurrency case: a single ``SimpleFunctionCallAssistant`` (a 3-node +workflow -- input LLM -> function call -> output LLM) is invoked several times at +once, each asking about a distinct postcode. Each response must reference its own +postcode and none of the others, proving the per-invocation isolation holds +across a multi-node, function-calling workflow (each run owns its own topic +queues, tracker, and stop flag). + +Requires ``OPENAI_API_KEY`` (loaded from ``.env`` at the repo root). +""" + +import asyncio +import os +import uuid + +from dotenv import load_dotenv + +from grafi.common.containers.container import container +from grafi.common.decorators.llm_function import llm_function +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from grafi.tools.function_calls.function_call_tool import FunctionCallTool +from tests_integration.function_call_assistant.simple_function_call_assistant import ( + SimpleFunctionCallAssistant, +) + +load_dotenv() + +event_store = container.event_store +api_key = os.getenv("OPENAI_API_KEY", "") + +# Distinct, non-overlapping postcodes so each answer is traceable to its request. +POSTCODES = ["11111", "22222", "33333"] + + +class WeatherMock(FunctionCallTool): + @llm_function + def get_weather_mock(self, postcode: str) -> str: + """Return a weather report for the given postcode. + + Args: + postcode (str): The postcode to report the weather for. + + Returns: + str: A weather report mentioning the postcode. + """ + return f"The weather of {postcode} is bad now." + + +def _invoke_context() -> InvokeContext: + request_id = uuid.uuid4().hex + return InvokeContext( + conversation_id=request_id, + invoke_id=request_id, + assistant_request_id=request_id, + ) + + +async def _run_once(assistant: SimpleFunctionCallAssistant, postcode: str) -> str: + contents = [] + event = PublishToTopicEvent( + invoke_context=_invoke_context(), + data=[Message(role="user", content=f"How's the weather in {postcode}?")], + ) + async for out in assistant.invoke(event, is_sequential=False): + for message in out.data: + if isinstance(message.content, str): + contents.append(message.content) + return " ".join(contents) + + +async def test_concurrent_function_call_invokes_are_isolated() -> None: + assert api_key, "OPENAI_API_KEY is not set (add it to .env at the repo root)" + + await event_store.clear_events() + assistant = ( + SimpleFunctionCallAssistant.builder() + .name("ConcurrentFunctionCallAssistant") + .api_key(api_key) + .function_tool(WeatherMock(name="WeatherMock")) + .build() + ) + + results = await asyncio.gather(*[_run_once(assistant, p) for p in POSTCODES]) + + for postcode, output in zip(POSTCODES, results): + others = [p for p in POSTCODES if p != postcode] + print(f"{postcode} -> {output!r}") + assert ( + postcode in output + ), f"response missing its postcode {postcode}: {output!r}" + assert not any( + other in output for other in others + ), f"cross-talk in response for {postcode}: {output!r}" + + print("Concurrent function-call invocation isolation: OK") + + +asyncio.run(test_concurrent_function_call_invokes_are_isolated()) diff --git a/tests_integration/simple_llm_assistant/concurrent_invocation_example.py b/tests_integration/simple_llm_assistant/concurrent_invocation_example.py new file mode 100644 index 0000000..946694f --- /dev/null +++ b/tests_integration/simple_llm_assistant/concurrent_invocation_example.py @@ -0,0 +1,101 @@ +"""Integration test: concurrent invocations on ONE assistant instance. + +Exercises the runtime isolation work (Defect #3 / runtime Gap 00) against a real +LLM. A single ``SimpleLLMAssistant`` instance is invoked several times +concurrently, each with a distinct one-word prompt. Each response must echo its +own word and none of the others -- proving the concurrent runs do not share or +corrupt each other's runtime state (topic queues, tracker, stop flag). + +API keys are read from the environment, loaded from a ``.env`` file at the repo +root (see ``.env.example``). Requires ``OPENAI_API_KEY``. +""" + +import asyncio +import os +import uuid + +from dotenv import load_dotenv + +from grafi.common.containers.container import container +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.instrumentations.tracing import TracingOptions +from grafi.common.instrumentations.tracing import setup_tracing +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from tests_integration.simple_llm_assistant.simple_llm_assistant import ( + SimpleLLMAssistant, +) + +load_dotenv() + +container.register_tracer(setup_tracing(tracing_options=TracingOptions.IN_MEMORY)) +event_store = container.event_store + +api_key = os.getenv("OPENAI_API_KEY", "") + +# Distinct, non-overlapping words so each response is unambiguously traceable to +# the prompt that produced it (none is a substring of another). +WORDS = ["ALPHA", "BETA", "GAMMA"] + +SYSTEM_MESSAGE = ( + "You are a word-echo service. Reply with ONLY the single uppercase word the " + "user sends, and nothing else." +) + + +def _invoke_context() -> InvokeContext: + request_id = uuid.uuid4().hex + return InvokeContext( + conversation_id=request_id, + invoke_id=request_id, + assistant_request_id=request_id, + ) + + +def _input(word: str) -> PublishToTopicEvent: + return PublishToTopicEvent( + invoke_context=_invoke_context(), + data=[Message(role="user", content=word)], + ) + + +async def _run_once(assistant: SimpleLLMAssistant, word: str) -> str: + """Drain one invocation's output into a single string.""" + contents = [] + async for event in assistant.invoke(_input(word), is_sequential=False): + for message in event.data: + if isinstance(message.content, str): + contents.append(message.content) + return " ".join(contents).upper() + + +async def test_concurrent_invokes_are_isolated() -> None: + """One assistant instance, several concurrent invokes; each response matches + its own prompt with no cross-talk.""" + assert api_key, "OPENAI_API_KEY is not set (add it to .env at the repo root)" + + await event_store.clear_events() + + # A single shared assistant instance drives every concurrent invocation. + assistant = ( + SimpleLLMAssistant.builder() + .name("ConcurrentAssistant") + .api_key(api_key) + .system_message(SYSTEM_MESSAGE) + .build() + ) + + results = await asyncio.gather(*[_run_once(assistant, w) for w in WORDS]) + + for word, output in zip(WORDS, results): + others = [w for w in WORDS if w != word] + print(f"{word} -> {output!r}") + assert word in output, f"response for {word!r} missing its word: {output!r}" + assert not any( + other in output for other in others + ), f"cross-talk in response for {word!r}: {output!r}" + + print("Concurrent invocation isolation: OK") + + +asyncio.run(test_concurrent_invokes_are_isolated()) diff --git a/tests_integration/simple_llm_assistant/concurrent_sequential_invocation_example.py b/tests_integration/simple_llm_assistant/concurrent_sequential_invocation_example.py new file mode 100644 index 0000000..f10162c --- /dev/null +++ b/tests_integration/simple_llm_assistant/concurrent_sequential_invocation_example.py @@ -0,0 +1,85 @@ +"""Integration test: concurrent SEQUENTIAL invocations on ONE assistant instance. + +Companion to ``concurrent_invocation_example.py`` (which covers the parallel +engine). Here several ``invoke(..., is_sequential=True)`` calls run concurrently +on a single ``SimpleLLMAssistant``; each must echo its own word with no +cross-talk, proving the sequential engine is per-invocation isolated too (its +ready-queue and topic queues are not shared across runs). + +Requires ``OPENAI_API_KEY`` (loaded from ``.env`` at the repo root). +""" + +import asyncio +import os +import uuid + +from dotenv import load_dotenv + +from grafi.common.containers.container import container +from grafi.common.events.topic_events.publish_to_topic_event import PublishToTopicEvent +from grafi.common.models.invoke_context import InvokeContext +from grafi.common.models.message import Message +from tests_integration.simple_llm_assistant.simple_llm_assistant import ( + SimpleLLMAssistant, +) + +load_dotenv() + +event_store = container.event_store +api_key = os.getenv("OPENAI_API_KEY", "") + +WORDS = ["ALPHA", "BETA", "GAMMA"] +SYSTEM_MESSAGE = ( + "You are a word-echo service. Reply with ONLY the single uppercase word the " + "user sends, and nothing else." +) + + +def _invoke_context() -> InvokeContext: + request_id = uuid.uuid4().hex + return InvokeContext( + conversation_id=request_id, + invoke_id=request_id, + assistant_request_id=request_id, + ) + + +async def _run_once(assistant: SimpleLLMAssistant, word: str) -> str: + contents = [] + event = PublishToTopicEvent( + invoke_context=_invoke_context(), + data=[Message(role="user", content=word)], + ) + async for out in assistant.invoke(event, is_sequential=True): + for message in out.data: + if isinstance(message.content, str): + contents.append(message.content) + return " ".join(contents).upper() + + +async def test_concurrent_sequential_invokes_are_isolated() -> None: + assert api_key, "OPENAI_API_KEY is not set (add it to .env at the repo root)" + + await event_store.clear_events() + assistant = ( + SimpleLLMAssistant.builder() + .name("ConcurrentSequentialAssistant") + .api_key(api_key) + .system_message(SYSTEM_MESSAGE) + .build() + ) + + results = await asyncio.gather(*[_run_once(assistant, w) for w in WORDS]) + + for word, output in zip(WORDS, results): + others = [w for w in WORDS if w != word] + print(f"{word} -> {output!r}") + assert word in output, f"response for {word!r} missing its word: {output!r}" + assert not any( + other in output for other in others + ), f"cross-talk in response for {word!r}: {output!r}" + + print("Concurrent sequential invocation isolation: OK") + + +asyncio.run(test_concurrent_sequential_invokes_are_isolated())