diff --git a/lib/crewai/src/crewai/agents/agent_builder/base_agent.py b/lib/crewai/src/crewai/agents/agent_builder/base_agent.py index c55f5e417..9d37373e8 100644 --- a/lib/crewai/src/crewai/agents/agent_builder/base_agent.py +++ b/lib/crewai/src/crewai/agents/agent_builder/base_agent.py @@ -51,7 +51,6 @@ from crewai.utilities.string_utils import interpolate_only if TYPE_CHECKING: from crewai.context import ExecutionContext from crewai.crew import Crew - from crewai.state.provider.core import BaseProvider def _validate_crew_ref(value: Any) -> Any: @@ -338,19 +337,16 @@ class BaseAgent(BaseModel, ABC, metaclass=AgentMeta): execution_context: ExecutionContext | None = Field(default=None) @classmethod - def from_checkpoint( - cls, path: str, *, provider: BaseProvider | None = None - ) -> Self: - """Restore an Agent from a checkpoint file.""" + def from_checkpoint(cls, config: CheckpointConfig) -> Self: + """Restore an Agent from a checkpoint. + + Args: + config: Checkpoint configuration with ``restore_from`` set. + """ from crewai.context import apply_execution_context - from crewai.state.provider.json_provider import JsonProvider from crewai.state.runtime import RuntimeState - state = RuntimeState.from_checkpoint( - path, - provider=provider or JsonProvider(), - context={"from_checkpoint": True}, - ) + state = RuntimeState.from_checkpoint(config, context={"from_checkpoint": True}) for entity in state.root: if isinstance(entity, cls): if entity.execution_context is not None: @@ -359,7 +355,9 @@ class BaseAgent(BaseModel, ABC, metaclass=AgentMeta): entity.agent_executor.agent = entity entity.agent_executor._resuming = True return entity - raise ValueError(f"No {cls.__name__} found in checkpoint: {path}") + raise ValueError( + f"No {cls.__name__} found in checkpoint: {config.restore_from}" + ) @model_validator(mode="before") @classmethod diff --git a/lib/crewai/src/crewai/cli/checkpoint_tui.py b/lib/crewai/src/crewai/cli/checkpoint_tui.py index 5f81d5fad..c58ed2ea6 100644 --- a/lib/crewai/src/crewai/cli/checkpoint_tui.py +++ b/lib/crewai/src/crewai/cli/checkpoint_tui.py @@ -353,8 +353,9 @@ async def _run_checkpoint_tui_async(location: str) -> None: click.echo(f"\nResuming from: {selected}\n") from crewai.crew import Crew + from crewai.state.checkpoint_config import CheckpointConfig - crew = Crew.from_checkpoint(selected) + crew = Crew.from_checkpoint(CheckpointConfig(restore_from=selected)) result = await crew.akickoff() click.echo(f"\nResult: {getattr(result, 'raw', result)}") diff --git a/lib/crewai/src/crewai/crew.py b/lib/crewai/src/crewai/crew.py index d279a265d..111babbf3 100644 --- a/lib/crewai/src/crewai/crew.py +++ b/lib/crewai/src/crewai/crew.py @@ -42,7 +42,6 @@ if TYPE_CHECKING: from opentelemetry.trace import Span from crewai.context import ExecutionContext - from crewai.state.provider.core import BaseProvider try: from crewai_files import get_supported_content_types @@ -104,7 +103,11 @@ from crewai.rag.types import SearchResult from crewai.security.fingerprint import Fingerprint from crewai.security.security_config import SecurityConfig from crewai.skills.models import Skill -from crewai.state.checkpoint_config import CheckpointConfig, _coerce_checkpoint +from crewai.state.checkpoint_config import ( + CheckpointConfig, + _coerce_checkpoint, + apply_checkpoint, +) from crewai.task import Task from crewai.tasks.conditional_task import ConditionalTask from crewai.tasks.task_output import TaskOutput @@ -365,32 +368,21 @@ class Crew(FlowTrackable, BaseModel): checkpoint_kickoff_event_id: str | None = Field(default=None) @classmethod - def from_checkpoint( - cls, path: str, *, provider: BaseProvider | None = None - ) -> Crew: - """Restore a Crew from a checkpoint file, ready to resume via kickoff(). + def from_checkpoint(cls, config: CheckpointConfig) -> Crew: + """Restore a Crew from a checkpoint, ready to resume via kickoff(). Args: - path: Path to a checkpoint JSON file. - provider: Storage backend to read from. Defaults to JsonProvider. + config: Checkpoint configuration with ``restore_from`` set to + the path of the checkpoint to load. Returns: A Crew instance. Call kickoff() to resume from the last completed task. """ from crewai.context import apply_execution_context from crewai.events.event_bus import crewai_event_bus - from crewai.state.provider.json_provider import JsonProvider - from crewai.state.provider.utils import detect_provider from crewai.state.runtime import RuntimeState - if provider is None: - provider = detect_provider(path) - - state = RuntimeState.from_checkpoint( - path, - provider=provider or JsonProvider(), - context={"from_checkpoint": True}, - ) + state = RuntimeState.from_checkpoint(config, context={"from_checkpoint": True}) crewai_event_bus.set_runtime_state(state) for entity in state.root: if isinstance(entity, cls): @@ -398,27 +390,24 @@ class Crew(FlowTrackable, BaseModel): apply_execution_context(entity.execution_context) entity._restore_runtime() return entity - raise ValueError(f"No Crew found in checkpoint: {path}") + raise ValueError(f"No Crew found in checkpoint: {config.restore_from}") @classmethod def fork( cls, - path: str, - *, + config: CheckpointConfig, branch: str | None = None, - provider: BaseProvider | None = None, ) -> Crew: """Fork a Crew from a checkpoint, creating a new execution branch. Args: - path: Path to a checkpoint file. + config: Checkpoint configuration with ``restore_from`` set. branch: Branch label for the fork. Auto-generated if not provided. - provider: Storage backend to read from. Defaults to auto-detect. Returns: A Crew instance on the new branch. Call kickoff() to run. """ - crew = cls.from_checkpoint(path, provider=provider) + crew = cls.from_checkpoint(config) state = crewai_event_bus._runtime_state if state is None: raise RuntimeError( @@ -882,16 +871,23 @@ class Crew(FlowTrackable, BaseModel): self, inputs: dict[str, Any] | None = None, input_files: dict[str, FileInput] | None = None, + from_checkpoint: CheckpointConfig | None = None, ) -> CrewOutput | CrewStreamingOutput: """Execute the crew's workflow. Args: inputs: Optional input dictionary for task interpolation. input_files: Optional dict of named file inputs for the crew. + from_checkpoint: Optional checkpoint config. If ``restore_from`` + is set, the crew resumes from that checkpoint. Remaining + config fields enable checkpointing for the run. Returns: CrewOutput or CrewStreamingOutput if streaming is enabled. """ + restored = apply_checkpoint(self, from_checkpoint) + if restored is not None: + return restored.kickoff(inputs=inputs, input_files=input_files) # type: ignore[no-any-return] get_env_context() if self.stream: enable_agent_streaming(self.agents) @@ -1004,12 +1000,15 @@ class Crew(FlowTrackable, BaseModel): self, inputs: dict[str, Any] | None = None, input_files: dict[str, FileInput] | None = None, + from_checkpoint: CheckpointConfig | None = None, ) -> CrewOutput | CrewStreamingOutput: """Asynchronous kickoff method to start the crew execution. Args: inputs: Optional input dictionary for task interpolation. input_files: Optional dict of named file inputs for the crew. + from_checkpoint: Optional checkpoint config. If ``restore_from`` + is set, the crew resumes from that checkpoint. Returns: CrewOutput or CrewStreamingOutput if streaming is enabled. @@ -1018,6 +1017,9 @@ class Crew(FlowTrackable, BaseModel): to get stream chunks. After iteration completes, access the final result via .result. """ + restored = apply_checkpoint(self, from_checkpoint) + if restored is not None: + return await restored.kickoff_async(inputs=inputs, input_files=input_files) # type: ignore[no-any-return] inputs = inputs or {} if self.stream: @@ -1078,6 +1080,7 @@ class Crew(FlowTrackable, BaseModel): self, inputs: dict[str, Any] | None = None, input_files: dict[str, FileInput] | None = None, + from_checkpoint: CheckpointConfig | None = None, ) -> CrewOutput | CrewStreamingOutput: """Native async kickoff method using async task execution throughout. @@ -1088,10 +1091,15 @@ class Crew(FlowTrackable, BaseModel): Args: inputs: Optional input dictionary for task interpolation. input_files: Optional dict of named file inputs for the crew. + from_checkpoint: Optional checkpoint config. If ``restore_from`` + is set, the crew resumes from that checkpoint. Returns: CrewOutput or CrewStreamingOutput if streaming is enabled. """ + restored = apply_checkpoint(self, from_checkpoint) + if restored is not None: + return await restored.akickoff(inputs=inputs, input_files=input_files) # type: ignore[no-any-return] if self.stream: enable_agent_streaming(self.agents) ctx = StreamingContext(use_async=True) diff --git a/lib/crewai/src/crewai/flow/flow.py b/lib/crewai/src/crewai/flow/flow.py index a85885bd9..057f60ffb 100644 --- a/lib/crewai/src/crewai/flow/flow.py +++ b/lib/crewai/src/crewai/flow/flow.py @@ -113,7 +113,11 @@ from crewai.flow.utils import ( ) from crewai.memory.memory_scope import MemoryScope, MemorySlice from crewai.memory.unified_memory import Memory -from crewai.state.checkpoint_config import CheckpointConfig, _coerce_checkpoint +from crewai.state.checkpoint_config import ( + CheckpointConfig, + _coerce_checkpoint, + apply_checkpoint, +) if TYPE_CHECKING: @@ -122,7 +126,6 @@ if TYPE_CHECKING: from crewai.context import ExecutionContext from crewai.flow.async_feedback.types import PendingFeedbackContext from crewai.llms.base_llm import BaseLLM - from crewai.state.provider.core import BaseProvider from crewai.flow.visualization import build_flow_structure, render_interactive from crewai.types.streaming import CrewStreamingOutput, FlowStreamingOutput @@ -928,20 +931,21 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): ] = Field(default=None) @classmethod - def from_checkpoint( - cls, path: str, *, provider: BaseProvider | None = None - ) -> Flow: # type: ignore[type-arg] - """Restore a Flow from a checkpoint file.""" + def from_checkpoint(cls, config: CheckpointConfig) -> Flow: # type: ignore[type-arg] + """Restore a Flow from a checkpoint. + + Args: + config: Checkpoint configuration with ``restore_from`` set to + the path of the checkpoint to load. + + Returns: + A Flow instance ready to resume. + """ from crewai.context import apply_execution_context from crewai.events.event_bus import crewai_event_bus - from crewai.state.provider.json_provider import JsonProvider from crewai.state.runtime import RuntimeState - state = RuntimeState.from_checkpoint( - path, - provider=provider or JsonProvider(), - context={"from_checkpoint": True}, - ) + state = RuntimeState.from_checkpoint(config, context={"from_checkpoint": True}) crewai_event_bus.set_runtime_state(state) for entity in state.root: if not isinstance(entity, Flow): @@ -958,27 +962,24 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): instance.checkpoint_state = entity.checkpoint_state instance._restore_from_checkpoint() return instance - raise ValueError(f"No Flow found in checkpoint: {path}") + raise ValueError(f"No Flow found in checkpoint: {config.restore_from}") @classmethod def fork( cls, - path: str, - *, + config: CheckpointConfig, branch: str | None = None, - provider: BaseProvider | None = None, ) -> Flow: # type: ignore[type-arg] """Fork a Flow from a checkpoint, creating a new execution branch. Args: - path: Path to a checkpoint file. + config: Checkpoint configuration with ``restore_from`` set. branch: Branch label for the fork. Auto-generated if not provided. - provider: Storage backend to read from. Defaults to auto-detect. Returns: A Flow instance on the new branch. Call kickoff() to run. """ - flow = cls.from_checkpoint(path, provider=provider) + flow = cls.from_checkpoint(config) state = crewai_event_bus._runtime_state if state is None: raise RuntimeError( @@ -1984,6 +1985,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): self, inputs: dict[str, Any] | None = None, input_files: dict[str, FileInput] | None = None, + from_checkpoint: CheckpointConfig | None = None, ) -> Any | FlowStreamingOutput: """Start the flow execution in a synchronous context. @@ -1993,10 +1995,15 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): Args: inputs: Optional dictionary containing input values and/or a state ID. input_files: Optional dict of named file inputs for the flow. + from_checkpoint: Optional checkpoint config. If ``restore_from`` + is set, the flow resumes from that checkpoint. Returns: The final output from the flow or FlowStreamingOutput if streaming. """ + restored = apply_checkpoint(self, from_checkpoint) + if restored is not None: + return restored.kickoff(inputs=inputs, input_files=input_files) get_env_context() if self.stream: result_holder: list[Any] = [] @@ -2053,6 +2060,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): self, inputs: dict[str, Any] | None = None, input_files: dict[str, FileInput] | None = None, + from_checkpoint: CheckpointConfig | None = None, ) -> Any | FlowStreamingOutput: """Start the flow execution asynchronously. @@ -2064,10 +2072,15 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): Args: inputs: Optional dictionary containing input values and/or a state ID for restoration. input_files: Optional dict of named file inputs for the flow. + from_checkpoint: Optional checkpoint config. If ``restore_from`` + is set, the flow resumes from that checkpoint. Returns: The final output from the flow, which is the result of the last executed method. """ + restored = apply_checkpoint(self, from_checkpoint) + if restored is not None: + return await restored.kickoff_async(inputs=inputs, input_files=input_files) if self.stream: result_holder: list[Any] = [] current_task_info: TaskInfo = { @@ -2326,17 +2339,20 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): self, inputs: dict[str, Any] | None = None, input_files: dict[str, FileInput] | None = None, + from_checkpoint: CheckpointConfig | None = None, ) -> Any | FlowStreamingOutput: """Native async method to start the flow execution. Alias for kickoff_async. Args: inputs: Optional dictionary containing input values and/or a state ID for restoration. input_files: Optional dict of named file inputs for the flow. + from_checkpoint: Optional checkpoint config. If ``restore_from`` + is set, the flow resumes from that checkpoint. Returns: The final output from the flow, which is the result of the last executed method. """ - return await self.kickoff_async(inputs, input_files) + return await self.kickoff_async(inputs, input_files, from_checkpoint) async def _execute_start_method(self, start_method_name: FlowMethodName) -> None: """Executes a flow's start method and its triggered listeners. diff --git a/lib/crewai/src/crewai/state/checkpoint_config.py b/lib/crewai/src/crewai/state/checkpoint_config.py index 38c6b0490..bd24ef8eb 100644 --- a/lib/crewai/src/crewai/state/checkpoint_config.py +++ b/lib/crewai/src/crewai/state/checkpoint_config.py @@ -2,6 +2,7 @@ from __future__ import annotations +from pathlib import Path from typing import Annotated, Any, Literal from pydantic import BaseModel, Field, model_validator @@ -201,6 +202,12 @@ class CheckpointConfig(BaseModel): description="Maximum checkpoints to keep. Oldest are pruned after " "each write. None means keep all.", ) + restore_from: Path | str | None = Field( + default=None, + description="Path or location of a checkpoint to restore from. " + "When passed via a kickoff method's from_checkpoint parameter, " + "the crew or flow resumes from this checkpoint.", + ) @model_validator(mode="after") def _register_handlers(self) -> CheckpointConfig: @@ -216,3 +223,25 @@ class CheckpointConfig(BaseModel): @property def trigger_events(self) -> set[str]: return set(self.on_events) + + +def apply_checkpoint(instance: Any, from_checkpoint: CheckpointConfig | None) -> Any: + """Handle checkpoint config for a kickoff method. + + If *from_checkpoint* carries a ``restore_from`` path, builds and returns a + restored instance (with ``restore_from`` cleared). The caller should + dispatch into its own kickoff variant on that restored instance. + + If *from_checkpoint* is present but has no ``restore_from``, sets + ``instance.checkpoint`` and returns ``None`` (proceed normally). + + If *from_checkpoint* is ``None``, returns ``None`` immediately. + """ + if from_checkpoint is None: + return None + if from_checkpoint.restore_from is not None: + restored = type(instance).from_checkpoint(from_checkpoint) + restored.checkpoint = from_checkpoint.model_copy(update={"restore_from": None}) + return restored + instance.checkpoint = from_checkpoint + return None diff --git a/lib/crewai/src/crewai/state/runtime.py b/lib/crewai/src/crewai/state/runtime.py index 79c836e1e..3f32457bb 100644 --- a/lib/crewai/src/crewai/state/runtime.py +++ b/lib/crewai/src/crewai/state/runtime.py @@ -23,6 +23,7 @@ from pydantic import ( ) from crewai.context import capture_execution_context +from crewai.state.checkpoint_config import CheckpointConfig from crewai.state.event_record import EventRecord from crewai.state.provider.core import BaseProvider from crewai.state.provider.json_provider import JsonProvider @@ -208,19 +209,22 @@ class RuntimeState(RootModel): # type: ignore[type-arg] self._branch = f"fork/{uuid.uuid4().hex[:8]}" @classmethod - def from_checkpoint( - cls, location: str, provider: BaseProvider, **kwargs: Any - ) -> RuntimeState: + def from_checkpoint(cls, config: CheckpointConfig, **kwargs: Any) -> RuntimeState: """Restore a RuntimeState from a checkpoint. Args: - location: The identifier returned by a previous ``checkpoint`` call. - provider: The storage backend to read from. + config: Checkpoint configuration with ``restore_from`` set. **kwargs: Passed to ``model_validate_json``. Returns: A restored RuntimeState. """ + from crewai.state.provider.utils import detect_provider + + if config.restore_from is None: + raise ValueError("CheckpointConfig.restore_from must be set") + location = str(config.restore_from) + provider = detect_provider(location) raw = provider.from_checkpoint(location) state = cls.model_validate_json(raw, **kwargs) checkpoint_id = provider.extract_id(location) @@ -230,18 +234,23 @@ class RuntimeState(RootModel): # type: ignore[type-arg] @classmethod async def afrom_checkpoint( - cls, location: str, provider: BaseProvider, **kwargs: Any + cls, config: CheckpointConfig, **kwargs: Any ) -> RuntimeState: """Async version of :meth:`from_checkpoint`. Args: - location: The identifier returned by a previous ``acheckpoint`` call. - provider: The storage backend to read from. + config: Checkpoint configuration with ``restore_from`` set. **kwargs: Passed to ``model_validate_json``. Returns: A restored RuntimeState. """ + from crewai.state.provider.utils import detect_provider + + if config.restore_from is None: + raise ValueError("CheckpointConfig.restore_from must be set") + location = str(config.restore_from) + provider = detect_provider(location) raw = await provider.afrom_checkpoint(location) state = cls.model_validate_json(raw, **kwargs) checkpoint_id = provider.extract_id(location) diff --git a/lib/crewai/tests/test_checkpoint.py b/lib/crewai/tests/test_checkpoint.py index 0a0c00027..1f1541790 100644 --- a/lib/crewai/tests/test_checkpoint.py +++ b/lib/crewai/tests/test_checkpoint.py @@ -172,6 +172,14 @@ class TestCheckpointConfig: cfg = CheckpointConfig(on_events=["*"]) assert cfg.trigger_all + def test_restore_from_field(self) -> None: + cfg = CheckpointConfig(restore_from="/path/to/checkpoint.json") + assert cfg.restore_from == "/path/to/checkpoint.json" + + def test_restore_from_default_none(self) -> None: + cfg = CheckpointConfig() + assert cfg.restore_from is None + def test_trigger_events(self) -> None: cfg = CheckpointConfig( on_events=["task_completed", "crew_kickoff_completed"] @@ -269,9 +277,9 @@ class TestRuntimeStateLineage: loc = state.checkpoint(d) written_id = state._checkpoint_id - provider = JsonProvider() + cfg = CheckpointConfig(restore_from=loc) restored = RuntimeState.from_checkpoint( - loc, provider, context={"from_checkpoint": True} + cfg, context={"from_checkpoint": True} ) assert restored._checkpoint_id == written_id assert restored._parent_id == written_id @@ -480,3 +488,51 @@ class TestSqliteProviderFork: agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini") crew = Crew(agents=[agent], tasks=[], verbose=False) return RuntimeState(root=[crew]) + + +# ---------- Kickoff from_checkpoint parameter ---------- + + +class TestKickoffFromCheckpoint: + def test_crew_kickoff_delegates_to_from_checkpoint(self) -> None: + mock_restored = MagicMock(spec=Crew) + mock_restored.kickoff.return_value = "result" + + cfg = CheckpointConfig(restore_from="/path/to/cp.json") + with patch.object(Crew, "from_checkpoint", return_value=mock_restored): + agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini") + crew = Crew(agents=[agent], tasks=[], verbose=False) + result = crew.kickoff(inputs={"k": "v"}, from_checkpoint=cfg) + + mock_restored.kickoff.assert_called_once_with( + inputs={"k": "v"}, input_files=None + ) + assert mock_restored.checkpoint.restore_from is None + assert result == "result" + + def test_crew_kickoff_config_only_sets_checkpoint(self) -> None: + cfg = CheckpointConfig(on_events=["task_completed"]) + agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini") + crew = Crew(agents=[agent], tasks=[], verbose=False) + assert crew.checkpoint is None + with patch("crewai.crew.get_env_context"), \ + patch("crewai.crew.prepare_kickoff", side_effect=RuntimeError("stop")): + with pytest.raises(RuntimeError, match="stop"): + crew.kickoff(from_checkpoint=cfg) + assert isinstance(crew.checkpoint, CheckpointConfig) + assert crew.checkpoint.on_events == ["task_completed"] + + def test_flow_kickoff_delegates_to_from_checkpoint(self) -> None: + mock_restored = MagicMock(spec=Flow) + mock_restored.kickoff.return_value = "flow_result" + + cfg = CheckpointConfig(restore_from="/path/to/flow_cp.json") + with patch.object(Flow, "from_checkpoint", return_value=mock_restored): + flow = Flow() + result = flow.kickoff(from_checkpoint=cfg) + + mock_restored.kickoff.assert_called_once_with( + inputs=None, input_files=None + ) + assert mock_restored.checkpoint.restore_from is None + assert result == "flow_result"