diff --git a/docs/wayflowcore/source/conf.py b/docs/wayflowcore/source/conf.py index e4065d942..8fb96405a 100644 --- a/docs/wayflowcore/source/conf.py +++ b/docs/wayflowcore/source/conf.py @@ -72,6 +72,10 @@ 'Cannot handle as a local function: "wayflowcore.agentspec.components.nodes.ExtendedLlmNode.check_either_prompt_str_or_object_is_used" (use @functools.wraps)', "py:class reference target not found: wayflowcore.mcp.mcphelpers.ToolOutuptTypeT", "py:class reference target not found: wayflowcore.mcp.mcphelpers.ContextType", + ( + 'Cannot resolve forward reference in type annotations of "wayflowcore.agentserver.serverstorageconfig.ServerStorageConfig": ' + "name 'Datastore' is not defined" + ), ] @@ -256,6 +260,13 @@ def filter(self, record): ("py:class", r"wayflowcore.executors._events.event.Event"), ("py:class", r"(?:wayflowcore\.executors\._executor\.)?ConversationExecutor"), ("py:class", r"(?:wayflowcore\.executors\._executionstate\.)?ConversationExecutionState"), + ("py:class", r"(?:wayflowcore\.executors\._agentconversation\.)?AgentConversation"), + ("py:class", r"(?:wayflowcore\.executors\._flowconversation\.)?FlowConversation"), + ("py:class", r"(?:wayflowcore\.executors\._a2aagentconversation\.)?A2AAgentConversation"), + ( + "py:class", + r"(?:wayflowcore\.executors\._managerworkersconversation\.)?ManagerWorkersConversation", + ), ("py:class", r"wayflowcore.executors._agentexecutor.AgentConversationExecutionState"), ("py:class", r"wayflowcore.executors._flowexecutor.FlowConversationExecutionState"), ("py:class", r"wayflowcore.agentserver.a2a._app.A2AApp"), diff --git a/docs/wayflowcore/source/core/api/checkpointing.rst b/docs/wayflowcore/source/core/api/checkpointing.rst new file mode 100644 index 000000000..a46df9b12 --- /dev/null +++ b/docs/wayflowcore/source/core/api/checkpointing.rst @@ -0,0 +1,34 @@ +.. _checkpointing: + +Checkpointing +============= + +Checkpointing APIs persist and restore conversation state across process restarts, +debugging sessions, and server requests. + +Core types +---------- + +.. autoclass:: wayflowcore.checkpointing.ConversationCheckpoint + +.. autoclass:: wayflowcore.checkpointing.CheckpointingInterval + +.. autoclass:: wayflowcore.checkpointing.Checkpointer + :members: load_latest, load, save, save_async, list_checkpoints, delete + +Storage configuration +--------------------- + +.. autoclass:: wayflowcore.checkpointing.StorageConfig + :members: to_schema + +Datastore-backed checkpointers +------------------------------ + +.. autoclass:: wayflowcore.checkpointing.DatastoreCheckpointer + +.. autoclass:: wayflowcore.checkpointing.InMemoryCheckpointer + +.. autoclass:: wayflowcore.checkpointing.PostgresCheckpointer + +.. autoclass:: wayflowcore.checkpointing.OracleDatabaseCheckpointer diff --git a/docs/wayflowcore/source/core/api/index.rst b/docs/wayflowcore/source/core/api/index.rst index eeb7b8f69..8b98986fb 100644 --- a/docs/wayflowcore/source/core/api/index.rst +++ b/docs/wayflowcore/source/core/api/index.rst @@ -25,6 +25,7 @@ API Reference :maxdepth: 3 Agent Spec Adapters + Checkpointing Conversations LLMs Events diff --git a/docs/wayflowcore/source/core/changelog.rst b/docs/wayflowcore/source/core/changelog.rst index d75c2c852..da6e9bccc 100644 --- a/docs/wayflowcore/source/core/changelog.rst +++ b/docs/wayflowcore/source/core/changelog.rst @@ -4,6 +4,22 @@ Changelog WayFlow |current_version| ------------------------- +New features +^^^^^^^^^^^^ + +* **First-class conversation checkpointing** + + Added shared conversation checkpointing for Agents, Flows, Swarms, ManagerWorkers, and A2A agents through + ``ConversationCheckpoint``, ``Checkpointer``, ``InMemoryCheckpointer``, ``PostgresCheckpointer``, and + ``OracleDatabaseCheckpointer``. Conversations can now resume from ``conversation_id``, load specific checkpoints for + time-travel debugging, and choose checkpoint save frequency with ``CheckpointingInterval``. + + The OpenAI Responses server path now uses this shared checkpointing subsystem as well, so persisted + ``previous_response_id`` and ``conversation`` behavior is handled through the same checkpoint model. + + For more information, see :doc:`how to checkpoint and resume conversations ` + and the :doc:`API reference on checkpointing `. + Improvements ^^^^^^^^^^^^ @@ -87,15 +103,12 @@ New features For more information read the :doc:`API Reference on LLM models ` and the guide on :doc:`how to use LLMs from different providers `. - * **Logprob support in `LlmGenerationConfig` and `PromptExecutionStep`** Add per-token log-probabilities support with the ``top_logprobs`` generation config parameter and support returning per-token log-probabilities in the ``PromptExecutionStep``. For more information please read the guide on :ref:`How to request per-token log-probabilities ` - - Improvements ^^^^^^^^^^^^ diff --git a/docs/wayflowcore/source/core/code_examples/howto_checkpointing.py b/docs/wayflowcore/source/core/code_examples/howto_checkpointing.py new file mode 100644 index 000000000..030d28a05 --- /dev/null +++ b/docs/wayflowcore/source/core/code_examples/howto_checkpointing.py @@ -0,0 +1,71 @@ +# Copyright © 2026 Oracle and/or its affiliates. +# +# This software is under the Apache License 2.0 +# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License +# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. + +# isort:skip_file +# fmt: off +# mypy: ignore-errors +# docs-title: Code Example - How to Checkpoint and Resume Conversations + +# .. start-##_Configure_your_LLM +from wayflowcore.models import VllmModel + +llm = VllmModel( + model_id="LLAMA_MODEL_ID", + host_port="LLAMA_API_URL", +) +# .. end-##_Configure_your_LLM + +llm: VllmModel # docs-skiprow +(llm,) = _update_globals(["llm_small"]) # docs-skiprow # type: ignore + +# .. start-##_Start_a_checkpointed_conversation +from wayflowcore import Agent +from wayflowcore.checkpointing import InMemoryCheckpointer + +agent = Agent(llm=llm) +checkpointer = InMemoryCheckpointer() +conversation_id = "support-conversation-1" + +conversation = agent.start_conversation( + conversation_id=conversation_id, + checkpointer=checkpointer, +) + +status = conversation.execute() +# .. end-##_Start_a_checkpointed_conversation + +# .. start-##_Resume_the_latest_checkpoint +restored_conversation = agent.start_conversation( + conversation_id=conversation_id, + checkpointer=checkpointer, +) + +restored_conversation.append_user_message("Continue from where you left off.") +status = restored_conversation.execute() +# .. end-##_Resume_the_latest_checkpoint + +# .. start-##_Load_a_specific_checkpoint +# Checkpoints are ordered oldest -> newest. +checkpoints = checkpointer.list_checkpoints(conversation_id) + +previous_checkpoint = checkpoints[-2] +rewound_conversation = agent.start_conversation( + conversation_id=conversation_id, + checkpoint_id=previous_checkpoint.checkpoint_id, + checkpointer=checkpointer, +) + +rewound_conversation.append_user_message("Try a different path from here.") +status = rewound_conversation.execute() +# .. end-##_Load_a_specific_checkpoint + +# .. start-##_Control_checkpoint_frequency +from wayflowcore.checkpointing import CheckpointingInterval, InMemoryCheckpointer + +checkpointer = InMemoryCheckpointer( + checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS, +) +# .. end-##_Control_checkpoint_frequency diff --git a/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst b/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst new file mode 100644 index 000000000..131750435 --- /dev/null +++ b/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst @@ -0,0 +1,151 @@ +.. _top-howtocheckpointing: + +========================================== +How to Checkpoint and Resume Conversations +========================================== + +.. admonition:: Prerequisites + + This guide assumes familiarity with: + + - :doc:`Agents <../tutorials/basic_agent>` + - :doc:`Flows <../tutorials/basic_flow>` + - :doc:`Serve Agents with WayFlow ` + +Checkpointing lets WayFlow save a conversation while it runs and load it again later using the +same ``conversation_id``. Use it when you want to: + +- continue after a process restart +- pause a long-running workflow and come back to it later +- inspect earlier checkpoints while debugging +- retry from an older checkpoint with different code or inputs + + +Choose a checkpointer +===================== + +The checkpointer is the object that reads and writes checkpoints. WayFlow includes: + +- ``InMemoryCheckpointer`` for tests and local experimentation +- ``PostgresCheckpointer`` for PostgreSQL-backed persistence +- ``OracleDatabaseCheckpointer`` for Oracle-backed persistence + +All checkpointers use the same methods for loading, listing, and deleting checkpoints. +WayFlow saves checkpoints automatically during conversation execution. + + +Start a checkpointed conversation +================================= + +Attach a checkpointer when you start the conversation. The ``conversation_id`` is the name WayFlow +uses to find that conversation again. + +.. literalinclude:: ../code_examples/howto_checkpointing.py + :language: python + :start-after: .. start-##_Start_a_checkpointed_conversation + :end-before: .. end-##_Start_a_checkpointed_conversation + +Once checkpointing is enabled, WayFlow saves the top-level conversation automatically at the +configured checkpoints. If an Agent or Flow starts child conversations internally, WayFlow keeps +them attached to the same saved conversation. Application code only needs to pass the public +``conversation_id`` shown above. + +For checkpointing, there are three useful identifiers: + +- ``conversation_id``: the durable conversation id used to resume and list checkpoints +- ``checkpoint_id``: the exact saved snapshot to reload +- ``conversation.id``: the id of one concrete ``Conversation`` within that conversation thread + +Each nested conversation gets its own ``conversation.id`` while inheriting the conversation thread's +``conversation_id``. Application code only supplies ``conversation_id``; child identities are +created and restored internally. + +.. warning:: + + Checkpoints contain the serialized conversation, including messages and intermediate state. + Treat checkpoint storage like other persisted user conversation data: protect access, choose an + appropriate retention policy, and avoid storing it in places meant only for test data. + + +Resume the latest checkpoint +============================ + +To resume a conversation, call ``start_conversation()`` again with the same ``conversation_id`` and +checkpointer. + +.. literalinclude:: ../code_examples/howto_checkpointing.py + :language: python + :start-after: .. start-##_Resume_the_latest_checkpoint + :end-before: .. end-##_Resume_the_latest_checkpoint + +If the checkpointer has no saved state for that id, WayFlow starts a new conversation. + + +Load a specific checkpoint +========================== + +You can also load an older checkpoint. This is useful when you want to replay part of a run or +compare what happens after changing a prompt, tool, or step. + +.. literalinclude:: ../code_examples/howto_checkpointing.py + :language: python + :start-after: .. start-##_Load_a_specific_checkpoint + :end-before: .. end-##_Load_a_specific_checkpoint + +``list_checkpoints()`` returns checkpoints ordered from oldest to newest, with the checkpoint id, +creation time, and metadata recorded when the checkpoint was saved. That means ``checkpoints[-1]`` +is the newest checkpoint and ``checkpoints[-2]`` is the one before it. + + +Control checkpoint frequency +============================ + +Use ``CheckpointingInterval`` to choose how often WayFlow should save state. + +.. literalinclude:: ../code_examples/howto_checkpointing.py + :language: python + :start-after: .. start-##_Control_checkpoint_frequency + :end-before: .. end-##_Control_checkpoint_frequency + +The available options are: + +- ``CONVERSATION_TURNS``: save after the main ``conversation.execute()`` call returns +- ``LLM_TURNS``: also save after internal turns that used an LLM +- ``ALL_INTERNAL_TURNS``: also save after each internal Agent or Flow turn + +Saving more often gives WayFlow a more recent place to resume from, but it also writes more rows to +the checkpoint store. + +WayFlow resumes from the last checkpoint it saved. It does not resume from the middle of a tool +call, LLM request, or step. If a crash happens after that checkpoint, code that ran after the +checkpoint may run again. Tools and steps with side effects should therefore be safe to retry. For +example, if a step writes to another service, sends a notification, or charges a payment method, +use your own idempotency key or tracking table so the side effect is not repeated. + +When multiple processes share a relational checkpoint store, use a single writer per +``conversation_id``. WayFlow updates the "latest checkpoint" marker in one transaction for normal +writes, but the table does not enforce uniqueness for that marker. If two writers save the same +conversation at the same time, the store can end up with more than one checkpoint marked as latest. +In that case, loading by ``conversation_id`` raises an error instead of choosing one at random. +When recovering that history, load a specific ``checkpoint_id``. + + +Use checkpointing with the OpenAI Responses server +================================================== + +The OpenAI Responses server uses the same checkpointing storage behind ``ServerStorageConfig``. +Existing OpenAI-compatible features such as ``previous_response_id``, ``conversation``, +``get_response()``, ``delete_response()``, and ``store=False`` continue to work through that shared +storage path. + +If you are serving agents, configure storage the same way as in +:doc:`Serve Agents with WayFlow `. The server creates and uses the matching +checkpointer internally. + + +Next steps +========== + +- :doc:`Serialize and Deserialize Conversations ` +- :doc:`Serve Agents with WayFlow ` +- :doc:`Build a Swarm of Agents ` diff --git a/docs/wayflowcore/source/core/howtoguides/index.rst b/docs/wayflowcore/source/core/howtoguides/index.rst index 234da7699..d8449b032 100644 --- a/docs/wayflowcore/source/core/howtoguides/index.rst +++ b/docs/wayflowcore/source/core/howtoguides/index.rst @@ -105,6 +105,7 @@ These guides demonstrate how to configure the components of assistants built wit :maxdepth: 1 Load and Execute an Agent Spec Configuration + Checkpoint and Resume Conversations Serialize and Deserialize Flows and Agents Serialize and Deserialize Conversations Build a New WayFlow Component diff --git a/wayflowcore/setup.cfg b/wayflowcore/setup.cfg index 02432bd2f..330462588 100644 --- a/wayflowcore/setup.cfg +++ b/wayflowcore/setup.cfg @@ -44,6 +44,8 @@ filterwarnings = # User Warning # oci genai models requires an explicit number of max_tokens ignore:Setting `max_tokens`:UserWarning + # In-memory storage is intentionally used by checkpointing integration tests. + ignore:InMemoryDatastore is for DEVELOPMENT and PROOF-OF-CONCEPT ONLY!:UserWarning # tests for interrupts ignore:The assistant is being executed without a time limit!:UserWarning ignore:Usage of non-encrypted requests to http urls is considered unsecure and strongly discouraged:UserWarning diff --git a/wayflowcore/src/wayflowcore/a2a/a2aagent.py b/wayflowcore/src/wayflowcore/a2a/a2aagent.py index 4dee6ae11..2c7f406a9 100644 --- a/wayflowcore/src/wayflowcore/a2a/a2aagent.py +++ b/wayflowcore/src/wayflowcore/a2a/a2aagent.py @@ -19,6 +19,8 @@ from wayflowcore.tools import Tool if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer + from wayflowcore.conversation import Conversation from wayflowcore.executors._a2aagentconversation import A2AAgentConversation logger = logging.getLogger(__name__) @@ -244,10 +246,17 @@ def __init__( __metadata_info__=__metadata_info__, ) + @property + def _supports_checkpointing(self) -> bool: + return True + def start_conversation( self, inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, Message, List[Message], MessageList] = None, + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, ) -> "A2AAgentConversation": """ Initiates a new conversation with the remote server agent. @@ -263,15 +272,53 @@ def start_conversation( messages: Optional initial message list for the conversation. Can be either a ``MessageList`` or a list of ``Message`` objects. Defaults to an empty ``MessageList`` if not provided. - + conversation_id: + Conversation id used for checkpointing the conversation and later resuming it. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires both ``checkpointer`` and + ``conversation_id``. Returns ------- Conversation: A new conversation object associated with this agent. """ + return self._start_conversation_impl( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=None, + ) + + def _start_conversation_impl( + self, + inputs: Optional[Dict[str, Any]] = None, + messages: Union[None, str, Message, List[Message], MessageList] = None, + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + parent_conversation: Optional["Conversation"] = None, + ) -> "A2AAgentConversation": from wayflowcore.executors._a2aagentconversation import A2AAgentConversation from wayflowcore.executors._a2aagentexecutor import A2AAgentState + restored_conversation, conversation_instance_id, conversation_thread_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + expected_conversation_type=A2AAgentConversation, + parent_conversation=parent_conversation, + ) + ) + if restored_conversation is not None: + return restored_conversation + if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) @@ -281,8 +328,10 @@ def start_conversation( inputs=inputs or {}, # Inputs are ignored in execution message_list=messages, status=None, - conversation_id=IdGenerator.get_or_generate_id(None), + id=conversation_instance_id, + checkpointer=checkpointer, name="a2a_conversation", + conversation_id=conversation_thread_id, __metadata_info__={}, ) diff --git a/wayflowcore/src/wayflowcore/agent.py b/wayflowcore/src/wayflowcore/agent.py index e19973cfc..9226263fc 100644 --- a/wayflowcore/src/wayflowcore/agent.py +++ b/wayflowcore/src/wayflowcore/agent.py @@ -26,7 +26,9 @@ from wayflowcore.transforms import MessageTransform if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.contextproviders import ContextProvider + from wayflowcore.conversation import Conversation from wayflowcore.executors._agentconversation import AgentConversation from wayflowcore.flow import Flow from wayflowcore.ociagent import OciAgent @@ -360,6 +362,15 @@ def __init__( self._all_variables = [] self._update_internal_state() + @property + def _supports_checkpointing(self) -> bool: + from wayflowcore.serialization.context import _get_nested_components + + return all( + nested_component._supports_checkpointing + for nested_component in _get_nested_components(self, only_conversational=True) + ) + @property def agent_id(self) -> str: return self.id @@ -393,6 +404,8 @@ def start_conversation( inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, ) -> "AgentConversation": """ Initializes a conversation with the agent. @@ -406,15 +419,53 @@ def start_conversation( Message list to which the agent will participate conversation_id: Conversation id of the parent conversation. - + It is used for resume, storage, and usage accounting. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires both ``checkpointer`` and + ``conversation_id``. Returns ------- Conversation: The conversation object of the agent. """ + return self._start_conversation_impl( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=None, + ) + + def _start_conversation_impl( + self, + inputs: Optional[Dict[str, Any]] = None, + messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + parent_conversation: Optional["Conversation"] = None, + ) -> "AgentConversation": from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event from wayflowcore.executors._agentconversation import AgentConversation + from wayflowcore.executors._agentexecutor import AgentConversationExecutionState + + restored_conversation, conversation_instance_id, conversation_thread_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + expected_conversation_type=AgentConversation, + parent_conversation=parent_conversation, + ) + ) + if restored_conversation is not None: + return restored_conversation if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) @@ -458,21 +509,21 @@ def start_conversation( conversational_component=self, inputs=inputs, messages=messages, - conversation_id=conversation_id, + conversation_id=conversation_instance_id, nesting_level=None, ) ) - from wayflowcore.executors._agentexecutor import AgentConversationExecutionState - return AgentConversation( component=self, message_list=messages, - conversation_id=IdGenerator.get_or_generate_id(conversation_id), + id=conversation_instance_id, + checkpointer=checkpointer, inputs=inputs or {}, name="agent_conversation", state=AgentConversationExecutionState(), status=None, + conversation_id=conversation_thread_id, __metadata_info__={}, ) diff --git a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py b/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py deleted file mode 100644 index 58b12d97d..000000000 --- a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright © 2025 Oracle and/or its affiliates. -# -# This software is under the Apache License 2.0 -# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License -# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. -import logging -from textwrap import dedent -from typing import Dict, Optional, cast - -from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig -from wayflowcore.component import Component -from wayflowcore.conversation import Conversation -from wayflowcore.datastore.oracle import OracleDatabaseConnectionConfig, _execute_query_on_oracle_db -from wayflowcore.datastore.postgres import ( - PostgresDatabaseConnectionConfig, - _execute_query_on_postgres_db, -) -from wayflowcore.serialization import autodeserialize -from wayflowcore.serialization.context import DeserializationContext -from wayflowcore.tools import Tool - -logger = logging.getLogger(__name__) - - -def _prepare_postgres_datastore( - connection_config: PostgresDatabaseConnectionConfig, storage_config: ServerStorageConfig -) -> None: - from sqlalchemy.exc import ProgrammingError - - create_table_query = dedent(f""" - CREATE TABLE {storage_config.table_name} ( - {storage_config.turn_id_column_name} VARCHAR(255) PRIMARY KEY, - {storage_config.agent_id_column_name} VARCHAR(255) NOT NULL, - {storage_config.conversation_id_column_name} VARCHAR(255) NOT NULL, - {storage_config.created_at_column_name} INTEGER NOT NULL, - {storage_config.conversation_turn_state_column_name} TEXT NOT NULL, - {storage_config.is_last_turn_column_name} INTEGER NOT NULL, - {storage_config.extra_metadata_column_name} TEXT NOT NULL - ); - """) - try: - _execute_query_on_postgres_db(connection_config, create_table_query) - except ProgrammingError as e: - if f'relation "{storage_config.table_name}" already exists' in str(e): - raise ValueError( - f'The datastore is already setup. Either delete the existing "{storage_config.table_name}" table or start the server with `--setup-datastore=no`.' - ) from e - else: - raise e - - -def _prepare_oracle_datastore( - connection_config: OracleDatabaseConnectionConfig, storage_config: ServerStorageConfig -) -> None: - create_table_query = dedent(f""" - CREATE TABLE {storage_config.table_name} ( - {storage_config.turn_id_column_name} VARCHAR2(255) PRIMARY KEY, - {storage_config.agent_id_column_name} VARCHAR2(255) NOT NULL, - {storage_config.conversation_id_column_name} VARCHAR2(255) NOT NULL, - {storage_config.created_at_column_name} INTEGER NOT NULL, - {storage_config.conversation_turn_state_column_name} CLOB NOT NULL, - {storage_config.is_last_turn_column_name} INTEGER NOT NULL, - {storage_config.extra_metadata_column_name} CLOB NOT NULL - ); - """) - try: - _execute_query_on_oracle_db(connection_config, query=create_table_query) - except Exception as e: - if "already exists" in str(e): - raise ValueError( - f'The datastore is already setup. Either delete the existing "{storage_config.table_name}" table or start the server with `--setup-datastore=no`.' - ) from e - else: - raise e - - -def _deserialize_conversation_safely( - serialized_state: str, - tool_registry: Optional[Dict[str, Tool]] = None, - component: Optional[Component] = None, -) -> Conversation: - """ - Tries to deserialize the conversation. If it does not work, try to deserialize it by considering - the component as a disaggregated component, and will use the already instantiated agent instead of deserializing - it from scratch. - """ - deserialization_context = DeserializationContext() - deserialization_context.registered_tools = tool_registry.copy() if tool_registry else {} - try: - conversation = autodeserialize( - serialized_state, deserialization_context=deserialization_context - ) - except (TypeError, ValueError) as e: - if component is None: - raise e - # we try adding the ref to the agent itself, so that we fall back if - # something went wrong during agent deserialization - logger.warning( - "Failed to deserialize conversation by itself: %s. Using a fallback approach that leverages the provided agent as a disaggregated one to deserialize the conversation.", - e, - ) - deserialization_context = DeserializationContext() - deserialization_context.registered_tools = tool_registry.copy() if tool_registry else {} - deserialization_context._add_component_to_context(component) - - conversation = autodeserialize( - serialized_state, deserialization_context=deserialization_context - ) - return cast(Conversation, conversation) diff --git a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py index 99c34ecc9..e9ff47c84 100644 --- a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py +++ b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py @@ -4,10 +4,9 @@ # (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License # (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. - -import json import logging import time +from collections import OrderedDict from typing import Any, AsyncIterable, Dict, List, Optional, Union, cast import anyio @@ -16,10 +15,14 @@ from fastapi import status as http_status_code from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig +from wayflowcore.checkpointing import ( + CheckpointRestoreCompatibilityError, + ConversationCheckpoint, + DatastoreCheckpointer, +) from wayflowcore.conversation import Conversation from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.datastore import Datastore, InMemoryDatastore -from wayflowcore.datastore._relational import RelationalDatastore from wayflowcore.events import register_event_listeners from wayflowcore.executors.executionstatus import ( ExecutionStatus, @@ -27,9 +30,7 @@ UserMessageRequestStatus, ) from wayflowcore.idgeneration import IdGenerator -from wayflowcore.serialization import serialize -from ..._storagehelpers import _deserialize_conversation_safely from ..models.openairesponsespydanticmodels import ( Conversation2, CreateResponse, @@ -62,6 +63,8 @@ class WayFlowOpenAIResponsesService(OpenAIResponsesService): + _RESPONSE_CONVERSATION_CACHE_MAX_SIZE = 1024 + def __init__( self, agents: Dict[str, ConversationalComponent], @@ -71,7 +74,15 @@ def __init__( self.agents = agents self.storage_config = storage_config or ServerStorageConfig() self.storage = storage or InMemoryDatastore(schema=self.storage_config.to_schema()) + self.checkpointer = DatastoreCheckpointer( + datastore=self.storage, + storage_config=self.storage_config, + ) self.created_at = int(time.time()) + # Process-local fast path from public OpenAI `response_id` to the owning + # WayFlow `conversation_id`, so we can fetch the exact saved checkpoint for + # that response via `checkpointer.load(conversation_id, response_id)`. + self._response_conversation_ids: OrderedDict[str, str] = OrderedDict() self.tool_registries = { agent_name: {t.name: t for t in agent._referenced_tools()} for agent_name, agent in self.agents.items() @@ -126,23 +137,23 @@ async def get_response( detail="Get endpoint for wayflow server only supports non-streaming requests", ) - try: - metadata = self._lookup_conversation( - where={self.storage_config.turn_id_column_name: response_id}, - what=self.storage_config.extra_metadata_column_name, + checkpoint = self._find_checkpoint_for_response_id(response_id) + if checkpoint is None: + raise HTTPException( + status_code=http_status_code.HTTP_404_NOT_FOUND, detail="Response not found" ) - except ValueError: + response_as_txt = checkpoint.metadata.get("response") + if not isinstance(response_as_txt, str): raise HTTPException( status_code=http_status_code.HTTP_404_NOT_FOUND, detail="Response not found" ) - response_as_txt = json.loads(metadata)["response"] return Response.model_validate_json(response_as_txt) async def delete_response(self, response_id: str) -> Optional[ResponseError]: - self.storage.delete( - collection_name=self.storage_config.table_name, - where={self.storage_config.turn_id_column_name: response_id}, - ) + checkpoint = self._find_checkpoint_for_response_id(response_id) + if checkpoint is not None: + self.checkpointer.delete(checkpoint.conversation_id, checkpoint.checkpoint_id) + self._response_conversation_ids.pop(response_id) return None async def cancel_response(self, response_id: str) -> Union[Response, ResponseError]: @@ -192,6 +203,23 @@ async def create_response(self, body: CreateResponse) -> AsyncIterable[ResponseS if conversation_id is not None and not isinstance(conversation_id, str): conversation_id = conversation_id.id + # Stored OpenAI responses are backed by durable WayFlow checkpoints. Components + # that cannot checkpoint may still serve non-stored, non-resumed responses. + should_store_response = body.store is not False + if not agent._supports_checkpointing: + if previous_response_id is not None or conversation_id is not None: + raise HTTPException( + status_code=http_status_code.HTTP_501_NOT_IMPLEMENTED, + detail=( + f"`{agent.__class__.__name__}` does not support stored response resume yet." + ), + ) + if should_store_response: + raise HTTPException( + status_code=http_status_code.HTTP_501_NOT_IMPLEMENTED, + detail=(f"`{agent.__class__.__name__}` does not support stored responses."), + ) + state = self._load_state( previous_response_id=previous_response_id, conversation_id=conversation_id, @@ -270,6 +298,7 @@ async def runner(conversation: Conversation) -> None: async for ev in receive_stream: # These events come from the synchronous callback yield ev + await receive_stream.aclose() if raised_exception: if "not a multimodal model" in str(raised_exception): @@ -309,11 +338,21 @@ async def runner(conversation: Conversation) -> None: token_usage_listener.usage ) - if body.store is None or body.store is True: - self._save_state( - state=state, - response=current_response, + # persists a completed OpenAI Responses response as a WayFlow checkpoint + if should_store_response: + self.checkpointer.save( + state, + # Use the public response id as the checkpoint id so GET/DELETE/resume + # can locate the exact checkpoint from only the OpenAI Responses API id. + checkpoint_id=current_response.id, + # The stored component id is the served model id that owns this + # persisted conversation state. Resume must use the same model id. + component_id=model, + metadata={ + "response": current_response.model_dump_json(), + }, ) + self._cache_response_conversation_id(current_response.id, state.id) if current_response.error is not None: yield ResponseFailedEvent( @@ -369,106 +408,96 @@ def _load_state( agent_id: str, ) -> Optional[Conversation]: if previous_response_id: - try: - serialized_conversation = self._lookup_conversation( - where={self.storage_config.turn_id_column_name: previous_response_id}, - what=self.storage_config.conversation_turn_state_column_name, - ) - except ValueError: + # previous_response_id resumes the exact checkpoint created for that public + # response id, not just the latest checkpoint in the conversation. + checkpoint = self._find_checkpoint_for_response_id(previous_response_id) + if checkpoint is None: raise HTTPException( status_code=http_status_code.HTTP_404_NOT_FOUND, detail=f"No previous response with id `{previous_response_id}` was found", ) + incompatible_detail = ( + f"Previous response `{previous_response_id}` is not compatible " + f"with model `{agent_id}`" + ) elif conversation_id: - try: - serialized_conversation = self._lookup_conversation( - where={ - self.storage_config.conversation_id_column_name: conversation_id, - self.storage_config.is_last_turn_column_name: 1, # only latest round - }, - what=self.storage_config.conversation_turn_state_column_name, - ) - except ValueError: + # conversation resume follows OpenAI Responses semantics: continue from the + # latest stored turn in that conversation rather than a specific response id. + checkpoint = self.checkpointer.load_latest(conversation_id) + if checkpoint is None: raise HTTPException( status_code=http_status_code.HTTP_404_NOT_FOUND, detail=f"No conversation with id `{conversation_id}` was found", ) + incompatible_detail = ( + f"Conversation `{conversation_id}` is not compatible with model `{agent_id}`" + ) else: return None + + checkpoint_model = checkpoint.component_id + if checkpoint_model != agent_id: + raise HTTPException( + status_code=http_status_code.HTTP_400_BAD_REQUEST, + detail=( + f"{incompatible_detail}: checkpoint belongs to model " + f"`{checkpoint_model}` and cannot be resumed with model `{agent_id}`." + ), + ) + + self._cache_response_conversation_id(checkpoint.checkpoint_id, checkpoint.conversation_id) + agent = self.agents[agent_id] try: - return _deserialize_conversation_safely( - serialized_state=serialized_conversation, - tool_registry=self.tool_registries[agent_id], - component=self.agents[agent_id], + # The service owns persistence for OpenAI Responses, so restore the + # exact checkpoint without attaching a checkpointer for execute(). + return agent._restore_checkpointed_conversation( + checkpoint=checkpoint, + expected_conversation_type=agent.conversation_class, + attached_checkpointer=None, ) + except CheckpointRestoreCompatibilityError as e: + raise HTTPException( + status_code=http_status_code.HTTP_400_BAD_REQUEST, + detail=f"{incompatible_detail}: {e}", + ) from e except (TypeError, ValueError) as e: raise HTTPException( status_code=http_status_code.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Conversation state is corrupted, it cannot be de-serialized: {e}", - ) - - def _save_state( - self, - response: Response, - state: Conversation, - ) -> None: - conversation_model = response.conversation - if conversation_model is None: - raise ValueError("Internal Error: Conversation should not be None") - conversation_id = conversation_model.id + ) from e + + def _find_checkpoint_for_response_id( + self, response_id: str + ) -> Optional[ConversationCheckpoint]: + """Resolve the checkpoint stored under a public OpenAI Responses `response_id`.""" + checkpoint = self._load_cached_checkpoint_for_response_id(response_id) + if checkpoint is not None: + return checkpoint + + checkpoint = self.checkpointer._find_checkpoint_by_id(response_id) + if checkpoint is not None: + self._cache_response_conversation_id(response_id, checkpoint.conversation_id) + return checkpoint + + def _load_cached_checkpoint_for_response_id( + self, response_id: str + ) -> Optional[ConversationCheckpoint]: + conversation_id = self._response_conversation_ids.get(response_id) if conversation_id is None: - raise ValueError("Internal Error: Conversation ID should not be None") - - updates = {self.storage_config.is_last_turn_column_name: 0} - updates_where = { - self.storage_config.conversation_id_column_name: conversation_id, - self.storage_config.is_last_turn_column_name: 1, - } - serialized_state = serialize(state) - new_entity = { - self.storage_config.agent_id_column_name: response.model, - self.storage_config.conversation_id_column_name: conversation_id, - self.storage_config.turn_id_column_name: response.id, - self.storage_config.created_at_column_name: int(time.time()), - self.storage_config.conversation_turn_state_column_name: serialized_state, - self.storage_config.is_last_turn_column_name: 1, - self.storage_config.extra_metadata_column_name: json.dumps( - {"response": response.model_dump_json()} - ), - } - if isinstance(self.storage, RelationalDatastore): - # for relational datastores, we prefer making a single - # transaction, to avoid corrupting the state of the DB - # if the process crashes between the update and the insert - data_table = self.storage.data_tables[self.storage_config.table_name] - sql_update_stmt = data_table._update_query( - where=updates_where, - update=updates, - ) - sql_create_stmt, new_entities = data_table._create_query([new_entity]) - with data_table.engine.connect() as connection: - connection.execute(sql_update_stmt) - connection.execute(sql_create_stmt, new_entities) - connection.commit() + return None - else: - self.storage.update( - collection_name=self.storage_config.table_name, - where=updates_where, - update=updates, - ) - self.storage.create( - collection_name=self.storage_config.table_name, - entities=[new_entity], - ) + self._response_conversation_ids.move_to_end(response_id) + try: + return self.checkpointer.load(conversation_id, response_id) + except ValueError: + self._response_conversation_ids.pop(response_id) + return None - def _lookup_conversation(self, where: Dict[str, Any], what: str) -> Any: - serialized_conversations = self.storage.list( - collection_name=self.storage_config.table_name, where=where - ) - if len(serialized_conversations) != 1: - raise ValueError(f"No conversation with: {where}") - return serialized_conversations[0][what] + def _cache_response_conversation_id(self, response_id: str, conversation_id: str) -> None: + self._response_conversation_ids[response_id] = conversation_id + self._response_conversation_ids.move_to_end(response_id) + if len(self._response_conversation_ids) > self._RESPONSE_CONVERSATION_CACHE_MAX_SIZE: + self._response_conversation_ids.popitem(last=False) async def _create_state( self, diff --git a/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py b/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py index 4a27344e7..b62114736 100644 --- a/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py +++ b/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py @@ -4,52 +4,14 @@ # (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License # (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. - from dataclasses import dataclass -from typing import Dict, Optional -from wayflowcore.datastore import Datastore, Entity -from wayflowcore.property import IntegerProperty, StringProperty +from wayflowcore.checkpointing import StorageConfig @dataclass -class ServerStorageConfig: +class ServerStorageConfig(StorageConfig): """Configuration for server storage management.""" - datastore: Optional[Datastore] = None - """Datastore to use for persistence""" - - table_name: str = "conversations" - """Name of the table in which the states are stored""" - agent_id_column_name: str = "agent_id" - """Name of the column where the agent id of the state is stored""" - conversation_id_column_name: str = "conversation_id" - """Name of the column where the id of the conversation is stored""" - turn_id_column_name: str = "turn_id" - """Name of the column where the turn id / response id is stored""" - created_at_column_name: str = "created_at" - """Name of the column where the creation timestamp is stored""" - conversation_turn_state_column_name: str = "conversation_turn_state" - """Name of the column where the serialized state of turn is store""" - is_last_turn_column_name: str = "is_last_turn" - """Name of the column where the marker for the most recent turn of a given conversation is stored""" - extra_metadata_column_name: str = "extra_metadata" - """Name of the column where the server stores its own attributes""" - - max_retention: Optional[int] = None - """Number of seconds for which to retain a conversation before discarding it""" - - def to_schema(self) -> Dict[str, Entity]: - return { - self.table_name: Entity( - properties={ - self.agent_id_column_name: StringProperty(), - self.conversation_id_column_name: StringProperty(), - self.turn_id_column_name: StringProperty(), - self.is_last_turn_column_name: IntegerProperty(), - self.conversation_turn_state_column_name: StringProperty(), - self.created_at_column_name: IntegerProperty(), - self.extra_metadata_column_name: StringProperty(), - } - ), - } + # this is kept as a backward compatibility alias, + # all parameters and functionalities are defined in StorageConfig diff --git a/wayflowcore/src/wayflowcore/checkpointing/__init__.py b/wayflowcore/src/wayflowcore/checkpointing/__init__.py new file mode 100644 index 000000000..023a2d960 --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/__init__.py @@ -0,0 +1,31 @@ +# Copyright © 2026 Oracle and/or its affiliates. +# +# This software is under the Apache License 2.0 +# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License +# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. + +from .checkpointer import ( + Checkpointer, + CheckpointingInterval, + CheckpointRestoreCompatibilityError, + ConversationCheckpoint, + StorageConfig, +) +from .datastorecheckpointer import ( + DatastoreCheckpointer, + InMemoryCheckpointer, + OracleDatabaseCheckpointer, + PostgresCheckpointer, +) + +__all__ = [ + "CheckpointingInterval", + "CheckpointRestoreCompatibilityError", + "Checkpointer", + "ConversationCheckpoint", + "DatastoreCheckpointer", + "InMemoryCheckpointer", + "OracleDatabaseCheckpointer", + "PostgresCheckpointer", + "StorageConfig", +] diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py new file mode 100644 index 000000000..2f648ad3b --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py @@ -0,0 +1,216 @@ +# Copyright © 2026 Oracle and/or its affiliates. +# +# This software is under the Apache License 2.0 +# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License +# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + + +class CheckpointRestoreCompatibilityError(ValueError): + """Raised when a checkpoint cannot be resumed against the current live graph.""" + + +if TYPE_CHECKING: + from wayflowcore.conversation import Conversation + from wayflowcore.datastore import Datastore + + +@dataclass(frozen=True) +class ConversationCheckpoint: + """Durable snapshot of a conversation at a checkpoint boundary.""" + + checkpoint_id: str + """ID of the checkpoint""" + conversation_id: str + """ID of the stored conversation""" + component_id: str + """ID of the component that created the conversation""" + created_at: int + """Checkpoint creation time in seconds since the Unix epoch.""" + state: str + """Serialized conversation state.""" + metadata: Dict[str, Any] = field(default_factory=dict) + """Auxiliary checkpoint metadata used for ordering and inspection.""" + + @property + def id(self) -> str: + return self.checkpoint_id + + +class CheckpointingInterval(Enum): + """Configure which completed execution boundary triggers a checkpoint save.""" + + # Save only after the outermost `Conversation.execute()` returns. + CONVERSATION_TURNS = "conversation_turns" + """Saves the state at the end of a component turn (before returning from conversation.execute()""" + # Save after completed internal turns that actually used an LLM. + LLM_TURNS = "llm_turns" + """Saves the state at the end of a turn that uses a LLM (agent llm turn, PromptExecutionNode, ...)""" + # Save after every completed internal agent/flow turn boundary. + ALL_INTERNAL_TURNS = "all_internal_turns" + """Saves the state at every internal turn (every step of a flow, every iteration of an agent, ...) recursively in all sub-components""" + + +@dataclass +class StorageConfig: + """Configuration for checkpoint storage.""" + + datastore: Optional["Datastore"] = None + """Datastore to use for persistence""" + table_name: str = "conversations" + """Name of the table in which the states are stored""" + agent_id_column_name: str = "agent_id" + """Name of the column where the agent id of the state is stored""" + conversation_id_column_name: str = "conversation_id" + """Name of the column where the id of the conversation is stored""" + turn_id_column_name: str = "turn_id" + """Name of the column where the turn id / response id is stored""" + created_at_column_name: str = "created_at" + """Name of the column where the creation timestamp is stored""" + remove_by_column_name: str = "remove_by" + """Name of the column where the retention deadline timestamp is stored""" + conversation_turn_state_column_name: str = "conversation_turn_state" + """Name of the column where the serialized state of turn is store""" + is_last_turn_column_name: str = "is_last_turn" + """Name of the column where the marker for the most recent turn of a given conversation is stored""" + extra_metadata_column_name: str = "extra_metadata" + """Name of the column where the server stores its own attributes""" + max_retention: Optional[int] = None + """Number of seconds for which to retain a conversation before discarding it""" + + def to_schema(self) -> Dict[str, Any]: + from wayflowcore.datastore import Entity, nullable + from wayflowcore.property import IntegerProperty, StringProperty + + properties = { + self.agent_id_column_name: StringProperty(), + self.conversation_id_column_name: StringProperty(), + self.turn_id_column_name: StringProperty(), + self.is_last_turn_column_name: IntegerProperty(), + self.conversation_turn_state_column_name: StringProperty(), + self.created_at_column_name: IntegerProperty(), + self.extra_metadata_column_name: StringProperty(), + } + if self.max_retention is not None: + properties[self.remove_by_column_name] = nullable(IntegerProperty()) + return { + self.table_name: Entity( + properties=properties, + ), + } + + +class Checkpointer(ABC): + """Backend that can persist and restore checkpoints for conversations.""" + + def __init__( + self, + checkpointing_interval: CheckpointingInterval = CheckpointingInterval.CONVERSATION_TURNS, + ) -> None: + self.checkpointing_interval = checkpointing_interval + self._save_sequence_by_conversation: Dict[str, int] = {} + + @abstractmethod + def load_latest(self, conversation_id: str) -> Optional[ConversationCheckpoint]: + raise NotImplementedError() + + @abstractmethod + def load(self, conversation_id: str, checkpoint_id: str) -> ConversationCheckpoint: + raise NotImplementedError() + + def save( + self, + checkpoint: Union["Conversation", ConversationCheckpoint], + *, + checkpoint_id: Optional[str] = None, + component_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> Optional[ConversationCheckpoint]: + """Persist a checkpoint or snapshot a live conversation before persisting it. + + Checkpoint identifiers, component identifiers, and metadata are only valid + when saving a live conversation. + """ + from wayflowcore.conversation import Conversation + + if isinstance(checkpoint, Conversation): + return self._save_conversation( + checkpoint, + checkpoint_id=checkpoint_id, + component_id=component_id, + metadata=metadata, + ) + if checkpoint_id is not None or component_id is not None or metadata is not None: + raise ValueError( + "`checkpoint_id`, `component_id`, and `metadata` can only be provided " + "when saving a live Conversation." + ) + self._save_checkpoint(checkpoint) + return None + + def _save_conversation( + self, + conversation: "Conversation", + *, + checkpoint_id: Optional[str] = None, + component_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> ConversationCheckpoint: + """Create and persist a checkpoint for a live conversation.""" + import time + + from wayflowcore.idgeneration import IdGenerator + from wayflowcore.serialization import serialize + from wayflowcore.serialization.context import SerializationContext + + if not conversation.component._supports_checkpointing: + raise NotImplementedError("Checkpointing this component is not supported yet.") + serialization_context = SerializationContext(root=conversation) + serialization_context._register_external_component_references(conversation.component) + checkpoint = ConversationCheckpoint( + checkpoint_id=checkpoint_id or IdGenerator.get_or_generate_id(), + conversation_id=conversation.conversation_id, + component_id=component_id or conversation.component.id, + created_at=int(time.time()), + state=serialize(conversation, serialization_context=serialization_context), + metadata=dict(metadata or {}), + ) + self._save_checkpoint(checkpoint) + conversation.checkpoint_id = checkpoint.checkpoint_id + return checkpoint + + @abstractmethod + def _save_checkpoint(self, checkpoint: ConversationCheckpoint) -> None: + """Persist an already-materialized checkpoint in the backend.""" + raise NotImplementedError() + + async def save_async( + self, + checkpoint: Union["Conversation", ConversationCheckpoint], + *, + checkpoint_id: Optional[str] = None, + component_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> Optional[ConversationCheckpoint]: + # Async persistence is not implemented yet; this preserves the async API + # contract while delegating to the synchronous backend implementation. + return self.save( + checkpoint, + checkpoint_id=checkpoint_id, + component_id=component_id, + metadata=metadata, + ) + + @abstractmethod + def list_checkpoints( + self, conversation_id: str, limit: Optional[int] = 50 + ) -> List[ConversationCheckpoint]: + raise NotImplementedError() + + @abstractmethod + def delete(self, conversation_id: str, checkpoint_id: str) -> None: + raise NotImplementedError() diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py new file mode 100644 index 000000000..dc7ddeb43 --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py @@ -0,0 +1,203 @@ +# Copyright © 2026 Oracle and/or its affiliates. +# +# This software is under the Apache License 2.0 +# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License +# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Dict, Generator, Optional + +from ..events import Event, EventListener +from ..events.event import ( + AgentExecutionIterationStartedEvent, + FlowExecutionIterationStartedEvent, + LlmGenerationResponseEvent, +) +from .checkpointer import Checkpointer, CheckpointingInterval + +if TYPE_CHECKING: + from wayflowcore.conversation import Conversation + +_CHECKPOINT_SAVE_REASON_CONVERSATION_TURN = "conversation_turn" +_CHECKPOINT_SAVE_REASON_INTERNAL_TURN_BOUNDARY = "internal_turn_boundary" +_IterationStartedEvent = AgentExecutionIterationStartedEvent | FlowExecutionIterationStartedEvent +_ITERATION_STARTED_EVENTS = ( + AgentExecutionIterationStartedEvent, + FlowExecutionIterationStartedEvent, +) + + +def _find_checkpointed_conversation( + conversation: "Conversation", + execution_state: object, +) -> Optional["Conversation"]: + if conversation.state is execution_state: + return conversation + + # Events can be emitted by temporary helper conversations that are not part of + # the root conversation and its nested subconversations. Only save boundaries whose execution state belongs + # to a conversation that will actually be serialized. + for checkpoint_conversation in conversation._get_all_sub_conversations_recursive(): + if checkpoint_conversation.state is execution_state: + return checkpoint_conversation + return None + + +def _is_agent_iteration_start_that_exits_immediately( + checkpointed_conversation: "Conversation", + event: _IterationStartedEvent, +) -> bool: + from wayflowcore.agent import Agent + + # The agent executor emits an iteration-start event at the top of every loop, + # including the final loop that immediately exits because there is no work left. + if not isinstance(event, AgentExecutionIterationStartedEvent): + return False + + checkpointed_agent = checkpointed_conversation.component + if not isinstance(checkpointed_agent, Agent): + return False + execution_state = event.execution_state + if execution_state.curr_iter < checkpointed_agent.max_iterations: + return False + has_no_pending_tool_requests = ( + execution_state.current_tool_request is None and not execution_state.tool_call_queue + ) + return has_no_pending_tool_requests + + +def _build_listener_checkpoint_metadata( + conversation: "Conversation", + save_reason: str, + event: Optional[_IterationStartedEvent] = None, + metadata: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """Build listener metadata describing why this checkpoint was saved.""" + checkpoint_metadata: Dict[str, Any] = { + "save_reason": save_reason, + "current_step_name": conversation.current_step_name, + "message_count": len(conversation.message_list.messages), + } + if conversation.status is not None: + checkpoint_metadata["status_type"] = type(conversation.status).__name__ + if event is not None: + checkpoint_metadata["event_type"] = event.__class__.__name__ + + if isinstance(event, AgentExecutionIterationStartedEvent): + checkpoint_metadata["agent_iteration"] = event.execution_state.curr_iter + elif isinstance(event, FlowExecutionIterationStartedEvent): + checkpoint_metadata["flow_step_name"] = event.execution_state.current_step_name + checkpoint_metadata["nesting_level"] = event.execution_state.nesting_level + if metadata: + checkpoint_metadata.update(metadata) + return checkpoint_metadata + + +class _ConversationCheckpointEventListener(EventListener): + """Translate execution events into checkpoint saves for one root conversation. + + The listener is intentionally attached only around the outermost execute() + call. Nested conversations contribute to the same root snapshot through the + serialized root conversation and its nested subconversations rather than + writing independent checkpoints. + """ + + def __init__(self, conversation: "Conversation", checkpointer: Checkpointer) -> None: + self.conversation = conversation + self.checkpointer = checkpointer + self._pending_llm_checkpoint = False + self._latest_checkpoint_boundary_event: Optional[_IterationStartedEvent] = None + + def __call__(self, event: Event) -> None: + # LLM_TURNS checkpoints are delayed until the next safe agent/flow boundary. + # This avoids saving in the middle of a model/tool iteration. + if isinstance(event, LlmGenerationResponseEvent): + self._pending_llm_checkpoint = True + return + + if not isinstance(event, _ITERATION_STARTED_EVENTS): + return + + self._latest_checkpoint_boundary_event = event + + checkpointed_conversation = _find_checkpointed_conversation( + self.conversation, event.execution_state + ) + if checkpointed_conversation is None: + return + + match self.checkpointer.checkpointing_interval: + case CheckpointingInterval.CONVERSATION_TURNS: + # Conversation-turn checkpoints are written once after execute() returns. + return + case CheckpointingInterval.ALL_INTERNAL_TURNS: + self._save_internal_turn_checkpoint(event) + case CheckpointingInterval.LLM_TURNS if self._pending_llm_checkpoint: + # Skip the synthetic "next loop started, then exited immediately" event. + # The final conversation-turn checkpoint will capture the same state. + if _is_agent_iteration_start_that_exits_immediately( + checkpointed_conversation, event + ): + self._pending_llm_checkpoint = False + return + self._save_internal_turn_checkpoint(event) + + def save_pending_llm_checkpoint(self) -> None: + if self.checkpointer.checkpointing_interval != CheckpointingInterval.LLM_TURNS: + return + if not self._pending_llm_checkpoint: + return + event = self._latest_checkpoint_boundary_event + if event is None: + return + if _find_checkpointed_conversation(self.conversation, event.execution_state) is None: + return + + # If execution returns immediately after an LLM turn, there may be no following + # iteration-start event. Flush that last completed LLM boundary before final save. + self._save_internal_turn_checkpoint(event) + + def _save_internal_turn_checkpoint(self, event: _IterationStartedEvent) -> None: + self.checkpointer.save( + self.conversation, + metadata=_build_listener_checkpoint_metadata( + self.conversation, + save_reason=_CHECKPOINT_SAVE_REASON_INTERNAL_TURN_BOUNDARY, + event=event, + ), + ) + self._pending_llm_checkpoint = False + + +@contextmanager +def get_conversation_checkpoint_execution_context( + conversation: "Conversation", + is_outermost_execution: bool, +) -> Generator[None, None, None]: + """Context manager that wraps one outermost execute() with checkpointing. + + The ordering is deliberate: + - listen to internal execution events while the execute() call runs + - flush any pending LLM-turn checkpoint before the final save + - write the conversation-turn checkpoint only if execution returned normally + """ + checkpointer = conversation.checkpointer + if checkpointer is None or not is_outermost_execution: + # Nested executes are captured through the root conversation snapshot. Saving each + # nested execute separately would create duplicate/incomplete root checkpoints. + yield + return + + from wayflowcore.events.eventlistener import register_event_listeners + + listener = _ConversationCheckpointEventListener(conversation, checkpointer) + with register_event_listeners([listener]): + yield + listener.save_pending_llm_checkpoint() + checkpointer.save( + conversation, + metadata=_build_listener_checkpoint_metadata( + conversation, + save_reason=_CHECKPOINT_SAVE_REASON_CONVERSATION_TURN, + ), + ) diff --git a/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py new file mode 100644 index 000000000..71842ffa0 --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py @@ -0,0 +1,508 @@ +# Copyright © 2026 Oracle and/or its affiliates. +# +# This software is under the Apache License 2.0 +# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License +# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. + +import hashlib +import json +import time +from textwrap import dedent +from typing import Any, Dict, List, Optional, Sequence + +from wayflowcore.datastore import ( + Datastore, + InMemoryDatastore, + OracleDatabaseConnectionConfig, + OracleDatabaseDatastore, + PostgresDatabaseConnectionConfig, + PostgresDatabaseDatastore, +) +from wayflowcore.datastore._relational import RelationalDatastore +from wayflowcore.datastore.oracle import _execute_query_on_oracle_db +from wayflowcore.datastore.postgres import _execute_query_on_postgres_db + +from .checkpointer import Checkpointer, CheckpointingInterval, ConversationCheckpoint, StorageConfig + + +def _build_checkpoint_create_table_columns( + storage_config: StorageConfig, + is_oracle: bool, +) -> List[str]: + text_type = "CLOB" if is_oracle else "TEXT" + varchar_type = "VARCHAR2(255)" if is_oracle else "VARCHAR(255)" + + columns = [ + f"{storage_config.turn_id_column_name} {varchar_type} PRIMARY KEY", + f"{storage_config.agent_id_column_name} {varchar_type} NOT NULL", + f"{storage_config.conversation_id_column_name} {varchar_type} NOT NULL", + f"{storage_config.created_at_column_name} INTEGER NOT NULL", + f"{storage_config.conversation_turn_state_column_name} {text_type} NOT NULL", + f"{storage_config.is_last_turn_column_name} INTEGER NOT NULL", + f"{storage_config.extra_metadata_column_name} {text_type} NOT NULL", + ] + if storage_config.max_retention is not None: + columns.append(f"{storage_config.remove_by_column_name} INTEGER") + return columns + + +def _build_checkpoint_latest_index_name(storage_config: StorageConfig) -> str: + # Keep generated index names portable across supported relational backends. + # Oracle identifiers are limited to 30 bytes, so long custom table names are + # shortened with a stable suffix instead of being passed through directly. + raw_index_name = f"{storage_config.table_name}_conv_latest_idx" + safe_index_name = "".join( + character if character.isalnum() or character == "_" else "_" + for character in raw_index_name + ) + if len(safe_index_name) <= 30: + return safe_index_name + + digest = hashlib.sha256(raw_index_name.encode("utf-8")).hexdigest()[:8] + prefix = safe_index_name[:21].rstrip("_") or "wf_checkpoint" + return f"{prefix}_{digest}" + + +def _build_checkpoint_latest_index_query(storage_config: StorageConfig) -> str: + # Composite index for the two durable read paths: + # load_latest() filters by conversation_id + is_last_turn, while list_checkpoints() + # filters by conversation_id. Keeping conversation_id first supports both. + return dedent( + f""" + CREATE INDEX {_build_checkpoint_latest_index_name(storage_config)} + ON {storage_config.table_name} ( + {storage_config.conversation_id_column_name}, + {storage_config.is_last_turn_column_name} + ); + """ + ) + + +def _prepare_postgres_checkpoint_datastore( + connection_config: PostgresDatabaseConnectionConfig, + storage_config: StorageConfig, +) -> None: + from sqlalchemy.exc import ProgrammingError + + create_table_query = dedent( + f""" + CREATE TABLE {storage_config.table_name} ( + {", ".join(_build_checkpoint_create_table_columns(storage_config, is_oracle=False))} + ); + """ + ) + try: + _execute_query_on_postgres_db(connection_config, create_table_query) + _execute_query_on_postgres_db( + connection_config, + _build_checkpoint_latest_index_query(storage_config), + ) + except ProgrammingError as e: + if f'relation "{storage_config.table_name}" already exists' in str(e): + raise ValueError( + f'The datastore is already setup. Either delete the existing "{storage_config.table_name}" table or start the server with `--setup-datastore=no`.' + ) from e + raise + + +def _prepare_oracle_checkpoint_datastore( + connection_config: OracleDatabaseConnectionConfig, + storage_config: StorageConfig, +) -> None: + create_table_query = dedent( + f""" + CREATE TABLE {storage_config.table_name} ( + {", ".join(_build_checkpoint_create_table_columns(storage_config, is_oracle=True))} + ); + """ + ) + try: + _execute_query_on_oracle_db(connection_config, query=create_table_query) + _execute_query_on_oracle_db( + connection_config, + query=_build_checkpoint_latest_index_query(storage_config), + ) + except Exception as e: + if "already exists" in str(e): + raise ValueError( + f'The datastore is already setup. Either delete the existing "{storage_config.table_name}" table or start the server with `--setup-datastore=no`.' + ) from e + raise + + +class DatastoreCheckpointer(Checkpointer): + """Checkpointer backed by a WayFlow datastore.""" + + def __init__( + self, + datastore: Datastore, + storage_config: Optional[StorageConfig] = None, + checkpointing_interval: CheckpointingInterval = CheckpointingInterval.CONVERSATION_TURNS, + ) -> None: + super().__init__(checkpointing_interval=checkpointing_interval) + self.datastore = datastore + self.storage_config = storage_config or StorageConfig() + + def _entity_to_checkpoint(self, entity: Dict[str, Any]) -> ConversationCheckpoint: + raw_metadata = entity.get(self.storage_config.extra_metadata_column_name, "{}") + metadata = raw_metadata if isinstance(raw_metadata, dict) else json.loads(raw_metadata) + return ConversationCheckpoint( + checkpoint_id=str(entity[self.storage_config.turn_id_column_name]), + conversation_id=str(entity[self.storage_config.conversation_id_column_name]), + component_id=str(entity[self.storage_config.agent_id_column_name]), + created_at=int(entity[self.storage_config.created_at_column_name]), + state=str(entity[self.storage_config.conversation_turn_state_column_name]), + metadata=metadata, + ) + + def _checkpoint_to_entity(self, checkpoint: ConversationCheckpoint) -> Dict[str, Any]: + entity = { + self.storage_config.agent_id_column_name: checkpoint.component_id, + self.storage_config.conversation_id_column_name: checkpoint.conversation_id, + self.storage_config.turn_id_column_name: checkpoint.checkpoint_id, + self.storage_config.created_at_column_name: checkpoint.created_at, + self.storage_config.conversation_turn_state_column_name: checkpoint.state, + self.storage_config.is_last_turn_column_name: 1, + self.storage_config.extra_metadata_column_name: json.dumps(checkpoint.metadata), + } + if self.storage_config.max_retention is not None: + # Retention is represented as metadata for backend/external cleanup. Reads do + # not filter by this value, so callers should not treat it as enforced expiry. + entity[self.storage_config.remove_by_column_name] = ( + checkpoint.created_at + self.storage_config.max_retention + ) + return entity + + @staticmethod + def _sort_checkpoints( + checkpoints: Sequence[ConversationCheckpoint], + ) -> List[ConversationCheckpoint]: + # `created_at` is second-resolution for schema compatibility. Prefer the + # nanosecond timestamp and save sequence when available so same-second saves + # have stable ordering in list/load_latest/delete promotion paths. + return sorted( + checkpoints, + key=lambda checkpoint: ( + checkpoint.created_at, + checkpoint.metadata.get("created_at_ns", checkpoint.created_at * 1_000_000_000), + checkpoint.metadata.get("save_sequence", -1), + checkpoint.id, + ), + ) + + def _find_checkpoint( + self, + conversation_id: str, + checkpoint_id: str, + ) -> Optional[ConversationCheckpoint]: + entities = self.datastore.list( + collection_name=self.storage_config.table_name, + where={ + self.storage_config.conversation_id_column_name: conversation_id, + self.storage_config.turn_id_column_name: checkpoint_id, + }, + limit=1, + ) + if len(entities) == 0: + return None + return self._entity_to_checkpoint(entities[0]) + + def _find_checkpoint_by_id(self, checkpoint_id: str) -> Optional[ConversationCheckpoint]: + entities = self.datastore.list( + collection_name=self.storage_config.table_name, + where={self.storage_config.turn_id_column_name: checkpoint_id}, + limit=1, + ) + if len(entities) == 0: + return None + return self._entity_to_checkpoint(entities[0]) + + def load_latest(self, conversation_id: str) -> Optional[ConversationCheckpoint]: + entities = self.datastore.list( + collection_name=self.storage_config.table_name, + where={ + self.storage_config.conversation_id_column_name: conversation_id, + self.storage_config.is_last_turn_column_name: 1, + }, + ) + if len(entities) == 0: + return None + if len(entities) > 1: + raise ValueError( + f"Multiple latest checkpoints found for conversation `{conversation_id}`. " + "This can happen if multiple writers checkpoint the same conversation " + "concurrently or if the checkpoint datastore was modified manually. Use a " + "single writer per conversation or load a specific checkpoint by checkpoint_id." + ) + checkpoints = self._sort_checkpoints( + [self._entity_to_checkpoint(entity) for entity in entities] + ) + return checkpoints[-1] + + def load(self, conversation_id: str, checkpoint_id: str) -> ConversationCheckpoint: + checkpoint = self._find_checkpoint( + conversation_id=conversation_id, checkpoint_id=checkpoint_id + ) + if checkpoint is None: + raise ValueError( + f"Checkpoint `{checkpoint_id}` was not found for conversation `{conversation_id}`." + ) + return checkpoint + + def _save_checkpoint(self, checkpoint: ConversationCheckpoint) -> None: + created_at_ns = time.time_ns() + next_save_sequence = ( + self._save_sequence_by_conversation.get(checkpoint.conversation_id, 0) + 1 + ) + self._save_sequence_by_conversation[checkpoint.conversation_id] = next_save_sequence + checkpoint = ConversationCheckpoint( + checkpoint_id=checkpoint.checkpoint_id, + conversation_id=checkpoint.conversation_id, + component_id=checkpoint.component_id, + created_at=checkpoint.created_at, + state=checkpoint.state, + metadata=checkpoint.metadata + | { + "created_at_ns": created_at_ns, + "save_sequence": next_save_sequence, + }, + ) + existing_checkpoint = self._find_checkpoint( + conversation_id=checkpoint.conversation_id, + checkpoint_id=checkpoint.checkpoint_id, + ) + if existing_checkpoint is not None: + checkpoint = ConversationCheckpoint( + checkpoint_id=checkpoint.checkpoint_id, + conversation_id=checkpoint.conversation_id, + component_id=checkpoint.component_id, + created_at=checkpoint.created_at, + state=checkpoint.state, + metadata=existing_checkpoint.metadata | checkpoint.metadata, + ) + + update_latest_where = { + self.storage_config.conversation_id_column_name: checkpoint.conversation_id, + self.storage_config.is_last_turn_column_name: 1, + } + update_latest_values = {self.storage_config.is_last_turn_column_name: 0} + entity = self._checkpoint_to_entity(checkpoint) + + if isinstance(self.datastore, RelationalDatastore): + data_table = self.datastore.data_tables[self.storage_config.table_name] + with data_table.engine.connect() as connection: + # Latest promotion is transactional for one writer, but the table does + # not define a uniqueness constraint on (conversation_id, is_last_turn). + # Shared relational deployments should use a single writer per conversation. + connection.execute( + data_table._update_query( + where=update_latest_where, + update=update_latest_values, + ) + ) + if existing_checkpoint is None: + sql_create_stmt, new_entities = data_table._create_query([entity]) + connection.execute(sql_create_stmt, new_entities) + else: + update_checkpoint_where = { + self.storage_config.conversation_id_column_name: checkpoint.conversation_id, + self.storage_config.turn_id_column_name: checkpoint.checkpoint_id, + } + update_checkpoint_values = { + self.storage_config.agent_id_column_name: checkpoint.component_id, + self.storage_config.created_at_column_name: checkpoint.created_at, + self.storage_config.conversation_turn_state_column_name: checkpoint.state, + self.storage_config.is_last_turn_column_name: 1, + self.storage_config.extra_metadata_column_name: json.dumps( + checkpoint.metadata + ), + } + if self.storage_config.max_retention is not None: + update_checkpoint_values[self.storage_config.remove_by_column_name] = ( + checkpoint.created_at + self.storage_config.max_retention + ) + connection.execute( + data_table._update_query( + where=update_checkpoint_where, + update=update_checkpoint_values, + ) + ) + connection.commit() + return + + # Non-relational datastores follow the same logical promotion sequence. + # They generally do not provide stronger concurrency guarantees than the + # backing datastore implementation. + self.datastore.update( + collection_name=self.storage_config.table_name, + where=update_latest_where, + update=update_latest_values, + ) + if existing_checkpoint is None: + self.datastore.create( + collection_name=self.storage_config.table_name, + entities=[entity], + ) + else: + update_checkpoint_values = { + self.storage_config.agent_id_column_name: checkpoint.component_id, + self.storage_config.created_at_column_name: checkpoint.created_at, + self.storage_config.conversation_turn_state_column_name: checkpoint.state, + self.storage_config.is_last_turn_column_name: 1, + self.storage_config.extra_metadata_column_name: json.dumps(checkpoint.metadata), + } + if self.storage_config.max_retention is not None: + update_checkpoint_values[self.storage_config.remove_by_column_name] = ( + checkpoint.created_at + self.storage_config.max_retention + ) + self.datastore.update( + collection_name=self.storage_config.table_name, + where={ + self.storage_config.conversation_id_column_name: checkpoint.conversation_id, + self.storage_config.turn_id_column_name: checkpoint.checkpoint_id, + }, + update=update_checkpoint_values, + ) + + def list_checkpoints( + self, conversation_id: str, limit: Optional[int] = 50 + ) -> List[ConversationCheckpoint]: + checkpoints = self._sort_checkpoints( + [ + self._entity_to_checkpoint(entity) + for entity in self.datastore.list( + collection_name=self.storage_config.table_name, + where={self.storage_config.conversation_id_column_name: conversation_id}, + ) + ] + ) + if limit is not None: + should_apply_limit = len(checkpoints) > limit + if should_apply_limit: + checkpoints = checkpoints[-limit:] + return checkpoints + + def delete(self, conversation_id: str, checkpoint_id: str) -> None: + latest_checkpoint = self.load_latest(conversation_id) + checkpoints = self.list_checkpoints(conversation_id, limit=None) + checkpoint_to_promote: Optional[ConversationCheckpoint] = None + is_deleting_latest_checkpoint = ( + latest_checkpoint is not None and latest_checkpoint.checkpoint_id == checkpoint_id + ) + if is_deleting_latest_checkpoint: + # Deleting the current latest checkpoint should leave the conversation + # resumable from the next-newest checkpoint when one exists. + remaining_checkpoints = [ + checkpoint + for checkpoint in checkpoints + if checkpoint.checkpoint_id != checkpoint_id + ] + checkpoint_to_promote = remaining_checkpoints[-1] if remaining_checkpoints else None + + delete_where = { + self.storage_config.conversation_id_column_name: conversation_id, + self.storage_config.turn_id_column_name: checkpoint_id, + } + + if isinstance(self.datastore, RelationalDatastore): + data_table = self.datastore.data_tables[self.storage_config.table_name] + with data_table.engine.connect() as connection: + # Keep deletion of the latest checkpoint and promotion of the previous + # checkpoint atomic. Otherwise a crash between datastore.delete() and + # datastore.update() could leave the conversation with no latest row. + # This bypasses datatable.delete() because that helper executes and + # commits immediately; here we need the DELETE and possible UPDATE to + # share the same transaction. + connection.execute(data_table._delete_query(where=delete_where)) + if checkpoint_to_promote is not None: + connection.execute( + data_table._update_query( + where={ + self.storage_config.conversation_id_column_name: ( + checkpoint_to_promote.conversation_id + ), + self.storage_config.turn_id_column_name: ( + checkpoint_to_promote.checkpoint_id + ), + }, + update={self.storage_config.is_last_turn_column_name: 1}, + ) + ) + connection.commit() + return + + self.datastore.delete( + collection_name=self.storage_config.table_name, + where=delete_where, + ) + + if checkpoint_to_promote is not None: + self.datastore.update( + collection_name=self.storage_config.table_name, + where={ + self.storage_config.conversation_id_column_name: conversation_id, + self.storage_config.turn_id_column_name: checkpoint_to_promote.checkpoint_id, + }, + update={self.storage_config.is_last_turn_column_name: 1}, + ) + + +class InMemoryCheckpointer(DatastoreCheckpointer): + """Checkpointer backed by an in-memory datastore.""" + + def __init__( + self, + storage_config: Optional[StorageConfig] = None, + checkpointing_interval: CheckpointingInterval = CheckpointingInterval.CONVERSATION_TURNS, + ) -> None: + resolved_storage_config = storage_config or StorageConfig() + datastore = InMemoryDatastore(schema=resolved_storage_config.to_schema()) + super().__init__( + datastore=datastore, + storage_config=resolved_storage_config, + checkpointing_interval=checkpointing_interval, + ) + + +class PostgresCheckpointer(DatastoreCheckpointer): + """Checkpointer backed by PostgreSQL.""" + + def __init__( + self, + connection_config: PostgresDatabaseConnectionConfig, + storage_config: Optional[StorageConfig] = None, + checkpointing_interval: CheckpointingInterval = CheckpointingInterval.CONVERSATION_TURNS, + ) -> None: + resolved_storage_config = storage_config or StorageConfig() + datastore = PostgresDatabaseDatastore( + schema=resolved_storage_config.to_schema(), + connection_config=connection_config, + ) + super().__init__( + datastore=datastore, + storage_config=resolved_storage_config, + checkpointing_interval=checkpointing_interval, + ) + self.connection_config = connection_config + + +class OracleDatabaseCheckpointer(DatastoreCheckpointer): + """Checkpointer backed by Oracle Database.""" + + def __init__( + self, + connection_config: OracleDatabaseConnectionConfig, + storage_config: Optional[StorageConfig] = None, + checkpointing_interval: CheckpointingInterval = CheckpointingInterval.CONVERSATION_TURNS, + ) -> None: + resolved_storage_config = storage_config or StorageConfig() + datastore = OracleDatabaseDatastore( + schema=resolved_storage_config.to_schema(), + connection_config=connection_config, + ) + super().__init__( + datastore=datastore, + storage_config=resolved_storage_config, + checkpointing_interval=checkpointing_interval, + ) + self.connection_config = connection_config diff --git a/wayflowcore/src/wayflowcore/cli/serve.py b/wayflowcore/src/wayflowcore/cli/serve.py index 66108ff46..74b9d54ef 100644 --- a/wayflowcore/src/wayflowcore/cli/serve.py +++ b/wayflowcore/src/wayflowcore/cli/serve.py @@ -17,12 +17,12 @@ import yaml from wayflowcore.agentserver import ServerStorageConfig -from wayflowcore.agentserver._storagehelpers import ( - _prepare_oracle_datastore, - _prepare_postgres_datastore, -) from wayflowcore.agentserver.app import create_server_app from wayflowcore.agentspec import AgentSpecLoader +from wayflowcore.checkpointing.datastorecheckpointer import ( + _prepare_oracle_checkpoint_datastore, + _prepare_postgres_checkpoint_datastore, +) from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.datastore import ( Datastore, @@ -305,7 +305,7 @@ def _get_persistence_arguments( "datastore_connection_config must be a PostgresDatabaseConnectionConfig instance." ) if setup_datastore: - _prepare_postgres_datastore(datastore_connection_config, storage_config) + _prepare_postgres_checkpoint_datastore(datastore_connection_config, storage_config) storage = PostgresDatabaseDatastore( schema=storage_schema, connection_config=datastore_connection_config, @@ -326,7 +326,7 @@ def _get_persistence_arguments( "datastore_connection_config must be an OracleDatabaseConnectionConfig instance." ) if setup_datastore: - _prepare_oracle_datastore(datastore_connection_config, storage_config) + _prepare_oracle_checkpoint_datastore(datastore_connection_config, storage_config) storage = OracleDatabaseDatastore( schema=storage_schema, connection_config=datastore_connection_config, @@ -401,6 +401,8 @@ def _load_datastore_connection_config( if storage_type == "oracle-db": if raw_type == "TlsOracleDatabaseConnectionConfig": return TlsOracleDatabaseConnectionConfig(**data) + elif raw_type == "MTlsOracleDatabaseConnectionConfig": + return MTlsOracleDatabaseConnectionConfig(**data) else: raise ValueError( f"For oracle db storage type `{raw_type}`, the connection config type should be either `TlsOracleDatabaseConnectionConfig` or `MTlsOracleDatabaseConnectionConfig` but got `{raw_type}`" diff --git a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py index c98d1db59..d4dbcc4db 100644 --- a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py +++ b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py @@ -92,7 +92,8 @@ async def call_async(self, conversation: "Conversation") -> Any: from wayflowcore.tracing.span import ContextProviderExecutionSpan with ContextProviderExecutionSpan(context_provider=self) as span: - conversation = self.flow.start_conversation( + conversation = self.flow._start_conversation_impl( + parent_conversation=conversation, inputs={}, messages=conversation.message_list, ) diff --git a/wayflowcore/src/wayflowcore/conversation.py b/wayflowcore/src/wayflowcore/conversation.py index 14fa440c2..594cfab11 100644 --- a/wayflowcore/src/wayflowcore/conversation.py +++ b/wayflowcore/src/wayflowcore/conversation.py @@ -24,7 +24,6 @@ from wayflowcore._utils.async_helpers import run_async_in_sync from wayflowcore.component import DataclassComponent -from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.executors._events.event import Event from wayflowcore.executors.executionstatus import ( ExecutionStatus, @@ -37,7 +36,9 @@ from wayflowcore.tokenusage import TokenUsage if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.contextproviders import ContextProvider + from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.executors._executionstate import ConversationExecutionState from wayflowcore.executors.interrupts.executioninterrupt import ExecutionInterrupt from wayflowcore.models._requesthelpers import TaggedMessageChunkType @@ -63,12 +64,21 @@ def _get_active_conversations(return_copy: bool = True) -> List["Conversation"]: def _get_current_conversation_id() -> Optional[str]: + """Return the instance id of the currently executing conversation object.""" active_conversations = _get_active_conversations(return_copy=True) if not active_conversations: return None return active_conversations[-1].id +def _get_current_conversation_thread_id() -> Optional[str]: + """Return the thread id shared by the current nested conversations.""" + active_conversations = _get_active_conversations(return_copy=True) + if not active_conversations: + return None + return active_conversations[-1].conversation_id + + @contextmanager def _register_conversation(conversation: "Conversation") -> Generator[None, Any, None]: try: @@ -85,14 +95,23 @@ def _register_conversation(conversation: "Conversation") -> Generator[None, Any, @dataclass class Conversation(DataclassComponent): - component: ConversationalComponent + component: "ConversationalComponent" state: "ConversationExecutionState" inputs: Dict[str, Any] message_list: MessageList status: Optional[ExecutionStatus] token_usage: TokenUsage = field(default_factory=TokenUsage, init=False) - conversation_id: str = "" # deprecated - + conversation_id: str = "" + """Conversation thread id used for checkpoint persistence and resume.""" + checkpointer: Optional["Checkpointer"] = field( + default=None, + repr=False, + compare=False, + metadata={"serialize": False}, + ) + """Optional checkpointer to save the conversation""" + checkpoint_id: Optional[str] = field(default=None, init=False, repr=False, compare=False) + """ID of the current checkpoint the conversation is stored with""" status_handled: bool = False """Whether the current status associated to this conversation was already handled or not (messages/tool results were added to the conversation)""" @@ -100,6 +119,8 @@ class Conversation(DataclassComponent): def __post_init__(self) -> None: if self.inputs is None: self.inputs = {} + if not self.conversation_id: + self.conversation_id = self.id @property def plan(self) -> Optional[ExecutionPlan]: @@ -138,11 +159,17 @@ async def execute_async( if self.status_handled is False: self._update_conversation_with_status() - with _register_conversation(self): - new_status = await self.component.runner.execute_async(self, execution_interrupts) + from wayflowcore.checkpointing.checkpointeventlistener import ( + get_conversation_checkpoint_execution_context, + ) - self.status = new_status - self.status_handled = False + with get_conversation_checkpoint_execution_context( + self, + is_outermost_execution=len(_get_active_conversations(return_copy=False)) == 0, + ): + with _register_conversation(self): + self.status = await self.component.runner.execute_async(self, execution_interrupts) + self.status_handled = False return self.status @property @@ -159,9 +186,17 @@ def _get_all_context_providers_from_parent_conversations( @abstractmethod def _get_all_sub_conversations(self) -> List["Conversation"]: - """Gathers all sub conversations""" + """Return direct child subconversations.""" raise NotImplementedError() + def _get_all_sub_conversations_recursive(self) -> List["Conversation"]: + """Return all nested child subconversations recursively.""" + all_sub_conversations: List["Conversation"] = [] + for child_conversation in self._get_all_sub_conversations(): + all_sub_conversations.append(child_conversation) + all_sub_conversations.extend(child_conversation._get_all_sub_conversations_recursive()) + return all_sub_conversations + @abstractmethod def __repr__(self) -> str: raise NotImplementedError() diff --git a/wayflowcore/src/wayflowcore/conversationalcomponent.py b/wayflowcore/src/wayflowcore/conversationalcomponent.py index f6eef9d00..8fce88b65 100644 --- a/wayflowcore/src/wayflowcore/conversationalcomponent.py +++ b/wayflowcore/src/wayflowcore/conversationalcomponent.py @@ -11,23 +11,27 @@ from wayflowcore._metadata import MetadataType from wayflowcore.componentwithio import ComponentWithInputsOutputs +from wayflowcore.idgeneration import IdGenerator +from wayflowcore.messagelist import Message from wayflowcore.property import Property logger = logging.getLogger(__name__) if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer + from wayflowcore.checkpointing.checkpointer import ConversationCheckpoint from wayflowcore.conversation import Conversation from wayflowcore.executors._executor import ConversationExecutor - from wayflowcore.messagelist import Message, MessageList + from wayflowcore.messagelist import MessageList from wayflowcore.models.llmmodel import LlmModel from wayflowcore.tools import Tool _HUMAN_ENTITY_ID = "human_user" +ConversationTypeT = TypeVar("ConversationTypeT", bound="Conversation") class ConversationalComponent(ComponentWithInputsOutputs, ABC): - def __init__( self, name: str, @@ -35,7 +39,7 @@ def __init__( input_descriptors: List["Property"], output_descriptors: List["Property"], runner: Type["ConversationExecutor"], - conversation_class: Any, + conversation_class: Type["Conversation"], id: Optional[str] = None, __metadata_info__: Optional[MetadataType] = None, ) -> None: @@ -66,13 +70,33 @@ def start_conversation( self, inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, ) -> "Conversation": pass + def _start_conversation_impl( + self, + inputs: Optional[Dict[str, Any]], + messages: Union[None, str, "Message", List["Message"], "MessageList"], + conversation_id: Optional[str], + checkpointer: Optional["Checkpointer"], + checkpoint_id: Optional[str], + parent_conversation: Optional["Conversation"] = None, + ) -> "Conversation": + raise NotImplementedError + @property def llms(self) -> List["LlmModel"]: raise NotImplementedError("to be implemented by child classes") + @property + @abstractmethod + def _supports_checkpointing(self) -> bool: + """Whether this component and its nested components support checkpointing.""" + raise NotImplementedError("to be implemented by child classes") + def _referenced_tools(self, recursive: bool = True) -> List["Tool"]: """ Returns a list of all tools that are present in this component's configuration, including tools @@ -115,6 +139,132 @@ def _update_internal_state(self) -> None: Method to update the attributes inside. """ + @staticmethod + def _messages_or_inputs_were_passed( + inputs: Optional[Dict[str, Any]], + messages: Union[None, str, "Message", List["Message"], "MessageList"], + ) -> bool: + if inputs: + return True + if messages is None: + return False + if isinstance(messages, str): + return len(messages) > 0 + if isinstance(messages, Message): + return True + return len(messages) > 0 + + def _prepare_conversation_start( + self, + inputs: Optional[Dict[str, Any]], + messages: Union[None, str, "Message", List["Message"], "MessageList"], + conversation_id: Optional[str], + checkpointer: Optional["Checkpointer"], + checkpoint_id: Optional[str], + expected_conversation_type: Type[ConversationTypeT], + parent_conversation: Optional["Conversation"], + ) -> tuple[Optional[ConversationTypeT], str, str]: + """Resolve whether start_conversation creates a fresh or restored conversation. + + Returns a tuple of: + - an already restored conversation, or ``None`` for a fresh start + - the generated id to assign to the concrete ``Conversation`` object + - the conversation thread id shared by nested conversations + + ``expected_conversation_type`` keeps the restored conversation typed for static typing; + runtime validation uses it too. + """ + if parent_conversation is not None: + if checkpoint_id is not None: + raise ValueError("Cannot restore a checkpoint as a subconversation.") + # Root checkpoints capture the complete child tree, so construct children + # normally instead of restoring them as independent conversations. + return None, IdGenerator.get_or_generate_id(), parent_conversation.conversation_id + + if checkpoint_id is not None and checkpointer is None: + raise ValueError("`checkpoint_id` requires a `checkpointer`.") + if checkpoint_id is not None and conversation_id is None: + raise ValueError("`checkpoint_id` requires a `conversation_id`.") + + # A root conversation is the first instance in its thread, so one ID serves both roles. + conversation_id = conversation_id or IdGenerator.get_or_generate_id() + + if checkpointer is None: + return None, conversation_id, conversation_id + + checkpoint = ( + checkpointer.load(conversation_id, checkpoint_id) + if checkpoint_id is not None + else checkpointer.load_latest(conversation_id) + ) + if checkpoint is None: + return None, conversation_id, conversation_id + + # Restoring and starting with new input are separate operations. + if self._messages_or_inputs_were_passed(inputs=inputs, messages=messages): + raise ValueError( + "Cannot restore a checkpoint while also passing new `inputs` or `messages`. " + "Load the conversation first, then append new user input explicitly." + ) + + conversation = self._restore_checkpointed_conversation( + checkpoint=checkpoint, + expected_conversation_type=expected_conversation_type, + attached_checkpointer=checkpointer, + ) + return conversation, conversation_id, conversation_id + + def _restore_checkpointed_conversation( + self, + checkpoint: "ConversationCheckpoint", + expected_conversation_type: Type[ConversationTypeT], + attached_checkpointer: Optional["Checkpointer"], + ) -> ConversationTypeT: + """Rehydrate a checkpoint against this component's live graph and tools. + + ``attached_checkpointer`` controls automatic checkpoint saves after restore. + It is separate from checkpoint loading, which has already completed. + """ + from wayflowcore.checkpointing import CheckpointRestoreCompatibilityError + from wayflowcore.exceptions import DataclassFieldDeserializationError + from wayflowcore.serialization import autodeserialize + from wayflowcore.serialization.context import ( + DeserializationContext, + _MissingDeserializationReferenceError, + ) + + # Rehydrate the stored conversation against this live component tree. + deserialization_context = DeserializationContext() + deserialization_context.registered_tools = { + tool.name: tool for tool in self._referenced_tools() + } + deserialization_context._register_external_component_references(self) + try: + conversation = autodeserialize( + checkpoint.state, + deserialization_context=deserialization_context, + ) + except (_MissingDeserializationReferenceError, DataclassFieldDeserializationError) as exc: + error: Optional[BaseException] = exc + while error is not None and not isinstance( + error, _MissingDeserializationReferenceError + ): + error = error.__cause__ + if error is None: + raise + raise CheckpointRestoreCompatibilityError( + "Cannot restore this checkpoint because the current component tree does not " + "match the serialized component ids. Restart-safe checkpoint restore requires stable component ids." + ) from exc + if not isinstance(conversation, expected_conversation_type): + raise CheckpointRestoreCompatibilityError( + "Cannot restore this checkpoint because this conversation was started with another " + f"component. Expected `{expected_conversation_type.__name__}`, got `{type(conversation).__name__}`." + ) + conversation.checkpoint_id = checkpoint.checkpoint_id + conversation.checkpointer = attached_checkpointer + return conversation + # Define a TypeVar that represents the component's type ConversationalComponentTypeT = TypeVar( diff --git a/wayflowcore/src/wayflowcore/datastore/_relational.py b/wayflowcore/src/wayflowcore/datastore/_relational.py index a973127cd..13aa4117c 100644 --- a/wayflowcore/src/wayflowcore/datastore/_relational.py +++ b/wayflowcore/src/wayflowcore/datastore/_relational.py @@ -331,8 +331,7 @@ def update(self, where: Dict[str, Any], update: EntityAsDictT) -> List[EntityAsD return result_as_dict def delete(self, where: Dict[str, Any]) -> None: - query = sqlalchemy.delete(self.sqlalchemy_table) - query = self._apply_where_clause(query, where) + query = self._delete_query(where) with self.engine.connect() as connection: result = connection.execute(query) if result.rowcount == 0: @@ -341,6 +340,12 @@ def delete(self, where: Dict[str, Any]) -> None: logger.info("Deleted %i entities", result.rowcount) connection.commit() + def _delete_query(self, where: Dict[str, Any]) -> "sqlalchemy.Delete": + # Used by DatastoreCheckpointer.delete() too, so it can compose the DELETE + # with a follow-up is_last_turn promotion UPDATE in one transaction. + query = sqlalchemy.delete(self.sqlalchemy_table) + return self._apply_where_clause(query, where) + class RelationalDatastore(Datastore, ABC): """A relational data store that supports querying data using diff --git a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py index 79fbc9ac3..49de04ff9 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py @@ -140,6 +140,14 @@ def _get_current_tool_request(self) -> None: if self.current_tool_request is None and self.tool_call_queue: self.current_tool_request = self.tool_call_queue.pop(0) + def _finalize_execute_turn(self) -> None: + # curr_iter is scoped to one top-level execute() call. Conversation-turn + # checkpoints and ordinary follow-up executes should restart with a fresh + # iteration budget, while internal-turn checkpoints preserve the live value + # captured before this finalization runs. + self.curr_iter = 0 + self.current_retrieved_tools = None + def _agent_as_client_tool(agent: Union[Agent, OciAgent]) -> ClientTool: agent_input_parameters: Dict[str, JsonSchemaParam] = {} @@ -461,8 +469,10 @@ def _get_or_create_expert_agent_subconversation( init_messages = MessageList.from_messages([]) init_messages.append_message(caller_request_message) - sub_agent_conversation = expert_agent.start_conversation( - messages=init_messages, inputs=inputs + sub_agent_conversation = expert_agent._start_conversation_impl( + parent_conversation=caller_conv, + messages=init_messages, + inputs=inputs, ) return sub_agent_conversation @@ -519,6 +529,7 @@ async def _execute_flow( messages: MessageList, flow: Flow, inputs: Dict[str, Any], + parent_conversation: "AgentConversation", ) -> Tuple[Any, str, ExecutionStatus]: """ Execute a flow and return its outputs and its execution status. @@ -528,7 +539,8 @@ async def _execute_flow( outputs: Any = None try: if state.current_flow_conversation is None: - state.current_flow_conversation = flow.start_conversation( + state.current_flow_conversation = flow._start_conversation_impl( + parent_conversation=parent_conversation, inputs=inputs, messages=messages, ) @@ -593,7 +605,7 @@ async def _execute_next_subcall( _descriptors_to_json_schema_map(flow.input_descriptors_dict.values()), ) return await AgentConversationExecutor._handle_flow_call( - config, state, flow, tool_request, messages + config, state, flow, tool_request, messages, conversation ) if state.current_retrieved_tools is None: @@ -746,6 +758,7 @@ async def _handle_flow_call( flow: Flow, tool_request: ToolRequest, messages: MessageList, + conversation: "AgentConversation", ) -> Optional[ExecutionStatus]: logger.debug( 'Agent executing flow "%s" (id=%s) with arguments: %s', @@ -754,7 +767,13 @@ async def _handle_flow_call( tool_request.args, ) output, serialized_output, flow_execution_status = ( - await AgentConversationExecutor._execute_flow(state, messages, flow, tool_request.args) + await AgentConversationExecutor._execute_flow( + state, + messages, + flow, + tool_request.args, + conversation, + ) ) logger.debug( @@ -1070,9 +1089,12 @@ async def _execute_agent( agent_state = conversation.state messages = conversation.message_list - agent_state.curr_iter = 0 should_yield = False + def _return_status(execution_status: ExecutionStatus) -> ExecutionStatus: + agent_state._finalize_execute_turn() + return execution_status + while not should_yield: record_event( @@ -1090,7 +1112,7 @@ async def _execute_agent( conversation=conversation, ) if execution_status is not None: - return execution_status + return _return_status(execution_status) agent_state.current_tool_request = None @@ -1108,8 +1130,10 @@ async def _execute_agent( output.name: default_outputs.get(output.name, output.default_value) for output in agent_config.output_descriptors } - return FinishedStatus( - output_values=default_outputs, _conversation_id=conversation.id + return _return_status( + FinishedStatus( + output_values=default_outputs, _conversation_id=conversation.id + ) ) break else: @@ -1118,7 +1142,7 @@ async def _execute_agent( config=agent_config, curr_iter=agent_state.curr_iter ) except AuthInterrupt as auth_interrupt: - return auth_interrupt.status + return _return_status(auth_interrupt.status) logger.debug("No open tool call, will decide next action by prompting the llm") agent_state.current_retrieved_tools = retrieved_tools @@ -1162,11 +1186,13 @@ async def _execute_agent( last_message = conversation.get_last_message() if last_message is None: raise ValueError("Something went wrong, should not happen") - return UserMessageRequestStatus( - message=last_message, - _conversation_id=conversation.id, + return _return_status( + UserMessageRequestStatus( + message=last_message, + _conversation_id=conversation.id, + ) ) - return FinishedStatus(output_values={}, _conversation_id=conversation.id) + return _return_status(FinishedStatus(output_values={}, _conversation_id=conversation.id)) @staticmethod def _get_submit_tool_name(agent_config: Agent) -> str: diff --git a/wayflowcore/src/wayflowcore/executors/_flowconversation.py b/wayflowcore/src/wayflowcore/executors/_flowconversation.py index de8399181..1f1c43a8a 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_flowconversation.py @@ -16,6 +16,7 @@ if TYPE_CHECKING: from wayflowcore.executors._flowexecutor import FlowConversationExecutor + from wayflowcore.serialization.context import DeserializationContext from wayflowcore.steps.step import Step @@ -24,6 +25,22 @@ class FlowConversation(Conversation): component: Flow state: FlowConversationExecutionState + @classmethod + def _deserialize_from_dict( + cls, input_dict: Dict[str, Any], deserialization_context: "DeserializationContext" + ) -> "FlowConversation": + conversation = cast( + FlowConversation, + super()._deserialize_from_dict(input_dict, deserialization_context), + ) + # FlowConversationExecutionState._serialize_context_value serializes + # _SUPER_CONVERSATION_KEY as None to avoid a parent -> child -> parent cycle, + # so restore direct child flow parents here. + for child_conversation in conversation._get_all_sub_conversations(): + if isinstance(child_conversation, FlowConversation): + child_conversation.state._register_super_conversation(conversation) + return conversation + def _gather_flow_outputs(self) -> Dict[str, Any]: return FlowConversationExecutor.gather_flow_outputs( state=self.state, @@ -184,14 +201,16 @@ def __str__(self) -> str: result = f"State: {self.state}\nList of messages:\n" for i, message in enumerate(self.message_list.messages): - message_str = dedent(""" + message_str = dedent( + """ Message #{} Message type: {} Message content:\n {}\n tool_requests: {} tool_results: {} - """).format( + """ + ).format( i, message.message_type, message.content, message.tool_requests, message.tool_result ) diff --git a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py index eb572ff22..222f54265 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py @@ -286,17 +286,18 @@ def create_sub_conversation( k: v for k, v in inputs.items() if k not in all_context_provider_keys } - sub_conversation = flow.start_conversation( - inputs_not_from_context_providers, - conversation_id=conversation.conversation_id, + resolved_sub_conversation_id = ( + sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY + ) + sub_conversation = flow._start_nested_flow_conversation( + parent_conversation=conversation, + inputs=inputs_not_from_context_providers, messages=conversation.message_list, nesting_level=conversation.state.nesting_level + 1, context_providers_from_parent_flow=all_context_provider_keys, ) - key = FlowConversationExecutor.make_key_for_step( - step, sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY - ) + key = FlowConversationExecutor.make_key_for_step(step, resolved_sub_conversation_id) conversation.state.internal_context_key_values[key] = sub_conversation sub_conversation._put_internal_context_key_value( FlowConversationExecutor._SUPER_CONVERSATION_KEY, conversation @@ -928,8 +929,11 @@ async def _execute_flow( def get_all_sub_conversations( state: FlowConversationExecutionState, ) -> List["Conversation"]: + from wayflowcore.conversation import Conversation + return [ conv for k, conv in state.internal_context_key_values.items() - if FlowConversationExecutor._SUB_CONVERSATION_KEY in k + if k != FlowConversationExecutor._SUPER_CONVERSATION_KEY + and isinstance(conv, Conversation) ] diff --git a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py index 980cebdff..d7dff8520 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py @@ -24,15 +24,16 @@ @dataclass class ManagerWorkersConversationExecutionState(ConversationExecutionState): - current_agent_name: str + current_agent_id: str subconversations: Dict[str, Union["AgentConversation", "ManagerWorkersConversation"]] def _create_subconversation_for_agent( - self, agent: Union[Agent, ManagerWorkers] + self, + agent: Union[Agent, ManagerWorkers], + parent_conversation: "ManagerWorkersConversation", ) -> Union["AgentConversation", "ManagerWorkersConversation"]: - subconv = agent.start_conversation() - self.subconversations[agent.name] = subconv - + subconv = agent._start_conversation_impl(parent_conversation=parent_conversation) + self.subconversations[agent.id] = subconv return subconv @@ -70,9 +71,9 @@ def __str__(self) -> str: return repr(self) def _get_agent_subconversation( - self, agent_name: str + self, agent: Union[Agent, ManagerWorkers] ) -> Optional[Union["AgentConversation", "ManagerWorkersConversation"]]: - return self.state.subconversations.get(agent_name) + return self.state.subconversations.get(agent.id) def _get_main_subconversation( self, @@ -80,7 +81,7 @@ def _get_main_subconversation( "Return subconversation between the manger agent and the user" from wayflowcore.agentconversation import AgentConversation - main_subconv = self._get_agent_subconversation(self.component.manager_agent.name) + main_subconv = self._get_agent_subconversation(self.component.manager_agent) if main_subconv is None: raise (ValueError(f"Internal error: Main subconversation is None")) @@ -90,7 +91,8 @@ def _get_main_subconversation( return main_subconv def append_tool_result(self, tool_result: ToolResult) -> None: - current_conv = self._get_agent_subconversation(self.state.current_agent_name) + current_agent = self.component._agent_by_id[self.state.current_agent_id] + current_conv = self._get_agent_subconversation(current_agent) if current_conv is None: raise (ValueError(f"Internal error: Current subconversation is None")) current_conv.append_tool_result(tool_result) diff --git a/wayflowcore/src/wayflowcore/executors/_managerworkersexecutor.py b/wayflowcore/src/wayflowcore/executors/_managerworkersexecutor.py index 50178f6cf..76f227c35 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersexecutor.py @@ -41,13 +41,16 @@ """ -def _create_manager_agent(group_manager: Union[Agent, LlmModel]) -> Agent: +def _create_manager_agent( + group_manager: Union[Agent, LlmModel], manager_agent_id: Optional[str] = None +) -> Agent: if isinstance(group_manager, LlmModel): manager_agent = Agent( name="manager_agent", description="Agent that can assign tasks to other agents.", llm=group_manager, custom_instruction=GROUP_MANAGER_CUSTOM_INSTRUCTION, + id=manager_agent_id, ) else: manager_agent = group_manager @@ -171,8 +174,9 @@ async def execute_async( managerworkers_config = conversation.component while True: - current_agent_name = conversation.state.current_agent_name - current_conversation = conversation._get_agent_subconversation(current_agent_name) + current_agent_id = conversation.state.current_agent_id + current_agent = managerworkers_config._agent_by_id[current_agent_id] + current_conversation = conversation._get_agent_subconversation(current_agent) if current_conversation is None: raise ValueError("Current conversation is None") @@ -180,11 +184,13 @@ async def execute_async( logger.info( "\n%s\nNew execution round. Current agent is %s\n%s\n", "-" * 30, - current_agent_name, + current_agent.name, "-" * 30, ) - if current_agent_name == managerworkers_config.manager_agent.name: - current_agent = managerworkers_config.manager_agent + if current_agent_id == managerworkers_config.manager_agent.id: + + if not isinstance(current_agent, Agent): + raise ValueError("Manager agent must be an Agent") if isinstance(current_conversation, ManagerWorkersConversation): raise ValueError( @@ -240,14 +246,14 @@ async def execute_async( if ( isinstance(status, ToolRequestStatus) - and current_agent_name == managerworkers_config.manager_agent.name + and current_agent_id == managerworkers_config.manager_agent.id ): # 1. current agent is the manager agent and is calling tools # These tool(s) will be handled in the next loop by checking the pending tools of the manager continue elif ( isinstance(status, UserMessageRequestStatus) - and current_agent_name == managerworkers_config.manager_agent.name + and current_agent_id == managerworkers_config.manager_agent.id ): # 2. current agent is the manager agent and is sending a message to the user logger.info( @@ -259,7 +265,7 @@ async def execute_async( # 3. current agent is a worker and is sending a message to the manager ManagerWorkersRunner._send_message_to_manager( message=_last_message, - manager_agent_name=managerworkers_config.manager_agent.name, + manager_agent=managerworkers_config.manager_agent, managerworkers_conversation=conversation, ) elif isinstance( @@ -327,7 +333,7 @@ def _handle_pending_tool_requests_of_manager( @staticmethod def _send_message_to_manager( message: "Message", - manager_agent_name: str, + manager_agent: Agent, managerworkers_conversation: "ManagerWorkersConversation", ) -> None: manager_subconversation = managerworkers_conversation._get_main_subconversation() @@ -352,7 +358,7 @@ def _send_message_to_manager( ) # Change current agent back to manager - managerworkers_conversation.state.current_agent_name = manager_agent_name + managerworkers_conversation.state.current_agent_id = manager_agent.id @staticmethod def _send_message_to_worker( @@ -379,14 +385,16 @@ def _send_message_to_worker( ) logger.debug("Failure when trying to call new agent: `%s`", error_message) else: + recipient_agent = managerworkers_config._agent_by_name[recipient_agent_name] worker_subconversation = managerworkers_conversation._get_agent_subconversation( - recipient_agent_name + recipient_agent ) if worker_subconversation is None: worker_subconversation = ( managerworkers_conversation.state._create_subconversation_for_agent( - managerworkers_config._agent_by_name[recipient_agent_name] + recipient_agent, + parent_conversation=managerworkers_conversation, ) ) logger.info( @@ -397,4 +405,4 @@ def _send_message_to_worker( logger.info("Calling agent %s with request `%s`", recipient_agent_name, message) # Change current agent for the next iteration - managerworkers_conversation.state.current_agent_name = recipient_agent_name + managerworkers_conversation.state.current_agent_id = recipient_agent.id diff --git a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py index 1a169f271..fc44ac75d 100644 --- a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py @@ -63,18 +63,14 @@ class SwarmConversationExecutionState(ConversationExecutionState): def __post_init__(self) -> None: self.current_thread = self.current_thread or self.main_thread - if not self.thread_subconversations: - self._create_subconversation_for_thread( - self.main_thread, inputs=self.inputs, message_list=self.messages - ) - def _create_subconversation_for_thread( self, thread: "SwarmThread", + parent_conversation: "SwarmConversation", inputs: Optional[Dict[str, Any]] = None, message_list: Optional[Union[MessageList, List[Message]]] = None, ) -> "AgentConversation": - thread_id = thread.identifier + thread_id = thread.id if thread_id in self.thread_subconversations: raise KeyError( f"Trying to create a new subconversation for thread {thread_id} but a conversation already exists" @@ -86,7 +82,8 @@ def _create_subconversation_for_thread( if isinstance(message_list, list) else message_list ) - conversation = thread.recipient_agent.start_conversation( + conversation = thread.recipient_agent._start_conversation_impl( + parent_conversation=parent_conversation, inputs=inputs, messages=thread.message_list, ) @@ -118,7 +115,7 @@ def _get_all_context_providers_from_parent_conversations( return [] def _get_all_sub_conversations(self) -> List["Conversation"]: - return [] + return list(self.state.thread_subconversations.values()) def __repr__(self) -> str: return f"{self.__class__.__name__}(state={self.state!r}, thread_subconversations={self.thread_subconversations!r})" @@ -163,9 +160,12 @@ def append_user_message(self, user_input: str | list[MessageContent]) -> None: def _get_subconversation_for_thread( self, thread: "SwarmThread" ) -> Optional["AgentConversation"]: - return self.thread_subconversations.get(thread.identifier) + return self.thread_subconversations.get(thread.id) def _get_recipient_names_for_agent(self, agent: Agent) -> List[str]: - if agent.name not in self.state.agents_and_threads: + if agent.id not in self.state.agents_and_threads: raise ValueError(f"Agent {agent} is not a sender of any thread") - return list(self.state.agents_and_threads[agent.name].keys()) + return [ + thread.recipient_agent.name + for thread in self.state.agents_and_threads[agent.id].values() + ] diff --git a/wayflowcore/src/wayflowcore/executors/_swarmexecutor.py b/wayflowcore/src/wayflowcore/executors/_swarmexecutor.py index 891d19c3c..8de3496e4 100644 --- a/wayflowcore/src/wayflowcore/executors/_swarmexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_swarmexecutor.py @@ -136,7 +136,8 @@ async def execute_async( agent_sub_conversation = conversation._get_subconversation_for_thread(current_thread) if agent_sub_conversation is None: agent_sub_conversation = conversation.state._create_subconversation_for_thread( - current_thread + current_thread, + parent_conversation=conversation, ) # Handle pending tool requests (if any) @@ -421,8 +422,9 @@ def _post_agent_message_to_next_thread( ) else: swarm_conversation.state.thread_stack.append(current_thread) - current_thread = swarm_conversation.state.agents_and_threads[current_agent.name][ - recipient_agent_name + recipient_agent = swarm_conversation.component._agent_by_name[recipient_agent_name] + current_thread = swarm_conversation.state.agents_and_threads[current_agent.id][ + recipient_agent.id ] current_thread.message_list.append_message( diff --git a/wayflowcore/src/wayflowcore/flow.py b/wayflowcore/src/wayflowcore/flow.py index 2c77ada85..9989ad4d9 100644 --- a/wayflowcore/src/wayflowcore/flow.py +++ b/wayflowcore/src/wayflowcore/flow.py @@ -45,6 +45,8 @@ from wayflowcore.tools import Tool if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer + from wayflowcore.conversation import Conversation from wayflowcore.executors._flowconversation import FlowConversation from wayflowcore.executors._flowexecutor import _IoKeyType from wayflowcore.messagelist import Message @@ -466,6 +468,15 @@ class Flow(ConversationalComponent, SerializableObject): _DEFAULT_STARTSTEP_NAME: ClassVar[str] = "__StartStep__" + @property + def _supports_checkpointing(self) -> bool: + from wayflowcore.serialization.context import _get_nested_components + + return all( + nested_component._supports_checkpointing + for nested_component in _get_nested_components(self, only_conversational=True) + ) + def __init__( self, steps: Optional[Union[Dict[str, "Step"], List["Step"]]] = None, @@ -1164,6 +1175,8 @@ def start_conversation( inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, nesting_level: int = 0, context_providers_from_parent_flow: Optional[Set[str]] = None, ) -> "FlowConversation": @@ -1176,23 +1189,78 @@ def start_conversation( Dictionary of inputs. Keys are the variable identifiers and values are the actual inputs to start the conversation. conversation_id: - Conversation id of the parent conversation. + Durable conversation id used for resume, storage, and usage accounting. messages: - List of messages (``MessageList`` object) before starting the conversation. - context_providers_from_parent_flow: - Context provider that don't need to be checked when validating existing inputs. + List of messages before starting the conversation. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. nesting_level: - Nesting level of the conversation. - - Returns - ------- - Conversation: - A Flow Conversation object. + Internal nesting level for flow execution. + context_providers_from_parent_flow: + Context providers inherited from a parent flow. """ + return self._start_conversation_impl( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + nesting_level=nesting_level, + context_providers_from_parent_flow=context_providers_from_parent_flow, + parent_conversation=None, + ) + + def _start_nested_flow_conversation( + self, + parent_conversation: "Conversation", + inputs: Optional[Dict[str, Any]] = None, + messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, + nesting_level: int = 0, + context_providers_from_parent_flow: Optional[Set[str]] = None, + ) -> "FlowConversation": + """Start a nested flow while preserving flow execution state from its parent.""" + return self._start_conversation_impl( + inputs=inputs, + messages=messages, + conversation_id=None, + checkpointer=None, + checkpoint_id=None, + nesting_level=nesting_level, + context_providers_from_parent_flow=context_providers_from_parent_flow, + parent_conversation=parent_conversation, + ) + + def _start_conversation_impl( + self, + inputs: Optional[Dict[str, Any]] = None, + messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + parent_conversation: Optional["Conversation"] = None, + nesting_level: int = 0, + context_providers_from_parent_flow: Optional[Set[str]] = None, + ) -> "FlowConversation": from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event from wayflowcore.executors._flowconversation import FlowConversation + restored_conversation, conversation_instance_id, conversation_thread_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + expected_conversation_type=FlowConversation, + parent_conversation=parent_conversation, + ) + ) + if restored_conversation is not None: + return restored_conversation + context_providers_from_parent_flow = context_providers_from_parent_flow or set() if inputs is None: inputs = {} @@ -1242,7 +1310,7 @@ def start_conversation( conversational_component=self, inputs=inputs, messages=messages, - conversation_id=conversation_id, + conversation_id=conversation_instance_id, nesting_level=nesting_level, ) ) @@ -1274,12 +1342,14 @@ def start_conversation( return FlowConversation( component=self, inputs=inputs, - conversation_id=IdGenerator.get_or_generate_id(conversation_id), + id=conversation_instance_id, + checkpointer=checkpointer, message_list=messages, __metadata_info__={}, status=None, name="flow_conversation", state=state, + conversation_id=conversation_thread_id, ) @property @@ -1633,6 +1703,7 @@ def from_steps( loop: bool = False, step_names: Optional[List[str]] = None, name: Optional[str] = None, + flow_id: Optional[str] = None, description: str = "", input_descriptors: Optional[List[Property]] = None, output_descriptors: Optional[List[Property]] = None, @@ -1656,6 +1727,8 @@ def from_steps( List of step names. Will default to "step_{idx}" if not passed. name: Name of the flow + flow_id: + Id of the flow description: Description of the flow input_descriptors: @@ -1718,6 +1791,7 @@ def from_steps( context_providers=context_providers, variables=variables, name=name, + flow_id=flow_id, description=description, output_descriptors=output_descriptors, input_descriptors=input_descriptors, diff --git a/wayflowcore/src/wayflowcore/flowbuilder.py b/wayflowcore/src/wayflowcore/flowbuilder.py index d7853b292..c89aeb461 100644 --- a/wayflowcore/src/wayflowcore/flowbuilder.py +++ b/wayflowcore/src/wayflowcore/flowbuilder.py @@ -312,7 +312,9 @@ def set_finish_points( self.add_edge(source_key, end_step_name) return self - def build(self, name: str = DEFAULT_FLOW_NAME, description: str = "") -> Flow: + def build( + self, name: str = DEFAULT_FLOW_NAME, flow_id: str | None = None, description: str = "" + ) -> Flow: """ Build the Flow. @@ -355,19 +357,23 @@ def build(self, name: str = DEFAULT_FLOW_NAME, description: str = "") -> Flow: control_flow_edges=self.control_flow_connections, data_flow_edges=self.data_flow_connections, name=name, + flow_id=flow_id, description=description, output_descriptors=self._output_descriptors, ) def build_agent_spec( - self, name: str = DEFAULT_FLOW_NAME, serialize_as: Literal["JSON", "YAML"] = "JSON" + self, + name: str = DEFAULT_FLOW_NAME, + flow_id: str | None = None, + serialize_as: Literal["JSON", "YAML"] = "JSON", ) -> str: """ Build the Flow and return its Agent Spec JSON or YAML configuration. Will raise errors if encountering any while building the Flow. """ - flow = self.build(name) + flow = self.build(name, flow_id=flow_id) if serialize_as == "JSON": return AgentSpecExporter().to_json(flow) elif serialize_as == "YAML": @@ -384,6 +390,7 @@ def build_linear_flow( cls, steps: list[Step], name: str = DEFAULT_FLOW_NAME, + flow_id: str | None = None, serialize_as: Literal[None] = None, data_flow_edges: list[DataFlowEdge] | None = None, input_descriptors: list[Property] | None = None, @@ -396,6 +403,7 @@ def build_linear_flow( cls, steps: list[Step], name: str = DEFAULT_FLOW_NAME, + flow_id: str | None = None, serialize_as: Literal["JSON", "YAML"] = "JSON", data_flow_edges: list[DataFlowEdge] | None = None, input_descriptors: list[Property] | None = None, @@ -407,6 +415,7 @@ def build_linear_flow( cls, steps: list[Step], name: str = DEFAULT_FLOW_NAME, + flow_id: str | None = None, serialize_as: Literal["JSON", "YAML"] | None = None, data_flow_edges: list[DataFlowEdge] | None = None, input_descriptors: list[Property] | None = None, @@ -435,6 +444,7 @@ def build_linear_flow( """ flow = Flow.from_steps( name=name, + flow_id=flow_id, steps=steps, data_flow_edges=data_flow_edges, input_descriptors=input_descriptors, diff --git a/wayflowcore/src/wayflowcore/flowhelpers.py b/wayflowcore/src/wayflowcore/flowhelpers.py index 5d956eea6..1bbd5444e 100644 --- a/wayflowcore/src/wayflowcore/flowhelpers.py +++ b/wayflowcore/src/wayflowcore/flowhelpers.py @@ -130,6 +130,7 @@ def create_single_step_flow( data_flow_edges: Optional[List[DataFlowEdge]] = None, variables: Optional[List["Variable"]] = None, flow_name: Optional[str] = None, + flow_id: Optional[str] = None, flow_description: str = "", ) -> Flow: """Create a flow that consist of one step only @@ -148,6 +149,8 @@ def create_single_step_flow( list of variables of the flow flow_name: optional name of the flow + flow_id: + optional id of the flow flow_description: optional description of the flow """ @@ -173,6 +176,7 @@ def create_single_step_flow( data_flow_edges=data_flow_edges, variables=variables, name=flow_name, + flow_id=flow_id, description=flow_description, ) diff --git a/wayflowcore/src/wayflowcore/managerworkers.py b/wayflowcore/src/wayflowcore/managerworkers.py index f59ace30b..051817583 100644 --- a/wayflowcore/src/wayflowcore/managerworkers.py +++ b/wayflowcore/src/wayflowcore/managerworkers.py @@ -6,7 +6,7 @@ import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union, cast from wayflowcore._metadata import MetadataType from wayflowcore.agent import Agent, CallerInputMode @@ -22,6 +22,8 @@ from wayflowcore.transforms import MessageTransform if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer + from wayflowcore.conversation import Conversation from wayflowcore.executors._managerworkersconversation import ManagerWorkersConversation from wayflowcore.messagelist import Message @@ -42,6 +44,15 @@ class ManagerWorkers(ConversationalComponent, SerializableDataclassMixin, Serial description: Optional[str] id: str + @property + def _supports_checkpointing(self) -> bool: + from wayflowcore.serialization.context import _get_nested_components + + return all( + nested_component._supports_checkpointing + for nested_component in _get_nested_components(self, only_conversational=True) + ) + def __init__( self, group_manager: Union[LlmModel, Agent], @@ -128,15 +139,23 @@ def __init__( if len(workers) == 0: raise ValueError("Cannot define a group with no worker agent.") + managerworkers_id = IdGenerator.get_or_generate_id(id) + self.group_manager = group_manager - self.manager_agent = _create_manager_agent(self.group_manager) + self.manager_agent = _create_manager_agent( + self.group_manager, manager_agent_id=f"{managerworkers_id}:manager" + ) self.workers = workers self.transforms = transforms or [] + agents = self.workers + [self.manager_agent] self._agent_by_name: Dict[str, Union["Agent", "ManagerWorkers"]] = _validate_agent_unicity( - self.workers + [self.manager_agent] + agents ) + self._agent_by_id: Dict[str, Union["Agent", "ManagerWorkers"]] = { + agent.id: agent for agent in agents + } # Create send message tools for the group manager self._manager_communication_tools = _create_communication_tools() @@ -153,7 +172,7 @@ def __init__( super().__init__( name=IdGenerator.get_or_generate_name(name, prefix="managerworkers_", length=8), description=description, - id=id, + id=managerworkers_id, input_descriptors=input_descriptors or [], output_descriptors=output_descriptors or [], runner=ManagerWorkersRunner, @@ -221,6 +240,8 @@ def start_conversation( inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, conversation_name: Optional[str] = None, ) -> "ManagerWorkersConversation": """ @@ -234,13 +255,38 @@ def start_conversation( messages: Message list of the manager agent and the end-user. conversation_id: - Conversation id of the main conversation. + Durable conversation id used for resume, storage, and usage accounting. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires both ``checkpointer`` and + ``conversation_id``. Returns ------- Conversation: The conversation object of the managerworkers. """ + return self._start_conversation_impl( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + conversation_name=conversation_name, + parent_conversation=None, + ) + + def _start_conversation_impl( + self, + inputs: Optional[Dict[str, Any]] = None, + messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + parent_conversation: Optional["Conversation"] = None, + conversation_name: Optional[str] = None, + ) -> "ManagerWorkersConversation": from wayflowcore.agentconversation import AgentConversation from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event @@ -249,43 +295,61 @@ def start_conversation( ManagerWorkersConversationExecutionState, ) + restored_conversation, conversation_instance_id, conversation_thread_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + expected_conversation_type=ManagerWorkersConversation, + parent_conversation=parent_conversation, + ) + ) + if restored_conversation is not None: + return restored_conversation + if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) - if conversation_id is None: - conversation_id = IdGenerator.get_or_generate_id(conversation_id) - record_event( ConversationCreatedEvent( conversational_component=self, inputs=inputs or {}, messages=messages, - conversation_id=conversation_id, + conversation_id=conversation_instance_id, nesting_level=None, ) ) subconversations: Dict[str, Union[AgentConversation, ManagerWorkersConversation]] = {} - subconversations[self.manager_agent.name] = self.manager_agent.start_conversation( - inputs=inputs, - messages=messages, - ) state = ManagerWorkersConversationExecutionState( - current_agent_name=self.manager_agent.name, + current_agent_id=self.manager_agent.id, subconversations=subconversations, ) - return ManagerWorkersConversation( + conversation = ManagerWorkersConversation( component=self, inputs={}, message_list=messages, + id=conversation_instance_id, name=conversation_name or "managerworkers_conversation", state=state, status=None, - conversation_id=conversation_id, + checkpointer=checkpointer, + conversation_id=conversation_thread_id, __metadata_info__={}, ) + subconversations[self.manager_agent.id] = cast( + "Union[AgentConversation, ManagerWorkersConversation]", + self.manager_agent._start_conversation_impl( + parent_conversation=conversation, + inputs=inputs, + messages=messages, + ), + ) + return conversation def _referenced_tools_dict_inner( self, recursive: bool, visited_set: Set[str] @@ -311,7 +375,11 @@ def _update_internal_state(self) -> None: _validate_agent_unicity, ) - self.manager_agent = _create_manager_agent(self.group_manager) - self._agent_by_name = _validate_agent_unicity(self.workers + [self.manager_agent]) + self.manager_agent = _create_manager_agent( + self.group_manager, manager_agent_id=f"{self.id}:manager" + ) + agents = self.workers + [self.manager_agent] + self._agent_by_name = _validate_agent_unicity(agents) + self._agent_by_id = {agent.id: agent for agent in agents} self._manager_communication_tools = _create_communication_tools() self._runtime_managerworkers_template = self._compose_runtime_managerworkers_template() diff --git a/wayflowcore/src/wayflowcore/mcp/_session_persistence.py b/wayflowcore/src/wayflowcore/mcp/_session_persistence.py index 143b069b0..cfb8c4b01 100644 --- a/wayflowcore/src/wayflowcore/mcp/_session_persistence.py +++ b/wayflowcore/src/wayflowcore/mcp/_session_persistence.py @@ -51,9 +51,9 @@ def get_current_conv_id_or_default() -> str: - from wayflowcore.conversation import _get_current_conversation_id + from wayflowcore.conversation import _get_current_conversation_thread_id - return _get_current_conversation_id() or _DEFAULT_MCP_SESSION_CONTEXT_ID + return _get_current_conversation_thread_id() or _DEFAULT_MCP_SESSION_CONTEXT_ID async def _call_with_parent_span( @@ -486,10 +486,10 @@ def shutdown_mcp_async_runtime() -> None: def _get_oauth_flow_handler(client_transport: "ClientTransport") -> "OAuthFlowHandler": - from wayflowcore.conversation import _get_current_conversation_id + from wayflowcore.conversation import _get_current_conversation_thread_id runtime = get_mcp_async_runtime() - conversation_id = _get_current_conversation_id() or _DEFAULT_MCP_SESSION_CONTEXT_ID + conversation_id = _get_current_conversation_thread_id() or _DEFAULT_MCP_SESSION_CONTEXT_ID transport_id = client_transport.id handler = runtime._oauth_handlers.get(transport_id, {}).get(conversation_id) diff --git a/wayflowcore/src/wayflowcore/ociagent.py b/wayflowcore/src/wayflowcore/ociagent.py index 4bfa36e92..2423059e7 100644 --- a/wayflowcore/src/wayflowcore/ociagent.py +++ b/wayflowcore/src/wayflowcore/ociagent.py @@ -19,6 +19,7 @@ from wayflowcore.tools import Tool if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.conversation import Conversation @@ -102,27 +103,58 @@ def __init__( __metadata_info__=__metadata_info__, ) + @property + def _supports_checkpointing(self) -> bool: + return False + def start_conversation( self, inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, Message, List[Message], MessageList] = None, - ) -> "Conversation": + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + ) -> Conversation: """ Initializes a conversation with the agent. Parameters ---------- inputs: - This argument is not used. - It is included for compatibility with the Flow class. + This argument is not used. It is included for compatibility with the Flow class. messages: - Message list to which the agent will participate - + Message list to which the agent will participate. + conversation_id: + Durable conversation id used for resume, storage, and usage accounting. + checkpointer: + Optional checkpoint backend. ``OciAgent`` does not support checkpoint restore yet, + so passing this raises ``NotImplementedError``. + checkpoint_id: + Optional checkpoint identifier. ``OciAgent`` does not support checkpoint restore yet, + so passing this raises ``NotImplementedError``. Returns ------- Conversation: The conversation object of the agent. """ + return self._start_conversation_impl( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=None, + ) + + def _start_conversation_impl( + self, + inputs: Optional[Dict[str, Any]] = None, + messages: Union[None, str, Message, List[Message], MessageList] = None, + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + parent_conversation: Optional["Conversation"] = None, + ) -> "Conversation": from wayflowcore.executors._ociagentconversation import OciAgentConversation from wayflowcore.executors._ociagentexecutor import ( OciAgentState, @@ -130,9 +162,24 @@ def start_conversation( _init_oci_agent_session, ) + if any(value is not None for value in (checkpointer, checkpoint_id)): + raise NotImplementedError("`OciAgent` checkpoint restore is not supported yet.") + if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) + _restored_conversation, conversation_instance_id, conversation_thread_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=None, + checkpoint_id=None, + expected_conversation_type=OciAgentConversation, + parent_conversation=parent_conversation, + ) + ) + _client = _init_oci_agent_client(self) return OciAgentConversation( @@ -145,8 +192,9 @@ def start_conversation( inputs=inputs or {}, message_list=messages, status=None, - conversation_id=IdGenerator.get_or_generate_id(None), + id=conversation_instance_id, name="oci_conversation", + conversation_id=conversation_thread_id, __metadata_info__={}, ) diff --git a/wayflowcore/src/wayflowcore/serialization/context.py b/wayflowcore/src/wayflowcore/serialization/context.py index a86d6cf17..728c263fb 100644 --- a/wayflowcore/src/wayflowcore/serialization/context.py +++ b/wayflowcore/src/wayflowcore/serialization/context.py @@ -7,10 +7,11 @@ import warnings from copy import deepcopy from functools import cached_property -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast, overload if TYPE_CHECKING: from wayflowcore.component import Component + from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.serialization.plugins import ( WayflowDeserializationPlugin, WayflowSerializationPlugin, @@ -47,6 +48,7 @@ def __init__(self, root: Any, plugins: Optional[List["WayflowSerializationPlugin """ self.root = root self._serialized_objects: Dict[str, Any] = {} + self._external_references: set[str] = set() self._started_serialization: Dict[str, bool] = {} self.plugins = plugins or [] @@ -113,6 +115,14 @@ def record_obj_dict(self, obj: Any, obj_as_dict: Dict[Any, Any]) -> None: """ self._serialized_objects[self.get_reference(obj)] = obj_as_dict + def _register_external_component_references(self, component: "Component") -> None: + """ + Marks the current component and all its nested components as provided externally to the + serialized object graph. + """ + for nested_component in _get_nested_components(component, include_root=True): + self._external_references.add(self.get_reference(nested_component)) + def check_obj_is_already_serialized(self, obj: Any) -> bool: """ Returns True if the object has already been serialized @@ -122,7 +132,11 @@ def check_obj_is_already_serialized(self, obj: Any) -> bool: obj: The original, non-serialized object """ - return self._serialized_objects.get(self.get_reference(obj)) is not None + obj_ref = self.get_reference(obj) + return ( + obj_ref in self._external_references + or self._serialized_objects.get(obj_ref) is not None + ) def get_reference_dict(self, obj: Any) -> Dict[str, str]: """ @@ -229,7 +243,7 @@ def get_referenced_dict(self, object_reference: str) -> Dict[Any, Any]: The reference of the object being deserialized """ if object_reference not in self._referenced_objects: - raise ValueError( + raise _MissingDeserializationReferenceError( f"During deserialization, encountered reference {object_reference} that is missing " f"in the _referenced_objects of the serialized root object." ) @@ -305,31 +319,86 @@ def _builtins_deserialization_plugin(cls) -> "WayflowDeserializationPlugin": return WayflowBuiltinsDeserializationPlugin() - def _add_component_to_context(self, component: "Component") -> None: + def _register_external_component_references(self, component: "Component") -> None: """ Adds the current components and all its subcomponents to this deserialization context. """ - from wayflowcore.component import Component + for nested_component in _get_nested_components(component, include_root=True): + component_ref = SerializationContext.get_reference(nested_component) + self._deserialized_objects.setdefault(component_ref, nested_component) + + +class _MissingDeserializationReferenceError(ValueError): + """Raised when deserialization encounters a reference missing from the root object.""" + + +@overload +def _get_nested_components( + value: Any, + include_root: bool = False, + only_conversational: Literal[True] = True, +) -> List["ConversationalComponent"]: ... + + +@overload +def _get_nested_components( + value: Any, + include_root: bool = False, + only_conversational: Literal[False] = False, +) -> List["Component"]: ... + + +def _get_nested_components( + value: Any, + include_root: bool = False, + only_conversational: bool = False, +) -> List[Any]: + """Return ordered components reachable from ``value``. + + Parameters + ---------- + value: + Object whose component graph should be traversed. + include_root: + Whether to include ``value`` when it is a component. + only_conversational: + Whether to return only conversational components. + """ + from wayflowcore.component import Component - component_ref = SerializationContext.get_reference(component) + component_type: Optional[Type["Component"]] = None + if only_conversational: + from wayflowcore.conversationalcomponent import ConversationalComponent + + component_type = ConversationalComponent + + ordered_components: List["Component"] = [] + visited_component_refs: set[str] = set() + + def _collect_nested_components(current_value: Any) -> None: + if isinstance(current_value, Component): + component_ref = SerializationContext.get_reference(current_value) + if component_ref in visited_component_refs: + return + visited_component_refs.add(component_ref) + if (current_value is not value or include_root) and ( + component_type is None or isinstance(current_value, component_type) + ): + ordered_components.append(current_value) + for name, attr in vars(current_value).items(): + if not name.startswith("_"): + _collect_nested_components(attr) + return - if component_ref in self._deserialized_objects: + if isinstance(current_value, dict): + for nested_value in current_value.values(): + _collect_nested_components(nested_value) return - self._deserialized_objects[component_ref] = component - - all_public_attrs = { - name: value for name, value in vars(component).items() if not name.startswith("_") - } - for attr_name, attr in all_public_attrs.items(): - if isinstance(attr, Component): - self._add_component_to_context(attr) - if isinstance(attr, dict): - for value in attr.values(): - if isinstance(value, Component): - self._add_component_to_context(value) - if isinstance(attr, list): - for value in attr: - if isinstance(value, Component): - self._add_component_to_context(value) + if isinstance(current_value, (list, tuple, set)): + for nested_value in current_value: + _collect_nested_components(nested_value) + + _collect_nested_components(value) + return ordered_components diff --git a/wayflowcore/src/wayflowcore/serialization/serializer.py b/wayflowcore/src/wayflowcore/serialization/serializer.py index fcb770db3..83f76b43a 100644 --- a/wayflowcore/src/wayflowcore/serialization/serializer.py +++ b/wayflowcore/src/wayflowcore/serialization/serializer.py @@ -218,6 +218,8 @@ def _serialize_to_dict(self, serialization_context: "SerializationContext") -> D or k.name == "__metadata_info__" # except the metadata ) and k.init # not part of the dataclass __init__ -> would fail at deserialization + # Runtime-only dataclass fields can opt out of durable serialization. + and k.metadata.get("serialize", True) ) } diff --git a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py index 893877282..c0e69a776 100644 --- a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py +++ b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py @@ -293,8 +293,10 @@ def _get_or_create_agent_subconversation( init_messages = ( caller_conv.message_list if self._share_conversation else MessageList.from_messages([]) ) - agent_sub_conversation = self.agent.start_conversation( - inputs=inputs, messages=init_messages + agent_sub_conversation = self.agent._start_conversation_impl( + parent_conversation=caller_conv, + inputs=inputs, + messages=init_messages, ) return agent_sub_conversation diff --git a/wayflowcore/src/wayflowcore/steps/retrystep.py b/wayflowcore/src/wayflowcore/steps/retrystep.py index 22f8598a9..7f79ac7de 100644 --- a/wayflowcore/src/wayflowcore/steps/retrystep.py +++ b/wayflowcore/src/wayflowcore/steps/retrystep.py @@ -8,10 +8,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, cast from wayflowcore._metadata import MetadataType -from wayflowcore.executors._flowexecutor import ( - FlowConversationExecutionState, - FlowConversationExecutor, -) +from wayflowcore.executors._flowexecutor import FlowConversationExecutor from wayflowcore.executors.executionstatus import FinishedStatus from wayflowcore.property import BooleanProperty, IntegerProperty, Property from wayflowcore.steps import FlowExecutionStep @@ -20,6 +17,7 @@ if TYPE_CHECKING: from wayflowcore.executors._flowconversation import FlowConversation + from wayflowcore.executors._flowexecutor import FlowConversationExecutionState from wayflowcore.flow import Flow logger = logging.getLogger(__name__) @@ -39,6 +37,10 @@ class RetryStep(Step): BRANCH_FAILURE = "failure" """Name of the branch taken in case the condition is still not met after the maximum number of trials""" + # Stored through FlowConversationExecutor.make_key_for_step so generated step names can + # be remapped when a checkpoint is restored into a freshly-instantiated flow. + _RETRY_COUNTER_KEY = "retry_counter" + def __init__( self, flow: "Flow", @@ -255,11 +257,13 @@ def might_yield(self) -> bool: """ return self.flow.might_yield - def _retry_count(self, state: FlowConversationExecutionState) -> int: - return cast(int, state.internal_context_key_values.get(f"retry_counter_{id(self)}", 0)) + def _retry_count(self, state: "FlowConversationExecutionState") -> int: + key = FlowConversationExecutor.make_key_for_step(self, self._RETRY_COUNTER_KEY) + return cast(int, state.internal_context_key_values.get(key, 0)) - def _set_counter(self, state: FlowConversationExecutionState, value: int) -> None: - state.internal_context_key_values[f"retry_counter_{id(self)}"] = value + def _set_counter(self, state: "FlowConversationExecutionState", value: int) -> None: + key = FlowConversationExecutor.make_key_for_step(self, self._RETRY_COUNTER_KEY) + state.internal_context_key_values[key] = value async def _invoke_step_async( self, diff --git a/wayflowcore/src/wayflowcore/swarm.py b/wayflowcore/src/wayflowcore/swarm.py index b4852fc18..4e150ae69 100644 --- a/wayflowcore/src/wayflowcore/swarm.py +++ b/wayflowcore/src/wayflowcore/swarm.py @@ -23,6 +23,7 @@ from wayflowcore.transforms import MessageTransform if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.conversation import Conversation from wayflowcore.messagelist import Message @@ -92,6 +93,15 @@ class Swarm(ConversationalComponent, SerializableDataclassMixin, SerializableObj description: Optional[str] id: str + @property + def _supports_checkpointing(self) -> bool: + from wayflowcore.serialization.context import _get_nested_components + + return all( + nested_component._supports_checkpointing + for nested_component in _get_nested_components(self, only_conversational=True) + ) + def __init__( self, first_agent: Agent, @@ -311,6 +321,50 @@ def start_conversation( inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], MessageList] = None, conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + conversation_name: Optional[str] = None, + ) -> "Conversation": + """ + Initializes a conversation with the swarm. + + Parameters + ---------- + inputs: + Dictionary of inputs used to initialize the conversation. + messages: + Message list of the swarm and the end-user. + conversation_id: + Durable conversation id used for resume, storage, and usage accounting. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires both ``checkpointer`` and + ``conversation_id``. + + Returns + ------- + Conversation: + The conversation object of the swarm. + """ + return self._start_conversation_impl( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + conversation_name=conversation_name, + parent_conversation=None, + ) + + def _start_conversation_impl( + self, + inputs: Optional[Dict[str, Any]] = None, + messages: Union[None, str, "Message", List["Message"], MessageList] = None, + conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + parent_conversation: Optional["Conversation"] = None, conversation_name: Optional[str] = None, ) -> "Conversation": from wayflowcore.executors._swarmconversation import ( @@ -320,21 +374,32 @@ def start_conversation( SwarmUser, ) + restored_conversation, conversation_instance_id, conversation_thread_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + expected_conversation_type=SwarmConversation, + parent_conversation=parent_conversation, + ) + ) + if restored_conversation is not None: + return restored_conversation + if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) - if conversation_id is None: - conversation_id = IdGenerator.get_or_generate_id(conversation_id) - main_thread = SwarmThread( caller=SwarmUser(), recipient_agent=self.first_agent, is_main_thread=True ) agents_and_threads: Dict[str, Dict[str, SwarmThread]] = {} for caller_agent, recipient_agent in self.relationships: - if caller_agent.name not in agents_and_threads: - agents_and_threads[caller_agent.name] = {} + if caller_agent.id not in agents_and_threads: + agents_and_threads[caller_agent.id] = {} - agents_and_threads[caller_agent.name][recipient_agent.name] = SwarmThread( + agents_and_threads[caller_agent.id][recipient_agent.id] = SwarmThread( caller=caller_agent, recipient_agent=recipient_agent, ) @@ -345,16 +410,25 @@ def start_conversation( inputs=inputs, messages=messages, ) - return SwarmConversation( + conversation = SwarmConversation( component=self, inputs=inputs or {}, message_list=messages, + id=conversation_instance_id, name=conversation_name or "swarm_conversation", state=state, status=None, - conversation_id=conversation_id, + checkpointer=checkpointer, + conversation_id=conversation_thread_id, __metadata_info__={}, ) + state._create_subconversation_for_thread( + main_thread, + parent_conversation=conversation, + inputs=inputs, + message_list=messages, + ) + return conversation def _referenced_tools_dict_inner( self, recursive: bool, visited_set: Set[str] diff --git a/wayflowcore/src/wayflowcore/tools/servertools.py b/wayflowcore/src/wayflowcore/tools/servertools.py index d6216c97f..0b5863602 100644 --- a/wayflowcore/src/wayflowcore/tools/servertools.py +++ b/wayflowcore/src/wayflowcore/tools/servertools.py @@ -636,7 +636,8 @@ async def __call__(self, **inputs: Any) -> Any: conversation = self.flow.start_conversation(inputs) interrupts = [] else: - conversation = self.flow.start_conversation( + conversation = self.flow._start_conversation_impl( + parent_conversation=self._parent_conversation, inputs=inputs, messages=self._parent_conversation.message_list, ) diff --git a/wayflowcore/tests/a2a/test_a2aagent.py b/wayflowcore/tests/a2a/test_a2aagent.py index 3a17dffca..a7046423a 100644 --- a/wayflowcore/tests/a2a/test_a2aagent.py +++ b/wayflowcore/tests/a2a/test_a2aagent.py @@ -10,8 +10,10 @@ import pytest from wayflowcore.a2a.a2aagent import A2AAgent, A2AConnectionConfig, A2ASessionParameters +from wayflowcore.checkpointing import InMemoryCheckpointer from wayflowcore.executors._a2aagentconversation import A2AAgentConversation from wayflowcore.executors._a2aagentexecutor import DEFAULT_RESPONSE +from wayflowcore.executors.executionstatus import UserMessageRequestStatus from wayflowcore.messagelist import Message from ..testhelpers.testhelpers import retry_test @@ -119,6 +121,60 @@ def test_a2aagent_handles_single_message_conversation(a2a_agent): assert "160" in conversation.get_last_message().content +@retry_test(max_attempts=4) +def test_a2aagent_checkpointing_supports_resume_and_time_travel(a2a_agent: A2AAgent) -> None: + """ + Failure rate: 0 out of 20 + Observed on: 2026-03-23 + Average success time: 0.00 seconds per successful attempt + Average failure time: No time measurement + Max attempt: 4 + Justification: (0.05 ** 4) ~= 0.6 / 100'000 + """ + checkpointer = InMemoryCheckpointer() + + conversation = a2a_agent.start_conversation( + conversation_id="a2a-checkpoint", checkpointer=checkpointer + ) + conversation.append_user_message("What is 5+5? Just output the answer.") + first_status = conversation.execute() + + assert isinstance(first_status, UserMessageRequestStatus) + checkpoint = checkpointer.load_latest(conversation.conversation_id) + assert checkpoint is not None + first_checkpoint_id = checkpoint.checkpoint_id + first_message_count = len(conversation.get_messages()) + + restored_conversation = a2a_agent.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + assert len(restored_conversation.get_messages()) == first_message_count + restored_conversation.append_user_message( + "What if you replace 5 by 10? Just output the answer." + ) + restored_status = restored_conversation.execute() + + assert isinstance(restored_status, UserMessageRequestStatus) + assert len(restored_conversation.get_messages()) > first_message_count + assert restored_conversation.get_last_message() is not None + assert checkpointer.load_latest(conversation.conversation_id) is not None + + rewound_conversation = a2a_agent.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + checkpoint_id=first_checkpoint_id, + ) + assert len(rewound_conversation.get_messages()) == first_message_count + assert len(rewound_conversation.get_messages()) < len(restored_conversation.get_messages()) + rewound_conversation.append_user_message("What if you replace 5 by 7? Just output the answer.") + rewound_status = rewound_conversation.execute() + + assert isinstance(rewound_status, UserMessageRequestStatus) + assert len(rewound_conversation.get_messages()) > first_message_count + assert rewound_conversation.get_last_message() is not None + + @retry_test(max_attempts=4) def test_a2aagent_passing_single_message_in_conversation_directly(a2a_agent): """ diff --git a/wayflowcore/tests/agentserver/conftest.py b/wayflowcore/tests/agentserver/conftest.py index 9723369ae..df18c98ff 100644 --- a/wayflowcore/tests/agentserver/conftest.py +++ b/wayflowcore/tests/agentserver/conftest.py @@ -16,18 +16,19 @@ import pytest import yaml -from wayflowcore.datastore import OracleDatabaseConnectionConfig, TlsOracleDatabaseConnectionConfig +from wayflowcore.datastore import OracleDatabaseConnectionConfig from wayflowcore.datastore.oracle import _execute_query_on_oracle_db from wayflowcore.datastore.postgres import ( PostgresDatabaseConnectionConfig, _execute_query_on_postgres_db, ) -from ..datastores.conftest import ( - all_oracle_tls_connection_config_env_variables_are_specified, +from ..conftest import ( all_postgres_connection_config_env_variables_are_specified, + get_oracle_connection_config, get_postgres_connection_config, ) +from ..datastores.conftest import all_oracle_tls_connection_config_env_variables_are_specified from ..utils import LogTee, _terminate_process_tree, get_available_port from .datastore_agent_server import ORACLE_DB_CREATE_DDL, ORACLE_DB_DELETE_DDL @@ -242,15 +243,6 @@ def _fixture(request): ) -def get_oracle_connection_config(): - return TlsOracleDatabaseConnectionConfig( - user=os.environ["ADB_DB_USER"], - password=os.environ["ADB_DB_PASSWORD"], - dsn=os.environ["ADB_DSN"], - config_dir=os.environ.get("ADB_CONFIG_DIR", None), - ) - - wayflow_server_http_postgres = register_wayflow_server_fixture( name="wayflow_server_http_postgres", **HR_AGENT_PARAMS, diff --git a/wayflowcore/tests/agentserver/test_wayflow_server.py b/wayflowcore/tests/agentserver/test_wayflow_server.py index 8a4bf2e3b..ffda05b89 100644 --- a/wayflowcore/tests/agentserver/test_wayflow_server.py +++ b/wayflowcore/tests/agentserver/test_wayflow_server.py @@ -787,7 +787,8 @@ def test_agent_with_datastore_is_supported(datastore_agent_inmemory_server): follow_up = _create_response( base_url=datastore_agent_inmemory_server, input_value="and the biggest city?", - model="datastore-assistant", + # Resume must use the same served component/model as the checkpoint. + model="datastore-swarm", previous_response_id=response["id"], ) output = follow_up["output"][0]["content"][0]["text"] diff --git a/wayflowcore/tests/conftest.py b/wayflowcore/tests/conftest.py index f23067155..fdbb4e96c 100644 --- a/wayflowcore/tests/conftest.py +++ b/wayflowcore/tests/conftest.py @@ -24,6 +24,7 @@ from wayflowcore._threading import shutdown_threadpool from wayflowcore.datastore import MTlsOracleDatabaseConnectionConfig from wayflowcore.datastore.oracle import TlsOracleDatabaseConnectionConfig +from wayflowcore.datastore.postgres import TlsPostgresDatabaseConnectionConfig from wayflowcore.embeddingmodels import VllmEmbeddingModel from wayflowcore.flowhelpers import create_single_step_flow from wayflowcore.mcp import authless_mcp_enabled @@ -961,6 +962,29 @@ def get_oracle_connection_config(): ) +def all_postgres_connection_config_env_variables_are_specified(): + postgres_connection_args = [ + "POSTGRES_DB_USER", + "POSTGRES_DB_PASSWORD", + "POSTGRES_DB_URL", + ] + return all(arg in os.environ for arg in postgres_connection_args) + + +def get_postgres_connection_config(): + if all_postgres_connection_config_env_variables_are_specified(): + return TlsPostgresDatabaseConnectionConfig( + user=os.environ["POSTGRES_DB_USER"], + password=os.environ["POSTGRES_DB_PASSWORD"], + url=os.environ["POSTGRES_DB_URL"], + sslmode="disable", + ) + pytest.skip( + "No database connection arguments configured in environment. " + "Skipping Postgres DB tests..." + ) + + @contextlib.contextmanager def disable_streaming(): """Temporarily disable message streaming of LLMs""" diff --git a/wayflowcore/tests/datastores/conftest.py b/wayflowcore/tests/datastores/conftest.py index 755534788..bb5320f09 100644 --- a/wayflowcore/tests/datastores/conftest.py +++ b/wayflowcore/tests/datastores/conftest.py @@ -24,7 +24,11 @@ from wayflowcore.property import FloatProperty, IntegerProperty, Property, StringProperty from wayflowcore.steps.step import Step -from ..conftest import get_oracle_connection_config +from ..conftest import ( + all_postgres_connection_config_env_variables_are_specified, + get_oracle_connection_config, + get_postgres_connection_config, +) def get_basic_office_entities(): @@ -134,30 +138,6 @@ def all_oracle_mtls_connection_config_env_variables_are_specified(): return all([arg in os.environ for arg in mtls_connection_args]) -def all_postgres_connection_config_env_variables_are_specified(): - tls_connection_args = [ - "POSTGRES_DB_USER", - "POSTGRES_DB_PASSWORD", - "POSTGRES_DB_URL", - ] - return all([arg in os.environ for arg in tls_connection_args]) - - -def get_postgres_connection_config(): - if all_postgres_connection_config_env_variables_are_specified(): - return TlsPostgresDatabaseConnectionConfig( - user=os.environ["POSTGRES_DB_USER"], - password=os.environ["POSTGRES_DB_PASSWORD"], - url=os.environ["POSTGRES_DB_URL"], - sslmode="disable", - ) - else: - pytest.skip( - "No database connection arguments configured in environment. " - "Skipping Postgres DB tests..." - ) - - def get_tls_postgres_connection_config(): if all_postgres_connection_config_env_variables_are_specified(): return TlsPostgresDatabaseConnectionConfig( diff --git a/wayflowcore/tests/integration/test_checkpointing.py b/wayflowcore/tests/integration/test_checkpointing.py new file mode 100644 index 000000000..c92f6cf77 --- /dev/null +++ b/wayflowcore/tests/integration/test_checkpointing.py @@ -0,0 +1,537 @@ +# Copyright © 2026 Oracle and/or its affiliates. +# +# This software is under the Apache License 2.0 +# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License +# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. + +from typing import Any, Dict, Optional +from uuid import uuid4 + +import pytest + +from wayflowcore.agent import Agent +from wayflowcore.checkpointing import ( + CheckpointingInterval, + InMemoryCheckpointer, + OracleDatabaseCheckpointer, + PostgresCheckpointer, + StorageConfig, +) +from wayflowcore.checkpointing.datastorecheckpointer import ( + _prepare_oracle_checkpoint_datastore, + _prepare_postgres_checkpoint_datastore, +) +from wayflowcore.conversation import Conversation +from wayflowcore.datastore.oracle import _execute_query_on_oracle_db +from wayflowcore.datastore.postgres import _execute_query_on_postgres_db +from wayflowcore.executors._events.event import Event, EventType +from wayflowcore.executors._executionstate import ConversationExecutionState +from wayflowcore.executors.executionstatus import FinishedStatus, UserMessageRequestStatus +from wayflowcore.executors.interrupts.executioninterrupt import ( + FlexibleExecutionInterrupt, + FlowExecutionInterrupt, + InterruptedExecutionStatus, + _AllEventsInterruptMixin, +) +from wayflowcore.flow import Flow +from wayflowcore.flowbuilder import FlowBuilder +from wayflowcore.managerworkers import ManagerWorkers +from wayflowcore.models import LlmModel +from wayflowcore.serialization import deserialize, serialize +from wayflowcore.steps import FlowExecutionStep +from wayflowcore.swarm import Swarm +from wayflowcore.tools import ToolRequest + +from ..conftest import get_oracle_connection_config, get_postgres_connection_config +from ..serialization.test_assistant_serialization import create_flow +from ..test_managerworkers import _send_message +from ..test_swarm import _handoff_message +from ..testhelpers.dummy import DoNothingStep, DummyModel +from ..testhelpers.patching import patch_llm + + +@pytest.fixture(params=["in-memory", "postgres", "oracle"]) +def integration_checkpointer(request): + if request.param == "in-memory": + yield InMemoryCheckpointer() + return + + storage_config = StorageConfig(table_name=f"test_cp_{uuid4().hex[:20]}") + if request.param == "postgres": + connection_config = get_postgres_connection_config() + _prepare_postgres_checkpoint_datastore(connection_config, storage_config) + checkpointer = PostgresCheckpointer(connection_config, storage_config) + else: + connection_config = get_oracle_connection_config() + _prepare_oracle_checkpoint_datastore(connection_config, storage_config) + checkpointer = OracleDatabaseCheckpointer(connection_config, storage_config) + + try: + yield checkpointer + finally: + drop_query = f"DROP TABLE {storage_config.table_name}" + if request.param == "postgres": + _execute_query_on_postgres_db(connection_config, drop_query) + else: + _execute_query_on_oracle_db(connection_config, drop_query) + + +def _build_nested_flow( + child_first_step_name: str, + child_second_step_name: str, + child_flow_name: str, + parent_step_name: str, + parent_flow_name: str, +) -> Flow: + first_step = DoNothingStep(name=child_first_step_name) + second_step = DoNothingStep(name=child_second_step_name) + child_flow = FlowBuilder.build_linear_flow( + [first_step, second_step], + name=child_flow_name, + flow_id=child_flow_name, + ) + parent_step = FlowExecutionStep(child_flow, name=parent_step_name) + return FlowBuilder.build_linear_flow( + [parent_step], + flow_id=parent_flow_name, + name=parent_flow_name, + ) + + +class _OnStepStartExecutionInterrupt( + _AllEventsInterruptMixin, FlexibleExecutionInterrupt, FlowExecutionInterrupt +): + def __init__(self, step_name: str) -> None: + self.step_name = step_name + self.triggered = False + self.current_event: Optional[Event] = None + super().__init__() + + def _return_status_if_condition_is_met( + self, state: ConversationExecutionState, conversation: Conversation + ) -> Optional[InterruptedExecutionStatus]: + if ( + self.current_event is not None + and self.current_event.type == EventType.STEP_EXECUTION_START + and self.step_name == conversation.current_step_name + and not self.triggered + ): + self.triggered = True + return InterruptedExecutionStatus( + interrupter=self, + reason=f"Start {self.step_name}", + _conversation_id=conversation.id, + ) + return None + + def on_event( + self, event: Event, state: ConversationExecutionState, conversation: Conversation + ) -> Optional[InterruptedExecutionStatus]: + self.current_event = event + return super().on_event(event, state, conversation) + + def _serialize_to_dict(self, serialization_context) -> Dict[str, Any]: + return {"step_name": self.step_name} + + @classmethod + def _deserialize_from_dict(cls, input_dict: Dict[str, Any], deserialization_context): + return cls(step_name=input_dict["step_name"]) + + +def _build_checkpointable_agent( + name: str, + initial_message: str, + llm: Optional[LlmModel] = None, +) -> Agent: + llm = llm if llm is not None else DummyModel() + agent = Agent( + llm=llm, + name=name, + description=f"{name} description", + custom_instruction="Be helpful.", + initial_message=initial_message, + agent_id=name, + ) + return agent + + +def _build_checkpointable_swarm(llm: LlmModel) -> Swarm: + first_agent = _build_checkpointable_agent( + name="checkpoint_swarm_first_agent", + initial_message="Hello from the swarm.", + llm=llm, + ) + second_agent = Agent( + llm=llm, + name="checkpoint_swarm_second_agent", + description="Swarm helper", + custom_instruction="Help with delegated tasks.", + agent_id="checkpoint_swarm_second_agent", + ) + swarm = Swarm( + first_agent=first_agent, + relationships=[(first_agent, second_agent)], + name="checkpoint_swarm", + id="checkpoint_swarm", + ) + return swarm + + +def _build_checkpointable_managerworkers(llm: LlmModel) -> ManagerWorkers: + manager_agent = _build_checkpointable_agent( + name="checkpoint_manager_agent", + initial_message="Hello from the manager.", + llm=llm, + ) + worker_agent = Agent( + llm=llm, + name="checkpoint_worker_agent", + description="Worker agent", + custom_instruction="Help the manager.", + agent_id="checkpoint_worker_agent", + ) + managerworkers = ManagerWorkers( + group_manager=manager_agent, + workers=[worker_agent], + name="checkpoint_managerworkers", + id="checkpoint_managerworkers", + ) + return managerworkers + + +def test_flow_checkpoint_restore_preserves_nested_interrupt_inheritance( + integration_checkpointer, +) -> None: + def build_flow() -> Flow: + return _build_nested_flow( + child_first_step_name="child_first_step", + child_second_step_name="child_second_step", + child_flow_name="checkpoint_child_flow", + parent_step_name="parent_flow_step", + parent_flow_name="checkpoint_parent_flow", + ) + + original_flow = build_flow() + conversation = original_flow.start_conversation( + conversation_id="flow-parent-link-restore", + checkpointer=integration_checkpointer, + ) + status = conversation.execute( + execution_interrupts=[_OnStepStartExecutionInterrupt("child_first_step")] + ) + assert isinstance(status, InterruptedExecutionStatus) + assert integration_checkpointer.load_latest(conversation.conversation_id) is not None + + # Checkpoint restoration must work with a freshly constructed flow instance; + # the saved execution state is matched using stable component IDs, rather + # than requiring the original Python object instances. + # Hence we call build_flow() again, which by definition + # returns a flow object with predefined ids. + restarted_flow = deserialize(Flow, serialize(build_flow())) + restored_conversation = restarted_flow.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=integration_checkpointer, + ) + # a top-level root conversation has the same instance id and thread id + assert restored_conversation.id == restored_conversation.conversation_id + + restored_status = restored_conversation.execute( + execution_interrupts=[_OnStepStartExecutionInterrupt("child_second_step")] + ) + + assert isinstance(restored_status, InterruptedExecutionStatus) + assert restored_status.reason == "Start child_second_step" + + +@pytest.mark.parametrize( + ( + "component_class", + "component_builder", + "model_getter", + "conversation_id", + "continuation_message", + "expected_status", + ), + [ + ( + Flow, + lambda _: create_flow(), + lambda _: None, + "serialized-flow-checkpoint", + "continue", + FinishedStatus, + ), + ( + Agent, + lambda llm: _build_checkpointable_agent( + name="serialized_checkpoint_agent", + initial_message="Initial response.", + llm=llm, + ), + lambda component: component.llm, + "serialized-agent-checkpoint", + "Continue.", + UserMessageRequestStatus, + ), + ( + Swarm, + _build_checkpointable_swarm, + lambda component: component.first_agent.llm, + "serialized-swarm-checkpoint", + "Continue.", + UserMessageRequestStatus, + ), + ( + ManagerWorkers, + _build_checkpointable_managerworkers, + lambda component: component.group_manager.llm, + "serialized-managerworkers-checkpoint", + "Continue.", + UserMessageRequestStatus, + ), + ], + ids=["flow", "agent", "swarm", "managerworkers"], +) +def test_checkpoint_restore_with_serialized_component_graph_supports_time_travel( + component_class, + component_builder, + model_getter, + conversation_id: str, + continuation_message: str, + expected_status, + integration_checkpointer, + vllm_responses_llm, +) -> None: + component = component_builder(vllm_responses_llm) + + conversation = component.start_conversation( + conversation_id=conversation_id, + checkpointer=integration_checkpointer, + ) + first_status = conversation.execute() + assert isinstance(first_status, UserMessageRequestStatus) + first_checkpoint = integration_checkpointer.load_latest(conversation.conversation_id) + assert first_checkpoint is not None + first_checkpoint_id = first_checkpoint.checkpoint_id + + restored_component = deserialize(component_class, serialize(component)) + assert isinstance(restored_component, component_class) + assert restored_component.id == component.id + + restored_conversation = restored_component.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=integration_checkpointer, + ) + assert isinstance(restored_conversation.status, UserMessageRequestStatus) + restored_conversation.append_user_message(continuation_message) + with patch_llm( + model_getter(restored_component) or vllm_responses_llm, + outputs=["Resumed response."], + ): + restored_status = restored_conversation.execute() + + assert isinstance(restored_status, expected_status) + + rewound_conversation = restored_component.start_conversation( + conversation_id=conversation.conversation_id, + checkpoint_id=first_checkpoint_id, + checkpointer=integration_checkpointer, + ) + assert len(rewound_conversation.get_messages()) < len(restored_conversation.get_messages()) + rewound_conversation.append_user_message("Try again.") + with patch_llm( + model_getter(restored_component) or vllm_responses_llm, + outputs=["Rewound response."], + ): + rewound_status = rewound_conversation.execute() + + assert isinstance(rewound_status, expected_status) + + +def test_agent_checkpoint_restore_relinks_current_flow_parent(integration_checkpointer) -> None: + def build_agent() -> tuple[Agent, Flow]: + parent_flow = _build_nested_flow( + child_first_step_name="agent_child_first_step", + child_second_step_name="agent_child_second_step", + child_flow_name="checkpoint_agent_child_flow", + parent_step_name="agent_parent_flow_step", + parent_flow_name="checkpoint_agent_parent_flow", + ) + agent = Agent( + llm=DummyModel(fails_if_not_set=False), + name="checkpoint_flow_agent", + description="Agent with a nested flow", + custom_instruction="Use the nested flow.", + flows=[parent_flow], + initial_message=None, + agent_id="checkpoint_flow_agent", + ) + return agent, parent_flow + + original_agent, original_parent_flow = build_agent() + conversation = original_agent.start_conversation( + conversation_id="agent-current-flow-restore", + checkpointer=integration_checkpointer, + ) + with patch_llm( + original_agent.llm, + outputs=[ + [ + ToolRequest( + name=original_parent_flow.name, + args={}, + tool_request_id="execute_parent_flow", + ) + ] + ], + ): + status = conversation.execute( + execution_interrupts=[_OnStepStartExecutionInterrupt("agent_child_first_step")] + ) + + assert isinstance(status, InterruptedExecutionStatus) + integration_checkpointer.save(conversation) + + restarted_agent, restarted_parent_flow = build_agent() + restored_conversation = restarted_agent.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=integration_checkpointer, + ) + + restored_parent_flow_conversation = restored_conversation.state.current_flow_conversation + assert restored_parent_flow_conversation is not None + + +def test_swarm_checkpoint_restore_uses_generated_agent_ids_for_threads( + integration_checkpointer, +) -> None: + integration_checkpointer.checkpointing_interval = CheckpointingInterval.ALL_INTERNAL_TURNS + + def build_swarm() -> tuple[Swarm, Agent, Agent]: + first_agent = Agent( + llm=DummyModel(fails_if_not_set=False), + description="Swarm first agent", + custom_instruction="Route work to another agent.", + agent_id="checkpoint_generated_name_swarm_first_agent", + ) + second_agent = Agent( + llm=DummyModel(fails_if_not_set=False), + description="Swarm second agent", + custom_instruction="Handle delegated work.", + agent_id="checkpoint_generated_name_swarm_second_agent", + ) + return ( + Swarm( + first_agent=first_agent, + relationships=[(first_agent, second_agent)], + name="checkpoint_generated_name_swarm", + id="checkpoint_generated_name_swarm", + ), + first_agent, + second_agent, + ) + + original_swarm, original_first_agent, original_second_agent = build_swarm() + conversation = original_swarm.start_conversation( + conversation_id="swarm-generated-agent-name-restart", + checkpointer=integration_checkpointer, + ) + original_first_agent.llm.set_next_output(_handoff_message(original_second_agent)) + original_second_agent.llm.set_next_output(["Delegated work completed.", "More work completed."]) + conversation.append_user_message("Delegate this task.") + conversation.execute() + conversation.append_user_message("Please continue.") + conversation.execute() + + restarted_swarm, restarted_first_agent, restarted_second_agent = build_swarm() + + restored_conversation = restarted_swarm.start_conversation( + conversation_id=conversation.conversation_id, checkpointer=integration_checkpointer + ) + assert any( + subconversation.component is restarted_second_agent + for subconversation in restored_conversation.thread_subconversations.values() + ) + + +def test_managerworkers_checkpoint_restore_uses_nested_agent_ids( + integration_checkpointer, +) -> None: + def build_managerworkers() -> tuple[ManagerWorkers, Agent, ManagerWorkers, Agent, Agent]: + nested_manager_agent = Agent( + llm=DummyModel(fails_if_not_set=False), + description="Nested generated manager agent", + custom_instruction="Assign nested work to worker agents.", + agent_id="checkpoint_nested_generated_name_manager_agent", + ) + nested_worker_agent = Agent( + llm=DummyModel(fails_if_not_set=False), + description="Nested generated worker agent", + custom_instruction="Handle nested delegated work.", + agent_id="checkpoint_nested_generated_name_worker_agent", + ) + nested_managerworkers = ManagerWorkers( + group_manager=nested_manager_agent, + workers=[nested_worker_agent], + name="checkpoint_nested_generated_name_managerworkers", + id="checkpoint_nested_generated_name_managerworkers", + ) + outer_manager_agent = Agent( + llm=DummyModel(fails_if_not_set=False), + description="Outer generated manager agent", + custom_instruction="Assign work to nested manager-workers.", + agent_id="checkpoint_outer_generated_name_manager_agent", + ) + return ( + ManagerWorkers( + group_manager=outer_manager_agent, + workers=[nested_managerworkers], + name="checkpoint_outer_generated_name_managerworkers", + id="checkpoint_outer_generated_name_managerworkers", + ), + outer_manager_agent, + nested_managerworkers, + nested_manager_agent, + nested_worker_agent, + ) + + ( + original_managerworkers, + original_outer_manager_agent, + original_nested_managerworkers, + original_nested_manager_agent, + original_nested_worker_agent, + ) = build_managerworkers() + conversation = original_managerworkers.start_conversation( + conversation_id="managerworkers-nested-generated-agent-name-restart", + checkpointer=integration_checkpointer, + ) + conversation.append_user_message("Save this nested conversation.") + original_outer_manager_agent.llm.set_next_output( + [_send_message(original_nested_managerworkers), "Saved."] + ) + original_nested_manager_agent.llm.set_next_output( + [_send_message(original_nested_worker_agent), "Nested work saved."] + ) + original_nested_worker_agent.llm.set_next_output("Worker saved the nested work.") + conversation.execute() + + ( + restarted_managerworkers, + _restarted_outer_manager_agent, + restarted_nested_managerworkers, + _restarted_nested_manager_agent, + restarted_nested_worker_agent, + ) = build_managerworkers() + + restored_conversation = restarted_managerworkers.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=integration_checkpointer, + ) + restored_nested_conversation = restored_conversation.subconversations[ + restarted_nested_managerworkers.id + ] + assert ( + restored_nested_conversation.subconversations[restarted_nested_worker_agent.id].component + is restarted_nested_worker_agent + ) diff --git a/wayflowcore/tests/mcptools/test_mcp_tools.py b/wayflowcore/tests/mcptools/test_mcp_tools.py index 5015d262a..94c6dbc31 100644 --- a/wayflowcore/tests/mcptools/test_mcp_tools.py +++ b/wayflowcore/tests/mcptools/test_mcp_tools.py @@ -2026,8 +2026,8 @@ def test_oauth_works_on_swarm_when_subagent_uses_mcp_tool( assert isinstance(status, UserMessageRequestStatus) - thread = conv.state.agents_and_threads[manager_agent.name][sub_agent.name] - sub_conv = conv.state.thread_subconversations[thread.identifier] + thread = conv.state.agents_and_threads[manager_agent.id][sub_agent.id] + sub_conv = conv.state.thread_subconversations[thread.id] last_tool_result_message = sub_conv.get_messages()[-2] assert ( last_tool_result_message.tool_result is not None @@ -2108,7 +2108,7 @@ def test_oauth_works_on_managerworkers_when_worker_uses_mcp_tool( status = conv.execute() assert isinstance(status, UserMessageRequestStatus) - sub_conv = conv.state.subconversations[worker_agent.name] + sub_conv = conv.state.subconversations[worker_agent.id] last_tool_result_message = sub_conv.get_messages()[-2] assert ( last_tool_result_message.tool_result is not None diff --git a/wayflowcore/tests/serialization/test_managerworkers_serialization.py b/wayflowcore/tests/serialization/test_managerworkers_serialization.py index 323f0c4c3..e48e83c54 100644 --- a/wayflowcore/tests/serialization/test_managerworkers_serialization.py +++ b/wayflowcore/tests/serialization/test_managerworkers_serialization.py @@ -118,7 +118,7 @@ def assert_managerworkers_conversation_states_are_equal( old_state: ManagerWorkersConversationExecutionState, new_state: ManagerWorkersConversationExecutionState, ): - assert old_state.current_agent_name == new_state.current_agent_name + assert old_state.current_agent_id == new_state.current_agent_id assert old_state.subconversations.keys() == new_state.subconversations.keys() for key in old_state.subconversations.keys(): assert_agent_conversations_are_equal( @@ -201,7 +201,7 @@ def test_can_deserialize_a_serialized_conversation( simple_conversation: ManagerWorkersConversation, simple_math_agents_example ): addition_agent, _ = simple_math_agents_example - simple_conversation.subconversations[addition_agent.name] = addition_agent.start_conversation() + simple_conversation.subconversations[addition_agent.id] = addition_agent.start_conversation() new_conversation = deserialize(ManagerWorkersConversation, serialize(simple_conversation)) assert_managerworkers_conversations_are_equal(simple_conversation, new_conversation) @@ -415,7 +415,7 @@ def test_can_deserialize_a_multi_level_managerworkers_serialized_conversation( multi_level_managerworkers_conversation: ManagerWorkersConversation, simple_math_agents_example ): addition_agent, _ = simple_math_agents_example - multi_level_managerworkers_conversation.subconversations[addition_agent.name] = ( + multi_level_managerworkers_conversation.subconversations[addition_agent.id] = ( addition_agent.start_conversation() ) new_conversation = deserialize( diff --git a/wayflowcore/tests/test_checkpointing.py b/wayflowcore/tests/test_checkpointing.py new file mode 100644 index 000000000..11b73cb16 --- /dev/null +++ b/wayflowcore/tests/test_checkpointing.py @@ -0,0 +1,336 @@ +# Copyright © 2026 Oracle and/or its affiliates. +# +# This software is under the Apache License 2.0 +# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License +# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. + +from unittest.mock import AsyncMock + +import pytest + +from wayflowcore.agent import Agent +from wayflowcore.checkpointing import CheckpointingInterval, InMemoryCheckpointer +from wayflowcore.controlconnection import ControlFlowEdge +from wayflowcore.executors.executionstatus import FinishedStatus, UserMessageRequestStatus +from wayflowcore.flow import Flow +from wayflowcore.flowhelpers import create_single_step_flow +from wayflowcore.serialization import deserialize, serialize +from wayflowcore.steps import CompleteStep, ConstantValuesStep, OutputMessageStep, RetryStep +from wayflowcore.steps.promptexecutionstep import PromptExecutionStep +from wayflowcore.swarm import Swarm + +from .testhelpers.dummy import DummyModel + + +def test_checkpointer_save_snapshots_a_live_conversation_with_custom_metadata() -> None: + checkpointer = InMemoryCheckpointer() + flow = create_single_step_flow(OutputMessageStep(message_template="Hello from checkpointing.")) + conversation = flow.start_conversation(conversation_id="custom-checkpoint") + + saved_checkpoint = checkpointer.save( + conversation, + checkpoint_id="response_123", + component_id="served-model", + metadata={"response": "serialized response"}, + ) + + assert saved_checkpoint is not None + assert saved_checkpoint.checkpoint_id == "response_123" + assert conversation.checkpoint_id == "response_123" + checkpoint = checkpointer.load("custom-checkpoint", "response_123") + assert checkpoint.component_id == "served-model" + assert checkpoint.metadata["response"] == "serialized response" + + checkpointer.save(conversation, checkpoint_id="response_456") + + checkpoints = checkpointer.list_checkpoints("custom-checkpoint") + assert [checkpoint.checkpoint_id for checkpoint in checkpoints] == [ + "response_123", + "response_456", + ] + assert checkpoints[-1].metadata["save_sequence"] == 2 + assert checkpointer.load_latest("custom-checkpoint").checkpoint_id == "response_456" + + checkpointer.delete("custom-checkpoint", "response_456") + + assert checkpointer.load_latest("custom-checkpoint").checkpoint_id == "response_123" + + +def test_checkpoint_restore_requires_conversation_id_when_checkpoint_id_is_provided() -> None: + checkpointer = InMemoryCheckpointer() + flow = create_single_step_flow(OutputMessageStep(message_template="Hello from checkpointing.")) + + conversation = flow.start_conversation( + conversation_id="checkpoint-missing-conversation-id", + checkpointer=checkpointer, + ) + status = conversation.execute() + assert isinstance(status, FinishedStatus) + + checkpoint = checkpointer.load_latest(conversation.conversation_id) + assert checkpoint is not None + + with pytest.raises(ValueError, match="`checkpoint_id` requires a `conversation_id`\\."): + flow.start_conversation( + checkpointer=checkpointer, + checkpoint_id=checkpoint.checkpoint_id, + ) + + +def test_checkpoint_restore_resumes_retry_flow() -> None: + # An internal-turn checkpoint must preserve retry progress instead of + # restarting the retry sequence when the conversation is restored. + inner_flow = create_single_step_flow( + ConstantValuesStep(constant_values={"success": False}, name="retry_result") + ) + retry_step = RetryStep( + flow=inner_flow, + success_condition="success", + max_num_trials=2, + name="retry_step", + ) + success_step = CompleteStep(name="success") + failure_step = CompleteStep(name="failure") + flow = Flow( + begin_step=retry_step, + steps={"retry_step": retry_step, "success": success_step, "failure": failure_step}, + control_flow_edges=[ + ControlFlowEdge( + source_step=retry_step, + source_branch=RetryStep.BRANCH_NEXT, + destination_step=success_step, + ), + ControlFlowEdge( + source_step=retry_step, + source_branch=RetryStep.BRANCH_FAILURE, + destination_step=failure_step, + ), + ], + ) + checkpointer = InMemoryCheckpointer( + checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS + ) + + conversation = flow.start_conversation( + conversation_id="checkpoint-retry-counter", + checkpointer=checkpointer, + ) + assert isinstance(conversation.execute(), FinishedStatus) + + # Restore from an in-progress checkpoint rather than the final checkpoint. + checkpoint = next( + checkpoint + for checkpoint in checkpointer.list_checkpoints(conversation.conversation_id) + if checkpoint.metadata["save_reason"] == "internal_turn_boundary" + ) + + reloaded_flow = deserialize(Flow, serialize(flow)) + restored_conversation = reloaded_flow.start_conversation( + conversation_id=conversation.conversation_id, + checkpoint_id=checkpoint.checkpoint_id, + checkpointer=checkpointer, + ) + + restored_status = restored_conversation.execute() + # The restored retry has exhausted its attempts and takes the failure branch. + assert isinstance(restored_status, FinishedStatus) + assert restored_status.complete_step_name == "failure" + + +@pytest.mark.parametrize( + ( + "interval", + "conversation_id", + "step", + "expected_reasons", + "expected_event_types", + ), + [ + ( + CheckpointingInterval.CONVERSATION_TURNS, + "conversation-turn", + OutputMessageStep(message_template="Hello once."), + ["conversation_turn"], + [], + ), + ( + CheckpointingInterval.ALL_INTERNAL_TURNS, + "all-internal-turns", + OutputMessageStep(message_template="Hello internal turns."), + ["internal_turn_boundary", "internal_turn_boundary", "conversation_turn"], + ["FlowExecutionIterationStartedEvent", "FlowExecutionIterationStartedEvent"], + ), + ( + CheckpointingInterval.LLM_TURNS, + "llm-turns", + None, + ["internal_turn_boundary", "conversation_turn"], + ["FlowExecutionIterationStartedEvent"], + ), + ], + ids=["conversation_turns", "all_internal_turns", "llm_turns"], +) +def test_checkpoint_intervals_save_expected_flow_checkpoints( + interval: CheckpointingInterval, + conversation_id: str, + step, + expected_reasons: list[str], + expected_event_types: list[str], +) -> None: + checkpointer = InMemoryCheckpointer(checkpointing_interval=interval) + if step is None: + llm = DummyModel() + llm.set_next_output("Hello from the prompt step.") + step = PromptExecutionStep(llm=llm, prompt_template="Say hello.") + flow = create_single_step_flow(step) + + status = flow.start_conversation( + conversation_id=conversation_id, + checkpointer=checkpointer, + ).execute() + + assert isinstance(status, FinishedStatus) + checkpoints = checkpointer.list_checkpoints(conversation_id) + assert [checkpoint.metadata["save_reason"] for checkpoint in checkpoints] == expected_reasons + assert [checkpoint.metadata.get("event_type") for checkpoint in checkpoints[:-1]] == ( + expected_event_types + ) + assert checkpoints[-1].metadata["status_type"] == "FinishedStatus" + + +@pytest.mark.parametrize( + "interval", + list(CheckpointingInterval), + ids=["conversation_turns", "all_internal_turns", "llm_turns"], +) +def test_checkpoint_intervals_save_expected_agent_checkpoints( + interval: CheckpointingInterval, +) -> None: + checkpointer = InMemoryCheckpointer(checkpointing_interval=interval) + conversation_id = f"agent-{interval.name.lower()}" + agent = Agent( + llm=DummyModel(), + name="checkpoint_interval_agent", + initial_message="Hello from the agent.", + ) + + status = agent.start_conversation( + conversation_id=conversation_id, + checkpointer=checkpointer, + ).execute() + + assert isinstance(status, UserMessageRequestStatus) + checkpoints = checkpointer.list_checkpoints(conversation_id) + assert checkpoints[-1].metadata["save_reason"] == "conversation_turn" + assert checkpoints[-1].metadata["status_type"] == "UserMessageRequestStatus" + if interval is CheckpointingInterval.ALL_INTERNAL_TURNS: + assert any( + checkpoint.metadata.get("event_type") == "AgentExecutionIterationStartedEvent" + for checkpoint in checkpoints + ) + assert all( + checkpoint.metadata.get("agent_iteration") is not None + for checkpoint in checkpoints[:-1] + ) + + +def test_execute_async_does_not_save_final_checkpoint_when_execution_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + checkpointer = InMemoryCheckpointer() + flow = create_single_step_flow(OutputMessageStep(message_template="Hello failure.")) + + conversation = flow.start_conversation( + conversation_id="checkpoint-final-exception", + checkpointer=checkpointer, + ) + monkeypatch.setattr( + conversation.component.runner, + "execute_async", + AsyncMock(side_effect=RuntimeError("runner failed")), + ) + + with pytest.raises(RuntimeError, match="runner failed"): + conversation.execute() + + assert checkpointer.list_checkpoints("checkpoint-final-exception") == [] + + +def _checkpoint_restore_generated_agent_id_scenario( + checkpointer: InMemoryCheckpointer, +): + original_agent = Agent( + llm=DummyModel(), + name="generated_id_restart_agent", + description="generated_id_restart_agent description", + custom_instruction="Be helpful.", + initial_message="Hello from the original generated-id agent.", + ) + conversation = original_agent.start_conversation( + conversation_id="agent-generated-id-restart", + checkpointer=checkpointer, + ) + assert isinstance(conversation.execute(), UserMessageRequestStatus) + restarted_agent = Agent( + llm=DummyModel(), + name="generated_id_restart_agent", + description="generated_id_restart_agent description", + custom_instruction="Be helpful.", + initial_message="Hello from the original generated-id agent.", + ) + return conversation, restarted_agent + + +def _checkpoint_restore_generated_swarm_child_id_scenario( + checkpointer: InMemoryCheckpointer, +): + def build_swarm() -> Swarm: + first_agent = Agent( + llm=DummyModel(), + name="checkpoint_swarm_first_agent", + description="checkpoint_swarm_first_agent description", + custom_instruction="Be helpful.", + initial_message="Hello from the swarm.", + ) + second_agent = Agent( + llm=DummyModel(fails_if_not_set=False), + name="checkpoint_swarm_second_agent", + description="Swarm helper", + custom_instruction="Help with delegated tasks.", + ) + return Swarm( + first_agent=first_agent, + relationships=[(first_agent, second_agent)], + name="checkpoint_swarm", + ) + + original_swarm = build_swarm() + conversation = original_swarm.start_conversation( + conversation_id="swarm-generated-child-id-restart", + checkpointer=checkpointer, + ) + assert isinstance(conversation.execute(), UserMessageRequestStatus) + restarted_swarm = build_swarm() + return conversation, restarted_swarm + + +@pytest.mark.parametrize( + "scenario", + [ + _checkpoint_restore_generated_agent_id_scenario, + _checkpoint_restore_generated_swarm_child_id_scenario, + ], + ids=[ + "generated_agent_id", + "generated_swarm_child_id", + ], +) +def test_checkpoint_restore_requires_matching_component_ids(scenario) -> None: + checkpointer = InMemoryCheckpointer() + conversation, restarted_component = scenario(checkpointer) + + with pytest.raises(ValueError, match="stable component ids"): + restarted_component.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) diff --git a/wayflowcore/tests/test_flowbuilder.py b/wayflowcore/tests/test_flowbuilder.py index 113691512..cf8a8b19e 100644 --- a/wayflowcore/tests/test_flowbuilder.py +++ b/wayflowcore/tests/test_flowbuilder.py @@ -254,8 +254,9 @@ def test_build_linear_flow_returns_flow_and_json(): s2 = OutputMessageStep(name="s2", message_template="B") # Flow object - flow = FlowBuilder.build_linear_flow([s1, s2], name="MyFlow") + flow = FlowBuilder.build_linear_flow([s1, s2], name="MyFlow", flow_id="my-flow") assert flow.name == "MyFlow" + assert flow.flow_id == "my-flow" assert set(flow.steps.keys()) >= {"s1", "s2"} # JSON spec diff --git a/wayflowcore/tests/test_managerworkers.py b/wayflowcore/tests/test_managerworkers.py index f2e5d601d..c7c234b5b 100644 --- a/wayflowcore/tests/test_managerworkers.py +++ b/wayflowcore/tests/test_managerworkers.py @@ -119,7 +119,7 @@ def test_manager_can_send_message_to_worker_and_worker_can_reply(): conversation.execute() # subconversation of worker should contain the message from manager - worker1_sub_conv = conversation.subconversations[worker1.name] + worker1_sub_conv = conversation.subconversations[worker1.id] worker1_first_message = worker1_sub_conv.message_list.messages[0] assert worker1_first_message.content == "Hey worker 1" and worker1_first_message.role == "user" @@ -135,6 +135,28 @@ def test_manager_can_send_message_to_worker_and_worker_can_reply(): assert last_message.tool_result.content == "Hello manager!" and last_message.role == "assistant" +def test_managerworkers_worker_conversation_inherits_parent_thread_identity(): + manager_llm = DummyModel() + worker = Agent( + DummyModel(fails_if_not_set=False), + name="identity_worker", + description="identity worker", + ) + group = ManagerWorkers(group_manager=manager_llm, workers=[worker]) + + conversation = group.start_conversation(conversation_id="manager-conversation") + conversation.append_user_message("Delegate this task.") + manager_llm.set_next_output([_send_message(worker, message="Please help."), "done"]) + + conversation.execute() + + worker_conversation = conversation.subconversations[worker.id] + assert conversation.id == conversation.conversation_id + assert worker_conversation.conversation_id == conversation.conversation_id + assert worker_conversation.id != worker.id + assert worker_conversation.id != conversation.id + + @pytest.fixture def simple_math_agents_example(remote_gemma_llm) -> Tuple[Agent, Agent, Agent]: llm = remote_gemma_llm @@ -340,7 +362,7 @@ def multiply( conversation.execute() # Check tool request `multiply` existing in the message list of multiplication_agent - subconversation = conversation.state.subconversations[multiplication_agent.name] + subconversation = conversation.state.subconversations[multiplication_agent.id] assert any( message.tool_requests is not None for message in subconversation.message_list.messages @@ -802,7 +824,7 @@ def test_multiple_tool_calls_including_with_nonexistent_tools(vllm_responses_llm assert conv.get_last_message().content == "fooza answers to user" tool_result_messages = [ m - for m in conv.state.subconversations["bwip_agent"].get_messages() + for m in conv.state.subconversations[bwip_agent.id].get_messages() if m.tool_result is not None ] assert ( diff --git a/wayflowcore/tests/test_ociagent.py b/wayflowcore/tests/test_ociagent.py index 511e163c2..ced7d9a71 100644 --- a/wayflowcore/tests/test_ociagent.py +++ b/wayflowcore/tests/test_ociagent.py @@ -8,6 +8,8 @@ import pytest from wayflowcore import Agent, Flow +from wayflowcore.checkpointing import InMemoryCheckpointer +from wayflowcore.models.ociclientconfig import OCIClientConfigWithApiKey from wayflowcore.ociagent import OciAgent from wayflowcore.serialization.serializer import deserialize_from_dict, serialize_to_dict from wayflowcore.steps import AgentExecutionStep @@ -51,6 +53,17 @@ def test_oci_knowledge_agent_simple_rag_question(agent): assert "technology" in last_message.lower() +def test_ociagent_explicitly_rejects_checkpoint_restore_arguments() -> None: + oci_agent = OciAgent( + agent_endpoint_id="ocid1.test.oc1..example", + client_config=OCIClientConfigWithApiKey(service_endpoint="https://example.com"), + name="checkpoint_oci_agent", + ) + + with pytest.raises(NotImplementedError, match="checkpoint restore"): + oci_agent.start_conversation(checkpointer=InMemoryCheckpointer()) + + # oci agent performance seems to fluctuate, we increase max_attempts to ensure CI reliability @retry_test(max_attempts=6) def test_oci_knowledge_agent_continue_conversation(agent): diff --git a/wayflowcore/tests/test_swarm.py b/wayflowcore/tests/test_swarm.py index 263b73b17..f28a035cc 100644 --- a/wayflowcore/tests/test_swarm.py +++ b/wayflowcore/tests/test_swarm.py @@ -494,9 +494,7 @@ def test_swarm_warns_agent_on_sending_message_to_caller_instead_of_using_talk_to # controlled execution conv.execute() - agent1_agent2_message_list = conv.state.agents_and_threads[agent1.name][ - agent2.name - ].message_list + agent1_agent2_message_list = conv.state.agents_and_threads[agent1.id][agent2.id].message_list last_message = ( agent1_agent2_message_list.get_last_message() ) # Message warning the `agent2` about what it is doing wrong @@ -581,9 +579,7 @@ def test_circular_calling_warning_with_handoff(): # controlled execution conv.execute() - agent3_agent2_message_list = conv.state.agents_and_threads[agent3.name][ - agent2.name - ].message_list + agent3_agent2_message_list = conv.state.agents_and_threads[agent3.id][agent2.id].message_list last_message = ( agent3_agent2_message_list.get_last_message() ) # Message warning the `agent3` about what it is doing wrong