From 831f78394072778646f0862368772ba4d75a2bbf Mon Sep 17 00:00:00 2001 From: jschweiz Date: Fri, 17 Apr 2026 15:40:53 +0200 Subject: [PATCH 01/10] [feat]: add checkpointing APIs --- docs/wayflowcore/source/core/changelog.rst | 11 +- .../core/code_examples/howto_checkpointing.py | 69 +++++++ .../core/howtoguides/howto_checkpointing.rst | 119 +++++++++++ .../source/core/howtoguides/index.rst | 1 + .../agentserver/_storagehelpers.py | 94 ++------- .../services/wayflowservice.py | 189 ++++++++---------- .../agentserver/serverstorageconfig.py | 47 +---- .../contextproviders/flowcontextprovider.py | 1 + .../wayflowcore/conversationalcomponent.py | 104 ++++++++++ .../wayflowcore/executors/_agentexecutor.py | 17 +- .../wayflowcore/executors/_flowexecutor.py | 2 +- .../executors/_managerworkersconversation.py | 3 +- .../tokenlimitexecutioninterrupt.py | 2 +- .../src/wayflowcore/models/llmmodel.py | 4 +- .../src/wayflowcore/serialization/context.py | 29 ++- .../wayflowcore/steps/agentexecutionstep.py | 4 +- .../src/wayflowcore/tools/servertools.py | 1 + wayflowcore/tests/test_managerworkers.py | 4 +- wayflowcore/tests/test_swarm.py | 4 +- 19 files changed, 449 insertions(+), 256 deletions(-) create mode 100644 docs/wayflowcore/source/core/code_examples/howto_checkpointing.py create mode 100644 docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst diff --git a/docs/wayflowcore/source/core/changelog.rst b/docs/wayflowcore/source/core/changelog.rst index d75c2c852..9f57dacef 100644 --- a/docs/wayflowcore/source/core/changelog.rst +++ b/docs/wayflowcore/source/core/changelog.rst @@ -87,14 +87,23 @@ 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 ` +* **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 `. 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..bcfc88cc4 --- /dev/null +++ b/docs/wayflowcore/source/core/code_examples/howto_checkpointing.py @@ -0,0 +1,69 @@ +# Copyright © 2025, 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 = agent.start_conversation( + conversation_id="support-thread-1", + checkpointer=checkpointer, +) + +status = conversation.execute() +# .. end-##_Start_a_checkpointed_conversation + +# .. start-##_Resume_the_latest_checkpoint +restored_conversation = agent.start_conversation( + conversation_id="support-thread-1", + 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 = checkpointer.list_checkpoints("support-thread-1") + +previous_checkpoint = checkpoints[-2] +rewound_conversation = agent.start_conversation( + conversation_id="support-thread-1", + 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..e22d9a64e --- /dev/null +++ b/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst @@ -0,0 +1,119 @@ +.. _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 ` + +WayFlow can now checkpoint the runtime state of a conversation and restore it later by +conversation id. This is useful when you want to: + +- resume after a crash or restart +- pause and continue a long-running workflow +- inspect prior checkpoints for debugging +- reload an earlier state and branch from it + + +Choose a checkpointer +===================== + +WayFlow exposes a shared checkpointing subsystem in ``wayflowcore.checkpointing``. +You can use: + +- ``InMemoryCheckpointer`` for tests and local experimentation +- ``PostgresCheckpointer`` for PostgreSQL-backed persistence +- ``OracleDatabaseCheckpointer`` for Oracle-backed persistence + +All checkpointers share the same API for saving, loading, listing, and deleting checkpoints. + + +Start a checkpointed conversation +================================= + +Attach a checkpointer when you start the conversation. ``conversation_id`` becomes the durable key +used to look up the conversation later. + +.. 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 root conversation automatically at the configured +checkpoint boundaries. For nested execution lineage without checkpoint restore, pass +``root_conversation_id`` explicitly. + + +Resume the latest checkpoint +============================ + +To restore the latest saved state, 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 no checkpoint exists for that id, WayFlow creates a new conversation instead. + + +Load a specific checkpoint +========================== + +You can inspect checkpoint history and reload an older checkpoint for replay or time-travel +debugging. + +.. 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 ordered checkpoint metadata, including the checkpoint id, +creation timestamp, and save metadata recorded at the boundary. + + +Control checkpoint frequency +============================ + +Use ``CheckpointingInterval`` to decide how often WayFlow should persist 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 outermost ``conversation.execute()`` call returns +- ``LLM_TURNS``: also save at internal turn boundaries after turns that used an LLM +- ``ALL_INTERNAL_TURNS``: also save at every internal agent/flow turn boundary + +Saving more frequently improves restart fidelity, but it also increases write volume. + + +Use checkpointing with the OpenAI Responses server +================================================== + +The OpenAI Responses server path now uses the shared checkpointing subsystem behind +``ServerStorageConfig``. That means the existing OpenAI-compatible features such as +``previous_response_id``, ``conversation``, ``get_response()``, ``delete_response()``, +and ``store=False`` all run through the same shared checkpoint model. + +If you are serving agents, keep using :doc:`Serve Agents with WayFlow ` to +configure the storage backend. The server will use the matching shared 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/src/wayflowcore/agentserver/_storagehelpers.py b/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py index 58b12d97d..6a09bad18 100644 --- a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py +++ b/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py @@ -4,19 +4,18 @@ # (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 typing import Dict, Optional from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig +from wayflowcore.checkpointing.datastore import ( + _prepare_oracle_checkpoint_datastore, + _prepare_postgres_checkpoint_datastore, +) +from wayflowcore.checkpointing.serialization import _deserialize_conversation_checkpoint_state 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.datastore.oracle import OracleDatabaseConnectionConfig +from wayflowcore.datastore.postgres import PostgresDatabaseConnectionConfig from wayflowcore.tools import Tool logger = logging.getLogger(__name__) @@ -25,53 +24,13 @@ 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 + _prepare_postgres_checkpoint_datastore(connection_config, storage_config) 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 + _prepare_oracle_checkpoint_datastore(connection_config, storage_config) def _deserialize_conversation_safely( @@ -79,31 +38,8 @@ def _deserialize_conversation_safely( 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) + return _deserialize_conversation_checkpoint_state( + serialized_state, + tool_registry=tool_registry, + component=component, + ) diff --git a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py index 99c34ecc9..2c36761ce 100644 --- a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py +++ b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py @@ -4,8 +4,6 @@ # (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 typing import Any, AsyncIterable, Dict, List, Optional, Union, cast @@ -16,10 +14,14 @@ from fastapi import status as http_status_code from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig +from wayflowcore.checkpointing import ConversationCheckpoint, DatastoreCheckpointer +from wayflowcore.checkpointing.runtime import ( + _detach_checkpointer_from_conversation, + _set_conversation_final_checkpoint_overrides, +) 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 +29,7 @@ UserMessageRequestStatus, ) from wayflowcore.idgeneration import IdGenerator -from wayflowcore.serialization import serialize -from ..._storagehelpers import _deserialize_conversation_safely from ..models.openairesponsespydanticmodels import ( Conversation2, CreateResponse, @@ -71,7 +71,12 @@ 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()) + self._response_conversation_ids: Dict[str, str] = {} self.tool_registries = { agent_name: {t.name: t for t in agent._referenced_tools()} for agent_name, agent in self.agents.items() @@ -126,23 +131,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._lookup_checkpoint_by_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._lookup_checkpoint_by_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, None) return None async def cancel_response(self, response_id: str) -> Union[Response, ResponseError]: @@ -198,14 +203,19 @@ async def create_response(self, body: CreateResponse) -> AsyncIterable[ResponseS agent_id=model, ) + response_id = IdGenerator.get_or_generate_id() state = await self._create_state( agent=agent, state=state, request=body, ) + if body.store is None or body.store is True: + _set_conversation_final_checkpoint_overrides(state, checkpoint_id=response_id) + else: + _detach_checkpointer_from_conversation(state) current_response = Response( - id=IdGenerator.get_or_generate_id(), + id=response_id, created_at=int(time.time()), error=None, incomplete_details=None, @@ -309,11 +319,13 @@ 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, + if (body.store is None or body.store is True) and state.checkpointer is not None: + self.checkpointer.save_conversation( + state, + checkpoint_id=current_response.id, + metadata={"response": current_response.model_dump_json()}, ) + self._response_conversation_ids[current_response.id] = state.id if current_response.error is not None: yield ResponseFailedEvent( @@ -369,106 +381,58 @@ 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: + checkpoint = self._lookup_checkpoint_by_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", ) - elif conversation_id: + self._response_conversation_ids[checkpoint.checkpoint_id] = checkpoint.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, + return self.agents[agent_id].start_conversation( + conversation_id=checkpoint.conversation_id, + checkpoint_id=checkpoint.checkpoint_id, + checkpointer=self.checkpointer, ) - except ValueError: + 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}", + ) from e + elif conversation_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", ) + self._response_conversation_ids[checkpoint.checkpoint_id] = checkpoint.conversation_id + try: + return self.agents[agent_id].start_conversation( + conversation_id=conversation_id, + checkpointer=self.checkpointer, + ) + 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}", + ) from e else: return None - try: - return _deserialize_conversation_safely( - serialized_state=serialized_conversation, - tool_registry=self.tool_registries[agent_id], - component=self.agents[agent_id], - ) - 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 - 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() - - 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], - ) - 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 _lookup_checkpoint_by_response_id( + self, response_id: str + ) -> Optional[ConversationCheckpoint]: + conversation_id = self._response_conversation_ids.get(response_id) + if conversation_id is not None: + try: + return self.checkpointer.load(conversation_id, response_id) + except ValueError: + self._response_conversation_ids.pop(response_id, None) + checkpoint = self.checkpointer._find_checkpoint_by_id(response_id) + if checkpoint is not None: + self._response_conversation_ids[response_id] = checkpoint.conversation_id + return checkpoint async def _create_state( self, @@ -495,13 +459,22 @@ async def _create_state( detail="Agent should have an `instructions` input descriptor to be able to take instructions as input", ) inputs = {"instructions": instructions} - state = agent.start_conversation(inputs=inputs, messages=new_messages) + if request.store is None or request.store is True: + state = agent.start_conversation( + inputs=inputs, + messages=new_messages, + checkpointer=self.checkpointer, + ) + else: + state = agent.start_conversation(inputs=inputs, messages=new_messages) else: # later: implement context provider for custom instructions if instructions is not None: raise NotImplementedError( "Instructions are only supported when creating a conversation" ) + if request.store is False: + _detach_checkpointer_from_conversation(state) # Add the new messages to the conversation for message in new_messages: state.append_message(message) diff --git a/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py b/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py index 4a27344e7..f2701cacb 100644 --- a/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py +++ b/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py @@ -4,52 +4,11 @@ # (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: - """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(), - } - ), - } +class ServerStorageConfig(StorageConfig): + """Configuration for agent-server conversation storage.""" diff --git a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py index c98d1db59..bd622d0ca 100644 --- a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py +++ b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py @@ -95,6 +95,7 @@ async def call_async(self, conversation: "Conversation") -> Any: conversation = self.flow.start_conversation( inputs={}, messages=conversation.message_list, + root_conversation_id=conversation.root_conversation_id, ) status = await conversation.execute_async() if status._requires_yielding: diff --git a/wayflowcore/src/wayflowcore/conversationalcomponent.py b/wayflowcore/src/wayflowcore/conversationalcomponent.py index f6eef9d00..2d15b20c2 100644 --- a/wayflowcore/src/wayflowcore/conversationalcomponent.py +++ b/wayflowcore/src/wayflowcore/conversationalcomponent.py @@ -10,13 +10,17 @@ from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, Set, Type, TypeVar, Union from wayflowcore._metadata import MetadataType +from wayflowcore.checkpointing.runtime import _attach_checkpointer_to_conversation +from wayflowcore.checkpointing.serialization import _deserialize_conversation_checkpoint_state from wayflowcore.componentwithio import ComponentWithInputsOutputs +from wayflowcore.idgeneration import IdGenerator from wayflowcore.property import Property logger = logging.getLogger(__name__) if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.conversation import Conversation from wayflowcore.executors._executor import ConversationExecutor from wayflowcore.messagelist import Message, MessageList @@ -66,6 +70,11 @@ def start_conversation( self, inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, + conversation_id: Optional[str] = None, + *, + root_conversation_id: Optional[str] = None, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, ) -> "Conversation": pass @@ -115,6 +124,101 @@ 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(message) > 0 + + def _restore_or_prepare_checkpoint_conversation( + self, + *, + inputs: Optional[Dict[str, Any]], + messages: Union[None, str, "Message", List["Message"], "MessageList"], + conversation_id: Optional[str], + root_conversation_id: Optional[str], + checkpointer: Optional["Checkpointer"], + checkpoint_id: Optional[str], + ) -> tuple[Optional["Conversation"], Optional[str]]: + if checkpointer is None: + if checkpoint_id is not None: + raise ValueError("`checkpoint_id` requires a `checkpointer`.") + return None, conversation_id + + if ( + root_conversation_id is not None + and conversation_id is not None + and root_conversation_id != conversation_id + ): + raise ValueError( + "`root_conversation_id` and `conversation_id` cannot differ when checkpointing is enabled." + ) + + resolved_conversation_id = conversation_id or root_conversation_id + if resolved_conversation_id is None and checkpoint_id is not None: + raise ValueError("`checkpoint_id` requires a `conversation_id`.") + if resolved_conversation_id is None: + resolved_conversation_id = IdGenerator.get_or_generate_id() + + checkpoint = ( + checkpointer.load(resolved_conversation_id, checkpoint_id) + if checkpoint_id is not None + else checkpointer.load_latest(resolved_conversation_id) + ) + if checkpoint is None: + return None, resolved_conversation_id + + 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 = _deserialize_conversation_checkpoint_state( + checkpoint.state, + tool_registry={tool.name: tool for tool in self._referenced_tools()}, + component=self, + ) + return ( + _attach_checkpointer_to_conversation( + conversation, + checkpointer=checkpointer, + checkpoint_id=checkpoint.checkpoint_id, + ), + resolved_conversation_id, + ) + + @staticmethod + def _resolve_runtime_and_root_conversation_ids( + *, + conversation_id: Optional[str], + root_conversation_id: Optional[str], + checkpointer: Optional["Checkpointer"], + restored_conversation_id: Optional[str], + ) -> tuple[str, str]: + if checkpointer is not None: + runtime_conversation_id = restored_conversation_id or IdGenerator.get_or_generate_id( + conversation_id or root_conversation_id + ) + resolved_root_conversation_id = root_conversation_id or runtime_conversation_id + if resolved_root_conversation_id != runtime_conversation_id: + raise ValueError( + "`root_conversation_id` and `conversation_id` cannot differ when checkpointing is enabled." + ) + return runtime_conversation_id, runtime_conversation_id + + runtime_conversation_id = IdGenerator.get_or_generate_id(conversation_id) + return runtime_conversation_id, root_conversation_id or runtime_conversation_id + # Define a TypeVar that represents the component's type ConversationalComponentTypeT = TypeVar( diff --git a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py index 79fbc9ac3..86229d52d 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py @@ -462,7 +462,9 @@ def _get_or_create_expert_agent_subconversation( init_messages.append_message(caller_request_message) sub_agent_conversation = expert_agent.start_conversation( - messages=init_messages, inputs=inputs + messages=init_messages, + inputs=inputs, + root_conversation_id=caller_conv.root_conversation_id, ) return sub_agent_conversation @@ -519,6 +521,7 @@ async def _execute_flow( messages: MessageList, flow: Flow, inputs: Dict[str, Any], + root_conversation_id: Optional[str], ) -> Tuple[Any, str, ExecutionStatus]: """ Execute a flow and return its outputs and its execution status. @@ -531,6 +534,7 @@ async def _execute_flow( state.current_flow_conversation = flow.start_conversation( inputs=inputs, messages=messages, + root_conversation_id=root_conversation_id, ) messages.append_message( Message( @@ -593,7 +597,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 +750,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 +759,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.root_conversation_id, + ) ) logger.debug( diff --git a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py index eb572ff22..df3d79849 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py @@ -288,7 +288,7 @@ def create_sub_conversation( sub_conversation = flow.start_conversation( inputs_not_from_context_providers, - conversation_id=conversation.conversation_id, + root_conversation_id=conversation.root_conversation_id, messages=conversation.message_list, nesting_level=conversation.state.nesting_level + 1, context_providers_from_parent_flow=all_context_provider_keys, diff --git a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py index 980cebdff..a242caa65 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py @@ -26,11 +26,12 @@ class ManagerWorkersConversationExecutionState(ConversationExecutionState): current_agent_name: str subconversations: Dict[str, Union["AgentConversation", "ManagerWorkersConversation"]] + root_conversation_id: str = "" def _create_subconversation_for_agent( self, agent: Union[Agent, ManagerWorkers] ) -> Union["AgentConversation", "ManagerWorkersConversation"]: - subconv = agent.start_conversation() + subconv = agent.start_conversation(root_conversation_id=self.root_conversation_id or None) self.subconversations[agent.name] = subconv return subconv diff --git a/wayflowcore/src/wayflowcore/executors/interrupts/tokenlimitexecutioninterrupt.py b/wayflowcore/src/wayflowcore/executors/interrupts/tokenlimitexecutioninterrupt.py index e96035453..fa7ce580e 100644 --- a/wayflowcore/src/wayflowcore/executors/interrupts/tokenlimitexecutioninterrupt.py +++ b/wayflowcore/src/wayflowcore/executors/interrupts/tokenlimitexecutioninterrupt.py @@ -126,7 +126,7 @@ def _return_status_if_condition_is_met( self, state: ConversationExecutionState, conversation: "Conversation" ) -> Optional[InterruptedExecutionStatus]: - conversation_id = conversation.conversation_id + conversation_id = conversation.root_conversation_id # We first check the global token limit, then we go over the llm-wise limits # Note that we must do the checks separately, because the list of all models diff --git a/wayflowcore/src/wayflowcore/models/llmmodel.py b/wayflowcore/src/wayflowcore/models/llmmodel.py index 27010c45b..33d52bff5 100644 --- a/wayflowcore/src/wayflowcore/models/llmmodel.py +++ b/wayflowcore/src/wayflowcore/models/llmmodel.py @@ -359,11 +359,11 @@ def _update_token_usage( if isinstance(conversation, FlowConversation): # generate with flow - self.token_usages_flow[conversation.conversation_id][ + self.token_usages_flow[conversation.root_conversation_id][ conversation.current_step_name ] += token_usage else: - self.token_usages_flexible[conversation.conversation_id] += token_usage + self.token_usages_flexible[conversation.root_conversation_id] += token_usage def get_total_token_consumption(self, conversation_id: str) -> TokenUsage: """Calculate and return the total token consumption for a given conversation. diff --git a/wayflowcore/src/wayflowcore/serialization/context.py b/wayflowcore/src/wayflowcore/serialization/context.py index a86d6cf17..232215778 100644 --- a/wayflowcore/src/wayflowcore/serialization/context.py +++ b/wayflowcore/src/wayflowcore/serialization/context.py @@ -312,6 +312,21 @@ def _add_component_to_context(self, component: "Component") -> None: """ from wayflowcore.component import Component + def _iter_nested_components(value: Any) -> List["Component"]: + if isinstance(value, Component): + return [value] + if isinstance(value, dict): + nested_components: List["Component"] = [] + for nested_value in value.values(): + nested_components.extend(_iter_nested_components(nested_value)) + return nested_components + if isinstance(value, (list, tuple, set)): + nested_components = [] + for nested_value in value: + nested_components.extend(_iter_nested_components(nested_value)) + return nested_components + return [] + component_ref = SerializationContext.get_reference(component) if component_ref in self._deserialized_objects: @@ -322,14 +337,6 @@ def _add_component_to_context(self, component: "Component") -> None: 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) + for attr in all_public_attrs.values(): + for nested_component in _iter_nested_components(attr): + self._add_component_to_context(nested_component) diff --git a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py index 893877282..25c4efe58 100644 --- a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py +++ b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py @@ -294,7 +294,9 @@ def _get_or_create_agent_subconversation( caller_conv.message_list if self._share_conversation else MessageList.from_messages([]) ) agent_sub_conversation = self.agent.start_conversation( - inputs=inputs, messages=init_messages + inputs=inputs, + messages=init_messages, + root_conversation_id=caller_conv.root_conversation_id, ) return agent_sub_conversation diff --git a/wayflowcore/src/wayflowcore/tools/servertools.py b/wayflowcore/src/wayflowcore/tools/servertools.py index d6216c97f..3d6bd6daf 100644 --- a/wayflowcore/src/wayflowcore/tools/servertools.py +++ b/wayflowcore/src/wayflowcore/tools/servertools.py @@ -639,6 +639,7 @@ async def __call__(self, **inputs: Any) -> Any: conversation = self.flow.start_conversation( inputs=inputs, messages=self._parent_conversation.message_list, + root_conversation_id=self._parent_conversation.root_conversation_id, ) interrupts = self._parent_conversation._get_interrupts() diff --git a/wayflowcore/tests/test_managerworkers.py b/wayflowcore/tests/test_managerworkers.py index f2e5d601d..040933d3a 100644 --- a/wayflowcore/tests/test_managerworkers.py +++ b/wayflowcore/tests/test_managerworkers.py @@ -184,14 +184,14 @@ def test_managerworkers_can_execute_with_initial_params_passed_in_start_conversa conversation = group.start_conversation( messages=[Message(content="Please compute 3*4 + 2", message_type=MessageType.USER)], inputs={"USER": "Iris"}, - conversation_id="12345", + root_conversation_id="12345", ) conversation.execute() # The first message must be not the default message as the init messages are passed. assert conversation.get_last_message().content != DEFAULT_INITIAL_MESSAGE - assert conversation.conversation_id == "12345" + assert conversation.root_conversation_id == "12345" @retry_test(max_attempts=2) diff --git a/wayflowcore/tests/test_swarm.py b/wayflowcore/tests/test_swarm.py index 263b73b17..f52e28fca 100644 --- a/wayflowcore/tests/test_swarm.py +++ b/wayflowcore/tests/test_swarm.py @@ -226,14 +226,14 @@ def test_can_execute_swarm_with_initial_params_passed_in_start_conversation( ) ], inputs={"USER": "Iris"}, - conversation_id="12345", + root_conversation_id="12345", ) conversation.execute() # The first message must be not the default message as the init messages are passed. assert conversation.get_last_message().content != "Hi! How can I help you?" - assert conversation.conversation_id == "12345" + assert conversation.root_conversation_id == "12345" def test_can_create_swarm(example_medical_agents): From fbd9d1ef7914dac6e4500e8a9a7e4c83eff3f930 Mon Sep 17 00:00:00 2001 From: jschweiz Date: Tue, 21 Apr 2026 15:18:18 +0200 Subject: [PATCH 02/10] [fix]: fix tests --- wayflowcore/src/wayflowcore/a2a/a2aagent.py | 53 +- wayflowcore/src/wayflowcore/agent.py | 49 +- .../services/wayflowservice.py | 22 +- .../src/wayflowcore/checkpointing/__init__.py | 27 + .../wayflowcore/checkpointing/checkpointer.py | 155 +++++ .../checkpointing/checkpointeventlistener.py | 178 +++++ .../checkpointing/datastorecheckpointer.py | 399 +++++++++++ .../checkpointing/serialization.py | 143 ++++ .../contextproviders/flowcontextprovider.py | 2 +- wayflowcore/src/wayflowcore/conversation.py | 55 +- .../wayflowcore/conversationalcomponent.py | 132 ++-- wayflowcore/src/wayflowcore/events/event.py | 4 +- .../wayflowcore/executors/_agentexecutor.py | 4 +- .../wayflowcore/executors/_flowexecutor.py | 2 +- .../executors/_managerworkersconversation.py | 2 +- .../executors/_swarmconversation.py | 2 + wayflowcore/src/wayflowcore/flow.py | 58 +- wayflowcore/src/wayflowcore/managerworkers.py | 56 +- wayflowcore/src/wayflowcore/ociagent.py | 47 +- .../src/wayflowcore/serialization/context.py | 17 +- .../wayflowcore/serialization/serializer.py | 34 +- .../wayflowcore/steps/agentexecutionstep.py | 2 +- wayflowcore/src/wayflowcore/swarm.py | 57 +- .../src/wayflowcore/tools/servertools.py | 2 +- wayflowcore/src/wayflowcore/tracing/span.py | 2 +- .../steps/test_prompt_execution_step.py | 2 +- .../test_conversation_checkpointing.py | 654 ++++++++++++++++++ wayflowcore/tests/test_managerworkers.py | 2 +- wayflowcore/tests/test_swarm.py | 2 +- .../tracing/spans/test_conversation_span.py | 2 +- 30 files changed, 2015 insertions(+), 151 deletions(-) create mode 100644 wayflowcore/src/wayflowcore/checkpointing/__init__.py create mode 100644 wayflowcore/src/wayflowcore/checkpointing/checkpointer.py create mode 100644 wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py create mode 100644 wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py create mode 100644 wayflowcore/src/wayflowcore/checkpointing/serialization.py create mode 100644 wayflowcore/tests/serialization/test_conversation_checkpointing.py diff --git a/wayflowcore/src/wayflowcore/a2a/a2aagent.py b/wayflowcore/src/wayflowcore/a2a/a2aagent.py index 4dee6ae11..7a654ad3f 100644 --- a/wayflowcore/src/wayflowcore/a2a/a2aagent.py +++ b/wayflowcore/src/wayflowcore/a2a/a2aagent.py @@ -19,6 +19,7 @@ from wayflowcore.tools import Tool if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.executors._a2aagentconversation import A2AAgentConversation logger = logging.getLogger(__name__) @@ -248,43 +249,71 @@ 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, + _root_conversation_id: Optional[str] = None, + _attach_checkpointer: bool = True, ) -> "A2AAgentConversation": """ - Initiates a new conversation with the remote server agent. - - Creates and returns a conversation instance tied to this agent, optionally initialized - with input data and a message history. + Start a conversation with the remote A2A agent. Parameters ---------- inputs: - Optional dictionary of initial input data for the conversation. Defaults to an empty - dictionary if not provided. + Optional structured inputs stored on the conversation for interface compatibility. + The A2A runtime currently executes from messages rather than these inputs. 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. + Optional initial message history for the remote conversation. + conversation_id: + Optional identifier for this A2A conversation. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires ``checkpointer``. + _root_conversation_id: + Internal lineage identifier shared with nested or parent conversations. Returns ------- - Conversation: - A new conversation object associated with this agent. + A2AAgentConversation + A new or restored A2A agent conversation. """ from wayflowcore.executors._a2aagentconversation import A2AAgentConversation from wayflowcore.executors._a2aagentexecutor import A2AAgentState + restored_conversation, conversation_runtime_id, conversation_root_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + _root_conversation_id=_root_conversation_id, + expected_conversation_type=A2AAgentConversation, + attach_checkpointer=_attach_checkpointer, + ) + ) + if restored_conversation is not None: + return restored_conversation + if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) - return A2AAgentConversation( + conversation = A2AAgentConversation( component=self, state=A2AAgentState(last_message_idx=-1), inputs=inputs or {}, # Inputs are ignored in execution message_list=messages, status=None, - conversation_id=IdGenerator.get_or_generate_id(None), + id=conversation_runtime_id, + checkpointer=checkpointer, name="a2a_conversation", + root_conversation_id=conversation_root_id, __metadata_info__={}, ) + return conversation @property def agent_id(self) -> str: diff --git a/wayflowcore/src/wayflowcore/agent.py b/wayflowcore/src/wayflowcore/agent.py index e19973cfc..d7f05f452 100644 --- a/wayflowcore/src/wayflowcore/agent.py +++ b/wayflowcore/src/wayflowcore/agent.py @@ -26,6 +26,7 @@ from wayflowcore.transforms import MessageTransform if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.contextproviders import ContextProvider from wayflowcore.executors._agentconversation import AgentConversation from wayflowcore.flow import Flow @@ -393,29 +394,54 @@ 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, + _root_conversation_id: Optional[str] = None, + _attach_checkpointer: bool = True, ) -> "AgentConversation": """ - Initializes a conversation with the agent. + Start a conversation with the agent. Parameters ---------- inputs: - This argument is not used. - It is included for compatibility with the Flow class. + Optional input values for the agent's declared input descriptors. messages: - Message list to which the agent will participate + Optional message history for the conversation. conversation_id: - Conversation id of the parent conversation. + Optional identifier for this agent conversation. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires ``checkpointer``. + _root_conversation_id: + Internal lineage identifier shared with nested or parent conversations. Returns ------- - Conversation: - The conversation object of the agent. + AgentConversation + A new or restored agent conversation. """ from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event from wayflowcore.executors._agentconversation import AgentConversation + restored_conversation, conversation_runtime_id, conversation_root_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + _root_conversation_id=_root_conversation_id, + expected_conversation_type=AgentConversation, + attach_checkpointer=_attach_checkpointer, + ) + ) + if restored_conversation is not None: + return restored_conversation + if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) @@ -458,23 +484,26 @@ def start_conversation( conversational_component=self, inputs=inputs, messages=messages, - conversation_id=conversation_id, + conversation_id=conversation_runtime_id, nesting_level=None, ) ) from wayflowcore.executors._agentexecutor import AgentConversationExecutionState - return AgentConversation( + conversation = AgentConversation( component=self, message_list=messages, - conversation_id=IdGenerator.get_or_generate_id(conversation_id), + id=conversation_runtime_id, + checkpointer=checkpointer, inputs=inputs or {}, name="agent_conversation", state=AgentConversationExecutionState(), status=None, + root_conversation_id=conversation_root_id, __metadata_info__={}, ) + return conversation @property def llms(self) -> List["LlmModel"]: diff --git a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py index 2c36761ce..e71d41085 100644 --- a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py +++ b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py @@ -15,10 +15,6 @@ from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig from wayflowcore.checkpointing import ConversationCheckpoint, DatastoreCheckpointer -from wayflowcore.checkpointing.runtime import ( - _detach_checkpointer_from_conversation, - _set_conversation_final_checkpoint_overrides, -) from wayflowcore.conversation import Conversation from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.datastore import Datastore, InMemoryDatastore @@ -197,10 +193,13 @@ 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 + should_store_response = body.store is None or body.store is True + state = self._load_state( previous_response_id=previous_response_id, conversation_id=conversation_id, agent_id=model, + attach_checkpointer=should_store_response, ) response_id = IdGenerator.get_or_generate_id() @@ -209,10 +208,6 @@ async def create_response(self, body: CreateResponse) -> AsyncIterable[ResponseS state=state, request=body, ) - if body.store is None or body.store is True: - _set_conversation_final_checkpoint_overrides(state, checkpoint_id=response_id) - else: - _detach_checkpointer_from_conversation(state) current_response = Response( id=response_id, @@ -266,7 +261,9 @@ async def runner(conversation: Conversation) -> None: nonlocal status try: with register_event_listeners([token_usage_listener, yielding_listener]): - status = await conversation.execute_async() + status = await conversation.execute_async( + _final_checkpoint_id=response_id if should_store_response else None, + ) except Exception as e: nonlocal raised_exception raised_exception = e @@ -319,7 +316,7 @@ async def runner(conversation: Conversation) -> None: token_usage_listener.usage ) - if (body.store is None or body.store is True) and state.checkpointer is not None: + if should_store_response and state.checkpointer is not None: self.checkpointer.save_conversation( state, checkpoint_id=current_response.id, @@ -379,6 +376,7 @@ def _load_state( previous_response_id: Optional[str], conversation_id: Optional[str], agent_id: str, + attach_checkpointer: bool = True, ) -> Optional[Conversation]: if previous_response_id: checkpoint = self._lookup_checkpoint_by_response_id(previous_response_id) @@ -393,6 +391,7 @@ def _load_state( conversation_id=checkpoint.conversation_id, checkpoint_id=checkpoint.checkpoint_id, checkpointer=self.checkpointer, + _attach_checkpointer=attach_checkpointer, ) except (TypeError, ValueError) as e: raise HTTPException( @@ -411,6 +410,7 @@ def _load_state( return self.agents[agent_id].start_conversation( conversation_id=conversation_id, checkpointer=self.checkpointer, + _attach_checkpointer=attach_checkpointer, ) except (TypeError, ValueError) as e: raise HTTPException( @@ -473,8 +473,6 @@ async def _create_state( raise NotImplementedError( "Instructions are only supported when creating a conversation" ) - if request.store is False: - _detach_checkpointer_from_conversation(state) # Add the new messages to the conversation for message in new_messages: state.append_message(message) diff --git a/wayflowcore/src/wayflowcore/checkpointing/__init__.py b/wayflowcore/src/wayflowcore/checkpointing/__init__.py new file mode 100644 index 000000000..7db08002e --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/__init__.py @@ -0,0 +1,27 @@ +# Copyright © 2025, 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 importlib import import_module +from typing import Any + +from .checkpointer import Checkpointer, CheckpointingInterval, ConversationCheckpoint, StorageConfig +from .datastorecheckpointer import ( + DatastoreCheckpointer, + InMemoryCheckpointer, + OracleDatabaseCheckpointer, + PostgresCheckpointer, +) + +__all__ = [ + "CheckpointingInterval", + "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..0ecafb48f --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py @@ -0,0 +1,155 @@ +# Copyright © 2025, 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 time +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from wayflowcore.idgeneration import IdGenerator + +from .serialization import _serialize_conversation_checkpoint_state + +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 + conversation_id: str + component_id: str + created_at: int + state: str + metadata: Dict[str, Any] = field(default_factory=dict) + + @property + def id(self) -> str: + return self.checkpoint_id + + +class CheckpointingInterval(Enum): + """ + Configure when the conversation is saved during execution. + """ + + CONVERSATION_TURNS = "conversation_turns" + LLM_TURNS = "llm_turns" + ALL_INTERNAL_TURNS = "all_internal_turns" + + +@dataclass +class StorageConfig: + """Configuration for checkpoint storage.""" + + datastore: Optional["Datastore"] = None + table_name: str = "conversations" + agent_id_column_name: str = "agent_id" + conversation_id_column_name: str = "conversation_id" + turn_id_column_name: str = "turn_id" + created_at_column_name: str = "created_at" + remove_by_column_name: str = "remove_by" + conversation_turn_state_column_name: str = "conversation_turn_state" + is_last_turn_column_name: str = "is_last_turn" + extra_metadata_column_name: str = "extra_metadata" + max_retention: Optional[int] = None + + 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: Any) -> None: + from wayflowcore.conversation import Conversation + + if isinstance(checkpoint, Conversation): + self.save_conversation(checkpoint) + return + if not isinstance(checkpoint, ConversationCheckpoint): + raise TypeError( + f"Expected a Conversation or ConversationCheckpoint, got {type(checkpoint).__name__}." + ) + self._save_checkpoint(checkpoint) + + async def save_async(self, checkpoint: Any) -> None: + self.save(checkpoint) + + def save_conversation( + self, + conversation: "Conversation", + *, + checkpoint_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> ConversationCheckpoint: + next_save_sequence = self._save_sequence_by_conversation.get(conversation.id, 0) + 1 + self._save_sequence_by_conversation[conversation.id] = next_save_sequence + checkpoint_metadata = {"save_sequence": next_save_sequence} + if metadata: + checkpoint_metadata.update(metadata) + checkpoint = ConversationCheckpoint( + checkpoint_id=checkpoint_id or IdGenerator.get_or_generate_id(), + conversation_id=conversation.id, + component_id=conversation.component.id, + created_at=int(time.time()), + state=_serialize_conversation_checkpoint_state(conversation), + metadata=checkpoint_metadata, + ) + self._save_checkpoint(checkpoint) + conversation.checkpoint_id = checkpoint.checkpoint_id + return checkpoint + + @abstractmethod + def _save_checkpoint(self, checkpoint: ConversationCheckpoint) -> None: + raise NotImplementedError() + + @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..af750a3b2 --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py @@ -0,0 +1,178 @@ +# Copyright © 2025, 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, nullcontext +from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional + +from ..events import EventListener +from .checkpointer import CheckpointingInterval + +if TYPE_CHECKING: + from wayflowcore.conversation import Conversation + + +def _build_checkpoint_metadata( + conversation: "Conversation", + *, + save_reason: str, + event: Optional[Any] = None, + metadata: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + 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__ + execution_state = getattr(event, "execution_state", None) + if execution_state is not None: + if hasattr(execution_state, "curr_iter"): + checkpoint_metadata["agent_iteration"] = execution_state.curr_iter + if hasattr(execution_state, "current_step_name"): + checkpoint_metadata["flow_step_name"] = execution_state.current_step_name + if hasattr(execution_state, "nesting_level"): + checkpoint_metadata["nesting_level"] = execution_state.nesting_level + if metadata: + checkpoint_metadata.update(metadata) + return checkpoint_metadata + + +def _save_conversation_checkpoint( + conversation: "Conversation", + *, + save_reason: str, + event: Optional[Any] = None, + metadata: Optional[Dict[str, Any]] = None, + checkpoint_id: Optional[str] = None, +) -> None: + checkpointer = conversation.checkpointer + if checkpointer is None: + return + + checkpoint_metadata = _build_checkpoint_metadata( + conversation, + save_reason=save_reason, + event=event, + metadata=metadata, + ) + + checkpointer.save_conversation( + conversation, + checkpoint_id=checkpoint_id, + metadata=checkpoint_metadata, + ) + + +class _ConversationCheckpointEventListener(EventListener): + def __init__(self, conversation: "Conversation") -> None: + self.conversation = conversation + self._llm_was_used_since_last_internal_turn = False + self._last_internal_turn_start_event: Optional[Any] = None + + def __call__(self, event: Any) -> None: + from wayflowcore.events.event import ( + AgentExecutionIterationStartedEvent, + FlowExecutionIterationStartedEvent, + LlmGenerationResponseEvent, + ) + + checkpointer = self.conversation.checkpointer + if checkpointer is None: + return + + if isinstance(event, LlmGenerationResponseEvent): + self._llm_was_used_since_last_internal_turn = True + return + + if not isinstance( + event, (AgentExecutionIterationStartedEvent, FlowExecutionIterationStartedEvent) + ): + return + + checkpointing_interval = checkpointer.checkpointing_interval + + if checkpointing_interval == CheckpointingInterval.CONVERSATION_TURNS: + self._last_internal_turn_start_event = event + return + + if checkpointing_interval == CheckpointingInterval.ALL_INTERNAL_TURNS: + _save_conversation_checkpoint( + self.conversation, + save_reason="internal_turn_boundary", + event=event, + metadata={ + "llm_used_in_previous_turn": self._llm_was_used_since_last_internal_turn, + }, + ) + self._llm_was_used_since_last_internal_turn = False + self._last_internal_turn_start_event = event + return + + if self._llm_was_used_since_last_internal_turn: + _save_conversation_checkpoint( + self.conversation, + save_reason="internal_turn_boundary", + event=event, + metadata={ + "llm_used_in_previous_turn": self._llm_was_used_since_last_internal_turn, + }, + ) + self._llm_was_used_since_last_internal_turn = False + + self._last_internal_turn_start_event = event + + def flush_pending_checkpoint(self) -> None: + checkpointer = self.conversation.checkpointer + if checkpointer is None: + return + if checkpointer.checkpointing_interval != CheckpointingInterval.LLM_TURNS: + return + if not self._llm_was_used_since_last_internal_turn: + return + + _save_conversation_checkpoint( + self.conversation, + save_reason="internal_turn_boundary", + event=self._last_internal_turn_start_event, + metadata={ + "llm_used_in_previous_turn": True, + }, + ) + self._llm_was_used_since_last_internal_turn = False + + +@contextmanager +def get_conversation_checkpoint_execution_context( + conversation: "Conversation", + *, + is_outermost_execution: bool, + final_checkpoint_id: Optional[str] = None, + final_checkpoint_metadata: Optional[Dict[str, Any]] = None, +) -> Iterator[None]: + if conversation.checkpointer is None or not is_outermost_execution: + with nullcontext(): + yield + return + + from wayflowcore.events.eventlistener import register_event_listeners + + listener = _ConversationCheckpointEventListener(conversation) + with register_event_listeners([listener]): + try: + yield + except Exception: + raise + else: + listener.flush_pending_checkpoint() + _save_conversation_checkpoint( + conversation, + save_reason="conversation_turn", + checkpoint_id=final_checkpoint_id, + metadata=final_checkpoint_metadata, + ) diff --git a/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py new file mode 100644 index 000000000..27bb5a2bb --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py @@ -0,0 +1,399 @@ +# Copyright © 2025, 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 json +import warnings +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.inmemory import _INMEMORY_USER_WARNING +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 _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) + 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) + 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: + 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]: + return sorted( + checkpoints, + key=lambda checkpoint: ( + checkpoint.created_at, + 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 + 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: + 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: + 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() + else: + 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 and len(checkpoints) > 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 + if latest_checkpoint is not None and latest_checkpoint.checkpoint_id == checkpoint_id: + remaining_checkpoints = [ + checkpoint + for checkpoint in checkpoints + if checkpoint.checkpoint_id != checkpoint_id + ] + checkpoint_to_promote = remaining_checkpoints[-1] if remaining_checkpoints else None + + self.datastore.delete( + 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, + }, + ) + + 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() + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=f"{_INMEMORY_USER_WARNING}*") + 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/checkpointing/serialization.py b/wayflowcore/src/wayflowcore/checkpointing/serialization.py new file mode 100644 index 000000000..d757a3bc0 --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/serialization.py @@ -0,0 +1,143 @@ +# Copyright © 2025, 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 TYPE_CHECKING, Any, Dict, List, Optional, Sequence, cast + +import yaml + +from wayflowcore.serialization import autodeserialize, serialize_to_dict +from wayflowcore.serialization.context import DeserializationContext, SerializationContext +from wayflowcore.serialization.serializer import autodeserialize_from_dict + +if TYPE_CHECKING: + from wayflowcore.component import Component + from wayflowcore.conversation import Conversation + + +_CHECKPOINT_ENVELOPE_FORMAT = "wayflow-conversation-checkpoint" +_CHECKPOINT_ENVELOPE_VERSION = 1 + + +def _iter_conversation_graph(root_conversation: "Conversation") -> Sequence["Conversation"]: + visited_conversation_ids: set[str] = set() + queue: List["Conversation"] = [root_conversation] + ordered_conversations: List["Conversation"] = [] + + while queue: + conversation = queue.pop() + if conversation.id in visited_conversation_ids: + continue + visited_conversation_ids.add(conversation.id) + ordered_conversations.append(conversation) + queue.extend(conversation._get_all_sub_conversations()) + + return ordered_conversations + + +def _ensure_checkpointing_supported(conversation: "Conversation") -> None: + from wayflowcore.ociagent import OciAgent + + for sub_conversation in _iter_conversation_graph(conversation): + if isinstance(sub_conversation.component, OciAgent): + raise NotImplementedError( + "Checkpointing conversations that contain `OciAgent` is not supported yet." + ) + + +def _iter_component_tree(component: "Component") -> Sequence["Component"]: + from wayflowcore.component import Component + + def _iter_nested_components(value: Any) -> List["Component"]: + if isinstance(value, Component): + return [value] + if isinstance(value, dict): + nested_components: List["Component"] = [] + for nested_value in value.values(): + nested_components.extend(_iter_nested_components(nested_value)) + return nested_components + if isinstance(value, (list, tuple, set)): + nested_components = [] + for nested_value in value: + nested_components.extend(_iter_nested_components(nested_value)) + return nested_components + return [] + + visited_component_ids: set[str] = set() + ordered_components: List["Component"] = [] + queue: List["Component"] = [component] + + while queue: + current_component = queue.pop() + current_component_ref = SerializationContext.get_reference(current_component) + if current_component_ref in visited_component_ids: + continue + visited_component_ids.add(current_component_ref) + ordered_components.append(current_component) + + all_public_attrs = { + name: value + for name, value in vars(current_component).items() + if not name.startswith("_") + } + for attr in all_public_attrs.values(): + queue.extend(_iter_nested_components(attr)) + + return ordered_components + + +def _build_checkpoint_serialization_context(conversation: "Conversation") -> SerializationContext: + serialization_context = SerializationContext(root=conversation) + for component in _iter_component_tree(conversation.component): + serialization_context.register_external_reference(component) + return serialization_context + + +def _serialize_conversation_checkpoint_state(conversation: "Conversation") -> str: + _ensure_checkpointing_supported(conversation) + + serialized_conversation = serialize_to_dict( + conversation, + serialization_context=_build_checkpoint_serialization_context(conversation), + ) + + envelope = { + "checkpoint_format": _CHECKPOINT_ENVELOPE_FORMAT, + "version": _CHECKPOINT_ENVELOPE_VERSION, + "conversation": serialized_conversation, + } + return yaml.safe_dump(envelope) + + +def _deserialize_conversation_checkpoint_state( + serialized_state: str, + *, + tool_registry: Optional[Dict[str, Any]] = None, + component: Optional["Component"] = None, +) -> "Conversation": + deserialization_context = DeserializationContext() + deserialization_context.registered_tools = tool_registry.copy() if tool_registry else {} + + if component is not None: + deserialization_context._add_component_to_context(component) + + state_payload = yaml.safe_load(serialized_state) + if ( + isinstance(state_payload, dict) + and state_payload.get("checkpoint_format") == _CHECKPOINT_ENVELOPE_FORMAT + and state_payload.get("version") == _CHECKPOINT_ENVELOPE_VERSION + and "conversation" in state_payload + ): + conversation = autodeserialize_from_dict( + state_payload["conversation"], + deserialization_context=deserialization_context, + ) + else: + conversation = autodeserialize( + serialized_state, + deserialization_context=deserialization_context, + ) + + return cast("Conversation", conversation) diff --git a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py index bd622d0ca..0bc90f4d1 100644 --- a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py +++ b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py @@ -95,7 +95,7 @@ async def call_async(self, conversation: "Conversation") -> Any: conversation = self.flow.start_conversation( inputs={}, messages=conversation.message_list, - root_conversation_id=conversation.root_conversation_id, + _root_conversation_id=conversation.root_conversation_id, ) status = await conversation.execute_async() if status._requires_yielding: diff --git a/wayflowcore/src/wayflowcore/conversation.py b/wayflowcore/src/wayflowcore/conversation.py index 14fa440c2..c2c2dadd2 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 @@ -62,6 +63,10 @@ def _get_active_conversations(return_copy: bool = True) -> List["Conversation"]: return copy(active_conversations) if return_copy else active_conversations +def is_outermost_execution() -> bool: + return len(_get_active_conversations(return_copy=False)) == 0 + + def _get_current_conversation_id() -> Optional[str]: active_conversations = _get_active_conversations(return_copy=True) if not active_conversations: @@ -85,13 +90,20 @@ 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 + root_conversation_id: str = "" + checkpointer: Optional["Checkpointer"] = field( + default=None, + repr=False, + compare=False, + metadata={"serialize": False}, + ) + checkpoint_id: Optional[str] = field(default=None, init=False, repr=False, compare=False) status_handled: bool = False """Whether the current status associated to this conversation was already handled or not @@ -100,6 +112,8 @@ class Conversation(DataclassComponent): def __post_init__(self) -> None: if self.inputs is None: self.inputs = {} + if not self.root_conversation_id: + self.root_conversation_id = self.id @property def plan(self) -> Optional[ExecutionPlan]: @@ -114,6 +128,9 @@ def _register_event(self, event: Event) -> None: def execute( self, execution_interrupts: Optional[Sequence["ExecutionInterrupt"]] = None, + *, + _final_checkpoint_id: Optional[str] = None, + _final_checkpoint_metadata: Optional[Dict[str, Any]] = None, ) -> "ExecutionStatus": """ Execute the conversation and get its ``ExecutionStatus`` based on the outcome. @@ -121,13 +138,22 @@ def execute( The ``Execution`` status is returned by the Assistant and indicates if the assistant yielded, finished the conversation. """ - return run_async_in_sync( - self.execute_async, execution_interrupts, method_name="execute_async" - ) + + async def _execute_async_wrapper() -> "ExecutionStatus": + return await self.execute_async( + execution_interrupts, + _final_checkpoint_id=_final_checkpoint_id, + _final_checkpoint_metadata=_final_checkpoint_metadata, + ) + + return run_async_in_sync(_execute_async_wrapper, method_name="execute_async") async def execute_async( self, execution_interrupts: Optional[Sequence["ExecutionInterrupt"]] = None, + *, + _final_checkpoint_id: Optional[str] = None, + _final_checkpoint_metadata: Optional[Dict[str, Any]] = None, ) -> "ExecutionStatus": """ Execute the conversation and get its ``ExecutionStatus`` based on the outcome. @@ -138,11 +164,20 @@ 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=is_outermost_execution(), + final_checkpoint_id=_final_checkpoint_id, + final_checkpoint_metadata=_final_checkpoint_metadata, + ): + with _register_conversation(self): + new_status = await self.component.runner.execute_async(self, execution_interrupts) + self.status = new_status + self.status_handled = False return self.status @property diff --git a/wayflowcore/src/wayflowcore/conversationalcomponent.py b/wayflowcore/src/wayflowcore/conversationalcomponent.py index 2d15b20c2..ee4ae5105 100644 --- a/wayflowcore/src/wayflowcore/conversationalcomponent.py +++ b/wayflowcore/src/wayflowcore/conversationalcomponent.py @@ -7,11 +7,20 @@ import logging from abc import ABC, abstractmethod from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, Set, Type, TypeVar, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Generic, + List, + Optional, + Set, + Type, + TypeVar, + Union, +) from wayflowcore._metadata import MetadataType -from wayflowcore.checkpointing.runtime import _attach_checkpointer_to_conversation -from wayflowcore.checkpointing.serialization import _deserialize_conversation_checkpoint_state from wayflowcore.componentwithio import ComponentWithInputsOutputs from wayflowcore.idgeneration import IdGenerator from wayflowcore.property import Property @@ -21,6 +30,7 @@ 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 @@ -28,6 +38,7 @@ from wayflowcore.tools import Tool _HUMAN_ENTITY_ID = "human_user" +ConversationTypeT = TypeVar("ConversationTypeT", bound="Conversation") class ConversationalComponent(ComponentWithInputsOutputs, ABC): @@ -72,11 +83,36 @@ def start_conversation( messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, conversation_id: Optional[str] = None, *, - root_conversation_id: Optional[str] = None, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, + _root_conversation_id: Optional[str] = None, + _attach_checkpointer: bool = True, ) -> "Conversation": - pass + """ + Start a conversation for this component. + + Parameters + ---------- + inputs: + Optional structured inputs used to initialize the conversation. + messages: + Optional initial message history. Concrete implementations normalize this into a + ``MessageList`` when needed. + conversation_id: + Optional identifier for the concrete conversation instance. + checkpointer: + Optional checkpoint backend used to restore and persist conversation state. + checkpoint_id: + Optional checkpoint identifier to restore. Requires ``checkpointer``. + _root_conversation_id: + Internal lineage identifier shared by nested conversations for usage accounting, + execution limits, and checkpoint lineage. + + Returns + ------- + Conversation + A new or restored conversation instance ready for execution. + """ @property def llms(self) -> List["LlmModel"]: @@ -137,33 +173,37 @@ def _messages_or_inputs_were_passed( return len(messages) > 0 if isinstance(messages, Message): return True - return len(message) > 0 + return len(messages) > 0 - def _restore_or_prepare_checkpoint_conversation( + def _prepare_conversation_start( self, *, inputs: Optional[Dict[str, Any]], messages: Union[None, str, "Message", List["Message"], "MessageList"], conversation_id: Optional[str], - root_conversation_id: Optional[str], + _root_conversation_id: Optional[str], checkpointer: Optional["Checkpointer"], checkpoint_id: Optional[str], - ) -> tuple[Optional["Conversation"], Optional[str]]: + expected_conversation_type: Type[ConversationTypeT], + attach_checkpointer: bool, + ) -> tuple[Optional[ConversationTypeT], str, str]: if checkpointer is None: if checkpoint_id is not None: raise ValueError("`checkpoint_id` requires a `checkpointer`.") - return None, conversation_id + + runtime_conversation_id = IdGenerator.get_or_generate_id(conversation_id) + return None, runtime_conversation_id, _root_conversation_id or runtime_conversation_id if ( - root_conversation_id is not None + _root_conversation_id is not None and conversation_id is not None - and root_conversation_id != conversation_id + and _root_conversation_id != conversation_id ): raise ValueError( "`root_conversation_id` and `conversation_id` cannot differ when checkpointing is enabled." ) - resolved_conversation_id = conversation_id or root_conversation_id + resolved_conversation_id = conversation_id or _root_conversation_id if resolved_conversation_id is None and checkpoint_id is not None: raise ValueError("`checkpoint_id` requires a `conversation_id`.") if resolved_conversation_id is None: @@ -175,7 +215,7 @@ def _restore_or_prepare_checkpoint_conversation( else checkpointer.load_latest(resolved_conversation_id) ) if checkpoint is None: - return None, resolved_conversation_id + return None, resolved_conversation_id, resolved_conversation_id if self._messages_or_inputs_were_passed(inputs=inputs, messages=messages): raise ValueError( @@ -183,41 +223,47 @@ def _restore_or_prepare_checkpoint_conversation( "Load the conversation first, then append new user input explicitly." ) + conversation = self._restore_checkpointed_conversation( + checkpoint=checkpoint, + checkpointer=checkpointer, + expected_conversation_type=expected_conversation_type, + attach_checkpointer=attach_checkpointer, + ) + return conversation, resolved_conversation_id, resolved_conversation_id + + def _restore_checkpointed_conversation( + self, + *, + checkpoint: "ConversationCheckpoint", + checkpointer: "Checkpointer", + expected_conversation_type: Type[ConversationTypeT], + attach_checkpointer: bool, + ) -> ConversationTypeT: + from wayflowcore.checkpointing.serialization import ( + _deserialize_conversation_checkpoint_state, + ) + + if checkpoint.component_id != self.id: + raise ValueError( + "Cannot restore this checkpoint because this conversation was started with another " + f"component. Checkpoint component id: `{checkpoint.component_id}`. Current component id: `{self.id}`." + ) + conversation = _deserialize_conversation_checkpoint_state( checkpoint.state, tool_registry={tool.name: tool for tool in self._referenced_tools()}, component=self, ) - return ( - _attach_checkpointer_to_conversation( - conversation, - checkpointer=checkpointer, - checkpoint_id=checkpoint.checkpoint_id, - ), - resolved_conversation_id, - ) - - @staticmethod - def _resolve_runtime_and_root_conversation_ids( - *, - conversation_id: Optional[str], - root_conversation_id: Optional[str], - checkpointer: Optional["Checkpointer"], - restored_conversation_id: Optional[str], - ) -> tuple[str, str]: - if checkpointer is not None: - runtime_conversation_id = restored_conversation_id or IdGenerator.get_or_generate_id( - conversation_id or root_conversation_id + if not isinstance(conversation, expected_conversation_type): + raise ValueError( + "Cannot restore this checkpoint because this conversation was started with another " + f"component. Expected `{expected_conversation_type.__name__}`, got `{type(conversation).__name__}`." ) - resolved_root_conversation_id = root_conversation_id or runtime_conversation_id - if resolved_root_conversation_id != runtime_conversation_id: - raise ValueError( - "`root_conversation_id` and `conversation_id` cannot differ when checkpointing is enabled." - ) - return runtime_conversation_id, runtime_conversation_id - - runtime_conversation_id = IdGenerator.get_or_generate_id(conversation_id) - return runtime_conversation_id, root_conversation_id or runtime_conversation_id + + if attach_checkpointer: + conversation.checkpointer = checkpointer + conversation.checkpoint_id = checkpoint.checkpoint_id + return conversation # Define a TypeVar that represents the component's type diff --git a/wayflowcore/src/wayflowcore/events/event.py b/wayflowcore/src/wayflowcore/events/event.py index 4c9c39f1c..c6428cbfa 100644 --- a/wayflowcore/src/wayflowcore/events/event.py +++ b/wayflowcore/src/wayflowcore/events/event.py @@ -767,7 +767,7 @@ class ConversationExecutionStartedEvent(StartSpanEvent["ConversationSpan"]): def to_tracing_info(self, mask_sensitive_information: bool = True) -> Dict[str, Any]: return { **super().to_tracing_info(mask_sensitive_information=mask_sensitive_information), - "conversation.id": self.conversation.conversation_id, + "conversation.id": self.conversation.id, "conversation.name": self.conversation.name, } @@ -788,7 +788,7 @@ class ConversationExecutionFinishedEvent(EndSpanEvent["ConversationSpan"]): def to_tracing_info(self, mask_sensitive_information: bool = True) -> Dict[str, Any]: return { **super().to_tracing_info(mask_sensitive_information=mask_sensitive_information), - "conversation.id": self.conversation.conversation_id, + "conversation.id": self.conversation.id, "conversation.name": self.conversation.name, "execution_status": self.execution_status.__class__.__name__, } diff --git a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py index 86229d52d..48706fe8e 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py @@ -464,7 +464,7 @@ def _get_or_create_expert_agent_subconversation( sub_agent_conversation = expert_agent.start_conversation( messages=init_messages, inputs=inputs, - root_conversation_id=caller_conv.root_conversation_id, + _root_conversation_id=caller_conv.root_conversation_id, ) return sub_agent_conversation @@ -534,7 +534,7 @@ async def _execute_flow( state.current_flow_conversation = flow.start_conversation( inputs=inputs, messages=messages, - root_conversation_id=root_conversation_id, + _root_conversation_id=root_conversation_id, ) messages.append_message( Message( diff --git a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py index df3d79849..fa8d3bef8 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py @@ -288,7 +288,7 @@ def create_sub_conversation( sub_conversation = flow.start_conversation( inputs_not_from_context_providers, - root_conversation_id=conversation.root_conversation_id, + _root_conversation_id=conversation.root_conversation_id, messages=conversation.message_list, nesting_level=conversation.state.nesting_level + 1, context_providers_from_parent_flow=all_context_provider_keys, diff --git a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py index a242caa65..e298be4a7 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py @@ -31,7 +31,7 @@ class ManagerWorkersConversationExecutionState(ConversationExecutionState): def _create_subconversation_for_agent( self, agent: Union[Agent, ManagerWorkers] ) -> Union["AgentConversation", "ManagerWorkersConversation"]: - subconv = agent.start_conversation(root_conversation_id=self.root_conversation_id or None) + subconv = agent.start_conversation(_root_conversation_id=self.root_conversation_id or None) self.subconversations[agent.name] = subconv return subconv diff --git a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py index 1a169f271..234d2ebe1 100644 --- a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py @@ -55,6 +55,7 @@ class SwarmConversationExecutionState(ConversationExecutionState): main_thread: SwarmThread agents_and_threads: Dict[str, Dict[str, SwarmThread]] context_providers: List["ContextProvider"] + root_conversation_id: str = "" current_thread: Optional["SwarmThread"] = None thread_stack: List["SwarmThread"] = field(default_factory=list) @@ -89,6 +90,7 @@ def _create_subconversation_for_thread( conversation = thread.recipient_agent.start_conversation( inputs=inputs, messages=thread.message_list, + _root_conversation_id=self.root_conversation_id or None, ) self.thread_subconversations[thread_id] = conversation diff --git a/wayflowcore/src/wayflowcore/flow.py b/wayflowcore/src/wayflowcore/flow.py index 2c77ada85..9710ca7ce 100644 --- a/wayflowcore/src/wayflowcore/flow.py +++ b/wayflowcore/src/wayflowcore/flow.py @@ -45,6 +45,7 @@ from wayflowcore.tools import Tool if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.executors._flowconversation import FlowConversation from wayflowcore.executors._flowexecutor import _IoKeyType from wayflowcore.messagelist import Message @@ -1166,33 +1167,59 @@ def start_conversation( conversation_id: Optional[str] = None, nesting_level: int = 0, context_providers_from_parent_flow: Optional[Set[str]] = None, + *, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + _root_conversation_id: Optional[str] = None, + _attach_checkpointer: bool = True, ) -> "FlowConversation": """ - Start the conversation. + Start a conversation for this flow. Parameters ---------- inputs: - 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. + Optional input values used to initialize flow execution. 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. + Optional message history available to the flow at startup. + conversation_id: + Optional identifier for this flow conversation. nesting_level: - Nesting level of the conversation. + Nesting level of the flow execution. Nested subflows increase this value. + context_providers_from_parent_flow: + Names of inputs already provided by parent-flow context providers when validating + required inputs for nested execution. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires ``checkpointer``. + _root_conversation_id: + Internal lineage identifier shared with nested or parent conversations. Returns ------- - Conversation: - A Flow Conversation object. + FlowConversation + A new or restored flow conversation. """ from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event from wayflowcore.executors._flowconversation import FlowConversation + restored_conversation, conversation_runtime_id, conversation_root_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + _root_conversation_id=_root_conversation_id, + expected_conversation_type=FlowConversation, + attach_checkpointer=_attach_checkpointer, + ) + ) + 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 +1269,7 @@ def start_conversation( conversational_component=self, inputs=inputs, messages=messages, - conversation_id=conversation_id, + conversation_id=conversation_runtime_id, nesting_level=nesting_level, ) ) @@ -1271,16 +1298,19 @@ def start_conversation( nesting_level=nesting_level, ) - return FlowConversation( + conversation = FlowConversation( component=self, inputs=inputs, - conversation_id=IdGenerator.get_or_generate_id(conversation_id), + id=conversation_runtime_id, + checkpointer=checkpointer, message_list=messages, __metadata_info__={}, status=None, name="flow_conversation", state=state, + root_conversation_id=conversation_root_id, ) + return conversation @property def llms(self) -> List["LlmModel"]: diff --git a/wayflowcore/src/wayflowcore/managerworkers.py b/wayflowcore/src/wayflowcore/managerworkers.py index f59ace30b..a64c47fd5 100644 --- a/wayflowcore/src/wayflowcore/managerworkers.py +++ b/wayflowcore/src/wayflowcore/managerworkers.py @@ -22,6 +22,7 @@ from wayflowcore.transforms import MessageTransform if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer from wayflowcore.executors._managerworkersconversation import ManagerWorkersConversation from wayflowcore.messagelist import Message @@ -222,24 +223,36 @@ def start_conversation( messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, conversation_id: Optional[str] = None, conversation_name: Optional[str] = None, + *, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + _root_conversation_id: Optional[str] = None, + _attach_checkpointer: bool = True, ) -> "ManagerWorkersConversation": """ - Initializes a conversation with the managerworkers. + Start a conversation for the manager-workers group. Parameters ---------- inputs: - Dictionary of inputs. Keys are the variable identifiers and - values are the actual inputs to start the main conversation. + Optional input values passed to the manager's main conversation. messages: - Message list of the manager agent and the end-user. + Optional shared message history between the user and the manager. conversation_id: - Conversation id of the main conversation. + Optional identifier for this manager-workers conversation. + conversation_name: + Optional display name used for the created conversation object. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires ``checkpointer``. + _root_conversation_id: + Internal lineage identifier shared with nested or parent conversations. Returns ------- - Conversation: - The conversation object of the managerworkers. + ManagerWorkersConversation + A new or restored manager-workers conversation. """ from wayflowcore.agentconversation import AgentConversation from wayflowcore.events.event import ConversationCreatedEvent @@ -249,18 +262,30 @@ def start_conversation( ManagerWorkersConversationExecutionState, ) + restored_conversation, conversation_runtime_id, conversation_root_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + _root_conversation_id=_root_conversation_id, + expected_conversation_type=ManagerWorkersConversation, + attach_checkpointer=_attach_checkpointer, + ) + ) + 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_runtime_id, nesting_level=None, ) ) @@ -269,23 +294,28 @@ def start_conversation( subconversations[self.manager_agent.name] = self.manager_agent.start_conversation( inputs=inputs, messages=messages, + _root_conversation_id=conversation_root_id, ) state = ManagerWorkersConversationExecutionState( current_agent_name=self.manager_agent.name, subconversations=subconversations, + root_conversation_id=conversation_root_id, ) - return ManagerWorkersConversation( + conversation = ManagerWorkersConversation( component=self, inputs={}, message_list=messages, + id=conversation_runtime_id, name=conversation_name or "managerworkers_conversation", state=state, status=None, - conversation_id=conversation_id, + checkpointer=checkpointer, + root_conversation_id=conversation_root_id, __metadata_info__={}, ) + return conversation def _referenced_tools_dict_inner( self, recursive: bool, visited_set: Set[str] diff --git a/wayflowcore/src/wayflowcore/ociagent.py b/wayflowcore/src/wayflowcore/ociagent.py index 4bfa36e92..ad74efde0 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 @@ -106,22 +107,37 @@ 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, + _root_conversation_id: Optional[str] = None, + _attach_checkpointer: bool = True, ) -> "Conversation": """ - Initializes a conversation with the agent. + Start a conversation with the OCI agent. Parameters ---------- inputs: - This argument is not used. - It is included for compatibility with the Flow class. + Optional structured inputs stored on the conversation for interface compatibility. messages: - Message list to which the agent will participate + Optional initial message history for the OCI agent session. + conversation_id: + Optional identifier for this OCI agent conversation. + 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``. + _root_conversation_id: + Internal lineage identifier shared with nested or parent conversations. Returns ------- - Conversation: - The conversation object of the agent. + Conversation + A new OCI agent conversation. """ from wayflowcore.executors._ociagentconversation import OciAgentConversation from wayflowcore.executors._ociagentexecutor import ( @@ -130,9 +146,25 @@ 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_runtime_id, conversation_root_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=None, + checkpoint_id=None, + _root_conversation_id=_root_conversation_id, + expected_conversation_type=OciAgentConversation, + attach_checkpointer=_attach_checkpointer, + ) + ) + _client = _init_oci_agent_client(self) return OciAgentConversation( @@ -145,8 +177,9 @@ def start_conversation( inputs=inputs or {}, message_list=messages, status=None, - conversation_id=IdGenerator.get_or_generate_id(None), + id=conversation_runtime_id, name="oci_conversation", + root_conversation_id=conversation_root_id, __metadata_info__={}, ) diff --git a/wayflowcore/src/wayflowcore/serialization/context.py b/wayflowcore/src/wayflowcore/serialization/context.py index 232215778..186b3406d 100644 --- a/wayflowcore/src/wayflowcore/serialization/context.py +++ b/wayflowcore/src/wayflowcore/serialization/context.py @@ -47,6 +47,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 +114,16 @@ 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_reference(self, obj: Any) -> None: + """ + Registers an object as provided externally to the serialized payload. + + The serializer will emit a ``$ref`` for this object, but it will not add the object to + the root ``_referenced_objects`` section because the deserialization context is expected + to already contain it. + """ + self._external_references.add(self.get_reference(obj)) + def check_obj_is_already_serialized(self, obj: Any) -> bool: """ Returns True if the object has already been serialized @@ -122,7 +133,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]: """ diff --git a/wayflowcore/src/wayflowcore/serialization/serializer.py b/wayflowcore/src/wayflowcore/serialization/serializer.py index fcb770db3..085b2cc1f 100644 --- a/wayflowcore/src/wayflowcore/serialization/serializer.py +++ b/wayflowcore/src/wayflowcore/serialization/serializer.py @@ -140,13 +140,18 @@ class MyDataclass: type_3: "MySecondCustomAttr" <--- resolves the actual type of this kind of attribute """ dataclass_fields: Dict[str, Any] = { - param.name: param.type for param in fields(cls) if param.init + param.name: param.type for param in fields(cls) if _should_serialize_dataclass_field(param) } # we resolve the forwards references (e.g. dataclasses with type annotations specified "between quotes") if any(isinstance(t, str) for t in dataclass_fields.values()): try: - dataclass_fields = get_type_hints(cls) + resolved_type_hints = get_type_hints(cls) + dataclass_fields = { + field_name: resolved_type_hints[field_name] + for field_name in dataclass_fields + if field_name in resolved_type_hints + } except NameError as e: pass @@ -193,7 +198,14 @@ def _resolve_legacy_field_name(cls: type, field_name: str) -> str: } if cls in _CLS_TO_ATTRIBUTE_MAPPING: - return _CLS_TO_ATTRIBUTE_MAPPING[cls].get(field_name, field_name) + resolved_field_name = _CLS_TO_ATTRIBUTE_MAPPING[cls].get(field_name) + if resolved_field_name is not None: + return resolved_field_name + + if field_name == "root_conversation_id" and any( + base.__name__ == "Conversation" for base in cls.__mro__ + ): + return "conversation_id" return field_name @@ -207,18 +219,20 @@ def _resolve_legacy_configurations(serialized_config: str) -> str: ) +def _should_serialize_dataclass_field(dataclass_field: Any) -> bool: + return bool( + (not dataclass_field.name.startswith("_") or dataclass_field.name == "__metadata_info__") + and dataclass_field.init + and dataclass_field.metadata.get("serialize", True) + ) + + class SerializableDataclassMixin: def _serialize_to_dict(self, serialization_context: "SerializationContext") -> Dict[str, Any]: return { k.name: serialize_any_to_dict(getattr(self, k.name), serialization_context) for k in fields(self) # type: ignore - if ( - ( - not k.name.startswith("_") # we don't serialize private fields - or k.name == "__metadata_info__" # except the metadata - ) - and k.init # not part of the dataclass __init__ -> would fail at deserialization - ) + if _should_serialize_dataclass_field(k) } @classmethod diff --git a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py index 25c4efe58..a26329035 100644 --- a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py +++ b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py @@ -296,7 +296,7 @@ def _get_or_create_agent_subconversation( agent_sub_conversation = self.agent.start_conversation( inputs=inputs, messages=init_messages, - root_conversation_id=caller_conv.root_conversation_id, + _root_conversation_id=caller_conv.root_conversation_id, ) return agent_sub_conversation diff --git a/wayflowcore/src/wayflowcore/swarm.py b/wayflowcore/src/wayflowcore/swarm.py index b4852fc18..676fa98f3 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 @@ -312,7 +313,37 @@ def start_conversation( messages: Union[None, str, "Message", List["Message"], MessageList] = None, conversation_id: Optional[str] = None, conversation_name: Optional[str] = None, + *, + checkpointer: Optional["Checkpointer"] = None, + checkpoint_id: Optional[str] = None, + _root_conversation_id: Optional[str] = None, + _attach_checkpointer: bool = True, ) -> "Conversation": + """ + Start a conversation for the swarm. + + Parameters + ---------- + inputs: + Optional input values available to the swarm execution state. + messages: + Optional shared message history for the swarm conversation. + conversation_id: + Optional identifier for this swarm conversation. + conversation_name: + Optional display name used for the created conversation object. + checkpointer: + Optional checkpoint backend used to restore and persist this conversation. + checkpoint_id: + Optional checkpoint identifier to restore. Requires ``checkpointer``. + _root_conversation_id: + Internal lineage identifier shared with nested or parent conversations. + + Returns + ------- + Conversation + A new or restored swarm conversation. + """ from wayflowcore.executors._swarmconversation import ( SwarmConversation, SwarmConversationExecutionState, @@ -320,12 +351,24 @@ def start_conversation( SwarmUser, ) + restored_conversation, conversation_runtime_id, conversation_root_id = ( + self._prepare_conversation_start( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + _root_conversation_id=_root_conversation_id, + expected_conversation_type=SwarmConversation, + attach_checkpointer=_attach_checkpointer, + ) + ) + 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 ) @@ -344,17 +387,21 @@ def start_conversation( context_providers=[], inputs=inputs, messages=messages, + root_conversation_id=conversation_root_id, ) - return SwarmConversation( + conversation = SwarmConversation( component=self, inputs=inputs or {}, message_list=messages, + id=conversation_runtime_id, name=conversation_name or "swarm_conversation", state=state, status=None, - conversation_id=conversation_id, + checkpointer=checkpointer, + root_conversation_id=conversation_root_id, __metadata_info__={}, ) + 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 3d6bd6daf..9c85792e3 100644 --- a/wayflowcore/src/wayflowcore/tools/servertools.py +++ b/wayflowcore/src/wayflowcore/tools/servertools.py @@ -639,7 +639,7 @@ async def __call__(self, **inputs: Any) -> Any: conversation = self.flow.start_conversation( inputs=inputs, messages=self._parent_conversation.message_list, - root_conversation_id=self._parent_conversation.root_conversation_id, + _root_conversation_id=self._parent_conversation.root_conversation_id, ) interrupts = self._parent_conversation._get_interrupts() diff --git a/wayflowcore/src/wayflowcore/tracing/span.py b/wayflowcore/src/wayflowcore/tracing/span.py index 07952157c..22dbfbbdc 100644 --- a/wayflowcore/src/wayflowcore/tracing/span.py +++ b/wayflowcore/src/wayflowcore/tracing/span.py @@ -550,7 +550,7 @@ class ConversationSpan(Span): def to_tracing_info(self, mask_sensitive_information: bool = True) -> Dict[str, Any]: return { **super().to_tracing_info(mask_sensitive_information=mask_sensitive_information), - "conversation.id": self.conversation.conversation_id, + "conversation.id": self.conversation.id, "conversation.name": self.conversation.name, "conversational_component.type": self.conversation.component.__class__.__name__, "conversational_component.id": self.conversation.component.id, diff --git a/wayflowcore/tests/integration/steps/test_prompt_execution_step.py b/wayflowcore/tests/integration/steps/test_prompt_execution_step.py index c53d9cf1c..92442b5c9 100644 --- a/wayflowcore/tests/integration/steps/test_prompt_execution_step.py +++ b/wayflowcore/tests/integration/steps/test_prompt_execution_step.py @@ -468,7 +468,7 @@ def test_check_token_consumption(remotely_hosted_llm): assert isinstance(status, FinishedStatus) assert PromptExecutionStep.OUTPUT in status.output_values assert isinstance(status.output_values[PromptExecutionStep.OUTPUT], str) - token_consumption = remotely_hosted_llm.get_total_token_consumption(conv.conversation_id) + token_consumption = remotely_hosted_llm.get_total_token_consumption(conv.root_conversation_id) assert token_consumption.input_tokens == 39 assert token_consumption.output_tokens == 10 diff --git a/wayflowcore/tests/serialization/test_conversation_checkpointing.py b/wayflowcore/tests/serialization/test_conversation_checkpointing.py new file mode 100644 index 000000000..2d694a346 --- /dev/null +++ b/wayflowcore/tests/serialization/test_conversation_checkpointing.py @@ -0,0 +1,654 @@ +# Copyright © 2025, 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 warnings +from types import SimpleNamespace +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock + +import pytest + +import wayflowcore.checkpointing.serialization as checkpoint_serialization +from wayflowcore.a2a.a2aagent import A2AAgent, A2AConnectionConfig +from wayflowcore.agent import Agent +from wayflowcore.checkpointing import ( + CheckpointingInterval, + ConversationCheckpoint, + InMemoryCheckpointer, +) +from wayflowcore.checkpointing.checkpointeventlistener import _save_conversation_checkpoint +from wayflowcore.executors.executionstatus import FinishedStatus, UserMessageRequestStatus +from wayflowcore.flowhelpers import create_single_step_flow +from wayflowcore.managerworkers import ManagerWorkers +from wayflowcore.models.ociclientconfig import OCIClientConfigWithApiKey +from wayflowcore.ociagent import OciAgent +from wayflowcore.serialization.serializer import _resolve_legacy_field_name +from wayflowcore.steps import OutputMessageStep, PromptExecutionStep +from wayflowcore.swarm import Swarm + +from ..testhelpers.dummy import DummyModel +from ..testhelpers.testhelpers import retry_test +from .test_assistant_serialization import create_flow + + +class RecordingCheckpointer: + def __init__(self, *, fail_first_save: bool = False) -> None: + self.checkpointing_interval = CheckpointingInterval.CONVERSATION_TURNS + self.should_fail_next_save = fail_first_save + self.saved_checkpoints: Dict[tuple[str, str], Dict[str, Any]] = {} + + def save_conversation( + self, + conversation, + *, + checkpoint_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ): + if self.should_fail_next_save: + self.should_fail_next_save = False + raise RuntimeError("checkpoint save failed") + resolved_checkpoint_id = checkpoint_id or "generated-checkpoint-id" + self.saved_checkpoints[(conversation.id, resolved_checkpoint_id)] = dict(metadata or {}) + conversation.checkpoint_id = resolved_checkpoint_id + return SimpleNamespace( + checkpoint_id=resolved_checkpoint_id, + metadata=dict(metadata or {}), + ) + + def load(self, conversation_id: str, checkpoint_id: str) -> Any: + return SimpleNamespace(metadata=self.saved_checkpoints[(conversation_id, checkpoint_id)]) + + def load_latest(self, conversation_id: str) -> Any: + return None + + +class StaticLoadCheckpointer: + def __init__(self, checkpoint: ConversationCheckpoint) -> None: + self.checkpointing_interval = CheckpointingInterval.CONVERSATION_TURNS + self.checkpoint = checkpoint + + def load(self, conversation_id: str, checkpoint_id: str) -> ConversationCheckpoint: + return self.checkpoint + + def load_latest(self, conversation_id: str) -> Optional[ConversationCheckpoint]: + if conversation_id != self.checkpoint.conversation_id: + return None + return self.checkpoint + + +def _build_checkpointable_agent( + *, + name: str, + initial_message: str, +) -> tuple[Agent, DummyModel]: + llm = DummyModel() + agent = Agent( + llm=llm, + name=name, + description=f"{name} description", + custom_instruction="Be helpful.", + initial_message=initial_message, + ) + return agent, llm + + +def _build_checkpointable_swarm() -> tuple[Swarm, DummyModel]: + first_agent, first_agent_llm = _build_checkpointable_agent( + name="checkpoint_swarm_first_agent", + 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.", + ) + swarm = Swarm( + first_agent=first_agent, + relationships=[(first_agent, second_agent)], + name="checkpoint_swarm", + ) + return swarm, first_agent_llm + + +def _build_checkpointable_managerworkers() -> tuple[ManagerWorkers, DummyModel]: + manager_agent, manager_llm = _build_checkpointable_agent( + name="checkpoint_manager_agent", + initial_message="Hello from the manager.", + ) + worker_agent = Agent( + llm=DummyModel(fails_if_not_set=False), + name="checkpoint_worker_agent", + description="Worker agent", + custom_instruction="Help the manager.", + ) + managerworkers = ManagerWorkers( + group_manager=manager_agent, + workers=[worker_agent], + name="checkpoint_managerworkers", + ) + return managerworkers, manager_llm + + +@pytest.fixture(scope="session") +def connection_config_no_verify(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + return A2AConnectionConfig(verify=False) + + +@pytest.fixture +def a2a_agent(a2a_server, connection_config_no_verify): + return A2AAgent( + name="Checkpoint A2A Agent", + agent_url=a2a_server, + connection_config=connection_config_no_verify, + ) + + +def test_inmemory_checkpointer_can_save_load_list_and_delete_checkpoints() -> None: + checkpointer = InMemoryCheckpointer() + flow = create_single_step_flow(OutputMessageStep(message_template="Hello from checkpointing.")) + + conversation = flow.start_conversation( + conversation_id="checkpoint-lifecycle", checkpointer=checkpointer + ) + assert conversation.checkpointer is checkpointer + + status = conversation.execute() + + assert isinstance(status, FinishedStatus) + first_checkpoint_id = conversation.checkpoint_id + assert first_checkpoint_id is not None + + checkpointer.save(conversation) + second_checkpoint_id = conversation.checkpoint_id + assert second_checkpoint_id is not None + assert second_checkpoint_id != first_checkpoint_id + + checkpoints = checkpointer.list_checkpoints("checkpoint-lifecycle") + assert [checkpoint.checkpoint_id for checkpoint in checkpoints] == [ + first_checkpoint_id, + second_checkpoint_id, + ] + assert checkpoints[-1].metadata["save_sequence"] == 2 + + latest_checkpoint = checkpointer.load_latest("checkpoint-lifecycle") + assert latest_checkpoint is not None + assert latest_checkpoint.checkpoint_id == second_checkpoint_id + + restored_conversation = flow.start_conversation( + conversation_id="checkpoint-lifecycle", + checkpoint_id=first_checkpoint_id, + checkpointer=checkpointer, + ) + assert restored_conversation.checkpointer is checkpointer + assert restored_conversation.checkpoint_id == first_checkpoint_id + assert restored_conversation.get_last_message().content == "Hello from checkpointing." + + checkpointer.delete("checkpoint-lifecycle", second_checkpoint_id) + promoted_checkpoint = checkpointer.load_latest("checkpoint-lifecycle") + assert promoted_checkpoint is not None + assert promoted_checkpoint.checkpoint_id == first_checkpoint_id + assert [ + checkpoint.checkpoint_id + for checkpoint in checkpointer.list_checkpoints("checkpoint-lifecycle") + ] == [first_checkpoint_id] + + +def test_conversation_turns_checkpoint_interval_saves_once_after_outer_execute() -> None: + checkpointer = InMemoryCheckpointer( + checkpointing_interval=CheckpointingInterval.CONVERSATION_TURNS + ) + flow = create_single_step_flow(OutputMessageStep(message_template="Hello once.")) + + status = flow.start_conversation( + conversation_id="conversation-turn", checkpointer=checkpointer + ).execute() + + assert isinstance(status, FinishedStatus) + checkpoints = checkpointer.list_checkpoints("conversation-turn") + assert len(checkpoints) == 1 + assert checkpoints[0].metadata["save_reason"] == "conversation_turn" + assert checkpoints[0].metadata["status_type"] == "FinishedStatus" + + +def test_all_internal_turns_checkpoint_interval_saves_before_each_flow_turn() -> None: + checkpointer = InMemoryCheckpointer( + checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS + ) + flow = create_single_step_flow(OutputMessageStep(message_template="Hello internal turns.")) + + status = flow.start_conversation( + conversation_id="all-internal-turns", checkpointer=checkpointer + ).execute() + + assert isinstance(status, FinishedStatus) + checkpoints = checkpointer.list_checkpoints("all-internal-turns") + assert len(checkpoints) == 3 + assert [checkpoint.metadata["save_reason"] for checkpoint in checkpoints] == [ + "internal_turn_boundary", + "internal_turn_boundary", + "conversation_turn", + ] + assert [checkpoint.metadata.get("event_type") for checkpoint in checkpoints[:-1]] == [ + "FlowExecutionIterationStartedEvent", + "FlowExecutionIterationStartedEvent", + ] + assert checkpoints[-1].metadata["status_type"] == "FinishedStatus" + + +def test_llm_turns_checkpoint_interval_saves_only_after_llm_backed_turns() -> None: + checkpointer = InMemoryCheckpointer(checkpointing_interval=CheckpointingInterval.LLM_TURNS) + dummy_llm = DummyModel() + dummy_llm.set_next_output("Hello from the prompt step.") + flow = create_single_step_flow( + PromptExecutionStep( + llm=dummy_llm, + prompt_template="Say hello.", + ) + ) + + status = flow.start_conversation( + conversation_id="llm-turns", checkpointer=checkpointer + ).execute() + + assert isinstance(status, FinishedStatus) + checkpoints = checkpointer.list_checkpoints("llm-turns") + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["save_reason"] == "internal_turn_boundary" + assert checkpoints[0].metadata["event_type"] == "FlowExecutionIterationStartedEvent" + assert checkpoints[0].metadata["llm_used_in_previous_turn"] is True + assert checkpoints[1].metadata["save_reason"] == "conversation_turn" + + +def test_checkpoint_serialization_context_registers_component_tree_as_external_refs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeSerializationContext: + def __init__(self, root: Any) -> None: + self.root = root + self.external_refs: set[str] = set() + self.recorded_refs: Dict[str, Dict[str, Any]] = {} + + @staticmethod + def get_reference(obj: Any) -> str: + obj_id = getattr(obj, "id", id(obj)) + return f"{obj.__class__.__name__.lower()}/{obj_id}" + + def register_external_reference(self, obj: Any) -> None: + self.external_refs.add(self.get_reference(obj)) + + def record_obj_dict(self, obj: Any, obj_as_dict: Dict[str, Any]) -> None: + self.recorded_refs[self.get_reference(obj)] = obj_as_dict + + monkeypatch.setattr( + checkpoint_serialization, + "SerializationContext", + _FakeSerializationContext, + ) + + flow = create_single_step_flow(OutputMessageStep(message_template="Hello external refs.")) + conversation = flow.start_conversation(conversation_id="checkpoint-external-refs") + serialization_context = checkpoint_serialization._build_checkpoint_serialization_context( + conversation + ) + + expected_component_refs = { + _FakeSerializationContext.get_reference(component) + for component in checkpoint_serialization._iter_component_tree(conversation.component) + } + + assert serialization_context.external_refs == expected_component_refs + assert serialization_context.recorded_refs == {} + + +def test_explicit_final_checkpoint_parameters_can_be_retried_after_save_fails() -> None: + checkpointer = RecordingCheckpointer(fail_first_save=True) + flow = create_single_step_flow(OutputMessageStep(message_template="Hello overrides.")) + + conversation = flow.start_conversation( + conversation_id="checkpoint-final-overrides", checkpointer=checkpointer + ) + + with pytest.raises(RuntimeError, match="checkpoint save failed"): + _save_conversation_checkpoint( + conversation, + save_reason="conversation_turn", + checkpoint_id="final-checkpoint-id", + metadata={"response_id": "resp-123"}, + ) + + _save_conversation_checkpoint( + conversation, + save_reason="conversation_turn", + checkpoint_id="final-checkpoint-id", + metadata={"response_id": "resp-123"}, + ) + + assert conversation.checkpoint_id == "final-checkpoint-id" + checkpoint = checkpointer.load("checkpoint-final-overrides", "final-checkpoint-id") + assert checkpoint.metadata["response_id"] == "resp-123" + + +def test_execute_final_checkpoint_parameters_are_applied_to_final_save() -> None: + checkpointer = RecordingCheckpointer() + flow = create_single_step_flow(OutputMessageStep(message_template="Hello execute.")) + + conversation = flow.start_conversation( + conversation_id="checkpoint-final-execute", + checkpointer=checkpointer, + ) + status = conversation.execute( + _final_checkpoint_id="final-checkpoint-id", + _final_checkpoint_metadata={"response_id": "resp-123"}, + ) + + assert isinstance(status, FinishedStatus) + assert conversation.checkpoint_id == "final-checkpoint-id" + checkpoint = checkpointer.load("checkpoint-final-execute", "final-checkpoint-id") + assert checkpoint.metadata["response_id"] == "resp-123" + assert checkpoint.metadata["save_reason"] == "conversation_turn" + + +def test_execute_async_does_not_save_final_checkpoint_when_execution_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + checkpointer = RecordingCheckpointer() + 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 conversation.checkpoint_id is None + assert checkpointer.saved_checkpoints == {} + + +def test_checkpoint_restore_rejects_conversations_from_other_components() -> None: + original_flow = create_single_step_flow(OutputMessageStep(message_template="Hello original.")) + other_flow = create_single_step_flow(OutputMessageStep(message_template="Hello other.")) + checkpointer = StaticLoadCheckpointer( + ConversationCheckpoint( + checkpoint_id="checkpoint-1", + conversation_id="checkpoint-other-component", + component_id=original_flow.id, + created_at=0, + state="unused-because-component-mismatch-is-checked-first", + metadata={}, + ) + ) + + with pytest.raises(ValueError, match="started with another component"): + other_flow.start_conversation( + conversation_id="checkpoint-other-component", + checkpointer=checkpointer, + ) + + +def test_restore_can_skip_attaching_live_checkpointer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + flow = create_single_step_flow(OutputMessageStep(message_template="Hello restore.")) + checkpoint_id = "checkpoint-no-attach-id" + checkpointer = StaticLoadCheckpointer( + ConversationCheckpoint( + checkpoint_id=checkpoint_id, + conversation_id="checkpoint-no-attach", + component_id=flow.id, + created_at=0, + state="unused-because-deserialization-is-mocked", + metadata={}, + ) + ) + + monkeypatch.setattr( + checkpoint_serialization, + "_deserialize_conversation_checkpoint_state", + lambda *args, **kwargs: flow.start_conversation(conversation_id="checkpoint-no-attach"), + ) + + restored_conversation = flow.start_conversation( + conversation_id="checkpoint-no-attach", + checkpoint_id=checkpoint_id, + checkpointer=checkpointer, + _attach_checkpointer=False, + ) + + assert restored_conversation.checkpointer is None + assert restored_conversation.checkpoint_id == checkpoint_id + + +def test_legacy_serialized_conversation_id_restores_root_conversation_id() -> None: + agent, _ = _build_checkpointable_agent( + name="legacy_checkpoint_agent", + initial_message="Hello from the past.", + ) + conversation = agent.start_conversation(_root_conversation_id="legacy-root-conversation") + + assert conversation.root_conversation_id == "legacy-root-conversation" + assert ( + _resolve_legacy_field_name(type(conversation), "root_conversation_id") == "conversation_id" + ) + assert not hasattr(conversation, "conversation_id") + + +def test_flow_checkpointing_supports_resume_and_time_travel() -> None: + checkpointer = InMemoryCheckpointer() + flow = create_flow() + + conversation = flow.start_conversation( + conversation_id="flow-checkpoint", checkpointer=checkpointer + ) + first_status = conversation.execute() + + assert isinstance(first_status, UserMessageRequestStatus) + first_checkpoint_id = conversation.checkpoint_id + assert first_checkpoint_id is not None + + restored_conversation = flow.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + ) + assert restored_conversation.checkpoint_id == first_checkpoint_id + assert isinstance(restored_conversation.status, UserMessageRequestStatus) + + restored_conversation.append_user_message("continue") + restored_status = restored_conversation.execute() + assert isinstance(restored_status, FinishedStatus) + + rewound_conversation = flow.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + checkpoint_id=first_checkpoint_id, + ) + assert len(rewound_conversation.get_messages()) < len(restored_conversation.get_messages()) + rewound_conversation.append_user_message("rewind") + rewound_status = rewound_conversation.execute() + assert isinstance(rewound_status, FinishedStatus) + + +def test_agent_checkpointing_supports_resume_and_time_travel() -> None: + checkpointer = InMemoryCheckpointer() + agent, llm = _build_checkpointable_agent( + name="checkpoint_agent", + initial_message="Hello from the agent.", + ) + + conversation = agent.start_conversation( + conversation_id="agent-checkpoint", checkpointer=checkpointer + ) + first_status = conversation.execute() + + assert isinstance(first_status, UserMessageRequestStatus) + first_checkpoint_id = conversation.checkpoint_id + assert first_checkpoint_id is not None + + restored_conversation = agent.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + ) + llm.set_next_output("Agent resumed successfully.") + restored_conversation.append_user_message("Please continue.") + restored_status = restored_conversation.execute() + assert isinstance(restored_status, UserMessageRequestStatus) + assert restored_conversation.get_last_message().content == "Agent resumed successfully." + + rewound_conversation = agent.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + checkpoint_id=first_checkpoint_id, + ) + llm.set_next_output("Agent rewound successfully.") + rewound_conversation.append_user_message("Try again.") + rewound_status = rewound_conversation.execute() + assert isinstance(rewound_status, UserMessageRequestStatus) + assert rewound_conversation.get_last_message().content == "Agent rewound successfully." + + +def test_swarm_checkpointing_supports_resume_and_time_travel() -> None: + checkpointer = InMemoryCheckpointer() + swarm, llm = _build_checkpointable_swarm() + + conversation = swarm.start_conversation( + conversation_id="swarm-checkpoint", checkpointer=checkpointer + ) + first_status = conversation.execute() + + assert isinstance(first_status, UserMessageRequestStatus) + first_checkpoint_id = conversation.checkpoint_id + assert first_checkpoint_id is not None + + restored_conversation = swarm.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + ) + llm.set_next_output("Swarm resumed successfully.") + restored_conversation.append_user_message("Continue the swarm conversation.") + restored_status = restored_conversation.execute() + assert isinstance(restored_status, UserMessageRequestStatus) + assert restored_conversation.get_last_message().content == "Swarm resumed successfully." + + rewound_conversation = swarm.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + checkpoint_id=first_checkpoint_id, + ) + llm.set_next_output("Swarm rewound successfully.") + rewound_conversation.append_user_message("Try the swarm again.") + rewound_status = rewound_conversation.execute() + assert isinstance(rewound_status, UserMessageRequestStatus) + assert rewound_conversation.get_last_message().content == "Swarm rewound successfully." + + +def test_managerworkers_checkpointing_supports_resume_and_time_travel() -> None: + checkpointer = InMemoryCheckpointer() + managerworkers, llm = _build_checkpointable_managerworkers() + + conversation = managerworkers.start_conversation( + conversation_id="managerworkers-checkpoint", + checkpointer=checkpointer, + ) + first_status = conversation.execute() + + assert isinstance(first_status, UserMessageRequestStatus) + first_checkpoint_id = conversation.checkpoint_id + assert first_checkpoint_id is not None + + restored_conversation = managerworkers.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + ) + llm.set_next_output("Manager resumed successfully.") + restored_conversation.append_user_message("Continue the manager workflow.") + restored_status = restored_conversation.execute() + assert isinstance(restored_status, UserMessageRequestStatus) + assert restored_conversation.get_last_message().content == "Manager resumed successfully." + + rewound_conversation = managerworkers.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + checkpoint_id=first_checkpoint_id, + ) + llm.set_next_output("Manager rewound successfully.") + rewound_conversation.append_user_message("Try the manager workflow again.") + rewound_status = rewound_conversation.execute() + assert isinstance(rewound_status, UserMessageRequestStatus) + assert rewound_conversation.get_last_message().content == "Manager rewound successfully." + + +@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) + first_checkpoint_id = conversation.checkpoint_id + assert first_checkpoint_id is not None + first_message_count = len(conversation.get_messages()) + + restored_conversation = a2a_agent.start_conversation( + conversation_id=conversation.id, + checkpointer=checkpointer, + ) + assert restored_conversation.checkpoint_id == first_checkpoint_id + 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.id) is not None + + rewound_conversation = a2a_agent.start_conversation( + conversation_id=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 + + +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()) diff --git a/wayflowcore/tests/test_managerworkers.py b/wayflowcore/tests/test_managerworkers.py index 040933d3a..ee7929d02 100644 --- a/wayflowcore/tests/test_managerworkers.py +++ b/wayflowcore/tests/test_managerworkers.py @@ -184,7 +184,7 @@ def test_managerworkers_can_execute_with_initial_params_passed_in_start_conversa conversation = group.start_conversation( messages=[Message(content="Please compute 3*4 + 2", message_type=MessageType.USER)], inputs={"USER": "Iris"}, - root_conversation_id="12345", + _root_conversation_id="12345", ) conversation.execute() diff --git a/wayflowcore/tests/test_swarm.py b/wayflowcore/tests/test_swarm.py index f52e28fca..0c2d2d1a6 100644 --- a/wayflowcore/tests/test_swarm.py +++ b/wayflowcore/tests/test_swarm.py @@ -226,7 +226,7 @@ def test_can_execute_swarm_with_initial_params_passed_in_start_conversation( ) ], inputs={"USER": "Iris"}, - root_conversation_id="12345", + _root_conversation_id="12345", ) conversation.execute() diff --git a/wayflowcore/tests/tracing/spans/test_conversation_span.py b/wayflowcore/tests/tracing/spans/test_conversation_span.py index 24e6d7098..274d28efc 100644 --- a/wayflowcore/tests/tracing/spans/test_conversation_span.py +++ b/wayflowcore/tests/tracing/spans/test_conversation_span.py @@ -92,7 +92,7 @@ def test_span_serialization_format( assert serialized_span["span_type"] == str(span.__class__.__name__) for attribute_name in attributes_to_check: assert getattr(span, attribute_name) == serialized_span[attribute_name] - assert serialized_span["conversation.id"] == span.conversation.conversation_id + assert serialized_span["conversation.id"] == span.conversation.id assert serialized_span["conversation.name"] == span.conversation.name assert ( serialized_span["conversational_component.type"] == "Agent" From 684eaabf0c11d6278131e6321a9e46c01dee3ca9 Mon Sep 17 00:00:00 2001 From: jschweiz Date: Wed, 22 Apr 2026 11:12:56 +0200 Subject: [PATCH 03/10] [fix]: fix import --- wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py b/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py index 6a09bad18..a269babf9 100644 --- a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py +++ b/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py @@ -7,7 +7,7 @@ from typing import Dict, Optional from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig -from wayflowcore.checkpointing.datastore import ( +from wayflowcore.checkpointing.datastorecheckpointer import ( _prepare_oracle_checkpoint_datastore, _prepare_postgres_checkpoint_datastore, ) From 2f9929c22bd3be759a7887a00ce79736862b7557 Mon Sep 17 00:00:00 2001 From: Son Le Date: Thu, 18 Jun 2026 16:13:58 +0200 Subject: [PATCH 04/10] refactor --- docs/wayflowcore/source/conf.py | 11 + .../source/core/api/checkpointing.rst | 34 + docs/wayflowcore/source/core/api/index.rst | 1 + docs/wayflowcore/source/core/changelog.rst | 28 +- .../core/code_examples/howto_checkpointing.py | 12 +- .../core/howtoguides/howto_checkpointing.rst | 102 ++- wayflowcore/src/wayflowcore/a2a/a2aagent.py | 37 +- wayflowcore/src/wayflowcore/agent.py | 35 +- .../agentserver/_storagehelpers.py | 17 - .../services/wayflowservice.py | 179 +++-- .../agentserver/serverstorageconfig.py | 2 +- .../src/wayflowcore/checkpointing/__init__.py | 5 +- .../checkpointing/checkpoint_state.py | 251 +++++++ .../wayflowcore/checkpointing/checkpointer.py | 78 +-- .../checkpointing/checkpointeventlistener.py | 263 +++---- .../checkpointing/datastorecheckpointer.py | 189 ++++- .../checkpointing/serialization.py | 143 ---- .../contextproviders/flowcontextprovider.py | 2 +- wayflowcore/src/wayflowcore/conversation.py | 52 +- .../wayflowcore/conversationalcomponent.py | 122 ++-- .../src/wayflowcore/datastore/_relational.py | 9 +- wayflowcore/src/wayflowcore/events/event.py | 4 +- .../executors/_agentconversation.py | 17 +- .../wayflowcore/executors/_agentexecutor.py | 45 +- .../executors/_flowconversation.py | 43 +- .../wayflowcore/executors/_flowexecutor.py | 30 +- .../executors/_managerworkersconversation.py | 7 +- .../executors/_swarmconversation.py | 7 +- .../tokenlimitexecutioninterrupt.py | 2 +- wayflowcore/src/wayflowcore/flow.py | 51 +- wayflowcore/src/wayflowcore/flowbuilder.py | 16 +- wayflowcore/src/wayflowcore/flowhelpers.py | 4 + wayflowcore/src/wayflowcore/managerworkers.py | 38 +- .../wayflowcore/mcp/_session_persistence.py | 8 +- .../src/wayflowcore/models/llmmodel.py | 4 +- wayflowcore/src/wayflowcore/ociagent.py | 26 +- .../src/wayflowcore/serialization/context.py | 83 ++- .../wayflowcore/serialization/serializer.py | 36 +- .../wayflowcore/steps/agentexecutionstep.py | 3 +- .../steps/parallelflowexecutionstep.py | 4 +- .../src/wayflowcore/steps/retrystep.py | 11 +- wayflowcore/src/wayflowcore/swarm.py | 41 +- .../src/wayflowcore/tools/servertools.py | 2 +- wayflowcore/src/wayflowcore/tracing/span.py | 2 +- wayflowcore/tests/a2a/test_a2aagent.py | 56 ++ .../tests/agentserver/test_wayflow_server.py | 3 +- .../steps/test_prompt_execution_step.py | 2 +- .../tests/integration/test_checkpointing.py | 528 ++++++++++++++ .../test_conversation_checkpointing.py | 654 ------------------ wayflowcore/tests/test_checkpointing.py | 310 +++++++++ wayflowcore/tests/test_flowbuilder.py | 3 +- wayflowcore/tests/test_managerworkers.py | 4 +- wayflowcore/tests/test_ociagent.py | 13 + wayflowcore/tests/test_swarm.py | 4 +- .../tracing/spans/test_conversation_span.py | 2 +- 55 files changed, 2147 insertions(+), 1488 deletions(-) create mode 100644 docs/wayflowcore/source/core/api/checkpointing.rst create mode 100644 wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py delete mode 100644 wayflowcore/src/wayflowcore/checkpointing/serialization.py create mode 100644 wayflowcore/tests/integration/test_checkpointing.py delete mode 100644 wayflowcore/tests/serialization/test_conversation_checkpointing.py create mode 100644 wayflowcore/tests/test_checkpointing.py 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 9f57dacef..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 ^^^^^^^^^^^^ @@ -93,18 +109,6 @@ New features per-token log-probabilities in the ``PromptExecutionStep``. For more information please read the guide on :ref:`How to request per-token log-probabilities ` -* **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 `. - Improvements ^^^^^^^^^^^^ diff --git a/docs/wayflowcore/source/core/code_examples/howto_checkpointing.py b/docs/wayflowcore/source/core/code_examples/howto_checkpointing.py index bcfc88cc4..030d28a05 100644 --- a/docs/wayflowcore/source/core/code_examples/howto_checkpointing.py +++ b/docs/wayflowcore/source/core/code_examples/howto_checkpointing.py @@ -1,4 +1,4 @@ -# Copyright © 2025, 2026 Oracle and/or its affiliates. +# 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 @@ -27,9 +27,10 @@ agent = Agent(llm=llm) checkpointer = InMemoryCheckpointer() +conversation_id = "support-conversation-1" conversation = agent.start_conversation( - conversation_id="support-thread-1", + conversation_id=conversation_id, checkpointer=checkpointer, ) @@ -38,7 +39,7 @@ # .. start-##_Resume_the_latest_checkpoint restored_conversation = agent.start_conversation( - conversation_id="support-thread-1", + conversation_id=conversation_id, checkpointer=checkpointer, ) @@ -47,11 +48,12 @@ # .. end-##_Resume_the_latest_checkpoint # .. start-##_Load_a_specific_checkpoint -checkpoints = checkpointer.list_checkpoints("support-thread-1") +# Checkpoints are ordered oldest -> newest. +checkpoints = checkpointer.list_checkpoints(conversation_id) previous_checkpoint = checkpoints[-2] rewound_conversation = agent.start_conversation( - conversation_id="support-thread-1", + conversation_id=conversation_id, checkpoint_id=previous_checkpoint.checkpoint_id, checkpointer=checkpointer, ) diff --git a/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst b/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst index e22d9a64e..a38f8e604 100644 --- a/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst +++ b/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst @@ -1,8 +1,8 @@ .. _top-howtocheckpointing: -========================================= +========================================== How to Checkpoint and Resume Conversations -========================================= +========================================== .. admonition:: Prerequisites @@ -12,78 +12,95 @@ How to Checkpoint and Resume Conversations - :doc:`Flows <../tutorials/basic_flow>` - :doc:`Serve Agents with WayFlow ` -WayFlow can now checkpoint the runtime state of a conversation and restore it later by -conversation id. This is useful when you want to: +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: -- resume after a crash or restart -- pause and continue a long-running workflow -- inspect prior checkpoints for debugging -- reload an earlier state and branch from it +- 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 ===================== -WayFlow exposes a shared checkpointing subsystem in ``wayflowcore.checkpointing``. -You can use: +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 share the same API for saving, loading, listing, and deleting checkpoints. +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. ``conversation_id`` becomes the durable key -used to look up the conversation later. +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 root conversation automatically at the configured -checkpoint boundaries. For nested execution lineage without checkpoint restore, pass -``root_conversation_id`` explicitly. +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 runtime id of this concrete ``Conversation`` object + +For a fresh top-level conversation, ``conversation.id`` usually matches +``conversation.conversation_id``. They can differ for restored or nested conversations, so use +``conversation_id`` in checkpointing APIs when you mean "this logical conversation". + +.. 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 restore the latest saved state, call ``start_conversation()`` again with the same -``conversation_id`` -and checkpointer. +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 no checkpoint exists for that id, WayFlow creates a new conversation instead. +If the checkpointer has no saved state for that id, WayFlow starts a new conversation. Load a specific checkpoint ========================== -You can inspect checkpoint history and reload an older checkpoint for replay or time-travel -debugging. +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 ordered checkpoint metadata, including the checkpoint id, -creation timestamp, and save metadata recorded at the boundary. +``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 decide how often WayFlow should persist state. +Use ``CheckpointingInterval`` to choose how often WayFlow should save state. .. literalinclude:: ../code_examples/howto_checkpointing.py :language: python @@ -92,23 +109,38 @@ Use ``CheckpointingInterval`` to decide how often WayFlow should persist state. The available options are: -- ``CONVERSATION_TURNS``: save after the outermost ``conversation.execute()`` call returns -- ``LLM_TURNS``: also save at internal turn boundaries after turns that used an LLM -- ``ALL_INTERNAL_TURNS``: also save at every internal agent/flow turn boundary +- ``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. -Saving more frequently improves restart fidelity, but it also increases write volume. +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 path now uses the shared checkpointing subsystem behind -``ServerStorageConfig``. That means the existing OpenAI-compatible features such as -``previous_response_id``, ``conversation``, ``get_response()``, ``delete_response()``, -and ``store=False`` all run through the same shared checkpoint model. +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, keep using :doc:`Serve Agents with WayFlow ` to -configure the storage backend. The server will use the matching shared checkpointer internally. +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 diff --git a/wayflowcore/src/wayflowcore/a2a/a2aagent.py b/wayflowcore/src/wayflowcore/a2a/a2aagent.py index 7a654ad3f..6a28ec341 100644 --- a/wayflowcore/src/wayflowcore/a2a/a2aagent.py +++ b/wayflowcore/src/wayflowcore/a2a/a2aagent.py @@ -250,35 +250,41 @@ 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, - _root_conversation_id: Optional[str] = None, + _runtime_conversation_id: Optional[str] = None, _attach_checkpointer: bool = True, ) -> "A2AAgentConversation": """ - Start a conversation with the remote A2A agent. + Initiates a new conversation with the remote server agent. + + Creates and returns a conversation instance tied to this agent, optionally initialized + with input data and a message history. Parameters ---------- inputs: - Optional structured inputs stored on the conversation for interface compatibility. - The A2A runtime currently executes from messages rather than these inputs. + Optional dictionary of initial input data for the conversation. Defaults to an empty + dictionary if not provided. messages: - Optional initial message history for the remote conversation. + 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: - Optional identifier for this A2A 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 ``checkpointer``. - _root_conversation_id: - Internal lineage identifier shared with nested or parent conversations. + Optional checkpoint identifier to restore. Requires both ``checkpointer`` and + ``conversation_id``. + _runtime_conversation_id: + Internal runtime id for a fresh conversation. When provided, it becomes + the created conversation object's ``.id`` instead of defaulting to + ``conversation_id``. Returns ------- - A2AAgentConversation - A new or restored A2A agent conversation. + Conversation: + A new conversation object associated with this agent. """ from wayflowcore.executors._a2aagentconversation import A2AAgentConversation from wayflowcore.executors._a2aagentexecutor import A2AAgentState @@ -290,7 +296,7 @@ def start_conversation( conversation_id=conversation_id, checkpointer=checkpointer, checkpoint_id=checkpoint_id, - _root_conversation_id=_root_conversation_id, + _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=A2AAgentConversation, attach_checkpointer=_attach_checkpointer, ) @@ -301,7 +307,7 @@ def start_conversation( if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) - conversation = A2AAgentConversation( + return A2AAgentConversation( component=self, state=A2AAgentState(last_message_idx=-1), inputs=inputs or {}, # Inputs are ignored in execution @@ -310,10 +316,9 @@ def start_conversation( id=conversation_runtime_id, checkpointer=checkpointer, name="a2a_conversation", - root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, __metadata_info__={}, ) - return conversation @property def agent_id(self) -> str: diff --git a/wayflowcore/src/wayflowcore/agent.py b/wayflowcore/src/wayflowcore/agent.py index d7f05f452..228f96d22 100644 --- a/wayflowcore/src/wayflowcore/agent.py +++ b/wayflowcore/src/wayflowcore/agent.py @@ -394,38 +394,42 @@ 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, - _root_conversation_id: Optional[str] = None, + _runtime_conversation_id: Optional[str] = None, _attach_checkpointer: bool = True, ) -> "AgentConversation": """ - Start a conversation with the agent. + Initializes a conversation with the agent. Parameters ---------- inputs: - Optional input values for the agent's declared input descriptors. + This argument is not used. + It is included for compatibility with the Flow class. messages: - Optional message history for the conversation. + Message list to which the agent will participate conversation_id: - Optional identifier for this agent 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 ``checkpointer``. - _root_conversation_id: - Internal lineage identifier shared with nested or parent conversations. + Optional checkpoint identifier to restore. Requires both ``checkpointer`` and + ``conversation_id``. + _runtime_conversation_id: + Internal runtime id for a fresh conversation. When provided, it becomes + the created conversation object's ``.id`` instead of defaulting to + ``conversation_id``. Returns ------- - AgentConversation - A new or restored agent conversation. + Conversation: + The conversation object of the agent. """ 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_runtime_id, conversation_root_id = ( self._prepare_conversation_start( @@ -434,7 +438,7 @@ def start_conversation( conversation_id=conversation_id, checkpointer=checkpointer, checkpoint_id=checkpoint_id, - _root_conversation_id=_root_conversation_id, + _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=AgentConversation, attach_checkpointer=_attach_checkpointer, ) @@ -489,9 +493,7 @@ def start_conversation( ) ) - from wayflowcore.executors._agentexecutor import AgentConversationExecutionState - - conversation = AgentConversation( + return AgentConversation( component=self, message_list=messages, id=conversation_runtime_id, @@ -500,10 +502,9 @@ def start_conversation( name="agent_conversation", state=AgentConversationExecutionState(), status=None, - root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, __metadata_info__={}, ) - return conversation @property def llms(self) -> List["LlmModel"]: diff --git a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py b/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py index a269babf9..e9fd69996 100644 --- a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py +++ b/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py @@ -4,19 +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. import logging -from typing import Dict, Optional from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig from wayflowcore.checkpointing.datastorecheckpointer import ( _prepare_oracle_checkpoint_datastore, _prepare_postgres_checkpoint_datastore, ) -from wayflowcore.checkpointing.serialization import _deserialize_conversation_checkpoint_state -from wayflowcore.component import Component -from wayflowcore.conversation import Conversation from wayflowcore.datastore.oracle import OracleDatabaseConnectionConfig from wayflowcore.datastore.postgres import PostgresDatabaseConnectionConfig -from wayflowcore.tools import Tool logger = logging.getLogger(__name__) @@ -31,15 +26,3 @@ def _prepare_oracle_datastore( connection_config: OracleDatabaseConnectionConfig, storage_config: ServerStorageConfig ) -> None: _prepare_oracle_checkpoint_datastore(connection_config, storage_config) - - -def _deserialize_conversation_safely( - serialized_state: str, - tool_registry: Optional[Dict[str, Tool]] = None, - component: Optional[Component] = None, -) -> Conversation: - return _deserialize_conversation_checkpoint_state( - serialized_state, - tool_registry=tool_registry, - component=component, - ) diff --git a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py index e71d41085..38b91c78f 100644 --- a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py +++ b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py @@ -6,6 +6,7 @@ import logging import time +from collections import OrderedDict from typing import Any, AsyncIterable, Dict, List, Optional, Union, cast import anyio @@ -15,6 +16,11 @@ from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig from wayflowcore.checkpointing import ConversationCheckpoint, DatastoreCheckpointer +from wayflowcore.checkpointing.checkpoint_state import ( + _CheckpointRestoreCompatibilityError, + _save_live_conversation_checkpoint, + _supports_checkpointing, +) from wayflowcore.conversation import Conversation from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.datastore import Datastore, InMemoryDatastore @@ -58,6 +64,8 @@ class WayFlowOpenAIResponsesService(OpenAIResponsesService): + _RESPONSE_CONVERSATION_CACHE_MAX_SIZE = 1024 + def __init__( self, agents: Dict[str, ConversationalComponent], @@ -72,7 +80,10 @@ def __init__( storage_config=self.storage_config, ) self.created_at = int(time.time()) - self._response_conversation_ids: Dict[str, str] = {} + # 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() @@ -127,7 +138,7 @@ async def get_response( detail="Get endpoint for wayflow server only supports non-streaming requests", ) - checkpoint = self._lookup_checkpoint_by_response_id(response_id) + 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" @@ -140,10 +151,10 @@ async def get_response( return Response.model_validate_json(response_as_txt) async def delete_response(self, response_id: str) -> Optional[ResponseError]: - checkpoint = self._lookup_checkpoint_by_response_id(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, None) + self._response_conversation_ids.pop(response_id) return None async def cancel_response(self, response_id: str) -> Union[Response, ResponseError]: @@ -193,16 +204,29 @@ 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 - should_store_response = body.store is None or body.store is True + # 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 _supports_checkpointing(agent): + 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, agent_id=model, - attach_checkpointer=should_store_response, ) - response_id = IdGenerator.get_or_generate_id() state = await self._create_state( agent=agent, state=state, @@ -210,7 +234,7 @@ async def create_response(self, body: CreateResponse) -> AsyncIterable[ResponseS ) current_response = Response( - id=response_id, + id=IdGenerator.get_or_generate_id(), created_at=int(time.time()), error=None, incomplete_details=None, @@ -261,9 +285,7 @@ async def runner(conversation: Conversation) -> None: nonlocal status try: with register_event_listeners([token_usage_listener, yielding_listener]): - status = await conversation.execute_async( - _final_checkpoint_id=response_id if should_store_response else None, - ) + status = await conversation.execute_async() except Exception as e: nonlocal raised_exception raised_exception = e @@ -277,6 +299,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): @@ -316,13 +339,22 @@ async def runner(conversation: Conversation) -> None: token_usage_listener.usage ) - if should_store_response and state.checkpointer is not None: - self.checkpointer.save_conversation( + # persists a completed OpenAI Responses response as a WayFlow checkpoint + if should_store_response: + _save_live_conversation_checkpoint( + self.checkpointer, 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, - metadata={"response": current_response.model_dump_json()}, + # 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._response_conversation_ids[current_response.id] = state.id + self._cache_response_conversation_id(current_response.id, state.id) if current_response.error is not None: yield ResponseFailedEvent( @@ -376,64 +408,102 @@ def _load_state( previous_response_id: Optional[str], conversation_id: Optional[str], agent_id: str, - attach_checkpointer: bool = True, ) -> Optional[Conversation]: if previous_response_id: - checkpoint = self._lookup_checkpoint_by_response_id(previous_response_id) + # 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", ) - self._response_conversation_ids[checkpoint.checkpoint_id] = checkpoint.conversation_id - try: - return self.agents[agent_id].start_conversation( - conversation_id=checkpoint.conversation_id, - checkpoint_id=checkpoint.checkpoint_id, - checkpointer=self.checkpointer, - _attach_checkpointer=attach_checkpointer, - ) - 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}", - ) from e + incompatible_detail = ( + f"Previous response `{previous_response_id}` is not compatible " + f"with model `{agent_id}`" + ) elif conversation_id: + # 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", ) - self._response_conversation_ids[checkpoint.checkpoint_id] = checkpoint.conversation_id - try: - return self.agents[agent_id].start_conversation( - conversation_id=conversation_id, - checkpointer=self.checkpointer, - _attach_checkpointer=attach_checkpointer, - ) - 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}", - ) from e + incompatible_detail = ( + f"Conversation `{conversation_id}` is not compatible with model `{agent_id}`" + ) else: return None - def _lookup_checkpoint_by_response_id( + 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: + # Use start_conversation() only to restore the checkpoint identified by + # (conversation_id, checkpoint_id), then clear conversation.checkpointer + # so execute() resumes without adding automatic checkpoint saves. + conversation = agent.start_conversation( + conversation_id=checkpoint.conversation_id, + checkpointer=self.checkpointer, + checkpoint_id=checkpoint.checkpoint_id, + ) + conversation.checkpointer = None + return conversation + 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}", + ) from e + + def _find_checkpoint_for_response_id( self, response_id: str ) -> Optional[ConversationCheckpoint]: - conversation_id = self._response_conversation_ids.get(response_id) - if conversation_id is not None: - try: - return self.checkpointer.load(conversation_id, response_id) - except ValueError: - self._response_conversation_ids.pop(response_id, None) + """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._response_conversation_ids[response_id] = checkpoint.conversation_id + 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: + return None + + 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 _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, agent: ConversationalComponent, @@ -459,14 +529,7 @@ async def _create_state( detail="Agent should have an `instructions` input descriptor to be able to take instructions as input", ) inputs = {"instructions": instructions} - if request.store is None or request.store is True: - state = agent.start_conversation( - inputs=inputs, - messages=new_messages, - checkpointer=self.checkpointer, - ) - else: - state = agent.start_conversation(inputs=inputs, messages=new_messages) + state = agent.start_conversation(inputs=inputs, messages=new_messages) else: # later: implement context provider for custom instructions if instructions is not None: diff --git a/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py b/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py index f2701cacb..990f3c927 100644 --- a/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py +++ b/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py @@ -11,4 +11,4 @@ @dataclass class ServerStorageConfig(StorageConfig): - """Configuration for agent-server conversation storage.""" + """Configuration for server storage management.""" diff --git a/wayflowcore/src/wayflowcore/checkpointing/__init__.py b/wayflowcore/src/wayflowcore/checkpointing/__init__.py index 7db08002e..4458ca4a0 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/__init__.py +++ b/wayflowcore/src/wayflowcore/checkpointing/__init__.py @@ -1,12 +1,9 @@ -# Copyright © 2025, 2026 Oracle and/or its affiliates. +# 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 importlib import import_module -from typing import Any - from .checkpointer import Checkpointer, CheckpointingInterval, ConversationCheckpoint, StorageConfig from .datastorecheckpointer import ( DatastoreCheckpointer, diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py b/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py new file mode 100644 index 000000000..bd9ac6b43 --- /dev/null +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py @@ -0,0 +1,251 @@ +# 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 time +from itertools import chain +from typing import TYPE_CHECKING, Any, Dict, Optional, Type + +from wayflowcore.exceptions import DataclassFieldDeserializationError +from wayflowcore.idgeneration import IdGenerator +from wayflowcore.serialization import autodeserialize, serialize +from wayflowcore.serialization.context import ( + DeserializationContext, + MissingDeserializationReferenceError, + SerializationContext, + _iter_nested_components, +) + +if TYPE_CHECKING: + from wayflowcore.checkpointing import Checkpointer + from wayflowcore.checkpointing.checkpointer import ConversationCheckpoint + from wayflowcore.conversation import Conversation + from wayflowcore.conversationalcomponent import ConversationalComponent + from wayflowcore.executors._agentexecutor import AgentConversationExecutionState + from wayflowcore.executors._flowexecutor import FlowConversationExecutionState + from wayflowcore.executors._managerworkersconversation import ( + ManagerWorkersConversationExecutionState, + ) + from wayflowcore.executors._swarmconversation import SwarmConversationExecutionState + + +_COMPONENT_ID_ERROR_HINT = "Restart-safe checkpoint restore requires stable component ids." +_STEP_NAME_ERROR_HINT = "Restart-safe checkpoint restore requires stable step names." +_AGENT_NAME_ERROR_HINT = "Restart-safe checkpoint restore requires stable agent names." + + +class _CheckpointRestoreCompatibilityError(ValueError): + """Raised when a checkpoint cannot be resumed against the current live graph.""" + + +def _save_live_conversation_checkpoint( + checkpointer: "Checkpointer", + conversation: "Conversation", + checkpoint_id: Optional[str] = None, + component_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, +) -> "ConversationCheckpoint": + from wayflowcore.checkpointing.checkpointer import ConversationCheckpoint + + if not _supports_checkpointing(conversation.component): + raise NotImplementedError( + "Checkpointing conversations that contain `OciAgent` is not supported yet." + ) + serialization_context = SerializationContext(root=conversation) + serialization_context._add_component_to_context(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 {}), + ) + checkpointer.save(checkpoint) + conversation.checkpoint_id = checkpoint.checkpoint_id + return checkpoint + + +def _load_checkpointed_conversation( + checkpoint: "ConversationCheckpoint", + component: "ConversationalComponent", + expected_conversation_type: Type["Conversation"], + tool_registry: Optional[Dict[str, Any]] = None, + checkpointer: Optional["Checkpointer"] = None, + attach_checkpointer: bool = True, +) -> "Conversation": + """Restore a Conversation object from stored checkpoint state. + + The serialized conversation omits live component/tool objects and rebuilds them from + the current component tree. After deserialization, the helper also repairs + derived runtime-only state that is not stored directly in the checkpoint. + """ + deserialization_context = DeserializationContext() + deserialization_context.registered_tools = tool_registry.copy() if tool_registry else {} + deserialization_context._add_component_to_context(component) + try: + conversation = autodeserialize( + checkpoint.state, + deserialization_context=deserialization_context, + ) + except (MissingDeserializationReferenceError, DataclassFieldDeserializationError) as exc: + if not _contains_missing_reference_error(exc): + raise + raise _CheckpointRestoreCompatibilityError( + "Cannot restore this checkpoint because the current component tree does not " + "match the serialized component ids. " + f"{_COMPONENT_ID_ERROR_HINT}" + ) 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__}`." + ) + _prepare_restored_conversation_states(conversation) + conversation.checkpoint_id = checkpoint.checkpoint_id + + if attach_checkpointer: + conversation.checkpointer = checkpointer + return conversation + + +def _contains_missing_reference_error(error: Exception) -> bool: + """Return whether the exception chain contains a missing component/tool reference.""" + while error is not None: + if isinstance(error, MissingDeserializationReferenceError): + return True + error = error.__cause__ # type: ignore[assignment] + return False + + +def _prepare_restored_conversation_states(root_conversation: "Conversation") -> None: + """Repair executor-specific state after a conversation is restored.""" + from wayflowcore.executors._agentexecutor import AgentConversationExecutionState + from wayflowcore.executors._flowexecutor import FlowConversationExecutionState + from wayflowcore.executors._managerworkersconversation import ( + ManagerWorkersConversationExecutionState, + ) + from wayflowcore.executors._swarmconversation import SwarmConversationExecutionState + + for conversation in chain( + (root_conversation,), + root_conversation._get_all_sub_conversations_recursive(), + ): + state = conversation.state + match state: + case AgentConversationExecutionState(): + _prepare_restored_agent_conversation_state(state) + case FlowConversationExecutionState(): + _validate_restored_flow_conversation_state(state) + case SwarmConversationExecutionState(): + _prepare_restored_swarm_conversation_state(state) + case ManagerWorkersConversationExecutionState(): + _prepare_restored_managerworkers_conversation_state(state) + + +def _prepare_restored_agent_conversation_state(state: "AgentConversationExecutionState") -> None: + """Re-key Agent subconversations by the canonical runtime slot key.""" + from wayflowcore.executors._agentconversation import AgentConversation + + state.current_sub_component_conversations = { + AgentConversation._sub_component_conversation_key( + subconversation.component + ): subconversation + for subconversation in state.current_sub_component_conversations.values() + } + + +def _validate_restored_flow_conversation_state(state: "FlowConversationExecutionState") -> None: + """Fail fast if restored Flow state names steps that no longer exist live.""" + valid_step_names = set(state.flow.steps) + if state.flow.begin_step_name is not None: + valid_step_names.add(state.flow.begin_step_name) + + for step_name in chain( + (state.current_step_name,), + state.step_history, + (step_name for step_name, _output_name in state.input_output_key_values), + ): + if step_name is None or step_name in valid_step_names: + continue + raise _CheckpointRestoreCompatibilityError( + "Cannot restore this checkpoint because flow conversation state refers to " + f"step `{step_name}` which is not present in the current Flow. " + f"{_STEP_NAME_ERROR_HINT}" + ) + + +def _prepare_restored_swarm_conversation_state(state: "SwarmConversationExecutionState") -> None: + """Rebuild Swarm's derived indexes, then assert the active thread is still valid.""" + threads = [state.main_thread] + [ + thread + for recipients_and_threads in state.agents_and_threads.values() + for thread in recipients_and_threads.values() + ] + # Rebuild the recipient lookup from the live thread objects. The serialized + # dict keys may no longer line up after deserialization. + state.agents_and_threads = {} + for thread in threads: + if not thread.is_main_thread: + state.agents_and_threads.setdefault(thread.caller.name, {})[ + thread.recipient_agent.name + ] = thread + # Thread subconversations are matched by sharing the same message list object. + # The thread identifiers are the stable keys; message-list object identity lets + # us reconnect them after deserialization rebuilt the in-memory objects. + thread_ids_by_message_list = {id(thread.message_list): thread.identifier for thread in threads} + state.thread_subconversations = { + thread_id: subconversation + for subconversation in state.thread_subconversations.values() + if (thread_id := thread_ids_by_message_list.get(id(subconversation.message_list))) + is not None + } + if state.current_thread is None: + raise _CheckpointRestoreCompatibilityError( + "Cannot restore this checkpoint because Swarm conversation state does not " + f"have an active thread. {_AGENT_NAME_ERROR_HINT}" + ) + + valid_thread_ids = {thread.identifier for thread in threads} + if state.current_thread.identifier not in valid_thread_ids: + raise _CheckpointRestoreCompatibilityError( + "Cannot restore this checkpoint because Swarm conversation state refers to " + f"thread `{state.current_thread.identifier}` which is not present in the current " + f"Swarm topology. {_AGENT_NAME_ERROR_HINT}" + ) + + +def _prepare_restored_managerworkers_conversation_state( + state: "ManagerWorkersConversationExecutionState", +) -> None: + """Re-key ManagerWorkers subconversations, then assert the active agent is still valid.""" + # Subconversations are restored as objects first; rebuild the name-keyed lookup + # used by the runtime from those live conversation objects. + state.subconversations = { + subconversation.component.name: subconversation + for subconversation in state.subconversations.values() + } + if ( + state.current_agent_name is not None + and state.current_agent_name not in state.subconversations + ): + raise _CheckpointRestoreCompatibilityError( + "Cannot restore this checkpoint because ManagerWorkers conversation state " + f"refers to agent `{state.current_agent_name}` which is not present " + f"in the current component tree. {_AGENT_NAME_ERROR_HINT}" + ) + + +# Checkpoint eligibility + + +def _supports_checkpointing(component: "ConversationalComponent") -> bool: + from wayflowcore.ociagent import OciAgent + + return not any( + isinstance(nested_component, OciAgent) + for nested_component in _iter_nested_components(component) + ) diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py index 0ecafb48f..040ca64da 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py @@ -1,21 +1,15 @@ -# Copyright © 2025, 2026 Oracle and/or its affiliates. +# 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 time from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import Enum from typing import TYPE_CHECKING, Any, Dict, List, Optional -from wayflowcore.idgeneration import IdGenerator - -from .serialization import _serialize_conversation_checkpoint_state - if TYPE_CHECKING: - from wayflowcore.conversation import Conversation from wayflowcore.datastore import Datastore @@ -24,11 +18,17 @@ class ConversationCheckpoint: """Durable snapshot of a conversation at a checkpoint boundary.""" checkpoint_id: str + """External identifier of this saved checkpoint.""" conversation_id: str + """Durable conversation id that this checkpoint belongs to.""" component_id: str + """Best-effort root component id stored for diagnostics and storage queries.""" 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: @@ -36,12 +36,13 @@ def id(self) -> str: class CheckpointingInterval(Enum): - """ - Configure when the conversation is saved during execution. - """ + """Configure which completed execution boundary triggers a checkpoint save.""" + # Save only after the outermost `Conversation.execute()` returns. CONVERSATION_TURNS = "conversation_turns" + # Save after completed internal turns that actually used an LLM. LLM_TURNS = "llm_turns" + # Save after every completed internal agent/flow turn boundary. ALL_INTERNAL_TURNS = "all_internal_turns" @@ -50,16 +51,27 @@ 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 @@ -101,49 +113,15 @@ def load_latest(self, conversation_id: str) -> Optional[ConversationCheckpoint]: def load(self, conversation_id: str, checkpoint_id: str) -> ConversationCheckpoint: raise NotImplementedError() - def save(self, checkpoint: Any) -> None: - from wayflowcore.conversation import Conversation - - if isinstance(checkpoint, Conversation): - self.save_conversation(checkpoint) - return - if not isinstance(checkpoint, ConversationCheckpoint): - raise TypeError( - f"Expected a Conversation or ConversationCheckpoint, got {type(checkpoint).__name__}." - ) - self._save_checkpoint(checkpoint) - - async def save_async(self, checkpoint: Any) -> None: - self.save(checkpoint) - - def save_conversation( - self, - conversation: "Conversation", - *, - checkpoint_id: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> ConversationCheckpoint: - next_save_sequence = self._save_sequence_by_conversation.get(conversation.id, 0) + 1 - self._save_sequence_by_conversation[conversation.id] = next_save_sequence - checkpoint_metadata = {"save_sequence": next_save_sequence} - if metadata: - checkpoint_metadata.update(metadata) - checkpoint = ConversationCheckpoint( - checkpoint_id=checkpoint_id or IdGenerator.get_or_generate_id(), - conversation_id=conversation.id, - component_id=conversation.component.id, - created_at=int(time.time()), - state=_serialize_conversation_checkpoint_state(conversation), - metadata=checkpoint_metadata, - ) - self._save_checkpoint(checkpoint) - conversation.checkpoint_id = checkpoint.checkpoint_id - return checkpoint - @abstractmethod - def _save_checkpoint(self, checkpoint: ConversationCheckpoint) -> None: + def save(self, checkpoint: ConversationCheckpoint) -> None: raise NotImplementedError() + async def save_async(self, checkpoint: ConversationCheckpoint) -> None: + # Async persistence is not implemented yet; this preserves the async API + # contract while delegating to the synchronous backend implementation. + self.save(checkpoint) + @abstractmethod def list_checkpoints( self, conversation_id: str, limit: Optional[int] = 50 diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py index af750a3b2..84a9199d9 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py @@ -1,26 +1,79 @@ -# Copyright © 2025, 2026 Oracle and/or its affiliates. +# 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, nullcontext +from contextlib import contextmanager from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional -from ..events import EventListener -from .checkpointer import CheckpointingInterval +from ..events import Event, EventListener +from ..events.event import ( + AgentExecutionIterationStartedEvent, + FlowExecutionIterationStartedEvent, + LlmGenerationResponseEvent, +) +from .checkpoint_state import _save_live_conversation_checkpoint +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 _build_checkpoint_metadata( + +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[Any] = None, + 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, @@ -28,151 +81,125 @@ def _build_checkpoint_metadata( } 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__ - execution_state = getattr(event, "execution_state", None) - if execution_state is not None: - if hasattr(execution_state, "curr_iter"): - checkpoint_metadata["agent_iteration"] = execution_state.curr_iter - if hasattr(execution_state, "current_step_name"): - checkpoint_metadata["flow_step_name"] = execution_state.current_step_name - if hasattr(execution_state, "nesting_level"): - checkpoint_metadata["nesting_level"] = execution_state.nesting_level + if isinstance(event, AgentExecutionIterationStartedEvent): + checkpoint_metadata["event_type"] = AgentExecutionIterationStartedEvent.__name__ + checkpoint_metadata["agent_iteration"] = event.execution_state.curr_iter + elif isinstance(event, FlowExecutionIterationStartedEvent): + checkpoint_metadata["event_type"] = FlowExecutionIterationStartedEvent.__name__ + 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 -def _save_conversation_checkpoint( - conversation: "Conversation", - *, - save_reason: str, - event: Optional[Any] = None, - metadata: Optional[Dict[str, Any]] = None, - checkpoint_id: Optional[str] = None, -) -> None: - checkpointer = conversation.checkpointer - if checkpointer is None: - return - - checkpoint_metadata = _build_checkpoint_metadata( - conversation, - save_reason=save_reason, - event=event, - metadata=metadata, - ) - - checkpointer.save_conversation( - conversation, - checkpoint_id=checkpoint_id, - metadata=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. + """ -class _ConversationCheckpointEventListener(EventListener): - def __init__(self, conversation: "Conversation") -> None: + def __init__(self, conversation: "Conversation", checkpointer: Checkpointer) -> None: self.conversation = conversation - self._llm_was_used_since_last_internal_turn = False - self._last_internal_turn_start_event: Optional[Any] = None - - def __call__(self, event: Any) -> None: - from wayflowcore.events.event import ( - AgentExecutionIterationStartedEvent, - FlowExecutionIterationStartedEvent, - LlmGenerationResponseEvent, - ) - - checkpointer = self.conversation.checkpointer - if checkpointer is None: - return + 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._llm_was_used_since_last_internal_turn = True + self._pending_llm_checkpoint = True return - if not isinstance( - event, (AgentExecutionIterationStartedEvent, FlowExecutionIterationStartedEvent) - ): + if not isinstance(event, _ITERATION_STARTED_EVENTS): return - checkpointing_interval = checkpointer.checkpointing_interval + self._latest_checkpoint_boundary_event = event - if checkpointing_interval == CheckpointingInterval.CONVERSATION_TURNS: - self._last_internal_turn_start_event = event + checkpointed_conversation = _find_checkpointed_conversation( + self.conversation, event.execution_state + ) + if checkpointed_conversation is None: return - if checkpointing_interval == CheckpointingInterval.ALL_INTERNAL_TURNS: - _save_conversation_checkpoint( - self.conversation, - save_reason="internal_turn_boundary", - event=event, - metadata={ - "llm_used_in_previous_turn": self._llm_was_used_since_last_internal_turn, - }, - ) - self._llm_was_used_since_last_internal_turn = False - self._last_internal_turn_start_event = event + 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 self._llm_was_used_since_last_internal_turn: - _save_conversation_checkpoint( - self.conversation, - save_reason="internal_turn_boundary", - event=event, - metadata={ - "llm_used_in_previous_turn": self._llm_was_used_since_last_internal_turn, - }, - ) - self._llm_was_used_since_last_internal_turn = False - - self._last_internal_turn_start_event = event - - def flush_pending_checkpoint(self) -> None: - checkpointer = self.conversation.checkpointer - if checkpointer is None: + if not self._pending_llm_checkpoint: return - if checkpointer.checkpointing_interval != CheckpointingInterval.LLM_TURNS: + event = self._latest_checkpoint_boundary_event + if event is None: return - if not self._llm_was_used_since_last_internal_turn: + if _find_checkpointed_conversation(self.conversation, event.execution_state) is None: return - _save_conversation_checkpoint( + # 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: + _save_live_conversation_checkpoint( + self.checkpointer, self.conversation, - save_reason="internal_turn_boundary", - event=self._last_internal_turn_start_event, - metadata={ - "llm_used_in_previous_turn": True, - }, + metadata=_build_listener_checkpoint_metadata( + self.conversation, + save_reason=_CHECKPOINT_SAVE_REASON_INTERNAL_TURN_BOUNDARY, + event=event, + ), ) - self._llm_was_used_since_last_internal_turn = False + self._pending_llm_checkpoint = False @contextmanager def get_conversation_checkpoint_execution_context( conversation: "Conversation", - *, is_outermost_execution: bool, - final_checkpoint_id: Optional[str] = None, - final_checkpoint_metadata: Optional[Dict[str, Any]] = None, ) -> Iterator[None]: - if conversation.checkpointer is None or not is_outermost_execution: - with nullcontext(): - yield + """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) + listener = _ConversationCheckpointEventListener(conversation, checkpointer) with register_event_listeners([listener]): - try: - yield - except Exception: - raise - else: - listener.flush_pending_checkpoint() - _save_conversation_checkpoint( + yield + listener.save_pending_llm_checkpoint() + _save_live_conversation_checkpoint( + checkpointer, + conversation, + metadata=_build_listener_checkpoint_metadata( conversation, - save_reason="conversation_turn", - checkpoint_id=final_checkpoint_id, - metadata=final_checkpoint_metadata, - ) + save_reason=_CHECKPOINT_SAVE_REASON_CONVERSATION_TURN, + ), + ) diff --git a/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py index 27bb5a2bb..51a247cce 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py +++ b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py @@ -1,10 +1,12 @@ -# Copyright © 2025, 2026 Oracle and/or its affiliates. +# 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 import warnings from textwrap import dedent from typing import Any, Dict, List, Optional, Sequence @@ -27,7 +29,6 @@ def _build_checkpoint_create_table_columns( storage_config: StorageConfig, - *, is_oracle: bool, ) -> List[str]: text_type = "CLOB" if is_oracle else "TEXT" @@ -47,6 +48,38 @@ def _build_checkpoint_create_table_columns( 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, @@ -62,6 +95,10 @@ def _prepare_postgres_checkpoint_datastore( ) 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( @@ -83,6 +120,10 @@ def _prepare_oracle_checkpoint_datastore( ) 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( @@ -127,6 +168,8 @@ def _checkpoint_to_entity(self, checkpoint: ConversationCheckpoint) -> Dict[str, 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 ) @@ -136,10 +179,14 @@ def _checkpoint_to_entity(self, checkpoint: ConversationCheckpoint) -> Dict[str, 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, ), @@ -147,7 +194,6 @@ def _sort_checkpoints( def _find_checkpoint( self, - *, conversation_id: str, checkpoint_id: str, ) -> Optional[ConversationCheckpoint]: @@ -183,6 +229,13 @@ def load_latest(self, conversation_id: str) -> Optional[ConversationCheckpoint]: ) 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] ) @@ -198,7 +251,24 @@ def load(self, conversation_id: str, checkpoint_id: str) -> ConversationCheckpoi ) return checkpoint - def _save_checkpoint(self, checkpoint: ConversationCheckpoint) -> None: + def save(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, @@ -223,6 +293,9 @@ def _save_checkpoint(self, checkpoint: ConversationCheckpoint) -> None: 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, @@ -257,37 +330,41 @@ def _save_checkpoint(self, checkpoint: ConversationCheckpoint) -> None: ) ) 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=update_latest_where, - update=update_latest_values, + 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, ) - 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 @@ -301,15 +378,22 @@ def list_checkpoints( ) ] ) - if limit is not None and len(checkpoints) > limit: - checkpoints = checkpoints[-limit:] + 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 - if latest_checkpoint is not None and latest_checkpoint.checkpoint_id == checkpoint_id: + 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 @@ -317,12 +401,41 @@ def delete(self, conversation_id: str, checkpoint_id: str) -> None: ] 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={ - self.storage_config.conversation_id_column_name: conversation_id, - self.storage_config.turn_id_column_name: checkpoint_id, - }, + where=delete_where, ) if checkpoint_to_promote is not None: diff --git a/wayflowcore/src/wayflowcore/checkpointing/serialization.py b/wayflowcore/src/wayflowcore/checkpointing/serialization.py deleted file mode 100644 index d757a3bc0..000000000 --- a/wayflowcore/src/wayflowcore/checkpointing/serialization.py +++ /dev/null @@ -1,143 +0,0 @@ -# Copyright © 2025, 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 TYPE_CHECKING, Any, Dict, List, Optional, Sequence, cast - -import yaml - -from wayflowcore.serialization import autodeserialize, serialize_to_dict -from wayflowcore.serialization.context import DeserializationContext, SerializationContext -from wayflowcore.serialization.serializer import autodeserialize_from_dict - -if TYPE_CHECKING: - from wayflowcore.component import Component - from wayflowcore.conversation import Conversation - - -_CHECKPOINT_ENVELOPE_FORMAT = "wayflow-conversation-checkpoint" -_CHECKPOINT_ENVELOPE_VERSION = 1 - - -def _iter_conversation_graph(root_conversation: "Conversation") -> Sequence["Conversation"]: - visited_conversation_ids: set[str] = set() - queue: List["Conversation"] = [root_conversation] - ordered_conversations: List["Conversation"] = [] - - while queue: - conversation = queue.pop() - if conversation.id in visited_conversation_ids: - continue - visited_conversation_ids.add(conversation.id) - ordered_conversations.append(conversation) - queue.extend(conversation._get_all_sub_conversations()) - - return ordered_conversations - - -def _ensure_checkpointing_supported(conversation: "Conversation") -> None: - from wayflowcore.ociagent import OciAgent - - for sub_conversation in _iter_conversation_graph(conversation): - if isinstance(sub_conversation.component, OciAgent): - raise NotImplementedError( - "Checkpointing conversations that contain `OciAgent` is not supported yet." - ) - - -def _iter_component_tree(component: "Component") -> Sequence["Component"]: - from wayflowcore.component import Component - - def _iter_nested_components(value: Any) -> List["Component"]: - if isinstance(value, Component): - return [value] - if isinstance(value, dict): - nested_components: List["Component"] = [] - for nested_value in value.values(): - nested_components.extend(_iter_nested_components(nested_value)) - return nested_components - if isinstance(value, (list, tuple, set)): - nested_components = [] - for nested_value in value: - nested_components.extend(_iter_nested_components(nested_value)) - return nested_components - return [] - - visited_component_ids: set[str] = set() - ordered_components: List["Component"] = [] - queue: List["Component"] = [component] - - while queue: - current_component = queue.pop() - current_component_ref = SerializationContext.get_reference(current_component) - if current_component_ref in visited_component_ids: - continue - visited_component_ids.add(current_component_ref) - ordered_components.append(current_component) - - all_public_attrs = { - name: value - for name, value in vars(current_component).items() - if not name.startswith("_") - } - for attr in all_public_attrs.values(): - queue.extend(_iter_nested_components(attr)) - - return ordered_components - - -def _build_checkpoint_serialization_context(conversation: "Conversation") -> SerializationContext: - serialization_context = SerializationContext(root=conversation) - for component in _iter_component_tree(conversation.component): - serialization_context.register_external_reference(component) - return serialization_context - - -def _serialize_conversation_checkpoint_state(conversation: "Conversation") -> str: - _ensure_checkpointing_supported(conversation) - - serialized_conversation = serialize_to_dict( - conversation, - serialization_context=_build_checkpoint_serialization_context(conversation), - ) - - envelope = { - "checkpoint_format": _CHECKPOINT_ENVELOPE_FORMAT, - "version": _CHECKPOINT_ENVELOPE_VERSION, - "conversation": serialized_conversation, - } - return yaml.safe_dump(envelope) - - -def _deserialize_conversation_checkpoint_state( - serialized_state: str, - *, - tool_registry: Optional[Dict[str, Any]] = None, - component: Optional["Component"] = None, -) -> "Conversation": - deserialization_context = DeserializationContext() - deserialization_context.registered_tools = tool_registry.copy() if tool_registry else {} - - if component is not None: - deserialization_context._add_component_to_context(component) - - state_payload = yaml.safe_load(serialized_state) - if ( - isinstance(state_payload, dict) - and state_payload.get("checkpoint_format") == _CHECKPOINT_ENVELOPE_FORMAT - and state_payload.get("version") == _CHECKPOINT_ENVELOPE_VERSION - and "conversation" in state_payload - ): - conversation = autodeserialize_from_dict( - state_payload["conversation"], - deserialization_context=deserialization_context, - ) - else: - conversation = autodeserialize( - serialized_state, - deserialization_context=deserialization_context, - ) - - return cast("Conversation", conversation) diff --git a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py index 0bc90f4d1..9d83e4dc1 100644 --- a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py +++ b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py @@ -95,7 +95,7 @@ async def call_async(self, conversation: "Conversation") -> Any: conversation = self.flow.start_conversation( inputs={}, messages=conversation.message_list, - _root_conversation_id=conversation.root_conversation_id, + conversation_id=conversation.conversation_id, ) status = await conversation.execute_async() if status._requires_yielding: diff --git a/wayflowcore/src/wayflowcore/conversation.py b/wayflowcore/src/wayflowcore/conversation.py index c2c2dadd2..67202bc92 100644 --- a/wayflowcore/src/wayflowcore/conversation.py +++ b/wayflowcore/src/wayflowcore/conversation.py @@ -63,17 +63,22 @@ def _get_active_conversations(return_copy: bool = True) -> List["Conversation"]: return copy(active_conversations) if return_copy else active_conversations -def is_outermost_execution() -> bool: - return len(_get_active_conversations(return_copy=False)) == 0 - - def _get_current_conversation_id() -> Optional[str]: + """Return the runtime 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_root_conversation_id() -> Optional[str]: + """Return the root conversation id shared by 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: @@ -96,15 +101,17 @@ class Conversation(DataclassComponent): message_list: MessageList status: Optional[ExecutionStatus] token_usage: TokenUsage = field(default_factory=TokenUsage, init=False) - root_conversation_id: str = "" + conversation_id: str = "" + """Root conversation 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)""" @@ -112,8 +119,8 @@ class Conversation(DataclassComponent): def __post_init__(self) -> None: if self.inputs is None: self.inputs = {} - if not self.root_conversation_id: - self.root_conversation_id = self.id + if not self.conversation_id: + self.conversation_id = self.id @property def plan(self) -> Optional[ExecutionPlan]: @@ -128,9 +135,6 @@ def _register_event(self, event: Event) -> None: def execute( self, execution_interrupts: Optional[Sequence["ExecutionInterrupt"]] = None, - *, - _final_checkpoint_id: Optional[str] = None, - _final_checkpoint_metadata: Optional[Dict[str, Any]] = None, ) -> "ExecutionStatus": """ Execute the conversation and get its ``ExecutionStatus`` based on the outcome. @@ -140,20 +144,13 @@ def execute( """ async def _execute_async_wrapper() -> "ExecutionStatus": - return await self.execute_async( - execution_interrupts, - _final_checkpoint_id=_final_checkpoint_id, - _final_checkpoint_metadata=_final_checkpoint_metadata, - ) + return await self.execute_async(execution_interrupts) return run_async_in_sync(_execute_async_wrapper, method_name="execute_async") async def execute_async( self, execution_interrupts: Optional[Sequence["ExecutionInterrupt"]] = None, - *, - _final_checkpoint_id: Optional[str] = None, - _final_checkpoint_metadata: Optional[Dict[str, Any]] = None, ) -> "ExecutionStatus": """ Execute the conversation and get its ``ExecutionStatus`` based on the outcome. @@ -170,13 +167,10 @@ async def execute_async( with get_conversation_checkpoint_execution_context( self, - is_outermost_execution=is_outermost_execution(), - final_checkpoint_id=_final_checkpoint_id, - final_checkpoint_metadata=_final_checkpoint_metadata, + is_outermost_execution=len(_get_active_conversations(return_copy=False)) == 0, ): with _register_conversation(self): - new_status = await self.component.runner.execute_async(self, execution_interrupts) - self.status = new_status + self.status = await self.component.runner.execute_async(self, execution_interrupts) self.status_handled = False return self.status @@ -194,9 +188,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 ee4ae5105..80ca50de9 100644 --- a/wayflowcore/src/wayflowcore/conversationalcomponent.py +++ b/wayflowcore/src/wayflowcore/conversationalcomponent.py @@ -18,11 +18,13 @@ Type, TypeVar, Union, + cast, ) 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__) @@ -30,10 +32,9 @@ 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 @@ -42,7 +43,6 @@ class ConversationalComponent(ComponentWithInputsOutputs, ABC): - def __init__( self, name: str, @@ -82,36 +82,17 @@ 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, - _root_conversation_id: Optional[str] = None, + _runtime_conversation_id: Optional[str] = None, _attach_checkpointer: bool = True, ) -> "Conversation": - """ - Start a conversation for this component. - - Parameters - ---------- - inputs: - Optional structured inputs used to initialize the conversation. - messages: - Optional initial message history. Concrete implementations normalize this into a - ``MessageList`` when needed. - conversation_id: - Optional identifier for the concrete conversation instance. - checkpointer: - Optional checkpoint backend used to restore and persist conversation state. - checkpoint_id: - Optional checkpoint identifier to restore. Requires ``checkpointer``. - _root_conversation_id: - Internal lineage identifier shared by nested conversations for usage accounting, - execution limits, and checkpoint lineage. - - Returns - ------- - Conversation - A new or restored conversation instance ready for execution. + """Start a fresh conversation or restore one from checkpoint storage. + + ``conversation_id`` is the durable conversation identifier used for resume and + storage. For fresh conversations, ``_runtime_conversation_id`` controls the + concrete ``Conversation.id`` assigned to the created object; if omitted, the + runtime id defaults to the durable conversation id. """ @property @@ -177,33 +158,44 @@ def _messages_or_inputs_were_passed( def _prepare_conversation_start( self, - *, inputs: Optional[Dict[str, Any]], messages: Union[None, str, "Message", List["Message"], "MessageList"], conversation_id: Optional[str], - _root_conversation_id: Optional[str], + _runtime_conversation_id: Optional[str], checkpointer: Optional["Checkpointer"], checkpoint_id: Optional[str], expected_conversation_type: Type[ConversationTypeT], attach_checkpointer: bool, ) -> 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 runtime conversation id to assign to ``Conversation.id`` on the concrete object + - the root conversation id shared across nested fresh starts + + ``expected_conversation_type`` keeps the restored conversation typed for mypy; + runtime validation uses it too. + """ + # No checkpointer means this is a fresh conversation; just resolve ids. if checkpointer is None: if checkpoint_id is not None: raise ValueError("`checkpoint_id` requires a `checkpointer`.") - runtime_conversation_id = IdGenerator.get_or_generate_id(conversation_id) - return None, runtime_conversation_id, _root_conversation_id or runtime_conversation_id + root_conversation_id = conversation_id or IdGenerator.get_or_generate_id() + # This value becomes the fresh conversation object's `.id`. + runtime_conversation_id = IdGenerator.get_or_generate_id( + _runtime_conversation_id or root_conversation_id + ) + return None, runtime_conversation_id, root_conversation_id - if ( - _root_conversation_id is not None - and conversation_id is not None - and _root_conversation_id != conversation_id - ): + if _runtime_conversation_id is not None: raise ValueError( - "`root_conversation_id` and `conversation_id` cannot differ when checkpointing is enabled." + "`_runtime_conversation_id` is not supported when restoring checkpoints." ) - resolved_conversation_id = conversation_id or _root_conversation_id + # Checkpoint restore needs the root conversation id to locate stored state. + resolved_conversation_id = conversation_id if resolved_conversation_id is None and checkpoint_id is not None: raise ValueError("`checkpoint_id` requires a `conversation_id`.") if resolved_conversation_id is None: @@ -217,53 +209,29 @@ def _prepare_conversation_start( if checkpoint is None: return None, resolved_conversation_id, resolved_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, - checkpointer=checkpointer, + from wayflowcore.checkpointing.checkpoint_state import _load_checkpointed_conversation + + # Rehydrate the stored conversation against this live component tree. + conversation = _load_checkpointed_conversation( + checkpoint, + component=self, expected_conversation_type=expected_conversation_type, + tool_registry={tool.name: tool for tool in self._referenced_tools()}, + checkpointer=checkpointer, attach_checkpointer=attach_checkpointer, ) - return conversation, resolved_conversation_id, resolved_conversation_id - - def _restore_checkpointed_conversation( - self, - *, - checkpoint: "ConversationCheckpoint", - checkpointer: "Checkpointer", - expected_conversation_type: Type[ConversationTypeT], - attach_checkpointer: bool, - ) -> ConversationTypeT: - from wayflowcore.checkpointing.serialization import ( - _deserialize_conversation_checkpoint_state, - ) - - if checkpoint.component_id != self.id: - raise ValueError( - "Cannot restore this checkpoint because this conversation was started with another " - f"component. Checkpoint component id: `{checkpoint.component_id}`. Current component id: `{self.id}`." - ) - - conversation = _deserialize_conversation_checkpoint_state( - checkpoint.state, - tool_registry={tool.name: tool for tool in self._referenced_tools()}, - component=self, + return ( + cast(ConversationTypeT, conversation), + resolved_conversation_id, + resolved_conversation_id, ) - if not isinstance(conversation, expected_conversation_type): - raise ValueError( - "Cannot restore this checkpoint because this conversation was started with another " - f"component. Expected `{expected_conversation_type.__name__}`, got `{type(conversation).__name__}`." - ) - - if attach_checkpointer: - conversation.checkpointer = checkpointer - conversation.checkpoint_id = checkpoint.checkpoint_id - return conversation # Define a TypeVar that represents the component's type 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/events/event.py b/wayflowcore/src/wayflowcore/events/event.py index c6428cbfa..4c9c39f1c 100644 --- a/wayflowcore/src/wayflowcore/events/event.py +++ b/wayflowcore/src/wayflowcore/events/event.py @@ -767,7 +767,7 @@ class ConversationExecutionStartedEvent(StartSpanEvent["ConversationSpan"]): def to_tracing_info(self, mask_sensitive_information: bool = True) -> Dict[str, Any]: return { **super().to_tracing_info(mask_sensitive_information=mask_sensitive_information), - "conversation.id": self.conversation.id, + "conversation.id": self.conversation.conversation_id, "conversation.name": self.conversation.name, } @@ -788,7 +788,7 @@ class ConversationExecutionFinishedEvent(EndSpanEvent["ConversationSpan"]): def to_tracing_info(self, mask_sensitive_information: bool = True) -> Dict[str, Any]: return { **super().to_tracing_info(mask_sensitive_information=mask_sensitive_information), - "conversation.id": self.conversation.id, + "conversation.id": self.conversation.conversation_id, "conversation.name": self.conversation.name, "execution_status": self.execution_status.__class__.__name__, } diff --git a/wayflowcore/src/wayflowcore/executors/_agentconversation.py b/wayflowcore/src/wayflowcore/executors/_agentconversation.py index b8e976710..ca8528a74 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_agentconversation.py @@ -65,15 +65,22 @@ def _get_all_sub_conversations(self) -> List["Conversation"]: sub_conversations += [self.state.current_flow_conversation] return sub_conversations + @staticmethod + def _sub_component_conversation_key(component: ConversationalComponent) -> str: + """Return the canonical runtime key for a subcomponent conversation slot.""" + return f"{component.__class__.__name__}:{component.name}" + def _get_sub_component_conversation( self, component: ConversationalComponent ) -> Optional["Conversation"]: - return self.state.current_sub_component_conversations.get(component.id) + return self.state.current_sub_component_conversations.get( + self._sub_component_conversation_key(component) + ) def _set_sub_component_conversation( self, component: ConversationalComponent, conversation: Optional["Conversation"] ) -> None: - identifier = component.id + identifier = self._sub_component_conversation_key(component) component_conversations = self.state.current_sub_component_conversations if not conversation and identifier in component_conversations: component_conversations.pop(identifier) @@ -91,14 +98,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/_agentexecutor.py b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py index 48706fe8e..8b7005d7c 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] = {} @@ -464,7 +472,8 @@ def _get_or_create_expert_agent_subconversation( sub_agent_conversation = expert_agent.start_conversation( messages=init_messages, inputs=inputs, - _root_conversation_id=caller_conv.root_conversation_id, + conversation_id=caller_conv.conversation_id, + _runtime_conversation_id=caller_conv._sub_component_conversation_key(expert_agent), ) return sub_agent_conversation @@ -521,7 +530,8 @@ async def _execute_flow( messages: MessageList, flow: Flow, inputs: Dict[str, Any], - root_conversation_id: Optional[str], + conversation_id: Optional[str], + runtime_conversation_id: Optional[str], ) -> Tuple[Any, str, ExecutionStatus]: """ Execute a flow and return its outputs and its execution status. @@ -534,7 +544,8 @@ async def _execute_flow( state.current_flow_conversation = flow.start_conversation( inputs=inputs, messages=messages, - _root_conversation_id=root_conversation_id, + conversation_id=conversation_id, + _runtime_conversation_id=runtime_conversation_id, ) messages.append_message( Message( @@ -764,7 +775,8 @@ async def _handle_flow_call( messages, flow, tool_request.args, - conversation.root_conversation_id, + conversation.conversation_id, + conversation._sub_component_conversation_key(flow), ) ) @@ -1081,9 +1093,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( @@ -1101,7 +1116,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 @@ -1119,8 +1134,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: @@ -1129,7 +1146,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 @@ -1173,11 +1190,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..aeec574a7 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, @@ -33,8 +50,8 @@ def _gather_flow_outputs(self) -> Dict[str, Any]: def _get_internal_context_value_for_step(self, assistant_step: "Step", key: str) -> Any: from wayflowcore.executors._flowexecutor import FlowConversationExecutor - key = FlowConversationExecutor().make_key_for_step(assistant_step, key) - return self.state.internal_context_key_values.get(key, None) + context_key = FlowConversationExecutor().make_key_for_step(assistant_step, key) + return self.state.internal_context_key_values.get(context_key) def _put_internal_context_key_value(self, key: str, value: Any) -> None: self.state.internal_context_key_values[key] = value @@ -70,15 +87,17 @@ def _get_step(self, step_name: str) -> "Step": return self.state.flow.steps[step_name] def _get_current_sub_conversation( - self, step: "Step", sub_conversation_id: Optional[str] = None + self, + step: "Step", + sub_conversation_id: Optional[str] = None, ) -> Optional["Conversation"]: from wayflowcore.executors._flowexecutor import FlowConversationExecutor key = FlowConversationExecutor().make_key_for_step( - step, sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY + step, + sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY, ) - value = self.state.internal_context_key_values.get(key, None) - return cast(Optional["Conversation"], value) + return cast(Optional["Conversation"], self.state.internal_context_key_values.get(key)) def _update_sub_conversation( self, @@ -101,7 +120,8 @@ def _get_or_create_current_sub_conversation( sub_conversation_id: Optional[str] = None, ) -> "FlowConversation": sub_conversation = self._get_current_sub_conversation( - step=step, sub_conversation_id=sub_conversation_id + step=step, + sub_conversation_id=sub_conversation_id, ) if sub_conversation is None: sub_conversation = self._create_sub_conversation( @@ -141,7 +161,8 @@ def _cleanup_sub_conversation( from wayflowcore.executors._flowexecutor import FlowConversationExecutor key = FlowConversationExecutor().make_key_for_step( - step, sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY + step, + sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY, ) self.state.internal_context_key_values.pop(key, None) @@ -184,14 +205,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 fa8d3bef8..6cab1bfd2 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py @@ -246,7 +246,7 @@ class FlowConversationExecutor(ConversationExecutor): @staticmethod def make_key_for_step(assistant_step: Step, key: str) -> str: - return str(assistant_step.id) + "_" + key + return str(assistant_step.name) + "_" + key @staticmethod def get_parent_conversation(state: FlowConversationExecutionState) -> Optional[Conversation]: @@ -286,17 +286,21 @@ def create_sub_conversation( k: v for k, v in inputs.items() if k not in all_context_provider_keys } + resolved_sub_conversation_id = ( + sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY + ) sub_conversation = flow.start_conversation( inputs_not_from_context_providers, - _root_conversation_id=conversation.root_conversation_id, + conversation_id=conversation.conversation_id, + _runtime_conversation_id=FlowConversationExecutor.make_key_for_step( + step, resolved_sub_conversation_id + ), 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 @@ -308,12 +312,14 @@ def create_sub_conversation( def cleanup_sub_conversation( state: FlowConversationExecutionState, step: Step, + sub_conversation_id: Optional[str] = None, ) -> None: """ Remove a subconversation saved in the internal context of the flow execution state, to cleanup the state after the subconversation is finished. """ key = FlowConversationExecutor.make_key_for_step( - step, FlowConversationExecutor._SUB_CONVERSATION_KEY + step, + sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY, ) # We have to pop finished sub-conversations from the context store as other methods assume the conversation # (dict value) is not `None`. @@ -325,10 +331,11 @@ def get_current_sub_conversation( step: Step, ) -> Optional[Conversation]: """Get the current sub conversation of a given step""" - key1 = FlowConversationExecutor.make_key_for_step( - step, FlowConversationExecutor._SUB_CONVERSATION_KEY + key = FlowConversationExecutor.make_key_for_step( + step, + FlowConversationExecutor._SUB_CONVERSATION_KEY, ) - sub_conv = state.internal_context_key_values.get(key1, None) + sub_conv = state.internal_context_key_values.get(key, None) return cast(Conversation, sub_conv) if sub_conv is not None else None @staticmethod @@ -928,8 +935,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 e298be4a7..6dbfb90c3 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py @@ -26,12 +26,15 @@ class ManagerWorkersConversationExecutionState(ConversationExecutionState): current_agent_name: str subconversations: Dict[str, Union["AgentConversation", "ManagerWorkersConversation"]] - root_conversation_id: str = "" + conversation_id: str = "" def _create_subconversation_for_agent( self, agent: Union[Agent, ManagerWorkers] ) -> Union["AgentConversation", "ManagerWorkersConversation"]: - subconv = agent.start_conversation(_root_conversation_id=self.root_conversation_id or None) + subconv = agent.start_conversation( + conversation_id=self.conversation_id or None, + _runtime_conversation_id=agent.name, + ) self.subconversations[agent.name] = subconv return subconv diff --git a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py index 234d2ebe1..6e06a831d 100644 --- a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py @@ -55,7 +55,7 @@ class SwarmConversationExecutionState(ConversationExecutionState): main_thread: SwarmThread agents_and_threads: Dict[str, Dict[str, SwarmThread]] context_providers: List["ContextProvider"] - root_conversation_id: str = "" + conversation_id: str = "" current_thread: Optional["SwarmThread"] = None thread_stack: List["SwarmThread"] = field(default_factory=list) @@ -90,7 +90,8 @@ def _create_subconversation_for_thread( conversation = thread.recipient_agent.start_conversation( inputs=inputs, messages=thread.message_list, - _root_conversation_id=self.root_conversation_id or None, + conversation_id=self.conversation_id or None, + _runtime_conversation_id=thread_id, ) self.thread_subconversations[thread_id] = conversation @@ -120,7 +121,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})" diff --git a/wayflowcore/src/wayflowcore/executors/interrupts/tokenlimitexecutioninterrupt.py b/wayflowcore/src/wayflowcore/executors/interrupts/tokenlimitexecutioninterrupt.py index fa7ce580e..e96035453 100644 --- a/wayflowcore/src/wayflowcore/executors/interrupts/tokenlimitexecutioninterrupt.py +++ b/wayflowcore/src/wayflowcore/executors/interrupts/tokenlimitexecutioninterrupt.py @@ -126,7 +126,7 @@ def _return_status_if_condition_is_met( self, state: ConversationExecutionState, conversation: "Conversation" ) -> Optional[InterruptedExecutionStatus]: - conversation_id = conversation.root_conversation_id + conversation_id = conversation.conversation_id # We first check the global token limit, then we go over the llm-wise limits # Note that we must do the checks separately, because the list of all models diff --git a/wayflowcore/src/wayflowcore/flow.py b/wayflowcore/src/wayflowcore/flow.py index 9710ca7ce..f0b32b22e 100644 --- a/wayflowcore/src/wayflowcore/flow.py +++ b/wayflowcore/src/wayflowcore/flow.py @@ -1165,41 +1165,43 @@ def start_conversation( inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, conversation_id: Optional[str] = None, - nesting_level: int = 0, - context_providers_from_parent_flow: Optional[Set[str]] = None, - *, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _root_conversation_id: Optional[str] = None, + _runtime_conversation_id: Optional[str] = None, _attach_checkpointer: bool = True, + nesting_level: int = 0, + context_providers_from_parent_flow: Optional[Set[str]] = None, ) -> "FlowConversation": """ - Start a conversation for this flow. + Start the conversation. Parameters ---------- inputs: - Optional input values used to initialize flow execution. - messages: - Optional message history available to the flow at startup. + Dictionary of inputs. Keys are the variable identifiers and + values are the actual inputs to start the conversation. conversation_id: - Optional identifier for this flow conversation. - nesting_level: - Nesting level of the flow execution. Nested subflows increase this value. - context_providers_from_parent_flow: - Names of inputs already provided by parent-flow context providers when validating - required inputs for nested execution. + Durable conversation id used for resume, storage, and usage accounting. + messages: + List of messages (``MessageList`` object) before starting the conversation. checkpointer: Optional checkpoint backend used to restore and persist this conversation. checkpoint_id: - Optional checkpoint identifier to restore. Requires ``checkpointer``. - _root_conversation_id: - Internal lineage identifier shared with nested or parent conversations. + Optional checkpoint identifier to restore. Requires both ``checkpointer`` and + ``conversation_id``. + _runtime_conversation_id: + Internal runtime id for a fresh conversation. When provided, it becomes + the created conversation object's ``.id`` instead of defaulting to + ``conversation_id``. + context_providers_from_parent_flow: + Context provider that don't need to be checked when validating existing inputs. + nesting_level: + Nesting level of the conversation. Returns ------- - FlowConversation - A new or restored flow conversation. + Conversation: + A Flow Conversation object. """ from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event @@ -1212,7 +1214,7 @@ def start_conversation( conversation_id=conversation_id, checkpointer=checkpointer, checkpoint_id=checkpoint_id, - _root_conversation_id=_root_conversation_id, + _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=FlowConversation, attach_checkpointer=_attach_checkpointer, ) @@ -1298,7 +1300,7 @@ def start_conversation( nesting_level=nesting_level, ) - conversation = FlowConversation( + return FlowConversation( component=self, inputs=inputs, id=conversation_runtime_id, @@ -1308,9 +1310,8 @@ def start_conversation( status=None, name="flow_conversation", state=state, - root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, ) - return conversation @property def llms(self) -> List["LlmModel"]: @@ -1663,6 +1664,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, @@ -1686,6 +1688,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: @@ -1748,6 +1752,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 a64c47fd5..2428beea8 100644 --- a/wayflowcore/src/wayflowcore/managerworkers.py +++ b/wayflowcore/src/wayflowcore/managerworkers.py @@ -222,37 +222,38 @@ def start_conversation( inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, conversation_id: Optional[str] = None, - conversation_name: Optional[str] = None, - *, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _root_conversation_id: Optional[str] = None, + _runtime_conversation_id: Optional[str] = None, _attach_checkpointer: bool = True, + conversation_name: Optional[str] = None, ) -> "ManagerWorkersConversation": """ - Start a conversation for the manager-workers group. + Initializes a conversation with the managerworkers. Parameters ---------- inputs: - Optional input values passed to the manager's main conversation. + Dictionary of inputs. Keys are the variable identifiers and + values are the actual inputs to start the main conversation. messages: - Optional shared message history between the user and the manager. + Message list of the manager agent and the end-user. conversation_id: - Optional identifier for this manager-workers conversation. - conversation_name: - Optional display name used for the created conversation object. + 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 ``checkpointer``. - _root_conversation_id: - Internal lineage identifier shared with nested or parent conversations. + Optional checkpoint identifier to restore. Requires both ``checkpointer`` and + ``conversation_id``. + _runtime_conversation_id: + Internal runtime id for a fresh conversation. When provided, it becomes + the created conversation object's ``.id`` instead of defaulting to + ``conversation_id``. Returns ------- - ManagerWorkersConversation - A new or restored manager-workers conversation. + Conversation: + The conversation object of the managerworkers. """ from wayflowcore.agentconversation import AgentConversation from wayflowcore.events.event import ConversationCreatedEvent @@ -269,7 +270,7 @@ def start_conversation( conversation_id=conversation_id, checkpointer=checkpointer, checkpoint_id=checkpoint_id, - _root_conversation_id=_root_conversation_id, + _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=ManagerWorkersConversation, attach_checkpointer=_attach_checkpointer, ) @@ -294,13 +295,14 @@ def start_conversation( subconversations[self.manager_agent.name] = self.manager_agent.start_conversation( inputs=inputs, messages=messages, - _root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, + _runtime_conversation_id=self.manager_agent.name, ) state = ManagerWorkersConversationExecutionState( current_agent_name=self.manager_agent.name, subconversations=subconversations, - root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, ) conversation = ManagerWorkersConversation( @@ -312,7 +314,7 @@ def start_conversation( state=state, status=None, checkpointer=checkpointer, - root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, __metadata_info__={}, ) return conversation diff --git a/wayflowcore/src/wayflowcore/mcp/_session_persistence.py b/wayflowcore/src/wayflowcore/mcp/_session_persistence.py index 143b069b0..e13524301 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_root_conversation_id - return _get_current_conversation_id() or _DEFAULT_MCP_SESSION_CONTEXT_ID + return _get_current_root_conversation_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_root_conversation_id runtime = get_mcp_async_runtime() - conversation_id = _get_current_conversation_id() or _DEFAULT_MCP_SESSION_CONTEXT_ID + conversation_id = _get_current_root_conversation_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/models/llmmodel.py b/wayflowcore/src/wayflowcore/models/llmmodel.py index 33d52bff5..27010c45b 100644 --- a/wayflowcore/src/wayflowcore/models/llmmodel.py +++ b/wayflowcore/src/wayflowcore/models/llmmodel.py @@ -359,11 +359,11 @@ def _update_token_usage( if isinstance(conversation, FlowConversation): # generate with flow - self.token_usages_flow[conversation.root_conversation_id][ + self.token_usages_flow[conversation.conversation_id][ conversation.current_step_name ] += token_usage else: - self.token_usages_flexible[conversation.root_conversation_id] += token_usage + self.token_usages_flexible[conversation.conversation_id] += token_usage def get_total_token_consumption(self, conversation_id: str) -> TokenUsage: """Calculate and return the total token consumption for a given conversation. diff --git a/wayflowcore/src/wayflowcore/ociagent.py b/wayflowcore/src/wayflowcore/ociagent.py index ad74efde0..c8438e7d3 100644 --- a/wayflowcore/src/wayflowcore/ociagent.py +++ b/wayflowcore/src/wayflowcore/ociagent.py @@ -108,36 +108,38 @@ 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, - _root_conversation_id: Optional[str] = None, + _runtime_conversation_id: Optional[str] = None, _attach_checkpointer: bool = True, ) -> "Conversation": """ - Start a conversation with the OCI agent. + Initializes a conversation with the agent. Parameters ---------- inputs: - Optional structured inputs stored on the conversation for interface compatibility. + This argument is not used. + It is included for compatibility with the Flow class. messages: - Optional initial message history for the OCI agent session. + Message list to which the agent will participate conversation_id: - Optional identifier for this OCI agent conversation. + 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``. - _root_conversation_id: - Internal lineage identifier shared with nested or parent conversations. + _runtime_conversation_id: + Internal runtime id for a fresh conversation. When provided, it becomes + the created conversation object's ``.id`` instead of defaulting to + ``conversation_id``. Returns ------- - Conversation - A new OCI agent conversation. + Conversation: + The conversation object of the agent. """ from wayflowcore.executors._ociagentconversation import OciAgentConversation from wayflowcore.executors._ociagentexecutor import ( @@ -159,7 +161,7 @@ def start_conversation( conversation_id=conversation_id, checkpointer=None, checkpoint_id=None, - _root_conversation_id=_root_conversation_id, + _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=OciAgentConversation, attach_checkpointer=_attach_checkpointer, ) @@ -179,7 +181,7 @@ def start_conversation( status=None, id=conversation_runtime_id, name="oci_conversation", - root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, __metadata_info__={}, ) diff --git a/wayflowcore/src/wayflowcore/serialization/context.py b/wayflowcore/src/wayflowcore/serialization/context.py index 186b3406d..a6c698eb5 100644 --- a/wayflowcore/src/wayflowcore/serialization/context.py +++ b/wayflowcore/src/wayflowcore/serialization/context.py @@ -38,6 +38,42 @@ def _create_component_type_to_plugin_mapping( return component_types_to_plugins +class MissingDeserializationReferenceError(ValueError): + """Raised when deserialization encounters a reference missing from the root object.""" + + +def _iter_nested_components(value: Any) -> List["Component"]: + """Return one ordered pass over all public nested components reachable from `value`.""" + from wayflowcore.component import Component + + 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) + ordered_components.append(current_value) + for name, attr in vars(current_value).items(): + if not name.startswith("_"): + _collect_nested_components(attr) + return + + if isinstance(current_value, dict): + for nested_value in current_value.values(): + _collect_nested_components(nested_value) + return + + 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 + + class SerializationContext: def __init__(self, root: Any, plugins: Optional[List["WayflowSerializationPlugin"]] = None): @@ -114,15 +150,13 @@ 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_reference(self, obj: Any) -> None: + def _add_component_to_context(self, component: "Component") -> None: """ - Registers an object as provided externally to the serialized payload. - - The serializer will emit a ``$ref`` for this object, but it will not add the object to - the root ``_referenced_objects`` section because the deserialization context is expected - to already contain it. + Marks the current component and all its nested components as provided externally to the + serialized object graph. """ - self._external_references.add(self.get_reference(obj)) + for nested_component in _iter_nested_components(component): + self._external_references.add(self.get_reference(nested_component)) def check_obj_is_already_serialized(self, obj: Any) -> bool: """ @@ -244,7 +278,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." ) @@ -325,33 +359,6 @@ def _add_component_to_context(self, component: "Component") -> None: Adds the current components and all its subcomponents to this deserialization context. """ - from wayflowcore.component import Component - - def _iter_nested_components(value: Any) -> List["Component"]: - if isinstance(value, Component): - return [value] - if isinstance(value, dict): - nested_components: List["Component"] = [] - for nested_value in value.values(): - nested_components.extend(_iter_nested_components(nested_value)) - return nested_components - if isinstance(value, (list, tuple, set)): - nested_components = [] - for nested_value in value: - nested_components.extend(_iter_nested_components(nested_value)) - return nested_components - return [] - - component_ref = SerializationContext.get_reference(component) - - if component_ref in self._deserialized_objects: - 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 in all_public_attrs.values(): - for nested_component in _iter_nested_components(attr): - self._add_component_to_context(nested_component) + for nested_component in _iter_nested_components(component): + component_ref = SerializationContext.get_reference(nested_component) + self._deserialized_objects.setdefault(component_ref, nested_component) diff --git a/wayflowcore/src/wayflowcore/serialization/serializer.py b/wayflowcore/src/wayflowcore/serialization/serializer.py index 085b2cc1f..83f76b43a 100644 --- a/wayflowcore/src/wayflowcore/serialization/serializer.py +++ b/wayflowcore/src/wayflowcore/serialization/serializer.py @@ -140,18 +140,13 @@ class MyDataclass: type_3: "MySecondCustomAttr" <--- resolves the actual type of this kind of attribute """ dataclass_fields: Dict[str, Any] = { - param.name: param.type for param in fields(cls) if _should_serialize_dataclass_field(param) + param.name: param.type for param in fields(cls) if param.init } # we resolve the forwards references (e.g. dataclasses with type annotations specified "between quotes") if any(isinstance(t, str) for t in dataclass_fields.values()): try: - resolved_type_hints = get_type_hints(cls) - dataclass_fields = { - field_name: resolved_type_hints[field_name] - for field_name in dataclass_fields - if field_name in resolved_type_hints - } + dataclass_fields = get_type_hints(cls) except NameError as e: pass @@ -198,14 +193,7 @@ def _resolve_legacy_field_name(cls: type, field_name: str) -> str: } if cls in _CLS_TO_ATTRIBUTE_MAPPING: - resolved_field_name = _CLS_TO_ATTRIBUTE_MAPPING[cls].get(field_name) - if resolved_field_name is not None: - return resolved_field_name - - if field_name == "root_conversation_id" and any( - base.__name__ == "Conversation" for base in cls.__mro__ - ): - return "conversation_id" + return _CLS_TO_ATTRIBUTE_MAPPING[cls].get(field_name, field_name) return field_name @@ -219,20 +207,20 @@ def _resolve_legacy_configurations(serialized_config: str) -> str: ) -def _should_serialize_dataclass_field(dataclass_field: Any) -> bool: - return bool( - (not dataclass_field.name.startswith("_") or dataclass_field.name == "__metadata_info__") - and dataclass_field.init - and dataclass_field.metadata.get("serialize", True) - ) - - class SerializableDataclassMixin: def _serialize_to_dict(self, serialization_context: "SerializationContext") -> Dict[str, Any]: return { k.name: serialize_any_to_dict(getattr(self, k.name), serialization_context) for k in fields(self) # type: ignore - if _should_serialize_dataclass_field(k) + if ( + ( + not k.name.startswith("_") # we don't serialize private fields + 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) + ) } @classmethod diff --git a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py index a26329035..c83b27c62 100644 --- a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py +++ b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py @@ -296,7 +296,8 @@ def _get_or_create_agent_subconversation( agent_sub_conversation = self.agent.start_conversation( inputs=inputs, messages=init_messages, - _root_conversation_id=caller_conv.root_conversation_id, + conversation_id=caller_conv.conversation_id, + _runtime_conversation_id=sub_conversation_id, ) return agent_sub_conversation diff --git a/wayflowcore/src/wayflowcore/steps/parallelflowexecutionstep.py b/wayflowcore/src/wayflowcore/steps/parallelflowexecutionstep.py index 16bfc8fe1..b5e4c4af7 100644 --- a/wayflowcore/src/wayflowcore/steps/parallelflowexecutionstep.py +++ b/wayflowcore/src/wayflowcore/steps/parallelflowexecutionstep.py @@ -255,9 +255,9 @@ async def _invoke_step_async( if input_name in set(p.name for p in flow.input_descriptors) # We extract the inputs needed by this specific flow }, - sub_conversation_id=flow.id, + sub_conversation_id=f"parallel_subflow_{index}", ) - for flow in self.flows + for index, flow in enumerate(self.flows) ] # We run only the conversations that did not reach the end diff --git a/wayflowcore/src/wayflowcore/steps/retrystep.py b/wayflowcore/src/wayflowcore/steps/retrystep.py index 22f8598a9..c19c0c762 100644 --- a/wayflowcore/src/wayflowcore/steps/retrystep.py +++ b/wayflowcore/src/wayflowcore/steps/retrystep.py @@ -39,6 +39,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", @@ -256,10 +260,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)) + context_key = FlowConversationExecutor.make_key_for_step(self, self._RETRY_COUNTER_KEY) + return cast(int, state.internal_context_key_values.get(context_key, 0)) def _set_counter(self, state: FlowConversationExecutionState, value: int) -> None: - state.internal_context_key_values[f"retry_counter_{id(self)}"] = value + state.internal_context_key_values[ + FlowConversationExecutor.make_key_for_step(self, self._RETRY_COUNTER_KEY) + ] = value async def _invoke_step_async( self, diff --git a/wayflowcore/src/wayflowcore/swarm.py b/wayflowcore/src/wayflowcore/swarm.py index 676fa98f3..0297b401d 100644 --- a/wayflowcore/src/wayflowcore/swarm.py +++ b/wayflowcore/src/wayflowcore/swarm.py @@ -312,37 +312,17 @@ def start_conversation( inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], MessageList] = None, conversation_id: Optional[str] = None, - conversation_name: Optional[str] = None, - *, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _root_conversation_id: Optional[str] = None, + _runtime_conversation_id: Optional[str] = None, _attach_checkpointer: bool = True, + conversation_name: Optional[str] = None, ) -> "Conversation": - """ - Start a conversation for the swarm. + """Start a fresh swarm conversation or restore one from a checkpoint. - Parameters - ---------- - inputs: - Optional input values available to the swarm execution state. - messages: - Optional shared message history for the swarm conversation. - conversation_id: - Optional identifier for this swarm conversation. - conversation_name: - Optional display name used for the created conversation object. - checkpointer: - Optional checkpoint backend used to restore and persist this conversation. - checkpoint_id: - Optional checkpoint identifier to restore. Requires ``checkpointer``. - _root_conversation_id: - Internal lineage identifier shared with nested or parent conversations. - - Returns - ------- - Conversation - A new or restored swarm conversation. + ``conversation_id`` is the durable conversation id. For fresh conversations, + ``_runtime_conversation_id`` becomes the created conversation object's + ``.id`` when provided. """ from wayflowcore.executors._swarmconversation import ( SwarmConversation, @@ -358,7 +338,7 @@ def start_conversation( conversation_id=conversation_id, checkpointer=checkpointer, checkpoint_id=checkpoint_id, - _root_conversation_id=_root_conversation_id, + _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=SwarmConversation, attach_checkpointer=_attach_checkpointer, ) @@ -387,9 +367,9 @@ def start_conversation( context_providers=[], inputs=inputs, messages=messages, - root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, ) - conversation = SwarmConversation( + return SwarmConversation( component=self, inputs=inputs or {}, message_list=messages, @@ -398,10 +378,9 @@ def start_conversation( state=state, status=None, checkpointer=checkpointer, - root_conversation_id=conversation_root_id, + conversation_id=conversation_root_id, __metadata_info__={}, ) - 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 9c85792e3..5e98a3aa3 100644 --- a/wayflowcore/src/wayflowcore/tools/servertools.py +++ b/wayflowcore/src/wayflowcore/tools/servertools.py @@ -639,7 +639,7 @@ async def __call__(self, **inputs: Any) -> Any: conversation = self.flow.start_conversation( inputs=inputs, messages=self._parent_conversation.message_list, - _root_conversation_id=self._parent_conversation.root_conversation_id, + conversation_id=self._parent_conversation.conversation_id, ) interrupts = self._parent_conversation._get_interrupts() diff --git a/wayflowcore/src/wayflowcore/tracing/span.py b/wayflowcore/src/wayflowcore/tracing/span.py index 22dbfbbdc..07952157c 100644 --- a/wayflowcore/src/wayflowcore/tracing/span.py +++ b/wayflowcore/src/wayflowcore/tracing/span.py @@ -550,7 +550,7 @@ class ConversationSpan(Span): def to_tracing_info(self, mask_sensitive_information: bool = True) -> Dict[str, Any]: return { **super().to_tracing_info(mask_sensitive_information=mask_sensitive_information), - "conversation.id": self.conversation.id, + "conversation.id": self.conversation.conversation_id, "conversation.name": self.conversation.name, "conversational_component.type": self.conversation.component.__class__.__name__, "conversational_component.id": self.conversation.component.id, 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/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/integration/steps/test_prompt_execution_step.py b/wayflowcore/tests/integration/steps/test_prompt_execution_step.py index 92442b5c9..c53d9cf1c 100644 --- a/wayflowcore/tests/integration/steps/test_prompt_execution_step.py +++ b/wayflowcore/tests/integration/steps/test_prompt_execution_step.py @@ -468,7 +468,7 @@ def test_check_token_consumption(remotely_hosted_llm): assert isinstance(status, FinishedStatus) assert PromptExecutionStep.OUTPUT in status.output_values assert isinstance(status.output_values[PromptExecutionStep.OUTPUT], str) - token_consumption = remotely_hosted_llm.get_total_token_consumption(conv.root_conversation_id) + token_consumption = remotely_hosted_llm.get_total_token_consumption(conv.conversation_id) assert token_consumption.input_tokens == 39 assert token_consumption.output_tokens == 10 diff --git a/wayflowcore/tests/integration/test_checkpointing.py b/wayflowcore/tests/integration/test_checkpointing.py new file mode 100644 index 000000000..951f1883a --- /dev/null +++ b/wayflowcore/tests/integration/test_checkpointing.py @@ -0,0 +1,528 @@ +# 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 + +import pytest + +from wayflowcore.agent import Agent +from wayflowcore.checkpointing import InMemoryCheckpointer +from wayflowcore.checkpointing.checkpoint_state import _save_live_conversation_checkpoint +from wayflowcore.controlconnection import ControlFlowEdge +from wayflowcore.conversation import Conversation +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.idgeneration import IdGenerator +from wayflowcore.managerworkers import ManagerWorkers +from wayflowcore.steps import FlowExecutionStep, OutputMessageStep +from wayflowcore.swarm import Swarm + +from ..serialization.test_assistant_serialization import create_flow +from ..testhelpers.dummy import DoNothingStep, DummyModel + + +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, + ) + + +def _start_interrupted_flow_conversation( + flow: Flow, + *, + conversation_id: str, + interrupt_step_name: str, + checkpointer: InMemoryCheckpointer | None = None, +): + conversation = flow.start_conversation( + conversation_id=conversation_id, + checkpointer=checkpointer, + ) + status = conversation.execute( + execution_interrupts=[_OnStepStartExecutionInterrupt(interrupt_step_name)] + ) + assert isinstance(status, InterruptedExecutionStatus) + return conversation + + +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, +) -> tuple[Agent, DummyModel]: + llm = DummyModel() + agent = Agent( + llm=llm, + name=name, + description=f"{name} description", + custom_instruction="Be helpful.", + initial_message=initial_message, + agent_id=name, + ) + return agent, llm + + +def _build_checkpointable_agent_component() -> tuple[Agent, DummyModel]: + return _build_checkpointable_agent( + name="checkpoint_agent", + initial_message="Hello from the agent.", + ) + + +def _build_checkpointable_swarm() -> tuple[Swarm, DummyModel]: + first_agent, first_agent_llm = _build_checkpointable_agent( + name="checkpoint_swarm_first_agent", + 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.", + 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, first_agent_llm + + +def _build_checkpointable_managerworkers() -> tuple[ManagerWorkers, DummyModel]: + manager_agent, manager_llm = _build_checkpointable_agent( + name="checkpoint_manager_agent", + initial_message="Hello from the manager.", + ) + worker_agent = Agent( + llm=DummyModel(fails_if_not_set=False), + 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, manager_llm + + +def test_flow_checkpoint_restore_rejects_generated_step_name_reinstantiation() -> None: + checkpointer = InMemoryCheckpointer() + + def build_flow() -> tuple[Flow, OutputMessageStep]: + first_step = DoNothingStep() + second_step = OutputMessageStep(message_template="Generated step resumed.") + first_step.id = "generated_step_restore_flow_first_step" + second_step.id = "generated_step_restore_flow_second_step" + return ( + Flow( + begin_step=first_step, + steps=[first_step, second_step], + control_flow_edges=[ + ControlFlowEdge(source_step=first_step, destination_step=second_step), + ControlFlowEdge(source_step=second_step, destination_step=None), + ], + name="generated_step_restore_flow", + flow_id="generated_step_restore_flow", + ), + second_step, + ) + + original_flow, original_second_step = build_flow() + assert IdGenerator.is_auto_generated(original_second_step.name) + conversation = original_flow.start_conversation( + conversation_id="generated-step-name-restore", + checkpointer=checkpointer, + ) + first_status = conversation.execute( + execution_interrupts=[_OnStepStartExecutionInterrupt(original_second_step.name)] + ) + + assert isinstance(first_status, InterruptedExecutionStatus) + assert checkpointer.load_latest(conversation.conversation_id) is not None + + restarted_flow, _ = build_flow() + + with pytest.raises(ValueError, match="stable step names"): + restarted_flow.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + + +def test_flow_checkpoint_restore_relinks_nested_flow_parent_for_interrupt_inheritance() -> None: + checkpointer = InMemoryCheckpointer() + + 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 = _start_interrupted_flow_conversation( + original_flow, + conversation_id="flow-parent-link-restore", + interrupt_step_name="child_first_step", + checkpointer=checkpointer, + ) + assert checkpointer.load_latest(conversation.conversation_id) is not None + + restarted_flow = build_flow() + restored_conversation = restarted_flow.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + restored_child_conversation = restored_conversation._get_current_sub_conversation( + restarted_flow.steps["parent_flow_step"] + ) + assert restored_child_conversation is not None + assert restored_child_conversation._get_parent_conversation() is restored_conversation + assert restored_conversation.id == restored_conversation.conversation_id + assert restored_child_conversation.id != restored_conversation.id + assert restored_child_conversation.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" + + +def test_flow_checkpointing_supports_resume_and_time_travel() -> None: + checkpointer = InMemoryCheckpointer() + flow = create_flow() + + conversation = flow.start_conversation( + conversation_id="flow-checkpoint", checkpointer=checkpointer + ) + first_status = conversation.execute() + + assert isinstance(first_status, UserMessageRequestStatus) + first_checkpoint = checkpointer.load_latest(conversation.conversation_id) + assert first_checkpoint is not None + first_checkpoint_id = first_checkpoint.checkpoint_id + assert conversation.checkpoint_id == first_checkpoint_id + + restored_conversation = flow.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + assert restored_conversation.checkpoint_id == first_checkpoint_id + assert isinstance(restored_conversation.status, UserMessageRequestStatus) + + restored_conversation.append_user_message("continue") + restored_status = restored_conversation.execute() + assert isinstance(restored_status, FinishedStatus) + + rewound_conversation = flow.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + checkpoint_id=first_checkpoint_id, + ) + assert len(rewound_conversation.get_messages()) < len(restored_conversation.get_messages()) + rewound_conversation.append_user_message("rewind") + rewound_status = rewound_conversation.execute() + assert isinstance(rewound_status, FinishedStatus) + + +@pytest.mark.parametrize( + ("builder", "conversation_id"), + [ + (_build_checkpointable_agent_component, "agent-checkpoint"), + (_build_checkpointable_swarm, "swarm-checkpoint"), + (_build_checkpointable_managerworkers, "managerworkers-checkpoint"), + ], +) +def test_multi_agent_checkpointing_supports_resume_and_time_travel( + builder, + conversation_id: str, +) -> None: + checkpointer = InMemoryCheckpointer() + component, llm = builder() + resumed_output = "Checkpoint resumed successfully." + rewound_output = "Checkpoint rewound successfully." + + conversation = component.start_conversation( + conversation_id=conversation_id, + checkpointer=checkpointer, + ) + first_status = conversation.execute() + + assert isinstance(first_status, UserMessageRequestStatus) + first_checkpoint = checkpointer.load_latest(conversation.conversation_id) + assert first_checkpoint is not None + first_checkpoint_id = first_checkpoint.checkpoint_id + assert conversation.checkpoint_id == first_checkpoint_id + + restored_conversation = component.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + assert restored_conversation.checkpoint_id == first_checkpoint_id + llm.set_next_output(resumed_output) + restored_conversation.append_user_message("Please continue.") + restored_status = restored_conversation.execute() + assert isinstance(restored_status, UserMessageRequestStatus) + assert restored_conversation.get_last_message().content == resumed_output + + rewound_conversation = component.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + checkpoint_id=first_checkpoint_id, + ) + assert rewound_conversation.checkpoint_id == first_checkpoint_id + llm.set_next_output(rewound_output) + rewound_conversation.append_user_message("Try again.") + rewound_status = rewound_conversation.execute() + assert isinstance(rewound_status, UserMessageRequestStatus) + assert rewound_conversation.get_last_message().content == rewound_output + + +def test_swarm_checkpoint_restore_rebuilds_generated_agent_name_threads() -> None: + checkpointer = InMemoryCheckpointer() + + 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=checkpointer, + ) + original_thread = conversation.state.agents_and_threads[original_first_agent.name][ + original_second_agent.name + ] + conversation.state._create_subconversation_for_thread(original_thread) + conversation.state.current_thread = original_thread + _save_live_conversation_checkpoint(checkpointer, conversation) + + restarted_swarm, restarted_first_agent, restarted_second_agent = build_swarm() + + restored_conversation = restarted_swarm.start_conversation( + conversation_id=conversation.conversation_id, checkpointer=checkpointer + ) + assert restored_conversation.state.current_thread is not None + assert restored_conversation.state.current_thread.identifier == ( + f"{restarted_first_agent.name}#{restarted_second_agent.name}" + ) + + +def test_agent_checkpoint_restore_relinks_current_flow_parent() -> None: + checkpointer = InMemoryCheckpointer() + + 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=checkpointer, + ) + parent_flow_conversation = _start_interrupted_flow_conversation( + original_parent_flow, + conversation_id=conversation.conversation_id, + interrupt_step_name="agent_child_first_step", + ) + conversation.state.current_flow_conversation = parent_flow_conversation + _save_live_conversation_checkpoint(checkpointer, conversation) + + restarted_agent, restarted_parent_flow = build_agent() + restored_conversation = restarted_agent.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + + restored_parent_flow_conversation = restored_conversation.state.current_flow_conversation + assert restored_parent_flow_conversation is not None + restored_child_conversation = restored_parent_flow_conversation._get_current_sub_conversation( + restarted_parent_flow.steps["agent_parent_flow_step"] + ) + assert restored_child_conversation is not None + assert ( + restored_child_conversation._get_parent_conversation() is restored_parent_flow_conversation + ) + + +def test_managerworkers_checkpoint_restore_rejects_nested_generated_agent_names() -> None: + checkpointer = InMemoryCheckpointer() + + 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=checkpointer, + ) + nested_conversation = conversation.state._create_subconversation_for_agent( + original_nested_managerworkers + ) + nested_conversation.state._create_subconversation_for_agent(original_nested_worker_agent) + conversation.state.current_agent_name = original_nested_managerworkers.name + nested_conversation.state.current_agent_name = original_nested_worker_agent.name + _save_live_conversation_checkpoint(checkpointer, conversation) + + restarted_managerworkers, *_ = build_managerworkers() + + with pytest.raises(ValueError, match="stable agent names"): + restarted_managerworkers.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) diff --git a/wayflowcore/tests/serialization/test_conversation_checkpointing.py b/wayflowcore/tests/serialization/test_conversation_checkpointing.py deleted file mode 100644 index 2d694a346..000000000 --- a/wayflowcore/tests/serialization/test_conversation_checkpointing.py +++ /dev/null @@ -1,654 +0,0 @@ -# Copyright © 2025, 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 warnings -from types import SimpleNamespace -from typing import Any, Dict, Optional -from unittest.mock import AsyncMock - -import pytest - -import wayflowcore.checkpointing.serialization as checkpoint_serialization -from wayflowcore.a2a.a2aagent import A2AAgent, A2AConnectionConfig -from wayflowcore.agent import Agent -from wayflowcore.checkpointing import ( - CheckpointingInterval, - ConversationCheckpoint, - InMemoryCheckpointer, -) -from wayflowcore.checkpointing.checkpointeventlistener import _save_conversation_checkpoint -from wayflowcore.executors.executionstatus import FinishedStatus, UserMessageRequestStatus -from wayflowcore.flowhelpers import create_single_step_flow -from wayflowcore.managerworkers import ManagerWorkers -from wayflowcore.models.ociclientconfig import OCIClientConfigWithApiKey -from wayflowcore.ociagent import OciAgent -from wayflowcore.serialization.serializer import _resolve_legacy_field_name -from wayflowcore.steps import OutputMessageStep, PromptExecutionStep -from wayflowcore.swarm import Swarm - -from ..testhelpers.dummy import DummyModel -from ..testhelpers.testhelpers import retry_test -from .test_assistant_serialization import create_flow - - -class RecordingCheckpointer: - def __init__(self, *, fail_first_save: bool = False) -> None: - self.checkpointing_interval = CheckpointingInterval.CONVERSATION_TURNS - self.should_fail_next_save = fail_first_save - self.saved_checkpoints: Dict[tuple[str, str], Dict[str, Any]] = {} - - def save_conversation( - self, - conversation, - *, - checkpoint_id: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - ): - if self.should_fail_next_save: - self.should_fail_next_save = False - raise RuntimeError("checkpoint save failed") - resolved_checkpoint_id = checkpoint_id or "generated-checkpoint-id" - self.saved_checkpoints[(conversation.id, resolved_checkpoint_id)] = dict(metadata or {}) - conversation.checkpoint_id = resolved_checkpoint_id - return SimpleNamespace( - checkpoint_id=resolved_checkpoint_id, - metadata=dict(metadata or {}), - ) - - def load(self, conversation_id: str, checkpoint_id: str) -> Any: - return SimpleNamespace(metadata=self.saved_checkpoints[(conversation_id, checkpoint_id)]) - - def load_latest(self, conversation_id: str) -> Any: - return None - - -class StaticLoadCheckpointer: - def __init__(self, checkpoint: ConversationCheckpoint) -> None: - self.checkpointing_interval = CheckpointingInterval.CONVERSATION_TURNS - self.checkpoint = checkpoint - - def load(self, conversation_id: str, checkpoint_id: str) -> ConversationCheckpoint: - return self.checkpoint - - def load_latest(self, conversation_id: str) -> Optional[ConversationCheckpoint]: - if conversation_id != self.checkpoint.conversation_id: - return None - return self.checkpoint - - -def _build_checkpointable_agent( - *, - name: str, - initial_message: str, -) -> tuple[Agent, DummyModel]: - llm = DummyModel() - agent = Agent( - llm=llm, - name=name, - description=f"{name} description", - custom_instruction="Be helpful.", - initial_message=initial_message, - ) - return agent, llm - - -def _build_checkpointable_swarm() -> tuple[Swarm, DummyModel]: - first_agent, first_agent_llm = _build_checkpointable_agent( - name="checkpoint_swarm_first_agent", - 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.", - ) - swarm = Swarm( - first_agent=first_agent, - relationships=[(first_agent, second_agent)], - name="checkpoint_swarm", - ) - return swarm, first_agent_llm - - -def _build_checkpointable_managerworkers() -> tuple[ManagerWorkers, DummyModel]: - manager_agent, manager_llm = _build_checkpointable_agent( - name="checkpoint_manager_agent", - initial_message="Hello from the manager.", - ) - worker_agent = Agent( - llm=DummyModel(fails_if_not_set=False), - name="checkpoint_worker_agent", - description="Worker agent", - custom_instruction="Help the manager.", - ) - managerworkers = ManagerWorkers( - group_manager=manager_agent, - workers=[worker_agent], - name="checkpoint_managerworkers", - ) - return managerworkers, manager_llm - - -@pytest.fixture(scope="session") -def connection_config_no_verify(): - with warnings.catch_warnings(): - warnings.simplefilter("ignore", UserWarning) - return A2AConnectionConfig(verify=False) - - -@pytest.fixture -def a2a_agent(a2a_server, connection_config_no_verify): - return A2AAgent( - name="Checkpoint A2A Agent", - agent_url=a2a_server, - connection_config=connection_config_no_verify, - ) - - -def test_inmemory_checkpointer_can_save_load_list_and_delete_checkpoints() -> None: - checkpointer = InMemoryCheckpointer() - flow = create_single_step_flow(OutputMessageStep(message_template="Hello from checkpointing.")) - - conversation = flow.start_conversation( - conversation_id="checkpoint-lifecycle", checkpointer=checkpointer - ) - assert conversation.checkpointer is checkpointer - - status = conversation.execute() - - assert isinstance(status, FinishedStatus) - first_checkpoint_id = conversation.checkpoint_id - assert first_checkpoint_id is not None - - checkpointer.save(conversation) - second_checkpoint_id = conversation.checkpoint_id - assert second_checkpoint_id is not None - assert second_checkpoint_id != first_checkpoint_id - - checkpoints = checkpointer.list_checkpoints("checkpoint-lifecycle") - assert [checkpoint.checkpoint_id for checkpoint in checkpoints] == [ - first_checkpoint_id, - second_checkpoint_id, - ] - assert checkpoints[-1].metadata["save_sequence"] == 2 - - latest_checkpoint = checkpointer.load_latest("checkpoint-lifecycle") - assert latest_checkpoint is not None - assert latest_checkpoint.checkpoint_id == second_checkpoint_id - - restored_conversation = flow.start_conversation( - conversation_id="checkpoint-lifecycle", - checkpoint_id=first_checkpoint_id, - checkpointer=checkpointer, - ) - assert restored_conversation.checkpointer is checkpointer - assert restored_conversation.checkpoint_id == first_checkpoint_id - assert restored_conversation.get_last_message().content == "Hello from checkpointing." - - checkpointer.delete("checkpoint-lifecycle", second_checkpoint_id) - promoted_checkpoint = checkpointer.load_latest("checkpoint-lifecycle") - assert promoted_checkpoint is not None - assert promoted_checkpoint.checkpoint_id == first_checkpoint_id - assert [ - checkpoint.checkpoint_id - for checkpoint in checkpointer.list_checkpoints("checkpoint-lifecycle") - ] == [first_checkpoint_id] - - -def test_conversation_turns_checkpoint_interval_saves_once_after_outer_execute() -> None: - checkpointer = InMemoryCheckpointer( - checkpointing_interval=CheckpointingInterval.CONVERSATION_TURNS - ) - flow = create_single_step_flow(OutputMessageStep(message_template="Hello once.")) - - status = flow.start_conversation( - conversation_id="conversation-turn", checkpointer=checkpointer - ).execute() - - assert isinstance(status, FinishedStatus) - checkpoints = checkpointer.list_checkpoints("conversation-turn") - assert len(checkpoints) == 1 - assert checkpoints[0].metadata["save_reason"] == "conversation_turn" - assert checkpoints[0].metadata["status_type"] == "FinishedStatus" - - -def test_all_internal_turns_checkpoint_interval_saves_before_each_flow_turn() -> None: - checkpointer = InMemoryCheckpointer( - checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS - ) - flow = create_single_step_flow(OutputMessageStep(message_template="Hello internal turns.")) - - status = flow.start_conversation( - conversation_id="all-internal-turns", checkpointer=checkpointer - ).execute() - - assert isinstance(status, FinishedStatus) - checkpoints = checkpointer.list_checkpoints("all-internal-turns") - assert len(checkpoints) == 3 - assert [checkpoint.metadata["save_reason"] for checkpoint in checkpoints] == [ - "internal_turn_boundary", - "internal_turn_boundary", - "conversation_turn", - ] - assert [checkpoint.metadata.get("event_type") for checkpoint in checkpoints[:-1]] == [ - "FlowExecutionIterationStartedEvent", - "FlowExecutionIterationStartedEvent", - ] - assert checkpoints[-1].metadata["status_type"] == "FinishedStatus" - - -def test_llm_turns_checkpoint_interval_saves_only_after_llm_backed_turns() -> None: - checkpointer = InMemoryCheckpointer(checkpointing_interval=CheckpointingInterval.LLM_TURNS) - dummy_llm = DummyModel() - dummy_llm.set_next_output("Hello from the prompt step.") - flow = create_single_step_flow( - PromptExecutionStep( - llm=dummy_llm, - prompt_template="Say hello.", - ) - ) - - status = flow.start_conversation( - conversation_id="llm-turns", checkpointer=checkpointer - ).execute() - - assert isinstance(status, FinishedStatus) - checkpoints = checkpointer.list_checkpoints("llm-turns") - assert len(checkpoints) == 2 - assert checkpoints[0].metadata["save_reason"] == "internal_turn_boundary" - assert checkpoints[0].metadata["event_type"] == "FlowExecutionIterationStartedEvent" - assert checkpoints[0].metadata["llm_used_in_previous_turn"] is True - assert checkpoints[1].metadata["save_reason"] == "conversation_turn" - - -def test_checkpoint_serialization_context_registers_component_tree_as_external_refs( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class _FakeSerializationContext: - def __init__(self, root: Any) -> None: - self.root = root - self.external_refs: set[str] = set() - self.recorded_refs: Dict[str, Dict[str, Any]] = {} - - @staticmethod - def get_reference(obj: Any) -> str: - obj_id = getattr(obj, "id", id(obj)) - return f"{obj.__class__.__name__.lower()}/{obj_id}" - - def register_external_reference(self, obj: Any) -> None: - self.external_refs.add(self.get_reference(obj)) - - def record_obj_dict(self, obj: Any, obj_as_dict: Dict[str, Any]) -> None: - self.recorded_refs[self.get_reference(obj)] = obj_as_dict - - monkeypatch.setattr( - checkpoint_serialization, - "SerializationContext", - _FakeSerializationContext, - ) - - flow = create_single_step_flow(OutputMessageStep(message_template="Hello external refs.")) - conversation = flow.start_conversation(conversation_id="checkpoint-external-refs") - serialization_context = checkpoint_serialization._build_checkpoint_serialization_context( - conversation - ) - - expected_component_refs = { - _FakeSerializationContext.get_reference(component) - for component in checkpoint_serialization._iter_component_tree(conversation.component) - } - - assert serialization_context.external_refs == expected_component_refs - assert serialization_context.recorded_refs == {} - - -def test_explicit_final_checkpoint_parameters_can_be_retried_after_save_fails() -> None: - checkpointer = RecordingCheckpointer(fail_first_save=True) - flow = create_single_step_flow(OutputMessageStep(message_template="Hello overrides.")) - - conversation = flow.start_conversation( - conversation_id="checkpoint-final-overrides", checkpointer=checkpointer - ) - - with pytest.raises(RuntimeError, match="checkpoint save failed"): - _save_conversation_checkpoint( - conversation, - save_reason="conversation_turn", - checkpoint_id="final-checkpoint-id", - metadata={"response_id": "resp-123"}, - ) - - _save_conversation_checkpoint( - conversation, - save_reason="conversation_turn", - checkpoint_id="final-checkpoint-id", - metadata={"response_id": "resp-123"}, - ) - - assert conversation.checkpoint_id == "final-checkpoint-id" - checkpoint = checkpointer.load("checkpoint-final-overrides", "final-checkpoint-id") - assert checkpoint.metadata["response_id"] == "resp-123" - - -def test_execute_final_checkpoint_parameters_are_applied_to_final_save() -> None: - checkpointer = RecordingCheckpointer() - flow = create_single_step_flow(OutputMessageStep(message_template="Hello execute.")) - - conversation = flow.start_conversation( - conversation_id="checkpoint-final-execute", - checkpointer=checkpointer, - ) - status = conversation.execute( - _final_checkpoint_id="final-checkpoint-id", - _final_checkpoint_metadata={"response_id": "resp-123"}, - ) - - assert isinstance(status, FinishedStatus) - assert conversation.checkpoint_id == "final-checkpoint-id" - checkpoint = checkpointer.load("checkpoint-final-execute", "final-checkpoint-id") - assert checkpoint.metadata["response_id"] == "resp-123" - assert checkpoint.metadata["save_reason"] == "conversation_turn" - - -def test_execute_async_does_not_save_final_checkpoint_when_execution_fails( - monkeypatch: pytest.MonkeyPatch, -) -> None: - checkpointer = RecordingCheckpointer() - 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 conversation.checkpoint_id is None - assert checkpointer.saved_checkpoints == {} - - -def test_checkpoint_restore_rejects_conversations_from_other_components() -> None: - original_flow = create_single_step_flow(OutputMessageStep(message_template="Hello original.")) - other_flow = create_single_step_flow(OutputMessageStep(message_template="Hello other.")) - checkpointer = StaticLoadCheckpointer( - ConversationCheckpoint( - checkpoint_id="checkpoint-1", - conversation_id="checkpoint-other-component", - component_id=original_flow.id, - created_at=0, - state="unused-because-component-mismatch-is-checked-first", - metadata={}, - ) - ) - - with pytest.raises(ValueError, match="started with another component"): - other_flow.start_conversation( - conversation_id="checkpoint-other-component", - checkpointer=checkpointer, - ) - - -def test_restore_can_skip_attaching_live_checkpointer( - monkeypatch: pytest.MonkeyPatch, -) -> None: - flow = create_single_step_flow(OutputMessageStep(message_template="Hello restore.")) - checkpoint_id = "checkpoint-no-attach-id" - checkpointer = StaticLoadCheckpointer( - ConversationCheckpoint( - checkpoint_id=checkpoint_id, - conversation_id="checkpoint-no-attach", - component_id=flow.id, - created_at=0, - state="unused-because-deserialization-is-mocked", - metadata={}, - ) - ) - - monkeypatch.setattr( - checkpoint_serialization, - "_deserialize_conversation_checkpoint_state", - lambda *args, **kwargs: flow.start_conversation(conversation_id="checkpoint-no-attach"), - ) - - restored_conversation = flow.start_conversation( - conversation_id="checkpoint-no-attach", - checkpoint_id=checkpoint_id, - checkpointer=checkpointer, - _attach_checkpointer=False, - ) - - assert restored_conversation.checkpointer is None - assert restored_conversation.checkpoint_id == checkpoint_id - - -def test_legacy_serialized_conversation_id_restores_root_conversation_id() -> None: - agent, _ = _build_checkpointable_agent( - name="legacy_checkpoint_agent", - initial_message="Hello from the past.", - ) - conversation = agent.start_conversation(_root_conversation_id="legacy-root-conversation") - - assert conversation.root_conversation_id == "legacy-root-conversation" - assert ( - _resolve_legacy_field_name(type(conversation), "root_conversation_id") == "conversation_id" - ) - assert not hasattr(conversation, "conversation_id") - - -def test_flow_checkpointing_supports_resume_and_time_travel() -> None: - checkpointer = InMemoryCheckpointer() - flow = create_flow() - - conversation = flow.start_conversation( - conversation_id="flow-checkpoint", checkpointer=checkpointer - ) - first_status = conversation.execute() - - assert isinstance(first_status, UserMessageRequestStatus) - first_checkpoint_id = conversation.checkpoint_id - assert first_checkpoint_id is not None - - restored_conversation = flow.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - ) - assert restored_conversation.checkpoint_id == first_checkpoint_id - assert isinstance(restored_conversation.status, UserMessageRequestStatus) - - restored_conversation.append_user_message("continue") - restored_status = restored_conversation.execute() - assert isinstance(restored_status, FinishedStatus) - - rewound_conversation = flow.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - checkpoint_id=first_checkpoint_id, - ) - assert len(rewound_conversation.get_messages()) < len(restored_conversation.get_messages()) - rewound_conversation.append_user_message("rewind") - rewound_status = rewound_conversation.execute() - assert isinstance(rewound_status, FinishedStatus) - - -def test_agent_checkpointing_supports_resume_and_time_travel() -> None: - checkpointer = InMemoryCheckpointer() - agent, llm = _build_checkpointable_agent( - name="checkpoint_agent", - initial_message="Hello from the agent.", - ) - - conversation = agent.start_conversation( - conversation_id="agent-checkpoint", checkpointer=checkpointer - ) - first_status = conversation.execute() - - assert isinstance(first_status, UserMessageRequestStatus) - first_checkpoint_id = conversation.checkpoint_id - assert first_checkpoint_id is not None - - restored_conversation = agent.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - ) - llm.set_next_output("Agent resumed successfully.") - restored_conversation.append_user_message("Please continue.") - restored_status = restored_conversation.execute() - assert isinstance(restored_status, UserMessageRequestStatus) - assert restored_conversation.get_last_message().content == "Agent resumed successfully." - - rewound_conversation = agent.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - checkpoint_id=first_checkpoint_id, - ) - llm.set_next_output("Agent rewound successfully.") - rewound_conversation.append_user_message("Try again.") - rewound_status = rewound_conversation.execute() - assert isinstance(rewound_status, UserMessageRequestStatus) - assert rewound_conversation.get_last_message().content == "Agent rewound successfully." - - -def test_swarm_checkpointing_supports_resume_and_time_travel() -> None: - checkpointer = InMemoryCheckpointer() - swarm, llm = _build_checkpointable_swarm() - - conversation = swarm.start_conversation( - conversation_id="swarm-checkpoint", checkpointer=checkpointer - ) - first_status = conversation.execute() - - assert isinstance(first_status, UserMessageRequestStatus) - first_checkpoint_id = conversation.checkpoint_id - assert first_checkpoint_id is not None - - restored_conversation = swarm.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - ) - llm.set_next_output("Swarm resumed successfully.") - restored_conversation.append_user_message("Continue the swarm conversation.") - restored_status = restored_conversation.execute() - assert isinstance(restored_status, UserMessageRequestStatus) - assert restored_conversation.get_last_message().content == "Swarm resumed successfully." - - rewound_conversation = swarm.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - checkpoint_id=first_checkpoint_id, - ) - llm.set_next_output("Swarm rewound successfully.") - rewound_conversation.append_user_message("Try the swarm again.") - rewound_status = rewound_conversation.execute() - assert isinstance(rewound_status, UserMessageRequestStatus) - assert rewound_conversation.get_last_message().content == "Swarm rewound successfully." - - -def test_managerworkers_checkpointing_supports_resume_and_time_travel() -> None: - checkpointer = InMemoryCheckpointer() - managerworkers, llm = _build_checkpointable_managerworkers() - - conversation = managerworkers.start_conversation( - conversation_id="managerworkers-checkpoint", - checkpointer=checkpointer, - ) - first_status = conversation.execute() - - assert isinstance(first_status, UserMessageRequestStatus) - first_checkpoint_id = conversation.checkpoint_id - assert first_checkpoint_id is not None - - restored_conversation = managerworkers.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - ) - llm.set_next_output("Manager resumed successfully.") - restored_conversation.append_user_message("Continue the manager workflow.") - restored_status = restored_conversation.execute() - assert isinstance(restored_status, UserMessageRequestStatus) - assert restored_conversation.get_last_message().content == "Manager resumed successfully." - - rewound_conversation = managerworkers.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - checkpoint_id=first_checkpoint_id, - ) - llm.set_next_output("Manager rewound successfully.") - rewound_conversation.append_user_message("Try the manager workflow again.") - rewound_status = rewound_conversation.execute() - assert isinstance(rewound_status, UserMessageRequestStatus) - assert rewound_conversation.get_last_message().content == "Manager rewound successfully." - - -@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) - first_checkpoint_id = conversation.checkpoint_id - assert first_checkpoint_id is not None - first_message_count = len(conversation.get_messages()) - - restored_conversation = a2a_agent.start_conversation( - conversation_id=conversation.id, - checkpointer=checkpointer, - ) - assert restored_conversation.checkpoint_id == first_checkpoint_id - 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.id) is not None - - rewound_conversation = a2a_agent.start_conversation( - conversation_id=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 - - -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()) diff --git a/wayflowcore/tests/test_checkpointing.py b/wayflowcore/tests/test_checkpointing.py new file mode 100644 index 000000000..e6fa3d799 --- /dev/null +++ b/wayflowcore/tests/test_checkpointing.py @@ -0,0 +1,310 @@ +# 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.executors.executionstatus import FinishedStatus, UserMessageRequestStatus +from wayflowcore.flowhelpers import create_single_step_flow +from wayflowcore.steps import OutputMessageStep +from wayflowcore.steps.promptexecutionstep import PromptExecutionStep +from wayflowcore.swarm import Swarm + +from .serialization.test_assistant_serialization import create_flow +from .testhelpers.dummy import DummyModel + + +def _checkpoint_restore_wrong_component_type_scenario( + checkpointer: InMemoryCheckpointer, +): + flow = create_single_step_flow(OutputMessageStep(message_template="Hello original.")) + agent = Agent( + llm=DummyModel(), + name="checkpoint_wrong_type_agent", + description="checkpoint_wrong_type_agent description", + custom_instruction="Be helpful.", + initial_message="Hello from the wrong component type.", + agent_id="checkpoint_wrong_type_agent", + ) + conversation = flow.start_conversation( + conversation_id="checkpoint-other-component-type", + checkpointer=checkpointer, + ) + assert isinstance(conversation.execute(), FinishedStatus) + return conversation, agent + + +def _checkpoint_restore_wrong_component_identity_scenario( + checkpointer: InMemoryCheckpointer, +): + original_agent = Agent( + llm=DummyModel(), + name="checkpoint_owner_agent", + agent_id="checkpoint-owner-agent-a", + description="Original checkpoint owner", + custom_instruction="Be helpful.", + initial_message="Hello from checkpoint owner A.", + ) + other_agent = Agent( + llm=DummyModel(), + name="checkpoint_owner_agent", + agent_id="checkpoint-owner-agent-b", + description="Other checkpoint owner", + custom_instruction="Be helpful.", + initial_message="Hello from checkpoint owner B.", + ) + conversation = original_agent.start_conversation( + conversation_id="checkpoint-explicit-owner", + checkpointer=checkpointer, + ) + assert isinstance(conversation.execute(), UserMessageRequestStatus) + return conversation, other_agent + + +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 + + +def test_inmemory_checkpointer_can_save_load_list_and_delete_checkpoints() -> None: + checkpointer = InMemoryCheckpointer() + flow = create_flow() + + conversation = flow.start_conversation( + conversation_id="checkpoint-lifecycle", checkpointer=checkpointer + ) + assert conversation.checkpointer is checkpointer + + first_status = conversation.execute() + assert isinstance(first_status, UserMessageRequestStatus) + first_checkpoint = checkpointer.load_latest(conversation.conversation_id) + assert first_checkpoint is not None + first_checkpoint_id = first_checkpoint.checkpoint_id + assert conversation.checkpoint_id == first_checkpoint_id + + conversation.append_user_message("continue") + second_status = conversation.execute() + assert isinstance(second_status, FinishedStatus) + second_checkpoint = checkpointer.load_latest(conversation.conversation_id) + assert second_checkpoint is not None + second_checkpoint_id = second_checkpoint.checkpoint_id + assert second_checkpoint_id != first_checkpoint_id + assert conversation.checkpoint_id == second_checkpoint_id + + checkpoints = checkpointer.list_checkpoints("checkpoint-lifecycle") + assert [checkpoint.checkpoint_id for checkpoint in checkpoints] == [ + first_checkpoint_id, + second_checkpoint_id, + ] + assert checkpoints[-1].metadata["save_sequence"] == 2 + + latest_checkpoint = checkpointer.load_latest("checkpoint-lifecycle") + assert latest_checkpoint is not None + assert latest_checkpoint.checkpoint_id == second_checkpoint_id + + restored_conversation = flow.start_conversation( + conversation_id="checkpoint-lifecycle", + checkpoint_id=first_checkpoint_id, + checkpointer=checkpointer, + ) + assert restored_conversation.checkpointer is checkpointer + assert restored_conversation.checkpoint_id == first_checkpoint_id + assert isinstance(restored_conversation.status, UserMessageRequestStatus) + + checkpointer.delete("checkpoint-lifecycle", second_checkpoint_id) + promoted_checkpoint = checkpointer.load_latest("checkpoint-lifecycle") + assert promoted_checkpoint is not None + assert promoted_checkpoint.checkpoint_id == first_checkpoint_id + assert [ + checkpoint.checkpoint_id + for checkpoint in checkpointer.list_checkpoints("checkpoint-lifecycle") + ] == [first_checkpoint_id] + + +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, + ) + + +@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_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" + + +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") == [] + + +@pytest.mark.parametrize( + "scenario", + [ + _checkpoint_restore_wrong_component_type_scenario, + _checkpoint_restore_wrong_component_identity_scenario, + _checkpoint_restore_generated_agent_id_scenario, + _checkpoint_restore_generated_swarm_child_id_scenario, + ], + ids=[ + "wrong_component_type", + "wrong_component_identity", + "generated_agent_id", + "generated_swarm_child_id", + ], +) +def test_checkpoint_restore_rejects_component_id_mismatches(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 ee7929d02..f2e5d601d 100644 --- a/wayflowcore/tests/test_managerworkers.py +++ b/wayflowcore/tests/test_managerworkers.py @@ -184,14 +184,14 @@ def test_managerworkers_can_execute_with_initial_params_passed_in_start_conversa conversation = group.start_conversation( messages=[Message(content="Please compute 3*4 + 2", message_type=MessageType.USER)], inputs={"USER": "Iris"}, - _root_conversation_id="12345", + conversation_id="12345", ) conversation.execute() # The first message must be not the default message as the init messages are passed. assert conversation.get_last_message().content != DEFAULT_INITIAL_MESSAGE - assert conversation.root_conversation_id == "12345" + assert conversation.conversation_id == "12345" @retry_test(max_attempts=2) 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 0c2d2d1a6..263b73b17 100644 --- a/wayflowcore/tests/test_swarm.py +++ b/wayflowcore/tests/test_swarm.py @@ -226,14 +226,14 @@ def test_can_execute_swarm_with_initial_params_passed_in_start_conversation( ) ], inputs={"USER": "Iris"}, - _root_conversation_id="12345", + conversation_id="12345", ) conversation.execute() # The first message must be not the default message as the init messages are passed. assert conversation.get_last_message().content != "Hi! How can I help you?" - assert conversation.root_conversation_id == "12345" + assert conversation.conversation_id == "12345" def test_can_create_swarm(example_medical_agents): diff --git a/wayflowcore/tests/tracing/spans/test_conversation_span.py b/wayflowcore/tests/tracing/spans/test_conversation_span.py index 274d28efc..24e6d7098 100644 --- a/wayflowcore/tests/tracing/spans/test_conversation_span.py +++ b/wayflowcore/tests/tracing/spans/test_conversation_span.py @@ -92,7 +92,7 @@ def test_span_serialization_format( assert serialized_span["span_type"] == str(span.__class__.__name__) for attribute_name in attributes_to_check: assert getattr(span, attribute_name) == serialized_span[attribute_name] - assert serialized_span["conversation.id"] == span.conversation.id + assert serialized_span["conversation.id"] == span.conversation.conversation_id assert serialized_span["conversation.name"] == span.conversation.name assert ( serialized_span["conversational_component.type"] == "Agent" From 33109805f28ea005b961d97565ae1d124d9eb1c3 Mon Sep 17 00:00:00 2001 From: Son Le Date: Fri, 17 Jul 2026 19:16:21 +0200 Subject: [PATCH 05/10] revert back to id-based identity instead of names; refactor in MW/Swarm to use agent ids instead of names for routing/keying subagents, etc --- .../checkpointing/checkpoint_state.py | 135 +---------------- .../checkpointing/checkpointeventlistener.py | 4 +- .../executors/_agentconversation.py | 3 +- .../executors/_flowconversation.py | 10 +- .../wayflowcore/executors/_flowexecutor.py | 3 +- .../executors/_managerworkersconversation.py | 15 +- .../executors/_managerworkersexecutor.py | 35 +++-- .../executors/_swarmconversation.py | 11 +- .../wayflowcore/executors/_swarmexecutor.py | 5 +- wayflowcore/src/wayflowcore/managerworkers.py | 28 +++- .../src/wayflowcore/serialization/context.py | 4 +- .../src/wayflowcore/steps/retrystep.py | 19 +-- wayflowcore/src/wayflowcore/swarm.py | 6 +- .../tests/integration/test_checkpointing.py | 85 +++-------- wayflowcore/tests/mcptools/test_mcp_tools.py | 6 +- .../test_managerworkers_serialization.py | 6 +- wayflowcore/tests/test_checkpointing.py | 143 +++++++++++++++++- wayflowcore/tests/test_managerworkers.py | 26 +++- wayflowcore/tests/test_swarm.py | 8 +- 19 files changed, 282 insertions(+), 270 deletions(-) diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py b/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py index bd9ac6b43..10da70f0a 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py @@ -5,7 +5,6 @@ # (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option. import time -from itertools import chain from typing import TYPE_CHECKING, Any, Dict, Optional, Type from wayflowcore.exceptions import DataclassFieldDeserializationError @@ -23,17 +22,9 @@ from wayflowcore.checkpointing.checkpointer import ConversationCheckpoint from wayflowcore.conversation import Conversation from wayflowcore.conversationalcomponent import ConversationalComponent - from wayflowcore.executors._agentexecutor import AgentConversationExecutionState - from wayflowcore.executors._flowexecutor import FlowConversationExecutionState - from wayflowcore.executors._managerworkersconversation import ( - ManagerWorkersConversationExecutionState, - ) - from wayflowcore.executors._swarmconversation import SwarmConversationExecutionState _COMPONENT_ID_ERROR_HINT = "Restart-safe checkpoint restore requires stable component ids." -_STEP_NAME_ERROR_HINT = "Restart-safe checkpoint restore requires stable step names." -_AGENT_NAME_ERROR_HINT = "Restart-safe checkpoint restore requires stable agent names." class _CheckpointRestoreCompatibilityError(ValueError): @@ -54,7 +45,7 @@ def _save_live_conversation_checkpoint( "Checkpointing conversations that contain `OciAgent` is not supported yet." ) serialization_context = SerializationContext(root=conversation) - serialization_context._add_component_to_context(conversation.component) + serialization_context._register_external_component_references(conversation.component) checkpoint = ConversationCheckpoint( checkpoint_id=checkpoint_id or IdGenerator.get_or_generate_id(), @@ -80,12 +71,11 @@ def _load_checkpointed_conversation( """Restore a Conversation object from stored checkpoint state. The serialized conversation omits live component/tool objects and rebuilds them from - the current component tree. After deserialization, the helper also repairs - derived runtime-only state that is not stored directly in the checkpoint. + the current component tree. """ deserialization_context = DeserializationContext() deserialization_context.registered_tools = tool_registry.copy() if tool_registry else {} - deserialization_context._add_component_to_context(component) + deserialization_context._register_external_component_references(component) try: conversation = autodeserialize( checkpoint.state, @@ -104,7 +94,6 @@ def _load_checkpointed_conversation( "Cannot restore this checkpoint because this conversation was started with another " f"component. Expected `{expected_conversation_type.__name__}`, got `{type(conversation).__name__}`." ) - _prepare_restored_conversation_states(conversation) conversation.checkpoint_id = checkpoint.checkpoint_id if attach_checkpointer: @@ -121,124 +110,6 @@ def _contains_missing_reference_error(error: Exception) -> bool: return False -def _prepare_restored_conversation_states(root_conversation: "Conversation") -> None: - """Repair executor-specific state after a conversation is restored.""" - from wayflowcore.executors._agentexecutor import AgentConversationExecutionState - from wayflowcore.executors._flowexecutor import FlowConversationExecutionState - from wayflowcore.executors._managerworkersconversation import ( - ManagerWorkersConversationExecutionState, - ) - from wayflowcore.executors._swarmconversation import SwarmConversationExecutionState - - for conversation in chain( - (root_conversation,), - root_conversation._get_all_sub_conversations_recursive(), - ): - state = conversation.state - match state: - case AgentConversationExecutionState(): - _prepare_restored_agent_conversation_state(state) - case FlowConversationExecutionState(): - _validate_restored_flow_conversation_state(state) - case SwarmConversationExecutionState(): - _prepare_restored_swarm_conversation_state(state) - case ManagerWorkersConversationExecutionState(): - _prepare_restored_managerworkers_conversation_state(state) - - -def _prepare_restored_agent_conversation_state(state: "AgentConversationExecutionState") -> None: - """Re-key Agent subconversations by the canonical runtime slot key.""" - from wayflowcore.executors._agentconversation import AgentConversation - - state.current_sub_component_conversations = { - AgentConversation._sub_component_conversation_key( - subconversation.component - ): subconversation - for subconversation in state.current_sub_component_conversations.values() - } - - -def _validate_restored_flow_conversation_state(state: "FlowConversationExecutionState") -> None: - """Fail fast if restored Flow state names steps that no longer exist live.""" - valid_step_names = set(state.flow.steps) - if state.flow.begin_step_name is not None: - valid_step_names.add(state.flow.begin_step_name) - - for step_name in chain( - (state.current_step_name,), - state.step_history, - (step_name for step_name, _output_name in state.input_output_key_values), - ): - if step_name is None or step_name in valid_step_names: - continue - raise _CheckpointRestoreCompatibilityError( - "Cannot restore this checkpoint because flow conversation state refers to " - f"step `{step_name}` which is not present in the current Flow. " - f"{_STEP_NAME_ERROR_HINT}" - ) - - -def _prepare_restored_swarm_conversation_state(state: "SwarmConversationExecutionState") -> None: - """Rebuild Swarm's derived indexes, then assert the active thread is still valid.""" - threads = [state.main_thread] + [ - thread - for recipients_and_threads in state.agents_and_threads.values() - for thread in recipients_and_threads.values() - ] - # Rebuild the recipient lookup from the live thread objects. The serialized - # dict keys may no longer line up after deserialization. - state.agents_and_threads = {} - for thread in threads: - if not thread.is_main_thread: - state.agents_and_threads.setdefault(thread.caller.name, {})[ - thread.recipient_agent.name - ] = thread - # Thread subconversations are matched by sharing the same message list object. - # The thread identifiers are the stable keys; message-list object identity lets - # us reconnect them after deserialization rebuilt the in-memory objects. - thread_ids_by_message_list = {id(thread.message_list): thread.identifier for thread in threads} - state.thread_subconversations = { - thread_id: subconversation - for subconversation in state.thread_subconversations.values() - if (thread_id := thread_ids_by_message_list.get(id(subconversation.message_list))) - is not None - } - if state.current_thread is None: - raise _CheckpointRestoreCompatibilityError( - "Cannot restore this checkpoint because Swarm conversation state does not " - f"have an active thread. {_AGENT_NAME_ERROR_HINT}" - ) - - valid_thread_ids = {thread.identifier for thread in threads} - if state.current_thread.identifier not in valid_thread_ids: - raise _CheckpointRestoreCompatibilityError( - "Cannot restore this checkpoint because Swarm conversation state refers to " - f"thread `{state.current_thread.identifier}` which is not present in the current " - f"Swarm topology. {_AGENT_NAME_ERROR_HINT}" - ) - - -def _prepare_restored_managerworkers_conversation_state( - state: "ManagerWorkersConversationExecutionState", -) -> None: - """Re-key ManagerWorkers subconversations, then assert the active agent is still valid.""" - # Subconversations are restored as objects first; rebuild the name-keyed lookup - # used by the runtime from those live conversation objects. - state.subconversations = { - subconversation.component.name: subconversation - for subconversation in state.subconversations.values() - } - if ( - state.current_agent_name is not None - and state.current_agent_name not in state.subconversations - ): - raise _CheckpointRestoreCompatibilityError( - "Cannot restore this checkpoint because ManagerWorkers conversation state " - f"refers to agent `{state.current_agent_name}` which is not present " - f"in the current component tree. {_AGENT_NAME_ERROR_HINT}" - ) - - # Checkpoint eligibility diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py index 84a9199d9..4be7fb1d7 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py @@ -5,7 +5,7 @@ # (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, Iterator, Optional +from typing import TYPE_CHECKING, Any, Dict, Generator, Optional from ..events import Event, EventListener from ..events.event import ( @@ -174,7 +174,7 @@ def _save_internal_turn_checkpoint(self, event: _IterationStartedEvent) -> None: def get_conversation_checkpoint_execution_context( conversation: "Conversation", is_outermost_execution: bool, -) -> Iterator[None]: +) -> Generator[None, None, None]: """Context manager that wraps one outermost execute() with checkpointing. The ordering is deliberate: diff --git a/wayflowcore/src/wayflowcore/executors/_agentconversation.py b/wayflowcore/src/wayflowcore/executors/_agentconversation.py index ca8528a74..793aa8574 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_agentconversation.py @@ -68,7 +68,8 @@ def _get_all_sub_conversations(self) -> List["Conversation"]: @staticmethod def _sub_component_conversation_key(component: ConversationalComponent) -> str: """Return the canonical runtime key for a subcomponent conversation slot.""" - return f"{component.__class__.__name__}:{component.name}" + # return f"{component.__class__.__name__}:{component.name}" + return component.id def _get_sub_component_conversation( self, component: ConversationalComponent diff --git a/wayflowcore/src/wayflowcore/executors/_flowconversation.py b/wayflowcore/src/wayflowcore/executors/_flowconversation.py index aeec574a7..6d5479666 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_flowconversation.py @@ -87,17 +87,15 @@ def _get_step(self, step_name: str) -> "Step": return self.state.flow.steps[step_name] def _get_current_sub_conversation( - self, - step: "Step", - sub_conversation_id: Optional[str] = None, + self, step: "Step", sub_conversation_id: Optional[str] = None ) -> Optional["Conversation"]: from wayflowcore.executors._flowexecutor import FlowConversationExecutor key = FlowConversationExecutor().make_key_for_step( - step, - sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY, + step, sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY ) - return cast(Optional["Conversation"], self.state.internal_context_key_values.get(key)) + value = self.state.internal_context_key_values.get(key, None) + return cast(Optional["Conversation"], value) def _update_sub_conversation( self, diff --git a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py index 6cab1bfd2..cb8db6b3a 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py @@ -246,7 +246,8 @@ class FlowConversationExecutor(ConversationExecutor): @staticmethod def make_key_for_step(assistant_step: Step, key: str) -> str: - return str(assistant_step.name) + "_" + key + return str(assistant_step.id) + "_" + key + # return str(assistant_step.name) + "_" + key @staticmethod def get_parent_conversation(state: FlowConversationExecutionState) -> Optional[Conversation]: diff --git a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py index 6dbfb90c3..87b09feb2 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py @@ -24,7 +24,7 @@ @dataclass class ManagerWorkersConversationExecutionState(ConversationExecutionState): - current_agent_name: str + current_agent_id: str subconversations: Dict[str, Union["AgentConversation", "ManagerWorkersConversation"]] conversation_id: str = "" @@ -33,9 +33,9 @@ def _create_subconversation_for_agent( ) -> Union["AgentConversation", "ManagerWorkersConversation"]: subconv = agent.start_conversation( conversation_id=self.conversation_id or None, - _runtime_conversation_id=agent.name, + _runtime_conversation_id=agent.id, ) - self.subconversations[agent.name] = subconv + self.subconversations[agent.id] = subconv return subconv @@ -74,9 +74,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, @@ -84,7 +84,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")) @@ -94,7 +94,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..e1b8ec9fd 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,15 @@ 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 ) ) logger.info( @@ -397,4 +404,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 6e06a831d..09d6235ae 100644 --- a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py @@ -75,7 +75,7 @@ def _create_subconversation_for_thread( 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" @@ -166,9 +166,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..67dd17438 100644 --- a/wayflowcore/src/wayflowcore/executors/_swarmexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_swarmexecutor.py @@ -421,8 +421,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/managerworkers.py b/wayflowcore/src/wayflowcore/managerworkers.py index 2428beea8..1c7df58c2 100644 --- a/wayflowcore/src/wayflowcore/managerworkers.py +++ b/wayflowcore/src/wayflowcore/managerworkers.py @@ -129,15 +129,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() @@ -154,7 +162,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, @@ -292,15 +300,15 @@ def start_conversation( ) subconversations: Dict[str, Union[AgentConversation, ManagerWorkersConversation]] = {} - subconversations[self.manager_agent.name] = self.manager_agent.start_conversation( + subconversations[self.manager_agent.id] = self.manager_agent.start_conversation( inputs=inputs, messages=messages, conversation_id=conversation_root_id, - _runtime_conversation_id=self.manager_agent.name, + _runtime_conversation_id=self.manager_agent.id, ) state = ManagerWorkersConversationExecutionState( - current_agent_name=self.manager_agent.name, + current_agent_id=self.manager_agent.id, subconversations=subconversations, conversation_id=conversation_root_id, ) @@ -343,7 +351,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/serialization/context.py b/wayflowcore/src/wayflowcore/serialization/context.py index a6c698eb5..c459a6429 100644 --- a/wayflowcore/src/wayflowcore/serialization/context.py +++ b/wayflowcore/src/wayflowcore/serialization/context.py @@ -150,7 +150,7 @@ 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 _add_component_to_context(self, component: "Component") -> None: + 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. @@ -354,7 +354,7 @@ 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. diff --git a/wayflowcore/src/wayflowcore/steps/retrystep.py b/wayflowcore/src/wayflowcore/steps/retrystep.py index c19c0c762..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__) @@ -259,14 +257,13 @@ def might_yield(self) -> bool: """ return self.flow.might_yield - def _retry_count(self, state: FlowConversationExecutionState) -> int: - context_key = FlowConversationExecutor.make_key_for_step(self, self._RETRY_COUNTER_KEY) - return cast(int, state.internal_context_key_values.get(context_key, 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[ - FlowConversationExecutor.make_key_for_step(self, self._RETRY_COUNTER_KEY) - ] = 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 0297b401d..ce92878c4 100644 --- a/wayflowcore/src/wayflowcore/swarm.py +++ b/wayflowcore/src/wayflowcore/swarm.py @@ -354,10 +354,10 @@ def start_conversation( ) 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, ) diff --git a/wayflowcore/tests/integration/test_checkpointing.py b/wayflowcore/tests/integration/test_checkpointing.py index 951f1883a..69c6ce462 100644 --- a/wayflowcore/tests/integration/test_checkpointing.py +++ b/wayflowcore/tests/integration/test_checkpointing.py @@ -11,7 +11,6 @@ from wayflowcore.agent import Agent from wayflowcore.checkpointing import InMemoryCheckpointer from wayflowcore.checkpointing.checkpoint_state import _save_live_conversation_checkpoint -from wayflowcore.controlconnection import ControlFlowEdge from wayflowcore.conversation import Conversation from wayflowcore.executors._events.event import Event, EventType from wayflowcore.executors._executionstate import ConversationExecutionState @@ -24,9 +23,8 @@ ) from wayflowcore.flow import Flow from wayflowcore.flowbuilder import FlowBuilder -from wayflowcore.idgeneration import IdGenerator from wayflowcore.managerworkers import ManagerWorkers -from wayflowcore.steps import FlowExecutionStep, OutputMessageStep +from wayflowcore.steps import FlowExecutionStep from wayflowcore.swarm import Swarm from ..serialization.test_assistant_serialization import create_flow @@ -180,50 +178,6 @@ def _build_checkpointable_managerworkers() -> tuple[ManagerWorkers, DummyModel]: return managerworkers, manager_llm -def test_flow_checkpoint_restore_rejects_generated_step_name_reinstantiation() -> None: - checkpointer = InMemoryCheckpointer() - - def build_flow() -> tuple[Flow, OutputMessageStep]: - first_step = DoNothingStep() - second_step = OutputMessageStep(message_template="Generated step resumed.") - first_step.id = "generated_step_restore_flow_first_step" - second_step.id = "generated_step_restore_flow_second_step" - return ( - Flow( - begin_step=first_step, - steps=[first_step, second_step], - control_flow_edges=[ - ControlFlowEdge(source_step=first_step, destination_step=second_step), - ControlFlowEdge(source_step=second_step, destination_step=None), - ], - name="generated_step_restore_flow", - flow_id="generated_step_restore_flow", - ), - second_step, - ) - - original_flow, original_second_step = build_flow() - assert IdGenerator.is_auto_generated(original_second_step.name) - conversation = original_flow.start_conversation( - conversation_id="generated-step-name-restore", - checkpointer=checkpointer, - ) - first_status = conversation.execute( - execution_interrupts=[_OnStepStartExecutionInterrupt(original_second_step.name)] - ) - - assert isinstance(first_status, InterruptedExecutionStatus) - assert checkpointer.load_latest(conversation.conversation_id) is not None - - restarted_flow, _ = build_flow() - - with pytest.raises(ValueError, match="stable step names"): - restarted_flow.start_conversation( - conversation_id=conversation.conversation_id, - checkpointer=checkpointer, - ) - - def test_flow_checkpoint_restore_relinks_nested_flow_parent_for_interrupt_inheritance() -> None: checkpointer = InMemoryCheckpointer() @@ -251,7 +205,7 @@ def build_flow() -> Flow: checkpointer=checkpointer, ) restored_child_conversation = restored_conversation._get_current_sub_conversation( - restarted_flow.steps["parent_flow_step"] + original_flow.steps["parent_flow_step"] ) assert restored_child_conversation is not None assert restored_child_conversation._get_parent_conversation() is restored_conversation @@ -357,7 +311,7 @@ def test_multi_agent_checkpointing_supports_resume_and_time_travel( assert rewound_conversation.get_last_message().content == rewound_output -def test_swarm_checkpoint_restore_rebuilds_generated_agent_name_threads() -> None: +def test_swarm_checkpoint_restore_uses_generated_agent_ids_for_threads() -> None: checkpointer = InMemoryCheckpointer() def build_swarm() -> tuple[Swarm, Agent, Agent]: @@ -389,8 +343,8 @@ def build_swarm() -> tuple[Swarm, Agent, Agent]: conversation_id="swarm-generated-agent-name-restart", checkpointer=checkpointer, ) - original_thread = conversation.state.agents_and_threads[original_first_agent.name][ - original_second_agent.name + original_thread = conversation.state.agents_and_threads[original_first_agent.id][ + original_second_agent.id ] conversation.state._create_subconversation_for_thread(original_thread) conversation.state.current_thread = original_thread @@ -405,6 +359,10 @@ def build_swarm() -> tuple[Swarm, Agent, Agent]: assert restored_conversation.state.current_thread.identifier == ( f"{restarted_first_agent.name}#{restarted_second_agent.name}" ) + restored_thread = restored_conversation.state.current_thread + assert restored_conversation.thread_subconversations[restored_thread.id].component is ( + restarted_second_agent + ) def test_agent_checkpoint_restore_relinks_current_flow_parent() -> None: @@ -451,7 +409,7 @@ def build_agent() -> tuple[Agent, Flow]: restored_parent_flow_conversation = restored_conversation.state.current_flow_conversation assert restored_parent_flow_conversation is not None restored_child_conversation = restored_parent_flow_conversation._get_current_sub_conversation( - restarted_parent_flow.steps["agent_parent_flow_step"] + original_parent_flow.steps["agent_parent_flow_step"] ) assert restored_child_conversation is not None assert ( @@ -459,7 +417,7 @@ def build_agent() -> tuple[Agent, Flow]: ) -def test_managerworkers_checkpoint_restore_rejects_nested_generated_agent_names() -> None: +def test_managerworkers_checkpoint_restore_uses_nested_agent_ids() -> None: checkpointer = InMemoryCheckpointer() def build_managerworkers() -> tuple[ManagerWorkers, Agent, ManagerWorkers, Agent, Agent]: @@ -515,14 +473,19 @@ def build_managerworkers() -> tuple[ManagerWorkers, Agent, ManagerWorkers, Agent original_nested_managerworkers ) nested_conversation.state._create_subconversation_for_agent(original_nested_worker_agent) - conversation.state.current_agent_name = original_nested_managerworkers.name - nested_conversation.state.current_agent_name = original_nested_worker_agent.name - _save_live_conversation_checkpoint(checkpointer, conversation) + nested_conversation.state.current_agent_id = original_nested_worker_agent.id + conversation.append_user_message("Save this nested conversation.") + original_outer_manager_agent.llm.set_next_output("Saved.") + conversation.execute() restarted_managerworkers, *_ = build_managerworkers() - with pytest.raises(ValueError, match="stable agent names"): - restarted_managerworkers.start_conversation( - conversation_id=conversation.conversation_id, - checkpointer=checkpointer, - ) + restored_conversation = restarted_managerworkers.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + assert restored_conversation.state.current_agent_id == original_outer_manager_agent.id + restored_nested_conversation = restored_conversation.subconversations[ + original_nested_managerworkers.id + ] + assert restored_nested_conversation.state.current_agent_id == original_nested_worker_agent.id 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 index e6fa3d799..09906a3ae 100644 --- a/wayflowcore/tests/test_checkpointing.py +++ b/wayflowcore/tests/test_checkpointing.py @@ -10,13 +10,19 @@ from wayflowcore.agent import Agent from wayflowcore.checkpointing import CheckpointingInterval, InMemoryCheckpointer +from wayflowcore.controlconnection import ControlFlowEdge +from wayflowcore.executors._flowexecutor import FlowConversationExecutor from wayflowcore.executors.executionstatus import FinishedStatus, UserMessageRequestStatus +from wayflowcore.flow import Flow from wayflowcore.flowhelpers import create_single_step_flow -from wayflowcore.steps import OutputMessageStep +from wayflowcore.managerworkers import ManagerWorkers +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 .serialization.test_assistant_serialization import create_flow +from .test_managerworkers import _send_message from .testhelpers.dummy import DummyModel @@ -201,6 +207,141 @@ def test_checkpoint_restore_requires_conversation_id_when_checkpoint_id_is_provi ) +def test_checkpoint_restore_with_serialized_component_graph() -> None: + checkpointer = InMemoryCheckpointer() + original_flow = create_flow() + + conversation = original_flow.start_conversation( + conversation_id="checkpoint-serialized-component-graph", + checkpointer=checkpointer, + ) + first_status = conversation.execute() + assert isinstance(first_status, UserMessageRequestStatus) + + serialized_flow = serialize(original_flow) + reloaded_flow = deserialize(Flow, serialized_flow) + assert reloaded_flow.id == original_flow.id + assert {name: step.id for name, step in reloaded_flow.steps.items()} == { + name: step.id for name, step in original_flow.steps.items() + } + + restored_conversation = reloaded_flow.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + assert isinstance(restored_conversation.status, UserMessageRequestStatus) + + restored_conversation.append_user_message("continue") + assert isinstance(restored_conversation.execute(), FinishedStatus) + + +def test_checkpoint_restore_preserves_retry_counter_with_recreated_step() -> None: + 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) + + counter_key = FlowConversationExecutor.make_key_for_step( + retry_step, RetryStep._RETRY_COUNTER_KEY + ) + checkpoint = next( + checkpoint + for checkpoint in checkpointer.list_checkpoints(conversation.conversation_id) + if f"{counter_key}: 1" in checkpoint.state + ) + + reloaded_flow = deserialize(Flow, serialize(flow)) + reloaded_retry_step = reloaded_flow.steps["retry_step"] + reloaded_retry_step.name = "renamed_retry_step" + restored_conversation = reloaded_flow.start_conversation( + conversation_id=conversation.conversation_id, + checkpoint_id=checkpoint.checkpoint_id, + checkpointer=checkpointer, + ) + + restored_counter_key = FlowConversationExecutor.make_key_for_step( + reloaded_retry_step, RetryStep._RETRY_COUNTER_KEY + ) + assert restored_counter_key == counter_key + assert restored_conversation.state.internal_context_key_values[restored_counter_key] == 1 + restored_status = restored_conversation.execute() + assert isinstance(restored_status, FinishedStatus) + assert restored_status.complete_step_name == "failure" + + +def test_managerworkers_checkpoint_restore_preserves_worker_subconversation() -> None: + manager_llm = DummyModel() + worker = Agent( + llm=DummyModel(fails_if_not_set=False), + name="checkpoint_worker", + description="Checkpoint worker", + initial_message="Worker ready.", + ) + group = ManagerWorkers( + group_manager=manager_llm, + workers=[worker], + name="checkpoint_managerworkers", + id="checkpoint-managerworkers", + ) + checkpointer = InMemoryCheckpointer( + checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS + ) + + conversation = group.start_conversation( + conversation_id="checkpoint-managerworkers", + checkpointer=checkpointer, + ) + conversation.append_user_message("Delegate this task.") + manager_llm.set_next_output([_send_message(worker, message="Please help."), "All done."]) + + status = conversation.execute() + assert isinstance(status, UserMessageRequestStatus) + assert worker.id in conversation.subconversations + + checkpoint = checkpointer.load_latest(conversation.conversation_id) + assert checkpoint is not None + restored_conversation = group.start_conversation( + conversation_id=conversation.conversation_id, + checkpoint_id=checkpoint.checkpoint_id, + checkpointer=checkpointer, + ) + + assert worker.id in restored_conversation.subconversations + assert restored_conversation.subconversations[worker.id].component is worker + + @pytest.mark.parametrize( ( "interval", diff --git a/wayflowcore/tests/test_managerworkers.py b/wayflowcore/tests/test_managerworkers.py index f2e5d601d..f2b622bc3 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,26 @@ 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_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 worker_conversation.conversation_id == conversation.conversation_id + assert worker_conversation.id == worker.id + + @pytest.fixture def simple_math_agents_example(remote_gemma_llm) -> Tuple[Agent, Agent, Agent]: llm = remote_gemma_llm @@ -340,7 +360,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 +822,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_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 From c97d6e62560d54fda94f60ff7eea1ea53706c6d4 Mon Sep 17 00:00:00 2001 From: Son Le Date: Wed, 22 Jul 2026 15:32:35 +0200 Subject: [PATCH 06/10] remove runtime_conversation_id and use conversation_id or conversation.id only --- .../core/howtoguides/howto_checkpointing.rst | 8 +- wayflowcore/src/wayflowcore/a2a/a2aagent.py | 36 ++++++--- wayflowcore/src/wayflowcore/agent.py | 43 ++++++---- .../contextproviders/flowcontextprovider.py | 4 +- wayflowcore/src/wayflowcore/conversation.py | 8 +- .../wayflowcore/conversationalcomponent.py | 78 +++++++++++++------ .../wayflowcore/executors/_agentexecutor.py | 16 ++-- .../wayflowcore/executors/_flowexecutor.py | 9 +-- .../executors/_managerworkersconversation.py | 14 ++-- .../executors/_managerworkersexecutor.py | 3 +- .../executors/_swarmconversation.py | 21 +++-- .../wayflowcore/executors/_swarmexecutor.py | 3 +- wayflowcore/src/wayflowcore/flow.py | 77 +++++++++++++++--- wayflowcore/src/wayflowcore/managerworkers.py | 73 ++++++++++++----- .../wayflowcore/mcp/_session_persistence.py | 8 +- wayflowcore/src/wayflowcore/ociagent.py | 46 ++++++----- .../wayflowcore/serialization/serializer.py | 27 +++++-- .../wayflowcore/steps/agentexecutionstep.py | 5 +- .../steps/parallelflowexecutionstep.py | 4 +- wayflowcore/src/wayflowcore/swarm.py | 63 ++++++++++++--- .../src/wayflowcore/tools/servertools.py | 4 +- .../tests/integration/test_checkpointing.py | 73 ++++++++--------- wayflowcore/tests/test_checkpointing.py | 29 ++++++- wayflowcore/tests/test_managerworkers.py | 6 +- 24 files changed, 437 insertions(+), 221 deletions(-) diff --git a/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst b/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst index a38f8e604..131750435 100644 --- a/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst +++ b/docs/wayflowcore/source/core/howtoguides/howto_checkpointing.rst @@ -54,11 +54,11 @@ 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 runtime id of this concrete ``Conversation`` object +- ``conversation.id``: the id of one concrete ``Conversation`` within that conversation thread -For a fresh top-level conversation, ``conversation.id`` usually matches -``conversation.conversation_id``. They can differ for restored or nested conversations, so use -``conversation_id`` in checkpointing APIs when you mean "this logical conversation". +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:: diff --git a/wayflowcore/src/wayflowcore/a2a/a2aagent.py b/wayflowcore/src/wayflowcore/a2a/a2aagent.py index 6a28ec341..55fe45a0c 100644 --- a/wayflowcore/src/wayflowcore/a2a/a2aagent.py +++ b/wayflowcore/src/wayflowcore/a2a/a2aagent.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: from wayflowcore.checkpointing import Checkpointer + from wayflowcore.conversation import Conversation from wayflowcore.executors._a2aagentconversation import A2AAgentConversation logger = logging.getLogger(__name__) @@ -252,8 +253,6 @@ def start_conversation( conversation_id: Optional[str] = None, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _runtime_conversation_id: Optional[str] = None, - _attach_checkpointer: bool = True, ) -> "A2AAgentConversation": """ Initiates a new conversation with the remote server agent. @@ -276,29 +275,42 @@ def start_conversation( checkpoint_id: Optional checkpoint identifier to restore. Requires both ``checkpointer`` and ``conversation_id``. - _runtime_conversation_id: - Internal runtime id for a fresh conversation. When provided, it becomes - the created conversation object's ``.id`` instead of defaulting to - ``conversation_id``. - Returns ------- Conversation: A new conversation object associated with this agent. """ + return self._start_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=None, + ) + + 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, + parent_conversation: Optional["Conversation"] = None, + ) -> "A2AAgentConversation": + """Create the concrete A2A conversation, including nested conversations.""" from wayflowcore.executors._a2aagentconversation import A2AAgentConversation from wayflowcore.executors._a2aagentexecutor import A2AAgentState - restored_conversation, conversation_runtime_id, conversation_root_id = ( + 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, - _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=A2AAgentConversation, - attach_checkpointer=_attach_checkpointer, + parent_conversation=parent_conversation, ) ) if restored_conversation is not None: @@ -313,10 +325,10 @@ def start_conversation( inputs=inputs or {}, # Inputs are ignored in execution message_list=messages, status=None, - id=conversation_runtime_id, + id=conversation_instance_id, checkpointer=checkpointer, name="a2a_conversation", - conversation_id=conversation_root_id, + conversation_id=conversation_thread_id, __metadata_info__={}, ) diff --git a/wayflowcore/src/wayflowcore/agent.py b/wayflowcore/src/wayflowcore/agent.py index 228f96d22..e9adf1097 100644 --- a/wayflowcore/src/wayflowcore/agent.py +++ b/wayflowcore/src/wayflowcore/agent.py @@ -28,6 +28,7 @@ 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 @@ -396,8 +397,6 @@ def start_conversation( conversation_id: Optional[str] = None, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _runtime_conversation_id: Optional[str] = None, - _attach_checkpointer: bool = True, ) -> "AgentConversation": """ Initializes a conversation with the agent. @@ -405,10 +404,9 @@ def start_conversation( 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: @@ -416,31 +414,44 @@ def start_conversation( checkpoint_id: Optional checkpoint identifier to restore. Requires both ``checkpointer`` and ``conversation_id``. - _runtime_conversation_id: - Internal runtime id for a fresh conversation. When provided, it becomes - the created conversation object's ``.id`` instead of defaulting to - ``conversation_id``. - Returns ------- Conversation: The conversation object of the agent. """ + return self._start_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=None, + ) + + 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, + parent_conversation: Optional["Conversation"] = None, + ) -> "AgentConversation": + """Create the concrete agent conversation, including nested conversations.""" 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_runtime_id, conversation_root_id = ( + 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, - _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=AgentConversation, - attach_checkpointer=_attach_checkpointer, + parent_conversation=parent_conversation, ) ) if restored_conversation is not None: @@ -488,7 +499,7 @@ def start_conversation( conversational_component=self, inputs=inputs, messages=messages, - conversation_id=conversation_runtime_id, + conversation_id=conversation_instance_id, nesting_level=None, ) ) @@ -496,13 +507,13 @@ def start_conversation( return AgentConversation( component=self, message_list=messages, - id=conversation_runtime_id, + id=conversation_instance_id, checkpointer=checkpointer, inputs=inputs or {}, name="agent_conversation", state=AgentConversationExecutionState(), status=None, - conversation_id=conversation_root_id, + conversation_id=conversation_thread_id, __metadata_info__={}, ) diff --git a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py index 9d83e4dc1..2f7d8703e 100644 --- a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py +++ b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py @@ -92,10 +92,10 @@ 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_subconversation( + parent_conversation=conversation, inputs={}, messages=conversation.message_list, - conversation_id=conversation.conversation_id, ) status = await conversation.execute_async() if status._requires_yielding: diff --git a/wayflowcore/src/wayflowcore/conversation.py b/wayflowcore/src/wayflowcore/conversation.py index 67202bc92..d69a667c5 100644 --- a/wayflowcore/src/wayflowcore/conversation.py +++ b/wayflowcore/src/wayflowcore/conversation.py @@ -64,15 +64,15 @@ def _get_active_conversations(return_copy: bool = True) -> List["Conversation"]: def _get_current_conversation_id() -> Optional[str]: - """Return the runtime id of the currently executing conversation object.""" + """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_root_conversation_id() -> Optional[str]: - """Return the root conversation id shared by nested conversations.""" +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 @@ -102,7 +102,7 @@ class Conversation(DataclassComponent): status: Optional[ExecutionStatus] token_usage: TokenUsage = field(default_factory=TokenUsage, init=False) conversation_id: str = "" - """Root conversation id used for checkpoint persistence and resume.""" + """Conversation thread id used for checkpoint persistence and resume.""" checkpointer: Optional["Checkpointer"] = field( default=None, repr=False, diff --git a/wayflowcore/src/wayflowcore/conversationalcomponent.py b/wayflowcore/src/wayflowcore/conversationalcomponent.py index 80ca50de9..7ca955bee 100644 --- a/wayflowcore/src/wayflowcore/conversationalcomponent.py +++ b/wayflowcore/src/wayflowcore/conversationalcomponent.py @@ -76,7 +76,6 @@ def __init__( __metadata_info__=__metadata_info__, ) - @abstractmethod def start_conversation( self, inputs: Optional[Dict[str, Any]] = None, @@ -84,16 +83,33 @@ def start_conversation( conversation_id: Optional[str] = None, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _runtime_conversation_id: Optional[str] = None, - _attach_checkpointer: bool = True, ) -> "Conversation": """Start a fresh conversation or restore one from checkpoint storage. - ``conversation_id`` is the durable conversation identifier used for resume and - storage. For fresh conversations, ``_runtime_conversation_id`` controls the - concrete ``Conversation.id`` assigned to the created object; if omitted, the - runtime id defaults to the durable conversation id. + ``conversation_id`` identifies the complete conversation thread for checkpoint + storage and resume. Each concrete conversation in that thread has its own ``id``. """ + return self._start_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=None, + ) + + @abstractmethod + def _start_conversation( + 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": + """Create a concrete conversation for this component.""" + raise NotImplementedError @property def llms(self) -> List["LlmModel"]: @@ -161,40 +177,36 @@ def _prepare_conversation_start( inputs: Optional[Dict[str, Any]], messages: Union[None, str, "Message", List["Message"], "MessageList"], conversation_id: Optional[str], - _runtime_conversation_id: Optional[str], checkpointer: Optional["Checkpointer"], checkpoint_id: Optional[str], expected_conversation_type: Type[ConversationTypeT], - attach_checkpointer: bool, + 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 runtime conversation id to assign to ``Conversation.id`` on the concrete object - - the root conversation id shared across nested fresh starts + - 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 mypy; runtime validation uses it too. """ - # No checkpointer means this is a fresh conversation; just resolve ids. + if parent_conversation is not None: + if checkpoint_id is not None: + raise ValueError("Cannot restore a checkpoint as a subconversation.") + return None, IdGenerator.get_or_generate_id(), parent_conversation.conversation_id + + # No checkpointer means this is a fresh root conversation; just resolve ids. if checkpointer is None: if checkpoint_id is not None: raise ValueError("`checkpoint_id` requires a `checkpointer`.") - root_conversation_id = conversation_id or IdGenerator.get_or_generate_id() - # This value becomes the fresh conversation object's `.id`. - runtime_conversation_id = IdGenerator.get_or_generate_id( - _runtime_conversation_id or root_conversation_id - ) - return None, runtime_conversation_id, root_conversation_id - - if _runtime_conversation_id is not None: - raise ValueError( - "`_runtime_conversation_id` is not supported when restoring checkpoints." - ) + # 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() + return None, conversation_id, conversation_id - # Checkpoint restore needs the root conversation id to locate stored state. + # Checkpoint restore uses the thread id to locate stored state. resolved_conversation_id = conversation_id if resolved_conversation_id is None and checkpoint_id is not None: raise ValueError("`checkpoint_id` requires a `conversation_id`.") @@ -225,7 +237,7 @@ def _prepare_conversation_start( expected_conversation_type=expected_conversation_type, tool_registry={tool.name: tool for tool in self._referenced_tools()}, checkpointer=checkpointer, - attach_checkpointer=attach_checkpointer, + attach_checkpointer=True, ) return ( cast(ConversationTypeT, conversation), @@ -233,6 +245,22 @@ def _prepare_conversation_start( resolved_conversation_id, ) + def _start_subconversation( + self, + parent_conversation: "Conversation", + inputs: Optional[Dict[str, Any]] = None, + messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, + ) -> "Conversation": + """Start a child that inherits its parent's conversation thread.""" + return self._start_conversation( + inputs=inputs, + messages=messages, + conversation_id=None, + checkpointer=None, + checkpoint_id=None, + parent_conversation=parent_conversation, + ) + # Define a TypeVar that represents the component's type ConversationalComponentTypeT = TypeVar( diff --git a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py index 8b7005d7c..6f7d1f583 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py @@ -469,11 +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( + sub_agent_conversation = expert_agent._start_subconversation( + parent_conversation=caller_conv, messages=init_messages, inputs=inputs, - conversation_id=caller_conv.conversation_id, - _runtime_conversation_id=caller_conv._sub_component_conversation_key(expert_agent), ) return sub_agent_conversation @@ -530,8 +529,7 @@ async def _execute_flow( messages: MessageList, flow: Flow, inputs: Dict[str, Any], - conversation_id: Optional[str], - runtime_conversation_id: Optional[str], + parent_conversation: "AgentConversation", ) -> Tuple[Any, str, ExecutionStatus]: """ Execute a flow and return its outputs and its execution status. @@ -541,11 +539,10 @@ 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_subconversation( + parent_conversation=parent_conversation, inputs=inputs, messages=messages, - conversation_id=conversation_id, - _runtime_conversation_id=runtime_conversation_id, ) messages.append_message( Message( @@ -775,8 +772,7 @@ async def _handle_flow_call( messages, flow, tool_request.args, - conversation.conversation_id, - conversation._sub_component_conversation_key(flow), + conversation, ) ) diff --git a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py index cb8db6b3a..2fadc2ab2 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py @@ -290,12 +290,9 @@ def create_sub_conversation( resolved_sub_conversation_id = ( sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY ) - sub_conversation = flow.start_conversation( - inputs_not_from_context_providers, - conversation_id=conversation.conversation_id, - _runtime_conversation_id=FlowConversationExecutor.make_key_for_step( - step, resolved_sub_conversation_id - ), + sub_conversation = flow._start_subconversation( + 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, diff --git a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py index 87b09feb2..fc5db4bd4 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py @@ -6,7 +6,7 @@ import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Union, cast from wayflowcore.agent import Agent from wayflowcore.conversation import Conversation @@ -26,17 +26,17 @@ class ManagerWorkersConversationExecutionState(ConversationExecutionState): current_agent_id: str subconversations: Dict[str, Union["AgentConversation", "ManagerWorkersConversation"]] - conversation_id: str = "" 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( - conversation_id=self.conversation_id or None, - _runtime_conversation_id=agent.id, + subconv = cast( + "Union[AgentConversation, ManagerWorkersConversation]", + agent._start_subconversation(parent_conversation=parent_conversation), ) self.subconversations[agent.id] = subconv - return subconv diff --git a/wayflowcore/src/wayflowcore/executors/_managerworkersexecutor.py b/wayflowcore/src/wayflowcore/executors/_managerworkersexecutor.py index e1b8ec9fd..76f227c35 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersexecutor.py @@ -393,7 +393,8 @@ def _send_message_to_worker( if worker_subconversation is None: worker_subconversation = ( managerworkers_conversation.state._create_subconversation_for_agent( - recipient_agent + recipient_agent, + parent_conversation=managerworkers_conversation, ) ) logger.info( diff --git a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py index 09d6235ae..f6fbac334 100644 --- a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py @@ -6,7 +6,7 @@ import logging from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast from wayflowcore.agent import Agent from wayflowcore.conversation import Conversation @@ -55,7 +55,6 @@ class SwarmConversationExecutionState(ConversationExecutionState): main_thread: SwarmThread agents_and_threads: Dict[str, Dict[str, SwarmThread]] context_providers: List["ContextProvider"] - conversation_id: str = "" current_thread: Optional["SwarmThread"] = None thread_stack: List["SwarmThread"] = field(default_factory=list) @@ -64,14 +63,10 @@ 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": @@ -87,11 +82,13 @@ def _create_subconversation_for_thread( if isinstance(message_list, list) else message_list ) - conversation = thread.recipient_agent.start_conversation( - inputs=inputs, - messages=thread.message_list, - conversation_id=self.conversation_id or None, - _runtime_conversation_id=thread_id, + conversation = cast( + "AgentConversation", + thread.recipient_agent._start_subconversation( + parent_conversation=parent_conversation, + inputs=inputs, + messages=thread.message_list, + ), ) self.thread_subconversations[thread_id] = conversation diff --git a/wayflowcore/src/wayflowcore/executors/_swarmexecutor.py b/wayflowcore/src/wayflowcore/executors/_swarmexecutor.py index 67dd17438..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) diff --git a/wayflowcore/src/wayflowcore/flow.py b/wayflowcore/src/wayflowcore/flow.py index f0b32b22e..834a8d8b8 100644 --- a/wayflowcore/src/wayflowcore/flow.py +++ b/wayflowcore/src/wayflowcore/flow.py @@ -46,6 +46,7 @@ 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 @@ -1167,10 +1168,67 @@ def start_conversation( conversation_id: Optional[str] = None, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _runtime_conversation_id: Optional[str] = None, - _attach_checkpointer: bool = True, nesting_level: int = 0, context_providers_from_parent_flow: Optional[Set[str]] = None, + ) -> "FlowConversation": + return self._start_flow_conversation( + 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_subconversation( + 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": + return self._start_flow_conversation( + 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( + 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, + ) -> "FlowConversation": + return self._start_flow_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=parent_conversation, + ) + + def _start_flow_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, + nesting_level: int = 0, + context_providers_from_parent_flow: Optional[Set[str]] = None, + parent_conversation: Optional["Conversation"] = None, ) -> "FlowConversation": """ Start the conversation. @@ -1189,10 +1247,6 @@ def start_conversation( checkpoint_id: Optional checkpoint identifier to restore. Requires both ``checkpointer`` and ``conversation_id``. - _runtime_conversation_id: - Internal runtime id for a fresh conversation. When provided, it becomes - the created conversation object's ``.id`` instead of defaulting to - ``conversation_id``. context_providers_from_parent_flow: Context provider that don't need to be checked when validating existing inputs. nesting_level: @@ -1207,16 +1261,15 @@ def start_conversation( from wayflowcore.events.eventlistener import record_event from wayflowcore.executors._flowconversation import FlowConversation - restored_conversation, conversation_runtime_id, conversation_root_id = ( + 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, - _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=FlowConversation, - attach_checkpointer=_attach_checkpointer, + parent_conversation=parent_conversation, ) ) if restored_conversation is not None: @@ -1271,7 +1324,7 @@ def start_conversation( conversational_component=self, inputs=inputs, messages=messages, - conversation_id=conversation_runtime_id, + conversation_id=conversation_instance_id, nesting_level=nesting_level, ) ) @@ -1303,14 +1356,14 @@ def start_conversation( return FlowConversation( component=self, inputs=inputs, - id=conversation_runtime_id, + id=conversation_instance_id, checkpointer=checkpointer, message_list=messages, __metadata_info__={}, status=None, name="flow_conversation", state=state, - conversation_id=conversation_root_id, + conversation_id=conversation_thread_id, ) @property diff --git a/wayflowcore/src/wayflowcore/managerworkers.py b/wayflowcore/src/wayflowcore/managerworkers.py index 1c7df58c2..0bb8dd526 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 @@ -23,6 +23,7 @@ if TYPE_CHECKING: from wayflowcore.checkpointing import Checkpointer + from wayflowcore.conversation import Conversation from wayflowcore.executors._managerworkersconversation import ManagerWorkersConversation from wayflowcore.messagelist import Message @@ -232,9 +233,45 @@ def start_conversation( conversation_id: Optional[str] = None, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _runtime_conversation_id: Optional[str] = None, - _attach_checkpointer: bool = True, conversation_name: Optional[str] = None, + ) -> "ManagerWorkersConversation": + return self._start_managerworkers_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + conversation_name=conversation_name, + parent_conversation=None, + ) + + def _start_conversation( + 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, + ) -> "ManagerWorkersConversation": + return self._start_managerworkers_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=parent_conversation, + ) + + def _start_managerworkers_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_name: Optional[str] = None, + parent_conversation: Optional["Conversation"] = None, ) -> "ManagerWorkersConversation": """ Initializes a conversation with the managerworkers. @@ -253,10 +290,6 @@ def start_conversation( checkpoint_id: Optional checkpoint identifier to restore. Requires both ``checkpointer`` and ``conversation_id``. - _runtime_conversation_id: - Internal runtime id for a fresh conversation. When provided, it becomes - the created conversation object's ``.id`` instead of defaulting to - ``conversation_id``. Returns ------- @@ -271,16 +304,15 @@ def start_conversation( ManagerWorkersConversationExecutionState, ) - restored_conversation, conversation_runtime_id, conversation_root_id = ( + 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, - _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=ManagerWorkersConversation, - attach_checkpointer=_attach_checkpointer, + parent_conversation=parent_conversation, ) ) if restored_conversation is not None: @@ -294,37 +326,38 @@ def start_conversation( conversational_component=self, inputs=inputs or {}, messages=messages, - conversation_id=conversation_runtime_id, + conversation_id=conversation_instance_id, nesting_level=None, ) ) subconversations: Dict[str, Union[AgentConversation, ManagerWorkersConversation]] = {} - subconversations[self.manager_agent.id] = self.manager_agent.start_conversation( - inputs=inputs, - messages=messages, - conversation_id=conversation_root_id, - _runtime_conversation_id=self.manager_agent.id, - ) state = ManagerWorkersConversationExecutionState( current_agent_id=self.manager_agent.id, subconversations=subconversations, - conversation_id=conversation_root_id, ) conversation = ManagerWorkersConversation( component=self, inputs={}, message_list=messages, - id=conversation_runtime_id, + id=conversation_instance_id, name=conversation_name or "managerworkers_conversation", state=state, status=None, checkpointer=checkpointer, - conversation_id=conversation_root_id, + conversation_id=conversation_thread_id, __metadata_info__={}, ) + subconversations[self.manager_agent.id] = cast( + "Union[AgentConversation, ManagerWorkersConversation]", + self.manager_agent._start_subconversation( + parent_conversation=conversation, + inputs=inputs, + messages=messages, + ), + ) return conversation def _referenced_tools_dict_inner( diff --git a/wayflowcore/src/wayflowcore/mcp/_session_persistence.py b/wayflowcore/src/wayflowcore/mcp/_session_persistence.py index e13524301..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_root_conversation_id + from wayflowcore.conversation import _get_current_conversation_thread_id - return _get_current_root_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_root_conversation_id + from wayflowcore.conversation import _get_current_conversation_thread_id runtime = get_mcp_async_runtime() - conversation_id = _get_current_root_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 c8438e7d3..4bf63b373 100644 --- a/wayflowcore/src/wayflowcore/ociagent.py +++ b/wayflowcore/src/wayflowcore/ociagent.py @@ -110,37 +110,48 @@ def start_conversation( conversation_id: Optional[str] = None, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _runtime_conversation_id: Optional[str] = None, - _attach_checkpointer: bool = True, - ) -> "Conversation": + ) -> 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``. + 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``. - _runtime_conversation_id: - Internal runtime id for a fresh conversation. When provided, it becomes - the created conversation object's ``.id`` instead of defaulting to - ``conversation_id``. - Returns ------- Conversation: The conversation object of the agent. """ + return self._start_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=None, + ) + + 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, + parent_conversation: Optional["Conversation"] = None, + ) -> "Conversation": + """Create the concrete OCI conversation, including nested conversations.""" from wayflowcore.executors._ociagentconversation import OciAgentConversation from wayflowcore.executors._ociagentexecutor import ( OciAgentState, @@ -154,16 +165,15 @@ def start_conversation( if not isinstance(messages, MessageList): messages = MessageList.from_messages(messages=messages) - _restored_conversation, conversation_runtime_id, conversation_root_id = ( + _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, - _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=OciAgentConversation, - attach_checkpointer=_attach_checkpointer, + parent_conversation=parent_conversation, ) ) @@ -179,9 +189,9 @@ def start_conversation( inputs=inputs or {}, message_list=messages, status=None, - id=conversation_runtime_id, + id=conversation_instance_id, name="oci_conversation", - conversation_id=conversation_root_id, + conversation_id=conversation_thread_id, __metadata_info__={}, ) diff --git a/wayflowcore/src/wayflowcore/serialization/serializer.py b/wayflowcore/src/wayflowcore/serialization/serializer.py index 83f76b43a..330dc6328 100644 --- a/wayflowcore/src/wayflowcore/serialization/serializer.py +++ b/wayflowcore/src/wayflowcore/serialization/serializer.py @@ -19,6 +19,7 @@ ForwardRef, List, Optional, + Protocol, Type, TypeVar, cast, @@ -558,10 +559,19 @@ def serialize( T = TypeVar("T", bound=SerializableObject) +T_co = TypeVar("T_co", bound=SerializableObject, covariant=True) + + +class _SerializableType(Protocol[T_co]): + """Type token for a serializable object, including abstract base classes.""" + + __name__: str + + def __call__(self, *args: Any, **kwargs: Any) -> T_co: ... def deserialize_from_dict( - deserialization_type: Type[T], + deserialization_type: _SerializableType[T], obj_as_dict: Dict[str, Any], deserialization_context: Optional[DeserializationContext] = None, plugins: Optional[List["WayflowDeserializationPlugin"]] = None, @@ -598,6 +608,8 @@ def deserialize_from_dict( >>> new_assistant = deserialize_from_dict(Flow, serialized_assistant) """ + runtime_deserialization_type = cast(Type[T], deserialization_type) + if deserialization_context is None: deserialization_context = DeserializationContext(plugins=plugins) elif plugins is not None: @@ -616,7 +628,7 @@ def deserialize_from_dict( deserialized_obj: SerializableObject = deserialization_context.get_deserialized_object( object_reference ) - if not isinstance(deserialized_obj, deserialization_type): + if not isinstance(deserialized_obj, runtime_deserialization_type): raise ValueError( f"A referenced objects found of type {deserialized_obj.__class__.__name__} " f"which is not compatible with the expected deserialization type of " @@ -628,10 +640,10 @@ def deserialize_from_dict( obj_as_dict = deserialization_context.get_referenced_dict(object_reference) deserialization_plugin = deserialization_context.get_deserialization_plugin_for_object( - deserialization_type + runtime_deserialization_type ) deserialized_obj = deserialization_plugin.deserialize( - deserialization_type, obj_as_dict, deserialization_context + runtime_deserialization_type, obj_as_dict, deserialization_context ) if object_reference: deserialization_context.recorddeserialized_object(object_reference, deserialized_obj) @@ -650,7 +662,7 @@ def _set_component_id(component: ObjectWithMetadata, reference: str) -> None: def deserialize( - deserialization_type: Type[T], + deserialization_type: _SerializableType[T], obj: str, deserialization_context: Optional[DeserializationContext] = None, plugins: Optional[List["WayflowDeserializationPlugin"]] = None, @@ -775,7 +787,10 @@ def autodeserialize_from_dict( "Failure to deserialize due to missing `_component_type`: The following object " f"does not seem to be a valid WayFlow component to deserialize:\n{obj_as_dict}" ) - deserialization_type = SerializableObject.get_component(component_type) + deserialization_type = cast( + _SerializableType[SerializableObject], + SerializableObject.get_component(component_type), + ) if component_type is not None and component_type != deserialization_type.__name__: raise ValueError( diff --git a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py index c83b27c62..662d70561 100644 --- a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py +++ b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py @@ -293,11 +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( + agent_sub_conversation = self.agent._start_subconversation( + parent_conversation=caller_conv, inputs=inputs, messages=init_messages, - conversation_id=caller_conv.conversation_id, - _runtime_conversation_id=sub_conversation_id, ) return agent_sub_conversation diff --git a/wayflowcore/src/wayflowcore/steps/parallelflowexecutionstep.py b/wayflowcore/src/wayflowcore/steps/parallelflowexecutionstep.py index b5e4c4af7..16bfc8fe1 100644 --- a/wayflowcore/src/wayflowcore/steps/parallelflowexecutionstep.py +++ b/wayflowcore/src/wayflowcore/steps/parallelflowexecutionstep.py @@ -255,9 +255,9 @@ async def _invoke_step_async( if input_name in set(p.name for p in flow.input_descriptors) # We extract the inputs needed by this specific flow }, - sub_conversation_id=f"parallel_subflow_{index}", + sub_conversation_id=flow.id, ) - for index, flow in enumerate(self.flows) + for flow in self.flows ] # We run only the conversations that did not reach the end diff --git a/wayflowcore/src/wayflowcore/swarm.py b/wayflowcore/src/wayflowcore/swarm.py index ce92878c4..7684b06de 100644 --- a/wayflowcore/src/wayflowcore/swarm.py +++ b/wayflowcore/src/wayflowcore/swarm.py @@ -314,15 +314,49 @@ def start_conversation( conversation_id: Optional[str] = None, checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, - _runtime_conversation_id: Optional[str] = None, - _attach_checkpointer: bool = True, conversation_name: Optional[str] = None, + ) -> "Conversation": + return self._start_swarm_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + conversation_name=conversation_name, + parent_conversation=None, + ) + + def _start_conversation( + 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": + return self._start_swarm_conversation( + inputs=inputs, + messages=messages, + conversation_id=conversation_id, + checkpointer=checkpointer, + checkpoint_id=checkpoint_id, + parent_conversation=parent_conversation, + ) + + def _start_swarm_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_name: Optional[str] = None, + parent_conversation: Optional["Conversation"] = None, ) -> "Conversation": """Start a fresh swarm conversation or restore one from a checkpoint. - ``conversation_id`` is the durable conversation id. For fresh conversations, - ``_runtime_conversation_id`` becomes the created conversation object's - ``.id`` when provided. + ``conversation_id`` identifies the complete conversation thread. """ from wayflowcore.executors._swarmconversation import ( SwarmConversation, @@ -331,16 +365,15 @@ def start_conversation( SwarmUser, ) - restored_conversation, conversation_runtime_id, conversation_root_id = ( + 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, - _runtime_conversation_id=_runtime_conversation_id, expected_conversation_type=SwarmConversation, - attach_checkpointer=_attach_checkpointer, + parent_conversation=parent_conversation, ) ) if restored_conversation is not None: @@ -367,20 +400,26 @@ def start_conversation( context_providers=[], inputs=inputs, messages=messages, - conversation_id=conversation_root_id, ) - return SwarmConversation( + conversation = SwarmConversation( component=self, inputs=inputs or {}, message_list=messages, - id=conversation_runtime_id, + id=conversation_instance_id, name=conversation_name or "swarm_conversation", state=state, status=None, checkpointer=checkpointer, - conversation_id=conversation_root_id, + 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 5e98a3aa3..6a5c09827 100644 --- a/wayflowcore/src/wayflowcore/tools/servertools.py +++ b/wayflowcore/src/wayflowcore/tools/servertools.py @@ -636,10 +636,10 @@ async def __call__(self, **inputs: Any) -> Any: conversation = self.flow.start_conversation(inputs) interrupts = [] else: - conversation = self.flow.start_conversation( + conversation = self.flow._start_subconversation( + parent_conversation=self._parent_conversation, inputs=inputs, messages=self._parent_conversation.message_list, - conversation_id=self._parent_conversation.conversation_id, ) interrupts = self._parent_conversation._get_interrupts() diff --git a/wayflowcore/tests/integration/test_checkpointing.py b/wayflowcore/tests/integration/test_checkpointing.py index 69c6ce462..2ea3ffb9c 100644 --- a/wayflowcore/tests/integration/test_checkpointing.py +++ b/wayflowcore/tests/integration/test_checkpointing.py @@ -9,7 +9,7 @@ import pytest from wayflowcore.agent import Agent -from wayflowcore.checkpointing import InMemoryCheckpointer +from wayflowcore.checkpointing import CheckpointingInterval, InMemoryCheckpointer from wayflowcore.checkpointing.checkpoint_state import _save_live_conversation_checkpoint from wayflowcore.conversation import Conversation from wayflowcore.executors._events.event import Event, EventType @@ -28,6 +28,8 @@ from wayflowcore.swarm import Swarm 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 @@ -204,14 +206,7 @@ def build_flow() -> Flow: conversation_id=conversation.conversation_id, checkpointer=checkpointer, ) - restored_child_conversation = restored_conversation._get_current_sub_conversation( - original_flow.steps["parent_flow_step"] - ) - assert restored_child_conversation is not None - assert restored_child_conversation._get_parent_conversation() is restored_conversation assert restored_conversation.id == restored_conversation.conversation_id - assert restored_child_conversation.id != restored_conversation.id - assert restored_child_conversation.conversation_id == restored_conversation.conversation_id restored_status = restored_conversation.execute( execution_interrupts=[_OnStepStartExecutionInterrupt("child_second_step")] @@ -312,7 +307,9 @@ def test_multi_agent_checkpointing_supports_resume_and_time_travel( def test_swarm_checkpoint_restore_uses_generated_agent_ids_for_threads() -> None: - checkpointer = InMemoryCheckpointer() + checkpointer = InMemoryCheckpointer( + checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS + ) def build_swarm() -> tuple[Swarm, Agent, Agent]: first_agent = Agent( @@ -343,25 +340,21 @@ def build_swarm() -> tuple[Swarm, Agent, Agent]: conversation_id="swarm-generated-agent-name-restart", checkpointer=checkpointer, ) - original_thread = conversation.state.agents_and_threads[original_first_agent.id][ - original_second_agent.id - ] - conversation.state._create_subconversation_for_thread(original_thread) - conversation.state.current_thread = original_thread - _save_live_conversation_checkpoint(checkpointer, conversation) + 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=checkpointer ) - assert restored_conversation.state.current_thread is not None - assert restored_conversation.state.current_thread.identifier == ( - f"{restarted_first_agent.name}#{restarted_second_agent.name}" - ) - restored_thread = restored_conversation.state.current_thread - assert restored_conversation.thread_subconversations[restored_thread.id].component is ( - restarted_second_agent + assert any( + subconversation.component is restarted_second_agent + for subconversation in restored_conversation.thread_subconversations.values() ) @@ -408,13 +401,6 @@ def build_agent() -> tuple[Agent, Flow]: restored_parent_flow_conversation = restored_conversation.state.current_flow_conversation assert restored_parent_flow_conversation is not None - restored_child_conversation = restored_parent_flow_conversation._get_current_sub_conversation( - original_parent_flow.steps["agent_parent_flow_step"] - ) - assert restored_child_conversation is not None - assert ( - restored_child_conversation._get_parent_conversation() is restored_parent_flow_conversation - ) def test_managerworkers_checkpoint_restore_uses_nested_agent_ids() -> None: @@ -469,23 +455,32 @@ def build_managerworkers() -> tuple[ManagerWorkers, Agent, ManagerWorkers, Agent conversation_id="managerworkers-nested-generated-agent-name-restart", checkpointer=checkpointer, ) - nested_conversation = conversation.state._create_subconversation_for_agent( - original_nested_managerworkers - ) - nested_conversation.state._create_subconversation_for_agent(original_nested_worker_agent) - nested_conversation.state.current_agent_id = original_nested_worker_agent.id conversation.append_user_message("Save this nested conversation.") - original_outer_manager_agent.llm.set_next_output("Saved.") + 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, *_ = build_managerworkers() + ( + 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=checkpointer, ) - assert restored_conversation.state.current_agent_id == original_outer_manager_agent.id restored_nested_conversation = restored_conversation.subconversations[ - original_nested_managerworkers.id + restarted_nested_managerworkers.id ] - assert restored_nested_conversation.state.current_agent_id == original_nested_worker_agent.id + assert ( + restored_nested_conversation.subconversations[restarted_nested_worker_agent.id].component + is restarted_nested_worker_agent + ) diff --git a/wayflowcore/tests/test_checkpointing.py b/wayflowcore/tests/test_checkpointing.py index 09906a3ae..42e5cf583 100644 --- a/wayflowcore/tests/test_checkpointing.py +++ b/wayflowcore/tests/test_checkpointing.py @@ -26,6 +26,25 @@ from .testhelpers.dummy import DummyModel +def test_root_conversation_uses_one_identity_for_instance_and_thread() -> None: + agent = Agent(llm=DummyModel(), name="root_identity_agent") + + conversation = agent.start_conversation() + + assert conversation.id == conversation.conversation_id + + +def test_checkpointed_root_conversation_uses_one_identity_for_instance_and_thread() -> None: + agent = Agent(llm=DummyModel(), name="checkpointed_root_identity_agent") + + conversation = agent.start_conversation( + conversation_id="checkpointed-root-identity", + checkpointer=InMemoryCheckpointer(), + ) + + assert conversation.id == conversation.conversation_id + + def _checkpoint_restore_wrong_component_type_scenario( checkpointer: InMemoryCheckpointer, ): @@ -329,6 +348,11 @@ def test_managerworkers_checkpoint_restore_preserves_worker_subconversation() -> status = conversation.execute() assert isinstance(status, UserMessageRequestStatus) assert worker.id in conversation.subconversations + worker_conversation = conversation.subconversations[worker.id] + manager_conversation = conversation.subconversations[group.manager_agent.id] + assert worker_conversation.id != worker.id + assert worker_conversation.id != manager_conversation.id + assert worker_conversation.conversation_id == conversation.conversation_id checkpoint = checkpointer.load_latest(conversation.conversation_id) assert checkpoint is not None @@ -339,7 +363,10 @@ def test_managerworkers_checkpoint_restore_preserves_worker_subconversation() -> ) assert worker.id in restored_conversation.subconversations - assert restored_conversation.subconversations[worker.id].component is worker + restored_worker_conversation = restored_conversation.subconversations[worker.id] + assert restored_worker_conversation.component is worker + assert restored_worker_conversation.id == worker_conversation.id + assert restored_worker_conversation.conversation_id == restored_conversation.conversation_id @pytest.mark.parametrize( diff --git a/wayflowcore/tests/test_managerworkers.py b/wayflowcore/tests/test_managerworkers.py index f2b622bc3..c7c234b5b 100644 --- a/wayflowcore/tests/test_managerworkers.py +++ b/wayflowcore/tests/test_managerworkers.py @@ -135,7 +135,7 @@ 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_identity(): +def test_managerworkers_worker_conversation_inherits_parent_thread_identity(): manager_llm = DummyModel() worker = Agent( DummyModel(fails_if_not_set=False), @@ -151,8 +151,10 @@ def test_managerworkers_worker_conversation_inherits_parent_identity(): 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 != worker.id + assert worker_conversation.id != conversation.id @pytest.fixture From 3ac600a9d8bd06e6b0c91d82c90773f079581dc2 Mon Sep 17 00:00:00 2001 From: Son Le Date: Wed, 22 Jul 2026 17:10:38 +0200 Subject: [PATCH 07/10] address comments --- wayflowcore/setup.cfg | 2 + wayflowcore/src/wayflowcore/a2a/a2aagent.py | 6 +- wayflowcore/src/wayflowcore/agent.py | 12 +-- .../agentserver/_storagehelpers.py | 28 ------ .../services/wayflowservice.py | 9 +- .../agentserver/serverstorageconfig.py | 3 + .../src/wayflowcore/checkpointing/__init__.py | 2 + .../checkpointing/checkpoint_state.py | 26 +++--- .../wayflowcore/checkpointing/checkpointer.py | 9 +- .../checkpointing/checkpointeventlistener.py | 5 +- .../checkpointing/datastorecheckpointer.py | 6 +- wayflowcore/src/wayflowcore/cli/serve.py | 12 +-- wayflowcore/src/wayflowcore/conversation.py | 8 +- .../wayflowcore/conversationalcomponent.py | 24 ++--- .../executors/_agentconversation.py | 12 +-- .../executors/_flowconversation.py | 10 +-- .../wayflowcore/executors/_flowexecutor.py | 12 +-- wayflowcore/src/wayflowcore/flow.py | 47 +++++----- wayflowcore/src/wayflowcore/managerworkers.py | 47 +++++----- wayflowcore/src/wayflowcore/ociagent.py | 4 +- .../src/wayflowcore/serialization/context.py | 10 +-- wayflowcore/src/wayflowcore/swarm.py | 46 +++++----- .../tests/integration/test_checkpointing.py | 90 ++++++++++++++----- wayflowcore/tests/test_checkpointing.py | 40 +++++++++ 24 files changed, 260 insertions(+), 210 deletions(-) delete mode 100644 wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py 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 55fe45a0c..4e69fc7fd 100644 --- a/wayflowcore/src/wayflowcore/a2a/a2aagent.py +++ b/wayflowcore/src/wayflowcore/a2a/a2aagent.py @@ -269,7 +269,7 @@ def start_conversation( 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: - Durable conversation id used for resume, storage, and usage accounting. + Conversation id (of the parent conversation) used for resume, storage, and usage accounting. checkpointer: Optional checkpoint backend used to restore and persist this conversation. checkpoint_id: @@ -280,7 +280,7 @@ def start_conversation( Conversation: A new conversation object associated with this agent. """ - return self._start_conversation( + return self._start_conversation_impl( inputs=inputs, messages=messages, conversation_id=conversation_id, @@ -289,7 +289,7 @@ def start_conversation( parent_conversation=None, ) - def _start_conversation( + def _start_conversation_impl( self, inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, Message, List[Message], MessageList] = None, diff --git a/wayflowcore/src/wayflowcore/agent.py b/wayflowcore/src/wayflowcore/agent.py index e9adf1097..dd705abe7 100644 --- a/wayflowcore/src/wayflowcore/agent.py +++ b/wayflowcore/src/wayflowcore/agent.py @@ -404,11 +404,13 @@ def start_conversation( 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. + 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: @@ -419,7 +421,7 @@ def start_conversation( Conversation: The conversation object of the agent. """ - return self._start_conversation( + return self._start_conversation_impl( inputs=inputs, messages=messages, conversation_id=conversation_id, @@ -428,7 +430,7 @@ def start_conversation( parent_conversation=None, ) - def _start_conversation( + def _start_conversation_impl( self, inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, diff --git a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py b/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py deleted file mode 100644 index e9fd69996..000000000 --- a/wayflowcore/src/wayflowcore/agentserver/_storagehelpers.py +++ /dev/null @@ -1,28 +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 wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig -from wayflowcore.checkpointing.datastorecheckpointer import ( - _prepare_oracle_checkpoint_datastore, - _prepare_postgres_checkpoint_datastore, -) -from wayflowcore.datastore.oracle import OracleDatabaseConnectionConfig -from wayflowcore.datastore.postgres import PostgresDatabaseConnectionConfig - -logger = logging.getLogger(__name__) - - -def _prepare_postgres_datastore( - connection_config: PostgresDatabaseConnectionConfig, storage_config: ServerStorageConfig -) -> None: - _prepare_postgres_checkpoint_datastore(connection_config, storage_config) - - -def _prepare_oracle_datastore( - connection_config: OracleDatabaseConnectionConfig, storage_config: ServerStorageConfig -) -> None: - _prepare_oracle_checkpoint_datastore(connection_config, storage_config) diff --git a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py index 38b91c78f..b834c1c20 100644 --- a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py +++ b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py @@ -15,9 +15,12 @@ from fastapi import status as http_status_code from wayflowcore.agentserver.serverstorageconfig import ServerStorageConfig -from wayflowcore.checkpointing import ConversationCheckpoint, DatastoreCheckpointer +from wayflowcore.checkpointing import ( + CheckpointRestoreCompatibilityError, + ConversationCheckpoint, + DatastoreCheckpointer, +) from wayflowcore.checkpointing.checkpoint_state import ( - _CheckpointRestoreCompatibilityError, _save_live_conversation_checkpoint, _supports_checkpointing, ) @@ -460,7 +463,7 @@ def _load_state( ) conversation.checkpointer = None return conversation - except _CheckpointRestoreCompatibilityError as e: + except CheckpointRestoreCompatibilityError as e: raise HTTPException( status_code=http_status_code.HTTP_400_BAD_REQUEST, detail=f"{incompatible_detail}: {e}", diff --git a/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py b/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py index 990f3c927..b62114736 100644 --- a/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py +++ b/wayflowcore/src/wayflowcore/agentserver/serverstorageconfig.py @@ -12,3 +12,6 @@ @dataclass class ServerStorageConfig(StorageConfig): """Configuration for server storage management.""" + + # 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 index 4458ca4a0..b6cdd29af 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/__init__.py +++ b/wayflowcore/src/wayflowcore/checkpointing/__init__.py @@ -4,6 +4,7 @@ # (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 .checkpoint_state import CheckpointRestoreCompatibilityError from .checkpointer import Checkpointer, CheckpointingInterval, ConversationCheckpoint, StorageConfig from .datastorecheckpointer import ( DatastoreCheckpointer, @@ -14,6 +15,7 @@ __all__ = [ "CheckpointingInterval", + "CheckpointRestoreCompatibilityError", "Checkpointer", "ConversationCheckpoint", "DatastoreCheckpointer", diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py b/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py index 10da70f0a..f1fcacefc 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py @@ -12,9 +12,9 @@ from wayflowcore.serialization import autodeserialize, serialize from wayflowcore.serialization.context import ( DeserializationContext, - MissingDeserializationReferenceError, SerializationContext, - _iter_nested_components, + _get_nested_components, + _MissingDeserializationReferenceError, ) if TYPE_CHECKING: @@ -24,10 +24,7 @@ from wayflowcore.conversationalcomponent import ConversationalComponent -_COMPONENT_ID_ERROR_HINT = "Restart-safe checkpoint restore requires stable component ids." - - -class _CheckpointRestoreCompatibilityError(ValueError): +class CheckpointRestoreCompatibilityError(ValueError): """Raised when a checkpoint cannot be resumed against the current live graph.""" @@ -41,9 +38,7 @@ def _save_live_conversation_checkpoint( from wayflowcore.checkpointing.checkpointer import ConversationCheckpoint if not _supports_checkpointing(conversation.component): - raise NotImplementedError( - "Checkpointing conversations that contain `OciAgent` is not supported yet." - ) + raise NotImplementedError("Checkpointing this component is not supported yet.") serialization_context = SerializationContext(root=conversation) serialization_context._register_external_component_references(conversation.component) @@ -81,16 +76,15 @@ def _load_checkpointed_conversation( checkpoint.state, deserialization_context=deserialization_context, ) - except (MissingDeserializationReferenceError, DataclassFieldDeserializationError) as exc: + except (_MissingDeserializationReferenceError, DataclassFieldDeserializationError) as exc: if not _contains_missing_reference_error(exc): raise - raise _CheckpointRestoreCompatibilityError( + raise CheckpointRestoreCompatibilityError( "Cannot restore this checkpoint because the current component tree does not " - "match the serialized component ids. " - f"{_COMPONENT_ID_ERROR_HINT}" + "match the serialized component ids. Restart-safe checkpoint restore requires stable component ids." ) from exc if not isinstance(conversation, expected_conversation_type): - raise _CheckpointRestoreCompatibilityError( + raise CheckpointRestoreCompatibilityError( "Cannot restore this checkpoint because this conversation was started with another " f"component. Expected `{expected_conversation_type.__name__}`, got `{type(conversation).__name__}`." ) @@ -104,7 +98,7 @@ def _load_checkpointed_conversation( def _contains_missing_reference_error(error: Exception) -> bool: """Return whether the exception chain contains a missing component/tool reference.""" while error is not None: - if isinstance(error, MissingDeserializationReferenceError): + if isinstance(error, _MissingDeserializationReferenceError): return True error = error.__cause__ # type: ignore[assignment] return False @@ -118,5 +112,5 @@ def _supports_checkpointing(component: "ConversationalComponent") -> bool: return not any( isinstance(nested_component, OciAgent) - for nested_component in _iter_nested_components(component) + for nested_component in _get_nested_components(component) ) diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py index 040ca64da..fd34970fa 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py @@ -18,11 +18,11 @@ class ConversationCheckpoint: """Durable snapshot of a conversation at a checkpoint boundary.""" checkpoint_id: str - """External identifier of this saved checkpoint.""" + """ID of the checkpoint""" conversation_id: str - """Durable conversation id that this checkpoint belongs to.""" + """ID of the stored conversation""" component_id: str - """Best-effort root component id stored for diagnostics and storage queries.""" + """ID of the component that created the conversation""" created_at: int """Checkpoint creation time in seconds since the Unix epoch.""" state: str @@ -40,10 +40,13 @@ class CheckpointingInterval(Enum): # 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 diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py index 4be7fb1d7..cc2d2903a 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py @@ -81,11 +81,12 @@ def _build_listener_checkpoint_metadata( } 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["event_type"] = AgentExecutionIterationStartedEvent.__name__ checkpoint_metadata["agent_iteration"] = event.execution_state.curr_iter elif isinstance(event, FlowExecutionIterationStartedEvent): - checkpoint_metadata["event_type"] = FlowExecutionIterationStartedEvent.__name__ checkpoint_metadata["flow_step_name"] = event.execution_state.current_step_name checkpoint_metadata["nesting_level"] = event.execution_state.nesting_level if metadata: diff --git a/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py index 51a247cce..2bbebbb28 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py +++ b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py @@ -7,7 +7,6 @@ import hashlib import json import time -import warnings from textwrap import dedent from typing import Any, Dict, List, Optional, Sequence @@ -20,7 +19,6 @@ PostgresDatabaseDatastore, ) from wayflowcore.datastore._relational import RelationalDatastore -from wayflowcore.datastore.inmemory import _INMEMORY_USER_WARNING from wayflowcore.datastore.oracle import _execute_query_on_oracle_db from wayflowcore.datastore.postgres import _execute_query_on_postgres_db @@ -458,9 +456,7 @@ def __init__( checkpointing_interval: CheckpointingInterval = CheckpointingInterval.CONVERSATION_TURNS, ) -> None: resolved_storage_config = storage_config or StorageConfig() - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", message=f"{_INMEMORY_USER_WARNING}*") - datastore = InMemoryDatastore(schema=resolved_storage_config.to_schema()) + datastore = InMemoryDatastore(schema=resolved_storage_config.to_schema()) super().__init__( datastore=datastore, storage_config=resolved_storage_config, diff --git a/wayflowcore/src/wayflowcore/cli/serve.py b/wayflowcore/src/wayflowcore/cli/serve.py index 66108ff46..ee50e9ac6 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, diff --git a/wayflowcore/src/wayflowcore/conversation.py b/wayflowcore/src/wayflowcore/conversation.py index d69a667c5..594cfab11 100644 --- a/wayflowcore/src/wayflowcore/conversation.py +++ b/wayflowcore/src/wayflowcore/conversation.py @@ -142,11 +142,9 @@ def execute( The ``Execution`` status is returned by the Assistant and indicates if the assistant yielded, finished the conversation. """ - - async def _execute_async_wrapper() -> "ExecutionStatus": - return await self.execute_async(execution_interrupts) - - return run_async_in_sync(_execute_async_wrapper, method_name="execute_async") + return run_async_in_sync( + self.execute_async, execution_interrupts, method_name="execute_async" + ) async def execute_async( self, diff --git a/wayflowcore/src/wayflowcore/conversationalcomponent.py b/wayflowcore/src/wayflowcore/conversationalcomponent.py index 7ca955bee..906b85272 100644 --- a/wayflowcore/src/wayflowcore/conversationalcomponent.py +++ b/wayflowcore/src/wayflowcore/conversationalcomponent.py @@ -76,6 +76,7 @@ def __init__( __metadata_info__=__metadata_info__, ) + @abstractmethod def start_conversation( self, inputs: Optional[Dict[str, Any]] = None, @@ -84,22 +85,9 @@ def start_conversation( checkpointer: Optional["Checkpointer"] = None, checkpoint_id: Optional[str] = None, ) -> "Conversation": - """Start a fresh conversation or restore one from checkpoint storage. - - ``conversation_id`` identifies the complete conversation thread for checkpoint - storage and resume. Each concrete conversation in that thread has its own ``id``. - """ - return self._start_conversation( - inputs=inputs, - messages=messages, - conversation_id=conversation_id, - checkpointer=checkpointer, - checkpoint_id=checkpoint_id, - parent_conversation=None, - ) + pass - @abstractmethod - def _start_conversation( + def _start_conversation_impl( self, inputs: Optional[Dict[str, Any]], messages: Union[None, str, "Message", List["Message"], "MessageList"], @@ -108,7 +96,7 @@ def _start_conversation( checkpoint_id: Optional[str], parent_conversation: Optional["Conversation"] = None, ) -> "Conversation": - """Create a concrete conversation for this component.""" + """Create a concrete conversation for internal root/child entry points.""" raise NotImplementedError @property @@ -189,7 +177,7 @@ def _prepare_conversation_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 mypy; + ``expected_conversation_type`` keeps the restored conversation typed for static typing; runtime validation uses it too. """ if parent_conversation is not None: @@ -252,7 +240,7 @@ def _start_subconversation( messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, ) -> "Conversation": """Start a child that inherits its parent's conversation thread.""" - return self._start_conversation( + return self._start_conversation_impl( inputs=inputs, messages=messages, conversation_id=None, diff --git a/wayflowcore/src/wayflowcore/executors/_agentconversation.py b/wayflowcore/src/wayflowcore/executors/_agentconversation.py index 793aa8574..ccc0da58c 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_agentconversation.py @@ -65,23 +65,15 @@ def _get_all_sub_conversations(self) -> List["Conversation"]: sub_conversations += [self.state.current_flow_conversation] return sub_conversations - @staticmethod - def _sub_component_conversation_key(component: ConversationalComponent) -> str: - """Return the canonical runtime key for a subcomponent conversation slot.""" - # return f"{component.__class__.__name__}:{component.name}" - return component.id - def _get_sub_component_conversation( self, component: ConversationalComponent ) -> Optional["Conversation"]: - return self.state.current_sub_component_conversations.get( - self._sub_component_conversation_key(component) - ) + return self.state.current_sub_component_conversations.get(component.id) def _set_sub_component_conversation( self, component: ConversationalComponent, conversation: Optional["Conversation"] ) -> None: - identifier = self._sub_component_conversation_key(component) + identifier = component.id component_conversations = self.state.current_sub_component_conversations if not conversation and identifier in component_conversations: component_conversations.pop(identifier) diff --git a/wayflowcore/src/wayflowcore/executors/_flowconversation.py b/wayflowcore/src/wayflowcore/executors/_flowconversation.py index 6d5479666..1f1c43a8a 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_flowconversation.py @@ -50,8 +50,8 @@ def _gather_flow_outputs(self) -> Dict[str, Any]: def _get_internal_context_value_for_step(self, assistant_step: "Step", key: str) -> Any: from wayflowcore.executors._flowexecutor import FlowConversationExecutor - context_key = FlowConversationExecutor().make_key_for_step(assistant_step, key) - return self.state.internal_context_key_values.get(context_key) + key = FlowConversationExecutor().make_key_for_step(assistant_step, key) + return self.state.internal_context_key_values.get(key, None) def _put_internal_context_key_value(self, key: str, value: Any) -> None: self.state.internal_context_key_values[key] = value @@ -118,8 +118,7 @@ def _get_or_create_current_sub_conversation( sub_conversation_id: Optional[str] = None, ) -> "FlowConversation": sub_conversation = self._get_current_sub_conversation( - step=step, - sub_conversation_id=sub_conversation_id, + step=step, sub_conversation_id=sub_conversation_id ) if sub_conversation is None: sub_conversation = self._create_sub_conversation( @@ -159,8 +158,7 @@ def _cleanup_sub_conversation( from wayflowcore.executors._flowexecutor import FlowConversationExecutor key = FlowConversationExecutor().make_key_for_step( - step, - sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY, + step, sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY ) self.state.internal_context_key_values.pop(key, None) diff --git a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py index 2fadc2ab2..fcabd72aa 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py @@ -247,7 +247,6 @@ class FlowConversationExecutor(ConversationExecutor): @staticmethod def make_key_for_step(assistant_step: Step, key: str) -> str: return str(assistant_step.id) + "_" + key - # return str(assistant_step.name) + "_" + key @staticmethod def get_parent_conversation(state: FlowConversationExecutionState) -> Optional[Conversation]: @@ -310,14 +309,12 @@ def create_sub_conversation( def cleanup_sub_conversation( state: FlowConversationExecutionState, step: Step, - sub_conversation_id: Optional[str] = None, ) -> None: """ Remove a subconversation saved in the internal context of the flow execution state, to cleanup the state after the subconversation is finished. """ key = FlowConversationExecutor.make_key_for_step( - step, - sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY, + step, FlowConversationExecutor._SUB_CONVERSATION_KEY ) # We have to pop finished sub-conversations from the context store as other methods assume the conversation # (dict value) is not `None`. @@ -329,11 +326,10 @@ def get_current_sub_conversation( step: Step, ) -> Optional[Conversation]: """Get the current sub conversation of a given step""" - key = FlowConversationExecutor.make_key_for_step( - step, - FlowConversationExecutor._SUB_CONVERSATION_KEY, + key1 = FlowConversationExecutor.make_key_for_step( + step, FlowConversationExecutor._SUB_CONVERSATION_KEY ) - sub_conv = state.internal_context_key_values.get(key, None) + sub_conv = state.internal_context_key_values.get(key1, None) return cast(Conversation, sub_conv) if sub_conv is not None else None @staticmethod diff --git a/wayflowcore/src/wayflowcore/flow.py b/wayflowcore/src/wayflowcore/flow.py index 834a8d8b8..1a7973631 100644 --- a/wayflowcore/src/wayflowcore/flow.py +++ b/wayflowcore/src/wayflowcore/flow.py @@ -1171,7 +1171,28 @@ def start_conversation( nesting_level: int = 0, context_providers_from_parent_flow: Optional[Set[str]] = None, ) -> "FlowConversation": - return self._start_flow_conversation( + """ + Start the conversation. + + Parameters + ---------- + inputs: + Dictionary of inputs. Keys are the variable identifiers and + values are the actual inputs to start the conversation. + conversation_id: + Durable conversation id used for resume, storage, and usage accounting. + messages: + 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: + 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, @@ -1190,7 +1211,7 @@ def _start_subconversation( nesting_level: int = 0, context_providers_from_parent_flow: Optional[Set[str]] = None, ) -> "FlowConversation": - return self._start_flow_conversation( + return self._start_conversation_impl( inputs=inputs, messages=messages, conversation_id=None, @@ -1201,34 +1222,16 @@ def _start_subconversation( parent_conversation=parent_conversation, ) - def _start_conversation( - 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, - ) -> "FlowConversation": - return self._start_flow_conversation( - inputs=inputs, - messages=messages, - conversation_id=conversation_id, - checkpointer=checkpointer, - checkpoint_id=checkpoint_id, - parent_conversation=parent_conversation, - ) - - def _start_flow_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, - parent_conversation: Optional["Conversation"] = None, ) -> "FlowConversation": """ Start the conversation. diff --git a/wayflowcore/src/wayflowcore/managerworkers.py b/wayflowcore/src/wayflowcore/managerworkers.py index 0bb8dd526..8d01ed84f 100644 --- a/wayflowcore/src/wayflowcore/managerworkers.py +++ b/wayflowcore/src/wayflowcore/managerworkers.py @@ -235,7 +235,30 @@ def start_conversation( checkpoint_id: Optional[str] = None, conversation_name: Optional[str] = None, ) -> "ManagerWorkersConversation": - return self._start_managerworkers_conversation( + """ + Initializes a conversation with the managerworkers. + + Parameters + ---------- + inputs: + Dictionary of inputs. Keys are the variable identifiers and + values are the actual inputs to start the main conversation. + messages: + Message list of the manager agent 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 managerworkers. + """ + return self._start_conversation_impl( inputs=inputs, messages=messages, conversation_id=conversation_id, @@ -245,33 +268,15 @@ def start_conversation( parent_conversation=None, ) - def _start_conversation( - 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, - ) -> "ManagerWorkersConversation": - return self._start_managerworkers_conversation( - inputs=inputs, - messages=messages, - conversation_id=conversation_id, - checkpointer=checkpointer, - checkpoint_id=checkpoint_id, - parent_conversation=parent_conversation, - ) - - def _start_managerworkers_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, - conversation_name: Optional[str] = None, parent_conversation: Optional["Conversation"] = None, + conversation_name: Optional[str] = None, ) -> "ManagerWorkersConversation": """ Initializes a conversation with the managerworkers. diff --git a/wayflowcore/src/wayflowcore/ociagent.py b/wayflowcore/src/wayflowcore/ociagent.py index 4bf63b373..1f7407c84 100644 --- a/wayflowcore/src/wayflowcore/ociagent.py +++ b/wayflowcore/src/wayflowcore/ociagent.py @@ -133,7 +133,7 @@ def start_conversation( Conversation: The conversation object of the agent. """ - return self._start_conversation( + return self._start_conversation_impl( inputs=inputs, messages=messages, conversation_id=conversation_id, @@ -142,7 +142,7 @@ def start_conversation( parent_conversation=None, ) - def _start_conversation( + def _start_conversation_impl( self, inputs: Optional[Dict[str, Any]] = None, messages: Union[None, str, Message, List[Message], MessageList] = None, diff --git a/wayflowcore/src/wayflowcore/serialization/context.py b/wayflowcore/src/wayflowcore/serialization/context.py index c459a6429..455496995 100644 --- a/wayflowcore/src/wayflowcore/serialization/context.py +++ b/wayflowcore/src/wayflowcore/serialization/context.py @@ -38,11 +38,11 @@ def _create_component_type_to_plugin_mapping( return component_types_to_plugins -class MissingDeserializationReferenceError(ValueError): +class _MissingDeserializationReferenceError(ValueError): """Raised when deserialization encounters a reference missing from the root object.""" -def _iter_nested_components(value: Any) -> List["Component"]: +def _get_nested_components(value: Any) -> List["Component"]: """Return one ordered pass over all public nested components reachable from `value`.""" from wayflowcore.component import Component @@ -155,7 +155,7 @@ def _register_external_component_references(self, component: "Component") -> Non Marks the current component and all its nested components as provided externally to the serialized object graph. """ - for nested_component in _iter_nested_components(component): + for nested_component in _get_nested_components(component): self._external_references.add(self.get_reference(nested_component)) def check_obj_is_already_serialized(self, obj: Any) -> bool: @@ -278,7 +278,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 MissingDeserializationReferenceError( + raise _MissingDeserializationReferenceError( f"During deserialization, encountered reference {object_reference} that is missing " f"in the _referenced_objects of the serialized root object." ) @@ -359,6 +359,6 @@ def _register_external_component_references(self, component: "Component") -> Non Adds the current components and all its subcomponents to this deserialization context. """ - for nested_component in _iter_nested_components(component): + for nested_component in _get_nested_components(component): component_ref = SerializationContext.get_reference(nested_component) self._deserialized_objects.setdefault(component_ref, nested_component) diff --git a/wayflowcore/src/wayflowcore/swarm.py b/wayflowcore/src/wayflowcore/swarm.py index 7684b06de..a0f36c6f7 100644 --- a/wayflowcore/src/wayflowcore/swarm.py +++ b/wayflowcore/src/wayflowcore/swarm.py @@ -316,7 +316,29 @@ def start_conversation( checkpoint_id: Optional[str] = None, conversation_name: Optional[str] = None, ) -> "Conversation": - return self._start_swarm_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, @@ -326,33 +348,15 @@ def start_conversation( parent_conversation=None, ) - def _start_conversation( - 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": - return self._start_swarm_conversation( - inputs=inputs, - messages=messages, - conversation_id=conversation_id, - checkpointer=checkpointer, - checkpoint_id=checkpoint_id, - parent_conversation=parent_conversation, - ) - - def _start_swarm_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, - conversation_name: Optional[str] = None, parent_conversation: Optional["Conversation"] = None, + conversation_name: Optional[str] = None, ) -> "Conversation": """Start a fresh swarm conversation or restore one from a checkpoint. diff --git a/wayflowcore/tests/integration/test_checkpointing.py b/wayflowcore/tests/integration/test_checkpointing.py index 2ea3ffb9c..33fbe5319 100644 --- a/wayflowcore/tests/integration/test_checkpointing.py +++ b/wayflowcore/tests/integration/test_checkpointing.py @@ -4,6 +4,7 @@ # (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 functools import partial from typing import Any, Dict, Optional import pytest @@ -12,6 +13,7 @@ from wayflowcore.checkpointing import CheckpointingInterval, InMemoryCheckpointer from wayflowcore.checkpointing.checkpoint_state import _save_live_conversation_checkpoint from wayflowcore.conversation import Conversation +from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.executors._events.event import Event, EventType from wayflowcore.executors._executionstate import ConversationExecutionState from wayflowcore.executors.executionstatus import FinishedStatus, UserMessageRequestStatus @@ -24,6 +26,7 @@ from wayflowcore.flow import Flow from wayflowcore.flowbuilder import FlowBuilder from wayflowcore.managerworkers import ManagerWorkers +from wayflowcore.serialization import deserialize, serialize from wayflowcore.steps import FlowExecutionStep from wayflowcore.swarm import Swarm @@ -31,10 +34,10 @@ from ..test_managerworkers import _send_message from ..test_swarm import _handoff_message from ..testhelpers.dummy import DoNothingStep, DummyModel +from ..testhelpers.patching import patch_llm def _build_nested_flow( - *, child_first_step_name: str, child_second_step_name: str, child_flow_name: str, @@ -115,10 +118,9 @@ def _deserialize_from_dict(cls, input_dict: Dict[str, Any], deserialization_cont def _build_checkpointable_agent( - *, name: str, initial_message: str, -) -> tuple[Agent, DummyModel]: +) -> Agent: llm = DummyModel() agent = Agent( llm=llm, @@ -128,18 +130,11 @@ def _build_checkpointable_agent( initial_message=initial_message, agent_id=name, ) - return agent, llm - - -def _build_checkpointable_agent_component() -> tuple[Agent, DummyModel]: - return _build_checkpointable_agent( - name="checkpoint_agent", - initial_message="Hello from the agent.", - ) + return agent -def _build_checkpointable_swarm() -> tuple[Swarm, DummyModel]: - first_agent, first_agent_llm = _build_checkpointable_agent( +def _build_checkpointable_swarm() -> Swarm: + first_agent = _build_checkpointable_agent( name="checkpoint_swarm_first_agent", initial_message="Hello from the swarm.", ) @@ -156,11 +151,11 @@ def _build_checkpointable_swarm() -> tuple[Swarm, DummyModel]: name="checkpoint_swarm", id="checkpoint_swarm", ) - return swarm, first_agent_llm + return swarm -def _build_checkpointable_managerworkers() -> tuple[ManagerWorkers, DummyModel]: - manager_agent, manager_llm = _build_checkpointable_agent( +def _build_checkpointable_managerworkers() -> ManagerWorkers: + manager_agent = _build_checkpointable_agent( name="checkpoint_manager_agent", initial_message="Hello from the manager.", ) @@ -177,10 +172,23 @@ def _build_checkpointable_managerworkers() -> tuple[ManagerWorkers, DummyModel]: name="checkpoint_managerworkers", id="checkpoint_managerworkers", ) - return managerworkers, manager_llm + return managerworkers + +def _get_checkpoint_test_llm(component: ConversationalComponent) -> DummyModel: + if isinstance(component, Agent): + assert isinstance(component.llm, DummyModel) + return component.llm + if isinstance(component, Swarm): + assert isinstance(component.first_agent.llm, DummyModel) + return component.first_agent.llm + assert isinstance(component, ManagerWorkers) + assert isinstance(component.group_manager, Agent) + assert isinstance(component.group_manager.llm, DummyModel) + return component.group_manager.llm -def test_flow_checkpoint_restore_relinks_nested_flow_parent_for_interrupt_inheritance() -> None: + +def test_flow_checkpoint_restore_preserves_nested_interrupt_inheritance() -> None: checkpointer = InMemoryCheckpointer() def build_flow() -> Flow: @@ -201,7 +209,7 @@ def build_flow() -> Flow: ) assert checkpointer.load_latest(conversation.conversation_id) is not None - restarted_flow = build_flow() + restarted_flow = deserialize(Flow, serialize(build_flow())) restored_conversation = restarted_flow.start_conversation( conversation_id=conversation.conversation_id, checkpointer=checkpointer, @@ -256,7 +264,14 @@ def test_flow_checkpointing_supports_resume_and_time_travel() -> None: @pytest.mark.parametrize( ("builder", "conversation_id"), [ - (_build_checkpointable_agent_component, "agent-checkpoint"), + ( + partial( + _build_checkpointable_agent, + name="checkpoint_agent", + initial_message="Hello from the agent.", + ), + "agent-checkpoint", + ), (_build_checkpointable_swarm, "swarm-checkpoint"), (_build_checkpointable_managerworkers, "managerworkers-checkpoint"), ], @@ -266,7 +281,8 @@ def test_multi_agent_checkpointing_supports_resume_and_time_travel( conversation_id: str, ) -> None: checkpointer = InMemoryCheckpointer() - component, llm = builder() + component = builder() + llm = _get_checkpoint_test_llm(component) resumed_output = "Checkpoint resumed successfully." rewound_output = "Checkpoint rewound successfully." @@ -306,6 +322,38 @@ def test_multi_agent_checkpointing_supports_resume_and_time_travel( assert rewound_conversation.get_last_message().content == rewound_output +def test_agent_checkpoint_restore_after_serialization(vllm_responses_llm) -> None: + checkpointer = InMemoryCheckpointer() + agent = Agent( + llm=vllm_responses_llm, + name="serialized_checkpoint_agent", + description="Agent used for serialized checkpoint restoration.", + custom_instruction="Be helpful.", + agent_id="serialized_checkpoint_agent", + ) + + with patch_llm(vllm_responses_llm, outputs=["Initial response."]): + conversation = agent.start_conversation( + conversation_id="serialized-agent-checkpoint", + checkpointer=checkpointer, + ) + assert isinstance(conversation.execute(), UserMessageRequestStatus) + + restored_agent = deserialize(Agent, serialize(agent)) + restored_llm = restored_agent.llm + restored_conversation = restored_agent.start_conversation( + conversation_id=conversation.conversation_id, + checkpointer=checkpointer, + ) + + with patch_llm(restored_llm, outputs=["Resumed response."]): + restored_conversation.append_user_message("Continue.") + restored_status = restored_conversation.execute() + + assert isinstance(restored_status, UserMessageRequestStatus) + assert restored_conversation.get_last_message().content == "Resumed response." + + def test_swarm_checkpoint_restore_uses_generated_agent_ids_for_threads() -> None: checkpointer = InMemoryCheckpointer( checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS diff --git a/wayflowcore/tests/test_checkpointing.py b/wayflowcore/tests/test_checkpointing.py index 42e5cf583..7b680106b 100644 --- a/wayflowcore/tests/test_checkpointing.py +++ b/wayflowcore/tests/test_checkpointing.py @@ -25,6 +25,10 @@ from .test_managerworkers import _send_message from .testhelpers.dummy import DummyModel +pytestmark = pytest.mark.filterwarnings( + "ignore:InMemoryDatastore is for DEVELOPMENT and PROOF-OF-CONCEPT ONLY!" +) + def test_root_conversation_uses_one_identity_for_instance_and_thread() -> None: agent = Agent(llm=DummyModel(), name="root_identity_agent") @@ -430,6 +434,42 @@ def test_checkpoint_intervals_save_expected_checkpoints( assert checkpoints[-1].metadata["status_type"] == "FinishedStatus" +@pytest.mark.parametrize( + "interval", + list(CheckpointingInterval), + ids=lambda interval: interval.name.lower(), +) +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: From 74b6028c6266af37f543ac8794688b3ad9582ab5 Mon Sep 17 00:00:00 2001 From: Son Le Date: Thu, 23 Jul 2026 14:46:42 +0200 Subject: [PATCH 08/10] Refactor checkpoint snapshot and restoration handling - Move live conversation checkpoint creation into Checkpointer.save(). - Support saving either live conversations or materialized checkpoints, including custom checkpoint IDs, component IDs, and metadata. - Move checkpoint restoration and compatibility validation into ConversationalComponent. - Add explicit checkpoint support detection across conversational components, while marking OCI agents as unsupported. - Preserve root component references during serialization and deserialization for nested graph restoration. - Rename nested flow conversation startup to clarify state preservation. - Update checkpoint listeners and OpenAI Responses persistence to use the unified Checkpointer API. - Remove the legacy checkpoint_state module. - Add unit coverage for saving live conversations with custom metadata. --- wayflowcore/src/wayflowcore/a2a/a2aagent.py | 6 +- wayflowcore/src/wayflowcore/agent.py | 9 ++ .../services/wayflowservice.py | 9 +- .../src/wayflowcore/checkpointing/__init__.py | 9 +- .../checkpointing/checkpoint_state.py | 116 ----------------- .../wayflowcore/checkpointing/checkpointer.py | 88 ++++++++++++- .../checkpointing/checkpointeventlistener.py | 7 +- .../checkpointing/datastorecheckpointer.py | 2 +- .../wayflowcore/conversationalcomponent.py | 111 ++++++++++------ .../executors/_agentconversation.py | 6 +- .../wayflowcore/executors/_flowexecutor.py | 2 +- wayflowcore/src/wayflowcore/flow.py | 12 +- wayflowcore/src/wayflowcore/managerworkers.py | 9 ++ wayflowcore/src/wayflowcore/ociagent.py | 4 + .../src/wayflowcore/serialization/context.py | 118 ++++++++++++------ wayflowcore/src/wayflowcore/swarm.py | 9 ++ .../tests/integration/test_checkpointing.py | 3 +- wayflowcore/tests/test_checkpointing.py | 20 +++ 18 files changed, 316 insertions(+), 224 deletions(-) delete mode 100644 wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py diff --git a/wayflowcore/src/wayflowcore/a2a/a2aagent.py b/wayflowcore/src/wayflowcore/a2a/a2aagent.py index 4e69fc7fd..d7591a434 100644 --- a/wayflowcore/src/wayflowcore/a2a/a2aagent.py +++ b/wayflowcore/src/wayflowcore/a2a/a2aagent.py @@ -246,6 +246,10 @@ def __init__( __metadata_info__=__metadata_info__, ) + @property + def _supports_checkpointing(self) -> bool: + return True + def start_conversation( self, inputs: Optional[Dict[str, Any]] = None, @@ -269,7 +273,7 @@ def start_conversation( 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 (of the parent conversation) used for resume, storage, and usage accounting. + Conversation id used for checkpointing the conversation and later resuming it. checkpointer: Optional checkpoint backend used to restore and persist this conversation. checkpoint_id: diff --git a/wayflowcore/src/wayflowcore/agent.py b/wayflowcore/src/wayflowcore/agent.py index dd705abe7..f1a93fc5e 100644 --- a/wayflowcore/src/wayflowcore/agent.py +++ b/wayflowcore/src/wayflowcore/agent.py @@ -362,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 diff --git a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py index b834c1c20..34ef54b0f 100644 --- a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py +++ b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py @@ -20,10 +20,6 @@ ConversationCheckpoint, DatastoreCheckpointer, ) -from wayflowcore.checkpointing.checkpoint_state import ( - _save_live_conversation_checkpoint, - _supports_checkpointing, -) from wayflowcore.conversation import Conversation from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.datastore import Datastore, InMemoryDatastore @@ -210,7 +206,7 @@ async def create_response(self, body: CreateResponse) -> AsyncIterable[ResponseS # 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 _supports_checkpointing(agent): + 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, @@ -344,8 +340,7 @@ async def runner(conversation: Conversation) -> None: # persists a completed OpenAI Responses response as a WayFlow checkpoint if should_store_response: - _save_live_conversation_checkpoint( - self.checkpointer, + 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. diff --git a/wayflowcore/src/wayflowcore/checkpointing/__init__.py b/wayflowcore/src/wayflowcore/checkpointing/__init__.py index b6cdd29af..023a2d960 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/__init__.py +++ b/wayflowcore/src/wayflowcore/checkpointing/__init__.py @@ -4,8 +4,13 @@ # (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 .checkpoint_state import CheckpointRestoreCompatibilityError -from .checkpointer import Checkpointer, CheckpointingInterval, ConversationCheckpoint, StorageConfig +from .checkpointer import ( + Checkpointer, + CheckpointingInterval, + CheckpointRestoreCompatibilityError, + ConversationCheckpoint, + StorageConfig, +) from .datastorecheckpointer import ( DatastoreCheckpointer, InMemoryCheckpointer, diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py b/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py deleted file mode 100644 index f1fcacefc..000000000 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpoint_state.py +++ /dev/null @@ -1,116 +0,0 @@ -# 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 time -from typing import TYPE_CHECKING, Any, Dict, Optional, Type - -from wayflowcore.exceptions import DataclassFieldDeserializationError -from wayflowcore.idgeneration import IdGenerator -from wayflowcore.serialization import autodeserialize, serialize -from wayflowcore.serialization.context import ( - DeserializationContext, - SerializationContext, - _get_nested_components, - _MissingDeserializationReferenceError, -) - -if TYPE_CHECKING: - from wayflowcore.checkpointing import Checkpointer - from wayflowcore.checkpointing.checkpointer import ConversationCheckpoint - from wayflowcore.conversation import Conversation - from wayflowcore.conversationalcomponent import ConversationalComponent - - -class CheckpointRestoreCompatibilityError(ValueError): - """Raised when a checkpoint cannot be resumed against the current live graph.""" - - -def _save_live_conversation_checkpoint( - checkpointer: "Checkpointer", - conversation: "Conversation", - checkpoint_id: Optional[str] = None, - component_id: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, -) -> "ConversationCheckpoint": - from wayflowcore.checkpointing.checkpointer import ConversationCheckpoint - - if not _supports_checkpointing(conversation.component): - 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 {}), - ) - checkpointer.save(checkpoint) - conversation.checkpoint_id = checkpoint.checkpoint_id - return checkpoint - - -def _load_checkpointed_conversation( - checkpoint: "ConversationCheckpoint", - component: "ConversationalComponent", - expected_conversation_type: Type["Conversation"], - tool_registry: Optional[Dict[str, Any]] = None, - checkpointer: Optional["Checkpointer"] = None, - attach_checkpointer: bool = True, -) -> "Conversation": - """Restore a Conversation object from stored checkpoint state. - - The serialized conversation omits live component/tool objects and rebuilds them from - the current component tree. - """ - deserialization_context = DeserializationContext() - deserialization_context.registered_tools = tool_registry.copy() if tool_registry else {} - deserialization_context._register_external_component_references(component) - try: - conversation = autodeserialize( - checkpoint.state, - deserialization_context=deserialization_context, - ) - except (_MissingDeserializationReferenceError, DataclassFieldDeserializationError) as exc: - if not _contains_missing_reference_error(exc): - 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 - - if attach_checkpointer: - conversation.checkpointer = checkpointer - return conversation - - -def _contains_missing_reference_error(error: Exception) -> bool: - """Return whether the exception chain contains a missing component/tool reference.""" - while error is not None: - if isinstance(error, _MissingDeserializationReferenceError): - return True - error = error.__cause__ # type: ignore[assignment] - return False - - -# Checkpoint eligibility - - -def _supports_checkpointing(component: "ConversationalComponent") -> bool: - from wayflowcore.ociagent import OciAgent - - return not any( - isinstance(nested_component, OciAgent) - for nested_component in _get_nested_components(component) - ) diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py index fd34970fa..2f648ad3b 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointer.py @@ -7,9 +7,15 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, List, Optional +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 @@ -116,14 +122,88 @@ def load_latest(self, conversation_id: str) -> Optional[ConversationCheckpoint]: 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(self, checkpoint: ConversationCheckpoint) -> None: + def _save_checkpoint(self, checkpoint: ConversationCheckpoint) -> None: + """Persist an already-materialized checkpoint in the backend.""" raise NotImplementedError() - async def save_async(self, checkpoint: ConversationCheckpoint) -> None: + 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. - self.save(checkpoint) + return self.save( + checkpoint, + checkpoint_id=checkpoint_id, + component_id=component_id, + metadata=metadata, + ) @abstractmethod def list_checkpoints( diff --git a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py index cc2d2903a..dc7ddeb43 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py +++ b/wayflowcore/src/wayflowcore/checkpointing/checkpointeventlistener.py @@ -13,7 +13,6 @@ FlowExecutionIterationStartedEvent, LlmGenerationResponseEvent, ) -from .checkpoint_state import _save_live_conversation_checkpoint from .checkpointer import Checkpointer, CheckpointingInterval if TYPE_CHECKING: @@ -159,8 +158,7 @@ def save_pending_llm_checkpoint(self) -> None: self._save_internal_turn_checkpoint(event) def _save_internal_turn_checkpoint(self, event: _IterationStartedEvent) -> None: - _save_live_conversation_checkpoint( - self.checkpointer, + self.checkpointer.save( self.conversation, metadata=_build_listener_checkpoint_metadata( self.conversation, @@ -196,8 +194,7 @@ def get_conversation_checkpoint_execution_context( with register_event_listeners([listener]): yield listener.save_pending_llm_checkpoint() - _save_live_conversation_checkpoint( - checkpointer, + checkpointer.save( conversation, metadata=_build_listener_checkpoint_metadata( conversation, diff --git a/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py index 2bbebbb28..71842ffa0 100644 --- a/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py +++ b/wayflowcore/src/wayflowcore/checkpointing/datastorecheckpointer.py @@ -249,7 +249,7 @@ def load(self, conversation_id: str, checkpoint_id: str) -> ConversationCheckpoi ) return checkpoint - def save(self, checkpoint: ConversationCheckpoint) -> None: + 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 diff --git a/wayflowcore/src/wayflowcore/conversationalcomponent.py b/wayflowcore/src/wayflowcore/conversationalcomponent.py index 906b85272..549593279 100644 --- a/wayflowcore/src/wayflowcore/conversationalcomponent.py +++ b/wayflowcore/src/wayflowcore/conversationalcomponent.py @@ -7,19 +7,7 @@ import logging from abc import ABC, abstractmethod from pathlib import Path -from typing import ( - TYPE_CHECKING, - Any, - Dict, - Generic, - List, - Optional, - Set, - Type, - TypeVar, - Union, - cast, -) +from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, Set, Type, TypeVar, Union from wayflowcore._metadata import MetadataType from wayflowcore.componentwithio import ComponentWithInputsOutputs @@ -32,6 +20,7 @@ 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 MessageList @@ -103,6 +92,12 @@ def _start_conversation_impl( 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 @@ -183,31 +178,28 @@ def _prepare_conversation_start( 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 - # No checkpointer means this is a fresh root conversation; just resolve ids. - if checkpointer is None: - if checkpoint_id is not None: - raise ValueError("`checkpoint_id` requires a `checkpointer`.") + 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() - return None, conversation_id, 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() - # Checkpoint restore uses the thread id to locate stored state. - resolved_conversation_id = conversation_id - if resolved_conversation_id is None and checkpoint_id is not None: - raise ValueError("`checkpoint_id` requires a `conversation_id`.") - if resolved_conversation_id is None: - resolved_conversation_id = IdGenerator.get_or_generate_id() + if checkpointer is None: + return None, conversation_id, conversation_id checkpoint = ( - checkpointer.load(resolved_conversation_id, checkpoint_id) + checkpointer.load(conversation_id, checkpoint_id) if checkpoint_id is not None - else checkpointer.load_latest(resolved_conversation_id) + else checkpointer.load_latest(conversation_id) ) if checkpoint is None: - return None, resolved_conversation_id, resolved_conversation_id + 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): @@ -216,23 +208,60 @@ def _prepare_conversation_start( "Load the conversation first, then append new user input explicitly." ) - from wayflowcore.checkpointing.checkpoint_state import _load_checkpointed_conversation - - # Rehydrate the stored conversation against this live component tree. - conversation = _load_checkpointed_conversation( - checkpoint, - component=self, + conversation = self._restore_checkpointed_conversation( + checkpoint=checkpoint, expected_conversation_type=expected_conversation_type, - tool_registry={tool.name: tool for tool in self._referenced_tools()}, checkpointer=checkpointer, - attach_checkpointer=True, ) - return ( - cast(ConversationTypeT, conversation), - resolved_conversation_id, - resolved_conversation_id, + return conversation, conversation_id, conversation_id + + def _restore_checkpointed_conversation( + self, + checkpoint: "ConversationCheckpoint", + expected_conversation_type: Type[ConversationTypeT], + checkpointer: "Checkpointer", + ) -> ConversationTypeT: + """Rehydrate a checkpoint against this component's live graph and tools.""" + 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 = checkpointer + return conversation + def _start_subconversation( self, parent_conversation: "Conversation", diff --git a/wayflowcore/src/wayflowcore/executors/_agentconversation.py b/wayflowcore/src/wayflowcore/executors/_agentconversation.py index ccc0da58c..b8e976710 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_agentconversation.py @@ -91,16 +91,14 @@ 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 fcabd72aa..222f54265 100644 --- a/wayflowcore/src/wayflowcore/executors/_flowexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_flowexecutor.py @@ -289,7 +289,7 @@ def create_sub_conversation( resolved_sub_conversation_id = ( sub_conversation_id or FlowConversationExecutor._SUB_CONVERSATION_KEY ) - sub_conversation = flow._start_subconversation( + sub_conversation = flow._start_nested_flow_conversation( parent_conversation=conversation, inputs=inputs_not_from_context_providers, messages=conversation.message_list, diff --git a/wayflowcore/src/wayflowcore/flow.py b/wayflowcore/src/wayflowcore/flow.py index 1a7973631..8516348ee 100644 --- a/wayflowcore/src/wayflowcore/flow.py +++ b/wayflowcore/src/wayflowcore/flow.py @@ -468,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, @@ -1203,7 +1212,7 @@ def start_conversation( parent_conversation=None, ) - def _start_subconversation( + def _start_nested_flow_conversation( self, parent_conversation: "Conversation", inputs: Optional[Dict[str, Any]] = None, @@ -1211,6 +1220,7 @@ def _start_subconversation( 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, diff --git a/wayflowcore/src/wayflowcore/managerworkers.py b/wayflowcore/src/wayflowcore/managerworkers.py index 8d01ed84f..e6934b832 100644 --- a/wayflowcore/src/wayflowcore/managerworkers.py +++ b/wayflowcore/src/wayflowcore/managerworkers.py @@ -44,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], diff --git a/wayflowcore/src/wayflowcore/ociagent.py b/wayflowcore/src/wayflowcore/ociagent.py index 1f7407c84..0c6540c06 100644 --- a/wayflowcore/src/wayflowcore/ociagent.py +++ b/wayflowcore/src/wayflowcore/ociagent.py @@ -103,6 +103,10 @@ def __init__( __metadata_info__=__metadata_info__, ) + @property + def _supports_checkpointing(self) -> bool: + return False + def start_conversation( self, inputs: Optional[Dict[str, Any]] = None, diff --git a/wayflowcore/src/wayflowcore/serialization/context.py b/wayflowcore/src/wayflowcore/serialization/context.py index 455496995..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, @@ -38,42 +39,6 @@ def _create_component_type_to_plugin_mapping( return component_types_to_plugins -class _MissingDeserializationReferenceError(ValueError): - """Raised when deserialization encounters a reference missing from the root object.""" - - -def _get_nested_components(value: Any) -> List["Component"]: - """Return one ordered pass over all public nested components reachable from `value`.""" - from wayflowcore.component import Component - - 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) - ordered_components.append(current_value) - for name, attr in vars(current_value).items(): - if not name.startswith("_"): - _collect_nested_components(attr) - return - - if isinstance(current_value, dict): - for nested_value in current_value.values(): - _collect_nested_components(nested_value) - return - - 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 - - class SerializationContext: def __init__(self, root: Any, plugins: Optional[List["WayflowSerializationPlugin"]] = None): @@ -155,7 +120,7 @@ def _register_external_component_references(self, component: "Component") -> Non 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): + 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: @@ -359,6 +324,81 @@ def _register_external_component_references(self, component: "Component") -> Non Adds the current components and all its subcomponents to this deserialization context. """ - for nested_component in _get_nested_components(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_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 isinstance(current_value, dict): + for nested_value in current_value.values(): + _collect_nested_components(nested_value) + return + + 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/swarm.py b/wayflowcore/src/wayflowcore/swarm.py index a0f36c6f7..ca6834e6c 100644 --- a/wayflowcore/src/wayflowcore/swarm.py +++ b/wayflowcore/src/wayflowcore/swarm.py @@ -93,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, diff --git a/wayflowcore/tests/integration/test_checkpointing.py b/wayflowcore/tests/integration/test_checkpointing.py index 33fbe5319..02dc694f9 100644 --- a/wayflowcore/tests/integration/test_checkpointing.py +++ b/wayflowcore/tests/integration/test_checkpointing.py @@ -11,7 +11,6 @@ from wayflowcore.agent import Agent from wayflowcore.checkpointing import CheckpointingInterval, InMemoryCheckpointer -from wayflowcore.checkpointing.checkpoint_state import _save_live_conversation_checkpoint from wayflowcore.conversation import Conversation from wayflowcore.conversationalcomponent import ConversationalComponent from wayflowcore.executors._events.event import Event, EventType @@ -439,7 +438,7 @@ def build_agent() -> tuple[Agent, Flow]: interrupt_step_name="agent_child_first_step", ) conversation.state.current_flow_conversation = parent_flow_conversation - _save_live_conversation_checkpoint(checkpointer, conversation) + checkpointer.save(conversation) restarted_agent, restarted_parent_flow = build_agent() restored_conversation = restarted_agent.start_conversation( diff --git a/wayflowcore/tests/test_checkpointing.py b/wayflowcore/tests/test_checkpointing.py index 7b680106b..5a6780f96 100644 --- a/wayflowcore/tests/test_checkpointing.py +++ b/wayflowcore/tests/test_checkpointing.py @@ -209,6 +209,26 @@ def test_inmemory_checkpointer_can_save_load_list_and_delete_checkpoints() -> No ] == [first_checkpoint_id] +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" + + 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.")) From e829e441fe055801f11f73e5bdc0ecd2e27c3e3a Mon Sep 17 00:00:00 2001 From: Son Le Date: Thu, 23 Jul 2026 16:48:53 +0200 Subject: [PATCH 09/10] refactor tests --- wayflowcore/src/wayflowcore/a2a/a2aagent.py | 1 - wayflowcore/src/wayflowcore/agent.py | 1 - .../services/wayflowservice.py | 15 +- .../contextproviders/flowcontextprovider.py | 2 +- .../wayflowcore/conversationalcomponent.py | 31 +- .../wayflowcore/executors/_agentexecutor.py | 4 +- .../executors/_managerworkersconversation.py | 7 +- .../executors/_swarmconversation.py | 13 +- wayflowcore/src/wayflowcore/flow.py | 27 -- wayflowcore/src/wayflowcore/managerworkers.py | 25 +- wayflowcore/src/wayflowcore/ociagent.py | 1 - .../wayflowcore/steps/agentexecutionstep.py | 2 +- wayflowcore/src/wayflowcore/swarm.py | 4 - .../src/wayflowcore/tools/servertools.py | 2 +- wayflowcore/tests/agentserver/conftest.py | 16 +- wayflowcore/tests/conftest.py | 24 ++ wayflowcore/tests/datastores/conftest.py | 30 +- .../tests/integration/test_checkpointing.py | 404 +++++++++--------- wayflowcore/tests/test_checkpointing.py | 364 ++++------------ 19 files changed, 346 insertions(+), 627 deletions(-) diff --git a/wayflowcore/src/wayflowcore/a2a/a2aagent.py b/wayflowcore/src/wayflowcore/a2a/a2aagent.py index d7591a434..2c7f406a9 100644 --- a/wayflowcore/src/wayflowcore/a2a/a2aagent.py +++ b/wayflowcore/src/wayflowcore/a2a/a2aagent.py @@ -302,7 +302,6 @@ def _start_conversation_impl( checkpoint_id: Optional[str] = None, parent_conversation: Optional["Conversation"] = None, ) -> "A2AAgentConversation": - """Create the concrete A2A conversation, including nested conversations.""" from wayflowcore.executors._a2aagentconversation import A2AAgentConversation from wayflowcore.executors._a2aagentexecutor import A2AAgentState diff --git a/wayflowcore/src/wayflowcore/agent.py b/wayflowcore/src/wayflowcore/agent.py index f1a93fc5e..9226263fc 100644 --- a/wayflowcore/src/wayflowcore/agent.py +++ b/wayflowcore/src/wayflowcore/agent.py @@ -448,7 +448,6 @@ def _start_conversation_impl( checkpoint_id: Optional[str] = None, parent_conversation: Optional["Conversation"] = None, ) -> "AgentConversation": - """Create the concrete agent conversation, including nested conversations.""" from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event from wayflowcore.executors._agentconversation import AgentConversation diff --git a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py index 34ef54b0f..e9ff47c84 100644 --- a/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py +++ b/wayflowcore/src/wayflowcore/agentserver/openairesponses/services/wayflowservice.py @@ -448,16 +448,13 @@ def _load_state( self._cache_response_conversation_id(checkpoint.checkpoint_id, checkpoint.conversation_id) agent = self.agents[agent_id] try: - # Use start_conversation() only to restore the checkpoint identified by - # (conversation_id, checkpoint_id), then clear conversation.checkpointer - # so execute() resumes without adding automatic checkpoint saves. - conversation = agent.start_conversation( - conversation_id=checkpoint.conversation_id, - checkpointer=self.checkpointer, - checkpoint_id=checkpoint.checkpoint_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, ) - conversation.checkpointer = None - return conversation except CheckpointRestoreCompatibilityError as e: raise HTTPException( status_code=http_status_code.HTTP_400_BAD_REQUEST, diff --git a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py index 2f7d8703e..d4dbcc4db 100644 --- a/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py +++ b/wayflowcore/src/wayflowcore/contextproviders/flowcontextprovider.py @@ -92,7 +92,7 @@ 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_subconversation( + conversation = self.flow._start_conversation_impl( parent_conversation=conversation, inputs={}, messages=conversation.message_list, diff --git a/wayflowcore/src/wayflowcore/conversationalcomponent.py b/wayflowcore/src/wayflowcore/conversationalcomponent.py index 549593279..8fce88b65 100644 --- a/wayflowcore/src/wayflowcore/conversationalcomponent.py +++ b/wayflowcore/src/wayflowcore/conversationalcomponent.py @@ -39,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: @@ -85,7 +85,6 @@ def _start_conversation_impl( checkpoint_id: Optional[str], parent_conversation: Optional["Conversation"] = None, ) -> "Conversation": - """Create a concrete conversation for internal root/child entry points.""" raise NotImplementedError @property @@ -211,7 +210,7 @@ def _prepare_conversation_start( conversation = self._restore_checkpointed_conversation( checkpoint=checkpoint, expected_conversation_type=expected_conversation_type, - checkpointer=checkpointer, + attached_checkpointer=checkpointer, ) return conversation, conversation_id, conversation_id @@ -219,9 +218,13 @@ def _restore_checkpointed_conversation( self, checkpoint: "ConversationCheckpoint", expected_conversation_type: Type[ConversationTypeT], - checkpointer: "Checkpointer", + attached_checkpointer: Optional["Checkpointer"], ) -> ConversationTypeT: - """Rehydrate a checkpoint against this component's live graph and tools.""" + """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 @@ -259,25 +262,9 @@ def _restore_checkpointed_conversation( f"component. Expected `{expected_conversation_type.__name__}`, got `{type(conversation).__name__}`." ) conversation.checkpoint_id = checkpoint.checkpoint_id - conversation.checkpointer = checkpointer + conversation.checkpointer = attached_checkpointer return conversation - def _start_subconversation( - self, - parent_conversation: "Conversation", - inputs: Optional[Dict[str, Any]] = None, - messages: Union[None, str, "Message", List["Message"], "MessageList"] = None, - ) -> "Conversation": - """Start a child that inherits its parent's conversation thread.""" - return self._start_conversation_impl( - inputs=inputs, - messages=messages, - conversation_id=None, - checkpointer=None, - checkpoint_id=None, - parent_conversation=parent_conversation, - ) - # Define a TypeVar that represents the component's type ConversationalComponentTypeT = TypeVar( diff --git a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py index 6f7d1f583..49de04ff9 100644 --- a/wayflowcore/src/wayflowcore/executors/_agentexecutor.py +++ b/wayflowcore/src/wayflowcore/executors/_agentexecutor.py @@ -469,7 +469,7 @@ 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_subconversation( + sub_agent_conversation = expert_agent._start_conversation_impl( parent_conversation=caller_conv, messages=init_messages, inputs=inputs, @@ -539,7 +539,7 @@ async def _execute_flow( outputs: Any = None try: if state.current_flow_conversation is None: - state.current_flow_conversation = flow._start_subconversation( + state.current_flow_conversation = flow._start_conversation_impl( parent_conversation=parent_conversation, inputs=inputs, messages=messages, diff --git a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py index fc5db4bd4..d7dff8520 100644 --- a/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_managerworkersconversation.py @@ -6,7 +6,7 @@ import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, Optional, Union, cast +from typing import TYPE_CHECKING, Dict, List, Optional, Union from wayflowcore.agent import Agent from wayflowcore.conversation import Conversation @@ -32,10 +32,7 @@ def _create_subconversation_for_agent( agent: Union[Agent, ManagerWorkers], parent_conversation: "ManagerWorkersConversation", ) -> Union["AgentConversation", "ManagerWorkersConversation"]: - subconv = cast( - "Union[AgentConversation, ManagerWorkersConversation]", - agent._start_subconversation(parent_conversation=parent_conversation), - ) + subconv = agent._start_conversation_impl(parent_conversation=parent_conversation) self.subconversations[agent.id] = subconv return subconv diff --git a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py index f6fbac334..fc44ac75d 100644 --- a/wayflowcore/src/wayflowcore/executors/_swarmconversation.py +++ b/wayflowcore/src/wayflowcore/executors/_swarmconversation.py @@ -6,7 +6,7 @@ import logging from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from wayflowcore.agent import Agent from wayflowcore.conversation import Conversation @@ -82,13 +82,10 @@ def _create_subconversation_for_thread( if isinstance(message_list, list) else message_list ) - conversation = cast( - "AgentConversation", - thread.recipient_agent._start_subconversation( - parent_conversation=parent_conversation, - inputs=inputs, - messages=thread.message_list, - ), + conversation = thread.recipient_agent._start_conversation_impl( + parent_conversation=parent_conversation, + inputs=inputs, + messages=thread.message_list, ) self.thread_subconversations[thread_id] = conversation diff --git a/wayflowcore/src/wayflowcore/flow.py b/wayflowcore/src/wayflowcore/flow.py index 8516348ee..9989ad4d9 100644 --- a/wayflowcore/src/wayflowcore/flow.py +++ b/wayflowcore/src/wayflowcore/flow.py @@ -1243,33 +1243,6 @@ def _start_conversation_impl( nesting_level: int = 0, context_providers_from_parent_flow: Optional[Set[str]] = None, ) -> "FlowConversation": - """ - Start the conversation. - - Parameters - ---------- - inputs: - Dictionary of inputs. Keys are the variable identifiers and - values are the actual inputs to start the conversation. - conversation_id: - Durable conversation id used for resume, storage, and usage accounting. - messages: - List of messages (``MessageList`` object) before starting the conversation. - checkpointer: - Optional checkpoint backend used to restore and persist this conversation. - checkpoint_id: - Optional checkpoint identifier to restore. Requires both ``checkpointer`` and - ``conversation_id``. - context_providers_from_parent_flow: - Context provider that don't need to be checked when validating existing inputs. - nesting_level: - Nesting level of the conversation. - - Returns - ------- - Conversation: - A Flow Conversation object. - """ from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event from wayflowcore.executors._flowconversation import FlowConversation diff --git a/wayflowcore/src/wayflowcore/managerworkers.py b/wayflowcore/src/wayflowcore/managerworkers.py index e6934b832..051817583 100644 --- a/wayflowcore/src/wayflowcore/managerworkers.py +++ b/wayflowcore/src/wayflowcore/managerworkers.py @@ -287,29 +287,6 @@ def _start_conversation_impl( parent_conversation: Optional["Conversation"] = None, conversation_name: Optional[str] = None, ) -> "ManagerWorkersConversation": - """ - Initializes a conversation with the managerworkers. - - Parameters - ---------- - inputs: - Dictionary of inputs. Keys are the variable identifiers and - values are the actual inputs to start the main conversation. - messages: - Message list of the manager agent 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 managerworkers. - """ from wayflowcore.agentconversation import AgentConversation from wayflowcore.events.event import ConversationCreatedEvent from wayflowcore.events.eventlistener import record_event @@ -366,7 +343,7 @@ def _start_conversation_impl( ) subconversations[self.manager_agent.id] = cast( "Union[AgentConversation, ManagerWorkersConversation]", - self.manager_agent._start_subconversation( + self.manager_agent._start_conversation_impl( parent_conversation=conversation, inputs=inputs, messages=messages, diff --git a/wayflowcore/src/wayflowcore/ociagent.py b/wayflowcore/src/wayflowcore/ociagent.py index 0c6540c06..2423059e7 100644 --- a/wayflowcore/src/wayflowcore/ociagent.py +++ b/wayflowcore/src/wayflowcore/ociagent.py @@ -155,7 +155,6 @@ def _start_conversation_impl( checkpoint_id: Optional[str] = None, parent_conversation: Optional["Conversation"] = None, ) -> "Conversation": - """Create the concrete OCI conversation, including nested conversations.""" from wayflowcore.executors._ociagentconversation import OciAgentConversation from wayflowcore.executors._ociagentexecutor import ( OciAgentState, diff --git a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py index 662d70561..c0e69a776 100644 --- a/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py +++ b/wayflowcore/src/wayflowcore/steps/agentexecutionstep.py @@ -293,7 +293,7 @@ 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_subconversation( + agent_sub_conversation = self.agent._start_conversation_impl( parent_conversation=caller_conv, inputs=inputs, messages=init_messages, diff --git a/wayflowcore/src/wayflowcore/swarm.py b/wayflowcore/src/wayflowcore/swarm.py index ca6834e6c..4e150ae69 100644 --- a/wayflowcore/src/wayflowcore/swarm.py +++ b/wayflowcore/src/wayflowcore/swarm.py @@ -367,10 +367,6 @@ def _start_conversation_impl( parent_conversation: Optional["Conversation"] = None, conversation_name: Optional[str] = None, ) -> "Conversation": - """Start a fresh swarm conversation or restore one from a checkpoint. - - ``conversation_id`` identifies the complete conversation thread. - """ from wayflowcore.executors._swarmconversation import ( SwarmConversation, SwarmConversationExecutionState, diff --git a/wayflowcore/src/wayflowcore/tools/servertools.py b/wayflowcore/src/wayflowcore/tools/servertools.py index 6a5c09827..0b5863602 100644 --- a/wayflowcore/src/wayflowcore/tools/servertools.py +++ b/wayflowcore/src/wayflowcore/tools/servertools.py @@ -636,7 +636,7 @@ async def __call__(self, **inputs: Any) -> Any: conversation = self.flow.start_conversation(inputs) interrupts = [] else: - conversation = self.flow._start_subconversation( + conversation = self.flow._start_conversation_impl( parent_conversation=self._parent_conversation, inputs=inputs, messages=self._parent_conversation.message_list, 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/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 index 02dc694f9..c92f6cf77 100644 --- a/wayflowcore/tests/integration/test_checkpointing.py +++ b/wayflowcore/tests/integration/test_checkpointing.py @@ -4,15 +4,26 @@ # (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 functools import partial from typing import Any, Dict, Optional +from uuid import uuid4 import pytest from wayflowcore.agent import Agent -from wayflowcore.checkpointing import CheckpointingInterval, InMemoryCheckpointer +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.conversationalcomponent import ConversationalComponent +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 @@ -25,10 +36,13 @@ 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 @@ -36,6 +50,32 @@ 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, @@ -58,24 +98,6 @@ def _build_nested_flow( ) -def _start_interrupted_flow_conversation( - flow: Flow, - *, - conversation_id: str, - interrupt_step_name: str, - checkpointer: InMemoryCheckpointer | None = None, -): - conversation = flow.start_conversation( - conversation_id=conversation_id, - checkpointer=checkpointer, - ) - status = conversation.execute( - execution_interrupts=[_OnStepStartExecutionInterrupt(interrupt_step_name)] - ) - assert isinstance(status, InterruptedExecutionStatus) - return conversation - - class _OnStepStartExecutionInterrupt( _AllEventsInterruptMixin, FlexibleExecutionInterrupt, FlowExecutionInterrupt ): @@ -119,8 +141,9 @@ def _deserialize_from_dict(cls, input_dict: Dict[str, Any], deserialization_cont def _build_checkpointable_agent( name: str, initial_message: str, + llm: Optional[LlmModel] = None, ) -> Agent: - llm = DummyModel() + llm = llm if llm is not None else DummyModel() agent = Agent( llm=llm, name=name, @@ -132,13 +155,14 @@ def _build_checkpointable_agent( return agent -def _build_checkpointable_swarm() -> Swarm: +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=DummyModel(fails_if_not_set=False), + llm=llm, name="checkpoint_swarm_second_agent", description="Swarm helper", custom_instruction="Help with delegated tasks.", @@ -153,13 +177,14 @@ def _build_checkpointable_swarm() -> Swarm: return swarm -def _build_checkpointable_managerworkers() -> ManagerWorkers: +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=DummyModel(fails_if_not_set=False), + llm=llm, name="checkpoint_worker_agent", description="Worker agent", custom_instruction="Help the manager.", @@ -174,22 +199,9 @@ def _build_checkpointable_managerworkers() -> ManagerWorkers: return managerworkers -def _get_checkpoint_test_llm(component: ConversationalComponent) -> DummyModel: - if isinstance(component, Agent): - assert isinstance(component.llm, DummyModel) - return component.llm - if isinstance(component, Swarm): - assert isinstance(component.first_agent.llm, DummyModel) - return component.first_agent.llm - assert isinstance(component, ManagerWorkers) - assert isinstance(component.group_manager, Agent) - assert isinstance(component.group_manager.llm, DummyModel) - return component.group_manager.llm - - -def test_flow_checkpoint_restore_preserves_nested_interrupt_inheritance() -> None: - checkpointer = InMemoryCheckpointer() - +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", @@ -200,19 +212,27 @@ def build_flow() -> Flow: ) original_flow = build_flow() - conversation = _start_interrupted_flow_conversation( - original_flow, + conversation = original_flow.start_conversation( conversation_id="flow-parent-link-restore", - interrupt_step_name="child_first_step", - checkpointer=checkpointer, + checkpointer=integration_checkpointer, ) - assert checkpointer.load_latest(conversation.conversation_id) is not None + 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=checkpointer, + 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( @@ -223,140 +243,169 @@ def build_flow() -> Flow: assert restored_status.reason == "Start child_second_step" -def test_flow_checkpointing_supports_resume_and_time_travel() -> None: - checkpointer = InMemoryCheckpointer() - flow = create_flow() - - conversation = flow.start_conversation( - conversation_id="flow-checkpoint", checkpointer=checkpointer - ) - first_status = conversation.execute() - - assert isinstance(first_status, UserMessageRequestStatus) - first_checkpoint = checkpointer.load_latest(conversation.conversation_id) - assert first_checkpoint is not None - first_checkpoint_id = first_checkpoint.checkpoint_id - assert conversation.checkpoint_id == first_checkpoint_id - - restored_conversation = flow.start_conversation( - conversation_id=conversation.conversation_id, - checkpointer=checkpointer, - ) - assert restored_conversation.checkpoint_id == first_checkpoint_id - assert isinstance(restored_conversation.status, UserMessageRequestStatus) - - restored_conversation.append_user_message("continue") - restored_status = restored_conversation.execute() - assert isinstance(restored_status, FinishedStatus) - - rewound_conversation = flow.start_conversation( - conversation_id=conversation.conversation_id, - checkpointer=checkpointer, - checkpoint_id=first_checkpoint_id, - ) - assert len(rewound_conversation.get_messages()) < len(restored_conversation.get_messages()) - rewound_conversation.append_user_message("rewind") - rewound_status = rewound_conversation.execute() - assert isinstance(rewound_status, FinishedStatus) - - @pytest.mark.parametrize( - ("builder", "conversation_id"), + ( + "component_class", + "component_builder", + "model_getter", + "conversation_id", + "continuation_message", + "expected_status", + ), [ ( - partial( - _build_checkpointable_agent, - name="checkpoint_agent", - initial_message="Hello from the agent.", + 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, ), - "agent-checkpoint", + 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, ), - (_build_checkpointable_swarm, "swarm-checkpoint"), - (_build_checkpointable_managerworkers, "managerworkers-checkpoint"), ], + ids=["flow", "agent", "swarm", "managerworkers"], ) -def test_multi_agent_checkpointing_supports_resume_and_time_travel( - builder, +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: - checkpointer = InMemoryCheckpointer() - component = builder() - llm = _get_checkpoint_test_llm(component) - resumed_output = "Checkpoint resumed successfully." - rewound_output = "Checkpoint rewound successfully." + component = component_builder(vllm_responses_llm) conversation = component.start_conversation( conversation_id=conversation_id, - checkpointer=checkpointer, + checkpointer=integration_checkpointer, ) first_status = conversation.execute() - assert isinstance(first_status, UserMessageRequestStatus) - first_checkpoint = checkpointer.load_latest(conversation.conversation_id) + first_checkpoint = integration_checkpointer.load_latest(conversation.conversation_id) assert first_checkpoint is not None first_checkpoint_id = first_checkpoint.checkpoint_id - assert conversation.checkpoint_id == first_checkpoint_id - restored_conversation = component.start_conversation( + 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=checkpointer, + checkpointer=integration_checkpointer, ) - assert restored_conversation.checkpoint_id == first_checkpoint_id - llm.set_next_output(resumed_output) - restored_conversation.append_user_message("Please continue.") - restored_status = restored_conversation.execute() - assert isinstance(restored_status, UserMessageRequestStatus) - assert restored_conversation.get_last_message().content == resumed_output - - rewound_conversation = component.start_conversation( + 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, - checkpointer=checkpointer, checkpoint_id=first_checkpoint_id, + checkpointer=integration_checkpointer, ) - assert rewound_conversation.checkpoint_id == first_checkpoint_id - llm.set_next_output(rewound_output) + assert len(rewound_conversation.get_messages()) < len(restored_conversation.get_messages()) rewound_conversation.append_user_message("Try again.") - rewound_status = rewound_conversation.execute() - assert isinstance(rewound_status, UserMessageRequestStatus) - assert rewound_conversation.get_last_message().content == rewound_output + 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_after_serialization(vllm_responses_llm) -> None: - checkpointer = InMemoryCheckpointer() - agent = Agent( - llm=vllm_responses_llm, - name="serialized_checkpoint_agent", - description="Agent used for serialized checkpoint restoration.", - custom_instruction="Be helpful.", - agent_id="serialized_checkpoint_agent", - ) - with patch_llm(vllm_responses_llm, outputs=["Initial response."]): - conversation = agent.start_conversation( - conversation_id="serialized-agent-checkpoint", - checkpointer=checkpointer, +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(conversation.execute(), UserMessageRequestStatus) - restored_agent = deserialize(Agent, serialize(agent)) - restored_llm = restored_agent.llm - restored_conversation = restored_agent.start_conversation( + 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=checkpointer, + checkpointer=integration_checkpointer, ) - with patch_llm(restored_llm, outputs=["Resumed response."]): - restored_conversation.append_user_message("Continue.") - restored_status = restored_conversation.execute() - - assert isinstance(restored_status, UserMessageRequestStatus) - assert restored_conversation.get_last_message().content == "Resumed response." + 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() -> None: - checkpointer = InMemoryCheckpointer( - checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS - ) +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( @@ -385,7 +434,7 @@ def build_swarm() -> tuple[Swarm, Agent, Agent]: original_swarm, original_first_agent, original_second_agent = build_swarm() conversation = original_swarm.start_conversation( conversation_id="swarm-generated-agent-name-restart", - checkpointer=checkpointer, + 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."]) @@ -397,7 +446,7 @@ def build_swarm() -> tuple[Swarm, Agent, Agent]: restarted_swarm, restarted_first_agent, restarted_second_agent = build_swarm() restored_conversation = restarted_swarm.start_conversation( - conversation_id=conversation.conversation_id, checkpointer=checkpointer + conversation_id=conversation.conversation_id, checkpointer=integration_checkpointer ) assert any( subconversation.component is restarted_second_agent @@ -405,54 +454,9 @@ def build_swarm() -> tuple[Swarm, Agent, Agent]: ) -def test_agent_checkpoint_restore_relinks_current_flow_parent() -> None: - checkpointer = InMemoryCheckpointer() - - 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=checkpointer, - ) - parent_flow_conversation = _start_interrupted_flow_conversation( - original_parent_flow, - conversation_id=conversation.conversation_id, - interrupt_step_name="agent_child_first_step", - ) - conversation.state.current_flow_conversation = parent_flow_conversation - checkpointer.save(conversation) - - restarted_agent, restarted_parent_flow = build_agent() - restored_conversation = restarted_agent.start_conversation( - conversation_id=conversation.conversation_id, - checkpointer=checkpointer, - ) - - restored_parent_flow_conversation = restored_conversation.state.current_flow_conversation - assert restored_parent_flow_conversation is not None - - -def test_managerworkers_checkpoint_restore_uses_nested_agent_ids() -> None: - checkpointer = InMemoryCheckpointer() - +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), @@ -500,7 +504,7 @@ def build_managerworkers() -> tuple[ManagerWorkers, Agent, ManagerWorkers, Agent ) = build_managerworkers() conversation = original_managerworkers.start_conversation( conversation_id="managerworkers-nested-generated-agent-name-restart", - checkpointer=checkpointer, + checkpointer=integration_checkpointer, ) conversation.append_user_message("Save this nested conversation.") original_outer_manager_agent.llm.set_next_output( @@ -522,7 +526,7 @@ def build_managerworkers() -> tuple[ManagerWorkers, Agent, ManagerWorkers, Agent restored_conversation = restarted_managerworkers.start_conversation( conversation_id=conversation.conversation_id, - checkpointer=checkpointer, + checkpointer=integration_checkpointer, ) restored_nested_conversation = restored_conversation.subconversations[ restarted_nested_managerworkers.id diff --git a/wayflowcore/tests/test_checkpointing.py b/wayflowcore/tests/test_checkpointing.py index 5a6780f96..11b73cb16 100644 --- a/wayflowcore/tests/test_checkpointing.py +++ b/wayflowcore/tests/test_checkpointing.py @@ -11,203 +11,16 @@ from wayflowcore.agent import Agent from wayflowcore.checkpointing import CheckpointingInterval, InMemoryCheckpointer from wayflowcore.controlconnection import ControlFlowEdge -from wayflowcore.executors._flowexecutor import FlowConversationExecutor from wayflowcore.executors.executionstatus import FinishedStatus, UserMessageRequestStatus from wayflowcore.flow import Flow from wayflowcore.flowhelpers import create_single_step_flow -from wayflowcore.managerworkers import ManagerWorkers 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 .serialization.test_assistant_serialization import create_flow -from .test_managerworkers import _send_message from .testhelpers.dummy import DummyModel -pytestmark = pytest.mark.filterwarnings( - "ignore:InMemoryDatastore is for DEVELOPMENT and PROOF-OF-CONCEPT ONLY!" -) - - -def test_root_conversation_uses_one_identity_for_instance_and_thread() -> None: - agent = Agent(llm=DummyModel(), name="root_identity_agent") - - conversation = agent.start_conversation() - - assert conversation.id == conversation.conversation_id - - -def test_checkpointed_root_conversation_uses_one_identity_for_instance_and_thread() -> None: - agent = Agent(llm=DummyModel(), name="checkpointed_root_identity_agent") - - conversation = agent.start_conversation( - conversation_id="checkpointed-root-identity", - checkpointer=InMemoryCheckpointer(), - ) - - assert conversation.id == conversation.conversation_id - - -def _checkpoint_restore_wrong_component_type_scenario( - checkpointer: InMemoryCheckpointer, -): - flow = create_single_step_flow(OutputMessageStep(message_template="Hello original.")) - agent = Agent( - llm=DummyModel(), - name="checkpoint_wrong_type_agent", - description="checkpoint_wrong_type_agent description", - custom_instruction="Be helpful.", - initial_message="Hello from the wrong component type.", - agent_id="checkpoint_wrong_type_agent", - ) - conversation = flow.start_conversation( - conversation_id="checkpoint-other-component-type", - checkpointer=checkpointer, - ) - assert isinstance(conversation.execute(), FinishedStatus) - return conversation, agent - - -def _checkpoint_restore_wrong_component_identity_scenario( - checkpointer: InMemoryCheckpointer, -): - original_agent = Agent( - llm=DummyModel(), - name="checkpoint_owner_agent", - agent_id="checkpoint-owner-agent-a", - description="Original checkpoint owner", - custom_instruction="Be helpful.", - initial_message="Hello from checkpoint owner A.", - ) - other_agent = Agent( - llm=DummyModel(), - name="checkpoint_owner_agent", - agent_id="checkpoint-owner-agent-b", - description="Other checkpoint owner", - custom_instruction="Be helpful.", - initial_message="Hello from checkpoint owner B.", - ) - conversation = original_agent.start_conversation( - conversation_id="checkpoint-explicit-owner", - checkpointer=checkpointer, - ) - assert isinstance(conversation.execute(), UserMessageRequestStatus) - return conversation, other_agent - - -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 - - -def test_inmemory_checkpointer_can_save_load_list_and_delete_checkpoints() -> None: - checkpointer = InMemoryCheckpointer() - flow = create_flow() - - conversation = flow.start_conversation( - conversation_id="checkpoint-lifecycle", checkpointer=checkpointer - ) - assert conversation.checkpointer is checkpointer - - first_status = conversation.execute() - assert isinstance(first_status, UserMessageRequestStatus) - first_checkpoint = checkpointer.load_latest(conversation.conversation_id) - assert first_checkpoint is not None - first_checkpoint_id = first_checkpoint.checkpoint_id - assert conversation.checkpoint_id == first_checkpoint_id - - conversation.append_user_message("continue") - second_status = conversation.execute() - assert isinstance(second_status, FinishedStatus) - second_checkpoint = checkpointer.load_latest(conversation.conversation_id) - assert second_checkpoint is not None - second_checkpoint_id = second_checkpoint.checkpoint_id - assert second_checkpoint_id != first_checkpoint_id - assert conversation.checkpoint_id == second_checkpoint_id - - checkpoints = checkpointer.list_checkpoints("checkpoint-lifecycle") - assert [checkpoint.checkpoint_id for checkpoint in checkpoints] == [ - first_checkpoint_id, - second_checkpoint_id, - ] - assert checkpoints[-1].metadata["save_sequence"] == 2 - - latest_checkpoint = checkpointer.load_latest("checkpoint-lifecycle") - assert latest_checkpoint is not None - assert latest_checkpoint.checkpoint_id == second_checkpoint_id - - restored_conversation = flow.start_conversation( - conversation_id="checkpoint-lifecycle", - checkpoint_id=first_checkpoint_id, - checkpointer=checkpointer, - ) - assert restored_conversation.checkpointer is checkpointer - assert restored_conversation.checkpoint_id == first_checkpoint_id - assert isinstance(restored_conversation.status, UserMessageRequestStatus) - - checkpointer.delete("checkpoint-lifecycle", second_checkpoint_id) - promoted_checkpoint = checkpointer.load_latest("checkpoint-lifecycle") - assert promoted_checkpoint is not None - assert promoted_checkpoint.checkpoint_id == first_checkpoint_id - assert [ - checkpoint.checkpoint_id - for checkpoint in checkpointer.list_checkpoints("checkpoint-lifecycle") - ] == [first_checkpoint_id] - def test_checkpointer_save_snapshots_a_live_conversation_with_custom_metadata() -> None: checkpointer = InMemoryCheckpointer() @@ -228,6 +41,20 @@ def test_checkpointer_save_snapshots_a_live_conversation_with_custom_metadata() 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() @@ -250,35 +77,9 @@ def test_checkpoint_restore_requires_conversation_id_when_checkpoint_id_is_provi ) -def test_checkpoint_restore_with_serialized_component_graph() -> None: - checkpointer = InMemoryCheckpointer() - original_flow = create_flow() - - conversation = original_flow.start_conversation( - conversation_id="checkpoint-serialized-component-graph", - checkpointer=checkpointer, - ) - first_status = conversation.execute() - assert isinstance(first_status, UserMessageRequestStatus) - - serialized_flow = serialize(original_flow) - reloaded_flow = deserialize(Flow, serialized_flow) - assert reloaded_flow.id == original_flow.id - assert {name: step.id for name, step in reloaded_flow.steps.items()} == { - name: step.id for name, step in original_flow.steps.items() - } - - restored_conversation = reloaded_flow.start_conversation( - conversation_id=conversation.conversation_id, - checkpointer=checkpointer, - ) - assert isinstance(restored_conversation.status, UserMessageRequestStatus) - - restored_conversation.append_user_message("continue") - assert isinstance(restored_conversation.execute(), FinishedStatus) - - -def test_checkpoint_restore_preserves_retry_counter_with_recreated_step() -> None: +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") ) @@ -316,83 +117,26 @@ def test_checkpoint_restore_preserves_retry_counter_with_recreated_step() -> Non ) assert isinstance(conversation.execute(), FinishedStatus) - counter_key = FlowConversationExecutor.make_key_for_step( - retry_step, RetryStep._RETRY_COUNTER_KEY - ) + # Restore from an in-progress checkpoint rather than the final checkpoint. checkpoint = next( checkpoint for checkpoint in checkpointer.list_checkpoints(conversation.conversation_id) - if f"{counter_key}: 1" in checkpoint.state + if checkpoint.metadata["save_reason"] == "internal_turn_boundary" ) reloaded_flow = deserialize(Flow, serialize(flow)) - reloaded_retry_step = reloaded_flow.steps["retry_step"] - reloaded_retry_step.name = "renamed_retry_step" restored_conversation = reloaded_flow.start_conversation( conversation_id=conversation.conversation_id, checkpoint_id=checkpoint.checkpoint_id, checkpointer=checkpointer, ) - restored_counter_key = FlowConversationExecutor.make_key_for_step( - reloaded_retry_step, RetryStep._RETRY_COUNTER_KEY - ) - assert restored_counter_key == counter_key - assert restored_conversation.state.internal_context_key_values[restored_counter_key] == 1 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" -def test_managerworkers_checkpoint_restore_preserves_worker_subconversation() -> None: - manager_llm = DummyModel() - worker = Agent( - llm=DummyModel(fails_if_not_set=False), - name="checkpoint_worker", - description="Checkpoint worker", - initial_message="Worker ready.", - ) - group = ManagerWorkers( - group_manager=manager_llm, - workers=[worker], - name="checkpoint_managerworkers", - id="checkpoint-managerworkers", - ) - checkpointer = InMemoryCheckpointer( - checkpointing_interval=CheckpointingInterval.ALL_INTERNAL_TURNS - ) - - conversation = group.start_conversation( - conversation_id="checkpoint-managerworkers", - checkpointer=checkpointer, - ) - conversation.append_user_message("Delegate this task.") - manager_llm.set_next_output([_send_message(worker, message="Please help."), "All done."]) - - status = conversation.execute() - assert isinstance(status, UserMessageRequestStatus) - assert worker.id in conversation.subconversations - worker_conversation = conversation.subconversations[worker.id] - manager_conversation = conversation.subconversations[group.manager_agent.id] - assert worker_conversation.id != worker.id - assert worker_conversation.id != manager_conversation.id - assert worker_conversation.conversation_id == conversation.conversation_id - - checkpoint = checkpointer.load_latest(conversation.conversation_id) - assert checkpoint is not None - restored_conversation = group.start_conversation( - conversation_id=conversation.conversation_id, - checkpoint_id=checkpoint.checkpoint_id, - checkpointer=checkpointer, - ) - - assert worker.id in restored_conversation.subconversations - restored_worker_conversation = restored_conversation.subconversations[worker.id] - assert restored_worker_conversation.component is worker - assert restored_worker_conversation.id == worker_conversation.id - assert restored_worker_conversation.conversation_id == restored_conversation.conversation_id - - @pytest.mark.parametrize( ( "interval", @@ -426,7 +170,7 @@ def test_managerworkers_checkpoint_restore_preserves_worker_subconversation() -> ], ids=["conversation_turns", "all_internal_turns", "llm_turns"], ) -def test_checkpoint_intervals_save_expected_checkpoints( +def test_checkpoint_intervals_save_expected_flow_checkpoints( interval: CheckpointingInterval, conversation_id: str, step, @@ -457,7 +201,7 @@ def test_checkpoint_intervals_save_expected_checkpoints( @pytest.mark.parametrize( "interval", list(CheckpointingInterval), - ids=lambda interval: interval.name.lower(), + ids=["conversation_turns", "all_internal_turns", "llm_turns"], ) def test_checkpoint_intervals_save_expected_agent_checkpoints( interval: CheckpointingInterval, @@ -512,22 +256,76 @@ def test_execute_async_does_not_save_final_checkpoint_when_execution_fails( 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_wrong_component_type_scenario, - _checkpoint_restore_wrong_component_identity_scenario, _checkpoint_restore_generated_agent_id_scenario, _checkpoint_restore_generated_swarm_child_id_scenario, ], ids=[ - "wrong_component_type", - "wrong_component_identity", "generated_agent_id", "generated_swarm_child_id", ], ) -def test_checkpoint_restore_rejects_component_id_mismatches(scenario) -> None: +def test_checkpoint_restore_requires_matching_component_ids(scenario) -> None: checkpointer = InMemoryCheckpointer() conversation, restarted_component = scenario(checkpointer) From 9b6f87293b645d2bb0bfbb8c83443b162716aacd Mon Sep 17 00:00:00 2001 From: Son Le Date: Fri, 24 Jul 2026 15:02:15 +0200 Subject: [PATCH 10/10] fix oracle db tests --- wayflowcore/src/wayflowcore/cli/serve.py | 2 ++ .../wayflowcore/serialization/serializer.py | 27 +++++-------------- 2 files changed, 8 insertions(+), 21 deletions(-) diff --git a/wayflowcore/src/wayflowcore/cli/serve.py b/wayflowcore/src/wayflowcore/cli/serve.py index ee50e9ac6..74b9d54ef 100644 --- a/wayflowcore/src/wayflowcore/cli/serve.py +++ b/wayflowcore/src/wayflowcore/cli/serve.py @@ -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/serialization/serializer.py b/wayflowcore/src/wayflowcore/serialization/serializer.py index 330dc6328..83f76b43a 100644 --- a/wayflowcore/src/wayflowcore/serialization/serializer.py +++ b/wayflowcore/src/wayflowcore/serialization/serializer.py @@ -19,7 +19,6 @@ ForwardRef, List, Optional, - Protocol, Type, TypeVar, cast, @@ -559,19 +558,10 @@ def serialize( T = TypeVar("T", bound=SerializableObject) -T_co = TypeVar("T_co", bound=SerializableObject, covariant=True) - - -class _SerializableType(Protocol[T_co]): - """Type token for a serializable object, including abstract base classes.""" - - __name__: str - - def __call__(self, *args: Any, **kwargs: Any) -> T_co: ... def deserialize_from_dict( - deserialization_type: _SerializableType[T], + deserialization_type: Type[T], obj_as_dict: Dict[str, Any], deserialization_context: Optional[DeserializationContext] = None, plugins: Optional[List["WayflowDeserializationPlugin"]] = None, @@ -608,8 +598,6 @@ def deserialize_from_dict( >>> new_assistant = deserialize_from_dict(Flow, serialized_assistant) """ - runtime_deserialization_type = cast(Type[T], deserialization_type) - if deserialization_context is None: deserialization_context = DeserializationContext(plugins=plugins) elif plugins is not None: @@ -628,7 +616,7 @@ def deserialize_from_dict( deserialized_obj: SerializableObject = deserialization_context.get_deserialized_object( object_reference ) - if not isinstance(deserialized_obj, runtime_deserialization_type): + if not isinstance(deserialized_obj, deserialization_type): raise ValueError( f"A referenced objects found of type {deserialized_obj.__class__.__name__} " f"which is not compatible with the expected deserialization type of " @@ -640,10 +628,10 @@ def deserialize_from_dict( obj_as_dict = deserialization_context.get_referenced_dict(object_reference) deserialization_plugin = deserialization_context.get_deserialization_plugin_for_object( - runtime_deserialization_type + deserialization_type ) deserialized_obj = deserialization_plugin.deserialize( - runtime_deserialization_type, obj_as_dict, deserialization_context + deserialization_type, obj_as_dict, deserialization_context ) if object_reference: deserialization_context.recorddeserialized_object(object_reference, deserialized_obj) @@ -662,7 +650,7 @@ def _set_component_id(component: ObjectWithMetadata, reference: str) -> None: def deserialize( - deserialization_type: _SerializableType[T], + deserialization_type: Type[T], obj: str, deserialization_context: Optional[DeserializationContext] = None, plugins: Optional[List["WayflowDeserializationPlugin"]] = None, @@ -787,10 +775,7 @@ def autodeserialize_from_dict( "Failure to deserialize due to missing `_component_type`: The following object " f"does not seem to be a valid WayFlow component to deserialize:\n{obj_as_dict}" ) - deserialization_type = cast( - _SerializableType[SerializableObject], - SerializableObject.get_component(component_type), - ) + deserialization_type = SerializableObject.get_component(component_type) if component_type is not None and component_type != deserialization_type.__name__: raise ValueError(