mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-07 07:59:29 +00:00
feat: add checkpoint and fork support to standalone agents
Add fork classmethod, _restore_runtime, and _restore_event_scope to BaseAgent. Fix from_checkpoint to set runtime state on the event bus and restore event scopes. Store kickoff event ID across checkpoints to skip re-emission on resume. Handle agent entity type in checkpoint CLI and TUI.
This commit is contained in:
@@ -29,7 +29,7 @@ from pydantic import (
|
||||
model_validator,
|
||||
)
|
||||
from pydantic.functional_serializers import PlainSerializer
|
||||
from typing_extensions import Self
|
||||
from typing_extensions import Self, TypeIs
|
||||
|
||||
from crewai.agent.planning_config import PlanningConfig
|
||||
from crewai.agent.utils import (
|
||||
@@ -133,6 +133,13 @@ _EXECUTOR_CLASS_MAP: dict[str, type] = {
|
||||
}
|
||||
|
||||
|
||||
def _is_resuming_agent_executor(
|
||||
executor: CrewAgentExecutor | AgentExecutor | None,
|
||||
) -> TypeIs[AgentExecutor]:
|
||||
"""Type guard: True when the executor is resuming from a checkpoint."""
|
||||
return isinstance(executor, AgentExecutor) and executor._resuming
|
||||
|
||||
|
||||
def _validate_executor_class(value: Any) -> Any:
|
||||
if isinstance(value, str):
|
||||
cls = _EXECUTOR_CLASS_MAP.get(value)
|
||||
@@ -1366,24 +1373,42 @@ class Agent(BaseAgent):
|
||||
|
||||
prompt, stop_words, rpm_limit_fn = self._build_execution_prompt(raw_tools)
|
||||
|
||||
executor = AgentExecutor(
|
||||
llm=cast(BaseLLM, self.llm),
|
||||
agent=self,
|
||||
prompt=prompt,
|
||||
max_iter=self.max_iter,
|
||||
tools=parsed_tools,
|
||||
tools_names=get_tool_names(parsed_tools),
|
||||
stop_words=stop_words,
|
||||
tools_description=render_text_description_and_args(parsed_tools),
|
||||
tools_handler=self.tools_handler,
|
||||
original_tools=raw_tools,
|
||||
step_callback=self.step_callback,
|
||||
function_calling_llm=self.function_calling_llm,
|
||||
respect_context_window=self.respect_context_window,
|
||||
request_within_rpm_limit=rpm_limit_fn,
|
||||
callbacks=[TokenCalcHandler(self._token_process)],
|
||||
response_model=response_format,
|
||||
)
|
||||
if _is_resuming_agent_executor(self.agent_executor):
|
||||
executor = self.agent_executor
|
||||
executor.tools = parsed_tools
|
||||
executor.tools_names = get_tool_names(parsed_tools)
|
||||
executor.tools_description = render_text_description_and_args(parsed_tools)
|
||||
executor.original_tools = raw_tools
|
||||
executor.prompt = prompt
|
||||
executor.response_model = response_format
|
||||
executor.stop_words = stop_words
|
||||
executor.tools_handler = self.tools_handler
|
||||
executor.step_callback = self.step_callback
|
||||
executor.function_calling_llm = cast(
|
||||
BaseLLM | None, self.function_calling_llm
|
||||
)
|
||||
executor.respect_context_window = self.respect_context_window
|
||||
executor.request_within_rpm_limit = rpm_limit_fn
|
||||
executor.callbacks = [TokenCalcHandler(self._token_process)]
|
||||
else:
|
||||
executor = AgentExecutor(
|
||||
llm=cast(BaseLLM, self.llm),
|
||||
agent=self,
|
||||
prompt=prompt,
|
||||
max_iter=self.max_iter,
|
||||
tools=parsed_tools,
|
||||
tools_names=get_tool_names(parsed_tools),
|
||||
stop_words=stop_words,
|
||||
tools_description=render_text_description_and_args(parsed_tools),
|
||||
tools_handler=self.tools_handler,
|
||||
original_tools=raw_tools,
|
||||
step_callback=self.step_callback,
|
||||
function_calling_llm=self.function_calling_llm,
|
||||
respect_context_window=self.respect_context_window,
|
||||
request_within_rpm_limit=rpm_limit_fn,
|
||||
callbacks=[TokenCalcHandler(self._token_process)],
|
||||
response_model=response_format,
|
||||
)
|
||||
|
||||
all_files: dict[str, Any] = {}
|
||||
if isinstance(messages, str):
|
||||
@@ -1504,14 +1529,17 @@ class Agent(BaseAgent):
|
||||
)
|
||||
|
||||
try:
|
||||
crewai_event_bus.emit(
|
||||
self,
|
||||
event=LiteAgentExecutionStartedEvent(
|
||||
if self.checkpoint_kickoff_event_id is not None:
|
||||
self._kickoff_event_id = self.checkpoint_kickoff_event_id
|
||||
self.checkpoint_kickoff_event_id = None
|
||||
else:
|
||||
started_event = LiteAgentExecutionStartedEvent(
|
||||
agent_info=agent_info,
|
||||
tools=parsed_tools,
|
||||
messages=messages,
|
||||
),
|
||||
)
|
||||
)
|
||||
crewai_event_bus.emit(self, event=started_event)
|
||||
self._kickoff_event_id = started_event.event_id
|
||||
|
||||
output = self._execute_and_build_output(executor, inputs, response_format)
|
||||
return self._finalize_kickoff(
|
||||
@@ -1808,14 +1836,17 @@ class Agent(BaseAgent):
|
||||
)
|
||||
|
||||
try:
|
||||
crewai_event_bus.emit(
|
||||
self,
|
||||
event=LiteAgentExecutionStartedEvent(
|
||||
if self.checkpoint_kickoff_event_id is not None:
|
||||
self._kickoff_event_id = self.checkpoint_kickoff_event_id
|
||||
self.checkpoint_kickoff_event_id = None
|
||||
else:
|
||||
started_event = LiteAgentExecutionStartedEvent(
|
||||
agent_info=agent_info,
|
||||
tools=parsed_tools,
|
||||
messages=messages,
|
||||
),
|
||||
)
|
||||
)
|
||||
crewai_event_bus.emit(self, event=started_event)
|
||||
self._kickoff_event_id = started_event.event_id
|
||||
|
||||
output = await self._execute_and_build_output_async(
|
||||
executor, inputs, response_format
|
||||
|
||||
@@ -28,6 +28,9 @@ from crewai.agents.agent_builder.base_agent_executor import BaseAgentExecutor
|
||||
from crewai.agents.agent_builder.utilities.base_token_process import TokenProcess
|
||||
from crewai.agents.cache.cache_handler import CacheHandler
|
||||
from crewai.agents.tools_handler import ToolsHandler
|
||||
from crewai.events.base_events import set_emission_counter
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
from crewai.events.event_context import restore_event_scope, set_last_event_id
|
||||
from crewai.knowledge.knowledge import Knowledge
|
||||
from crewai.knowledge.knowledge_config import KnowledgeConfig
|
||||
from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource
|
||||
@@ -51,6 +54,7 @@ from crewai.utilities.string_utils import interpolate_only
|
||||
if TYPE_CHECKING:
|
||||
from crewai.context import ExecutionContext
|
||||
from crewai.crew import Crew
|
||||
from crewai.state.runtime import RuntimeState
|
||||
|
||||
|
||||
def _validate_crew_ref(value: Any) -> Any:
|
||||
@@ -219,6 +223,7 @@ class BaseAgent(BaseModel, ABC, metaclass=AgentMeta):
|
||||
_original_goal: str | None = PrivateAttr(default=None)
|
||||
_original_backstory: str | None = PrivateAttr(default=None)
|
||||
_token_process: TokenProcess = PrivateAttr(default_factory=TokenProcess)
|
||||
_kickoff_event_id: str | None = PrivateAttr(default=None)
|
||||
id: UUID4 = Field(default_factory=uuid.uuid4, frozen=True)
|
||||
role: str = Field(description="Role of the agent")
|
||||
goal: str = Field(description="Objective of the agent")
|
||||
@@ -335,30 +340,90 @@ class BaseAgent(BaseModel, ABC, metaclass=AgentMeta):
|
||||
min_length=1,
|
||||
)
|
||||
execution_context: ExecutionContext | None = Field(default=None)
|
||||
checkpoint_kickoff_event_id: str | None = Field(default=None)
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint(cls, config: CheckpointConfig) -> Self:
|
||||
"""Restore an Agent from a checkpoint.
|
||||
"""Restore an Agent from a checkpoint, ready to resume via kickoff().
|
||||
|
||||
Args:
|
||||
config: Checkpoint configuration with ``restore_from`` set.
|
||||
config: Checkpoint configuration with ``restore_from`` set to
|
||||
the path of the checkpoint to load.
|
||||
|
||||
Returns:
|
||||
An Agent instance. Call kickoff() to resume execution.
|
||||
"""
|
||||
from crewai.context import apply_execution_context
|
||||
from crewai.state.runtime import RuntimeState
|
||||
|
||||
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):
|
||||
if entity.execution_context is not None:
|
||||
apply_execution_context(entity.execution_context)
|
||||
if entity.agent_executor is not None:
|
||||
entity.agent_executor.agent = entity
|
||||
entity.agent_executor._resuming = True
|
||||
entity._restore_runtime(state)
|
||||
return entity
|
||||
raise ValueError(
|
||||
f"No {cls.__name__} found in checkpoint: {config.restore_from}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fork(cls, config: CheckpointConfig, branch: str | None = None) -> Self:
|
||||
"""Fork an Agent from a checkpoint, creating a new execution branch.
|
||||
|
||||
Args:
|
||||
config: Checkpoint configuration with ``restore_from`` set.
|
||||
branch: Branch label for the fork. Auto-generated if not provided.
|
||||
|
||||
Returns:
|
||||
An Agent instance on the new branch. Call kickoff() to run.
|
||||
"""
|
||||
agent = cls.from_checkpoint(config)
|
||||
state = crewai_event_bus._runtime_state
|
||||
if state is None:
|
||||
raise RuntimeError("Cannot fork: no runtime state on the event bus.")
|
||||
state.fork(branch)
|
||||
return agent
|
||||
|
||||
def _restore_runtime(self, state: RuntimeState) -> None:
|
||||
"""Re-create runtime objects after restoring from a checkpoint.
|
||||
|
||||
Args:
|
||||
state: The RuntimeState containing the event record.
|
||||
"""
|
||||
if self.agent_executor is not None:
|
||||
self.agent_executor.agent = self
|
||||
self.agent_executor._resuming = True
|
||||
if self.checkpoint_kickoff_event_id is not None:
|
||||
self._kickoff_event_id = self.checkpoint_kickoff_event_id
|
||||
self._restore_event_scope(state)
|
||||
|
||||
def _restore_event_scope(self, state: RuntimeState) -> None:
|
||||
"""Rebuild the event scope stack from the checkpoint's event record.
|
||||
|
||||
Args:
|
||||
state: The RuntimeState containing the event record.
|
||||
"""
|
||||
stack: list[tuple[str, str]] = []
|
||||
kickoff_id = self._kickoff_event_id
|
||||
if kickoff_id:
|
||||
stack.append((kickoff_id, "lite_agent_execution_started"))
|
||||
|
||||
restore_event_scope(tuple(stack))
|
||||
|
||||
last_event_id: str | None = None
|
||||
max_seq = 0
|
||||
for node in state.event_record.nodes.values():
|
||||
seq = node.event.emission_sequence or 0
|
||||
if seq > max_seq:
|
||||
max_seq = seq
|
||||
last_event_id = node.event.event_id
|
||||
if last_event_id is not None:
|
||||
set_last_event_id(last_event_id)
|
||||
if max_seq > 0:
|
||||
set_emission_counter(max_seq)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def process_model_config(cls, values: Any) -> dict[str, Any]:
|
||||
|
||||
@@ -472,6 +472,8 @@ def _entity_type_from_meta(meta: dict[str, Any]) -> str:
|
||||
for ent in meta.get("entities", []):
|
||||
if ent.get("type") == "flow":
|
||||
return "flow"
|
||||
if ent.get("type") == "agent":
|
||||
return "agent"
|
||||
return "crew"
|
||||
|
||||
|
||||
@@ -505,6 +507,11 @@ def resume_checkpoint(location: str, checkpoint_id: str | None) -> None:
|
||||
|
||||
flow = Flow.from_checkpoint(config)
|
||||
result = asyncio.run(flow.kickoff_async(inputs=inputs))
|
||||
elif entity_type == "agent":
|
||||
from crewai.agent import Agent
|
||||
|
||||
agent = Agent.from_checkpoint(config)
|
||||
result = asyncio.run(agent.akickoff(messages="Resume execution."))
|
||||
else:
|
||||
from crewai.crew import Crew
|
||||
|
||||
|
||||
@@ -123,7 +123,7 @@ _TuiResult = (
|
||||
str,
|
||||
dict[str, Any] | None,
|
||||
dict[int, str] | None,
|
||||
Literal["crew", "flow"],
|
||||
Literal["crew", "flow", "agent"],
|
||||
]
|
||||
| None
|
||||
)
|
||||
@@ -682,10 +682,14 @@ class CheckpointTUI(App[_TuiResult]):
|
||||
overrides[task_idx] = editor.text
|
||||
return overrides or None
|
||||
|
||||
def _detect_entity_type(self, entry: dict[str, Any]) -> Literal["crew", "flow"]:
|
||||
def _detect_entity_type(
|
||||
self, entry: dict[str, Any]
|
||||
) -> Literal["crew", "flow", "agent"]:
|
||||
for ent in entry.get("entities", []):
|
||||
if ent.get("type") == "flow":
|
||||
return "flow"
|
||||
if ent.get("type") == "agent":
|
||||
return "agent"
|
||||
return "crew"
|
||||
|
||||
def _resolve_location(self, entry: dict[str, Any]) -> str:
|
||||
@@ -829,6 +833,21 @@ async def _run_checkpoint_tui_async(location: str) -> None:
|
||||
click.echo(f"\nResult: {getattr(result, 'raw', result)}")
|
||||
return
|
||||
|
||||
if entity_type == "agent":
|
||||
from crewai.agent import Agent
|
||||
|
||||
if action == "fork":
|
||||
click.echo(f"\nForking agent from: {selected}\n")
|
||||
agent = Agent.fork(config)
|
||||
else:
|
||||
click.echo(f"\nResuming agent from: {selected}\n")
|
||||
agent = Agent.from_checkpoint(config)
|
||||
|
||||
click.echo()
|
||||
result = await agent.akickoff(messages="Resume execution.")
|
||||
click.echo(f"\nResult: {getattr(result, 'raw', result)}")
|
||||
return
|
||||
|
||||
from crewai.crew import Crew
|
||||
|
||||
if action == "fork":
|
||||
|
||||
@@ -44,9 +44,12 @@ def _sync_checkpoint_fields(entity: object) -> None:
|
||||
entity: The entity whose private runtime attributes will be
|
||||
copied into its public checkpoint fields.
|
||||
"""
|
||||
from crewai.agents.agent_builder.base_agent import BaseAgent
|
||||
from crewai.crew import Crew
|
||||
from crewai.flow.flow import Flow
|
||||
|
||||
if isinstance(entity, BaseAgent):
|
||||
entity.checkpoint_kickoff_event_id = entity._kickoff_event_id
|
||||
if isinstance(entity, Flow):
|
||||
entity.checkpoint_completed_methods = (
|
||||
set(entity._completed_methods) if entity._completed_methods else None
|
||||
|
||||
@@ -562,3 +562,75 @@ class TestKickoffFromCheckpoint:
|
||||
)
|
||||
assert mock_restored.checkpoint.restore_from is None
|
||||
assert result == "flow_result"
|
||||
|
||||
|
||||
# ---------- Agent checkpoint/fork ----------
|
||||
|
||||
|
||||
class TestAgentCheckpoint:
|
||||
def _make_agent_state(self) -> RuntimeState:
|
||||
agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini")
|
||||
return RuntimeState(root=[agent])
|
||||
|
||||
def test_agent_from_checkpoint_sets_runtime_state(self) -> None:
|
||||
state = self._make_agent_state()
|
||||
state._provider = JsonProvider()
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
loc = state.checkpoint(d)
|
||||
cfg = CheckpointConfig(restore_from=loc)
|
||||
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
|
||||
crewai_event_bus._runtime_state = None
|
||||
Agent.from_checkpoint(cfg)
|
||||
assert crewai_event_bus._runtime_state is not None
|
||||
|
||||
def test_agent_fork_sets_branch(self) -> None:
|
||||
state = self._make_agent_state()
|
||||
state._provider = JsonProvider()
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
loc = state.checkpoint(d)
|
||||
cfg = CheckpointConfig(restore_from=loc)
|
||||
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
|
||||
Agent.fork(cfg, branch="agent-experiment")
|
||||
rt = crewai_event_bus._runtime_state
|
||||
assert rt is not None
|
||||
assert rt._branch == "agent-experiment"
|
||||
|
||||
def test_agent_fork_auto_branch(self) -> None:
|
||||
state = self._make_agent_state()
|
||||
state._provider = JsonProvider()
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
loc = state.checkpoint(d)
|
||||
cfg = CheckpointConfig(restore_from=loc)
|
||||
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
|
||||
Agent.fork(cfg)
|
||||
rt = crewai_event_bus._runtime_state
|
||||
assert rt is not None
|
||||
assert rt._branch.startswith("fork/")
|
||||
|
||||
def test_sync_checkpoint_fields_agent(self) -> None:
|
||||
from crewai.state.runtime import _sync_checkpoint_fields
|
||||
|
||||
agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini")
|
||||
agent._kickoff_event_id = "evt-123"
|
||||
_sync_checkpoint_fields(agent)
|
||||
assert agent.checkpoint_kickoff_event_id == "evt-123"
|
||||
|
||||
def test_agent_restore_kickoff_event_id(self) -> None:
|
||||
agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini")
|
||||
agent._kickoff_event_id = "evt-456"
|
||||
state = RuntimeState(root=[agent])
|
||||
state._provider = JsonProvider()
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
from crewai.state.runtime import _prepare_entities
|
||||
|
||||
_prepare_entities(state.root)
|
||||
loc = state.checkpoint(d)
|
||||
cfg = CheckpointConfig(restore_from=loc)
|
||||
restored = Agent.from_checkpoint(cfg)
|
||||
assert restored._kickoff_event_id == "evt-456"
|
||||
|
||||
Reference in New Issue
Block a user