mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-01 13:18:10 +00:00
6ad821b157
* Add expressions to FlowDefinition actions
Let definitions compute values without Python. A new `call: expression`
action evaluates a Common Expression Language (CEL) expression, and tool
`with:` blocks now render `${...}` CEL templates.
Example 1:
```yaml
decide:
do:
call: expression
expr: "state.score >= 80 ? 'qualified' : 'nurture'"
router: true
emit: [qualified, nurture]
```
Example 2:
```yaml
search:
do:
call: tool
ref: my.pkg:SearchTool
with:
search_query: "${outputs.build_query.query + ' news'}"
max_results: "${state.limit}"
```
* Address code review comments
* Address code review comments
* Fix linting offenses
* Address code review comments
* Fix scrapgraph issue
807 lines
29 KiB
Python
807 lines
29 KiB
Python
"""Tests for CheckpointConfig, checkpoint listener, pruning, and forking."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import tempfile
|
|
import time
|
|
from typing import Any
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
from crewai.agent.core import Agent
|
|
from crewai.agents.agent_builder.base_agent import BaseAgent
|
|
from crewai.crew import Crew
|
|
from crewai.llms.base_llm import BaseLLM
|
|
from crewai.flow.flow import _INITIAL_STATE_CLASS_MARKER, Flow, start
|
|
from crewai.state.checkpoint_config import CheckpointConfig
|
|
from crewai.state.checkpoint_listener import (
|
|
_find_checkpoint,
|
|
_resolve,
|
|
_SENTINEL,
|
|
)
|
|
from crewai.state.provider.json_provider import JsonProvider
|
|
from crewai.state.provider.sqlite_provider import SqliteProvider
|
|
from crewai.state.runtime import RuntimeState
|
|
from crewai.task import Task
|
|
|
|
|
|
|
|
|
|
class TestResolve:
|
|
def test_none_returns_none(self) -> None:
|
|
assert _resolve(None) is None
|
|
|
|
def test_false_returns_sentinel(self) -> None:
|
|
assert _resolve(False) is _SENTINEL
|
|
|
|
def test_true_returns_config(self) -> None:
|
|
result = _resolve(True)
|
|
assert isinstance(result, CheckpointConfig)
|
|
assert result.location == "./.checkpoints"
|
|
|
|
def test_config_returns_config(self) -> None:
|
|
cfg = CheckpointConfig(location="/tmp/cp")
|
|
assert _resolve(cfg) is cfg
|
|
|
|
|
|
|
|
|
|
class TestFindCheckpoint:
|
|
def _make_agent(self, checkpoint: Any = None) -> Agent:
|
|
return Agent(role="r", goal="g", backstory="b", checkpoint=checkpoint)
|
|
|
|
def _make_crew(
|
|
self, agents: list[Agent], checkpoint: Any = None
|
|
) -> Crew:
|
|
crew = Crew(agents=agents, tasks=[], checkpoint=checkpoint)
|
|
for a in agents:
|
|
a.crew = crew
|
|
return crew
|
|
|
|
def test_crew_true(self) -> None:
|
|
a = self._make_agent()
|
|
self._make_crew([a], checkpoint=True)
|
|
cfg = _find_checkpoint(a)
|
|
assert isinstance(cfg, CheckpointConfig)
|
|
|
|
def test_crew_true_agent_false_opts_out(self) -> None:
|
|
a = self._make_agent(checkpoint=False)
|
|
self._make_crew([a], checkpoint=True)
|
|
assert _find_checkpoint(a) is None
|
|
|
|
def test_crew_none_agent_none(self) -> None:
|
|
a = self._make_agent()
|
|
self._make_crew([a])
|
|
assert _find_checkpoint(a) is None
|
|
|
|
def test_agent_config_overrides_crew(self) -> None:
|
|
a = self._make_agent(
|
|
checkpoint=CheckpointConfig(location="/agent_cp")
|
|
)
|
|
self._make_crew([a], checkpoint=True)
|
|
cfg = _find_checkpoint(a)
|
|
assert isinstance(cfg, CheckpointConfig)
|
|
assert cfg.location == "/agent_cp"
|
|
|
|
def test_task_inherits_from_crew(self) -> None:
|
|
a = self._make_agent()
|
|
self._make_crew([a], checkpoint=True)
|
|
task = Task(description="d", expected_output="e", agent=a)
|
|
cfg = _find_checkpoint(task)
|
|
assert isinstance(cfg, CheckpointConfig)
|
|
|
|
def test_task_agent_false_blocks(self) -> None:
|
|
a = self._make_agent(checkpoint=False)
|
|
self._make_crew([a], checkpoint=True)
|
|
task = Task(description="d", expected_output="e", agent=a)
|
|
assert _find_checkpoint(task) is None
|
|
|
|
def test_flow_direct(self) -> None:
|
|
flow = Flow(checkpoint=True)
|
|
cfg = _find_checkpoint(flow)
|
|
assert isinstance(cfg, CheckpointConfig)
|
|
|
|
def test_flow_none(self) -> None:
|
|
flow = Flow()
|
|
assert _find_checkpoint(flow) is None
|
|
|
|
def test_unknown_source(self) -> None:
|
|
assert _find_checkpoint("random") is None
|
|
|
|
|
|
|
|
|
|
class TestPrune:
|
|
def test_prune_keeps_newest(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
branch_dir = os.path.join(d, "main")
|
|
os.makedirs(branch_dir)
|
|
for i in range(5):
|
|
path = os.path.join(branch_dir, f"cp_{i}.json")
|
|
with open(path, "w") as f:
|
|
f.write("{}")
|
|
# Ensure distinct mtime
|
|
time.sleep(0.01)
|
|
|
|
JsonProvider().prune(d, max_keep=2, branch="main")
|
|
remaining = os.listdir(branch_dir)
|
|
assert len(remaining) == 2
|
|
assert "cp_3.json" in remaining
|
|
assert "cp_4.json" in remaining
|
|
|
|
def test_prune_zero_removes_all(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
branch_dir = os.path.join(d, "main")
|
|
os.makedirs(branch_dir)
|
|
for i in range(3):
|
|
with open(os.path.join(branch_dir, f"cp_{i}.json"), "w") as f:
|
|
f.write("{}")
|
|
|
|
JsonProvider().prune(d, max_keep=0, branch="main")
|
|
assert os.listdir(branch_dir) == []
|
|
|
|
def test_prune_more_than_existing(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
branch_dir = os.path.join(d, "main")
|
|
os.makedirs(branch_dir)
|
|
with open(os.path.join(branch_dir, "cp.json"), "w") as f:
|
|
f.write("{}")
|
|
|
|
JsonProvider().prune(d, max_keep=10, branch="main")
|
|
assert len(os.listdir(branch_dir)) == 1
|
|
|
|
|
|
|
|
|
|
class TestCheckpointConfig:
|
|
def test_defaults(self) -> None:
|
|
cfg = CheckpointConfig()
|
|
assert cfg.location == "./.checkpoints"
|
|
assert cfg.on_events == ["task_completed"]
|
|
assert cfg.max_checkpoints is None
|
|
assert not cfg.trigger_all
|
|
|
|
def test_trigger_all(self) -> None:
|
|
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"]
|
|
)
|
|
assert cfg.trigger_events == {"task_completed", "crew_kickoff_completed"}
|
|
|
|
|
|
|
|
|
|
class TestRuntimeStateLineage:
|
|
def _make_state(self) -> RuntimeState:
|
|
from crewai import Agent, Crew
|
|
|
|
agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini")
|
|
crew = Crew(agents=[agent], tasks=[], verbose=False)
|
|
return RuntimeState(root=[crew])
|
|
|
|
def test_default_lineage_fields(self) -> None:
|
|
state = self._make_state()
|
|
assert state._checkpoint_id is None
|
|
assert state._parent_id is None
|
|
assert state._branch == "main"
|
|
|
|
def test_serialize_includes_version(self) -> None:
|
|
from crewai_core.version import get_crewai_version
|
|
|
|
state = self._make_state()
|
|
dumped = json.loads(state.model_dump_json())
|
|
assert dumped["crewai_version"] == get_crewai_version()
|
|
|
|
def test_deserialize_migrates_on_version_mismatch(self, caplog: Any) -> None:
|
|
import logging
|
|
|
|
state = self._make_state()
|
|
raw = state.model_dump_json()
|
|
data = json.loads(raw)
|
|
data["crewai_version"] = "0.1.0"
|
|
with caplog.at_level(logging.DEBUG):
|
|
RuntimeState.model_validate_json(
|
|
json.dumps(data), context={"from_checkpoint": True}
|
|
)
|
|
assert "Migrating checkpoint from crewAI 0.1.0" in caplog.text
|
|
|
|
def test_deserialize_warns_on_missing_version(self, caplog: Any) -> None:
|
|
import logging
|
|
|
|
state = self._make_state()
|
|
raw = state.model_dump_json()
|
|
data = json.loads(raw)
|
|
data.pop("crewai_version", None)
|
|
with caplog.at_level(logging.WARNING):
|
|
RuntimeState.model_validate_json(
|
|
json.dumps(data), context={"from_checkpoint": True}
|
|
)
|
|
assert "treating as 0.0.0" in caplog.text
|
|
|
|
def test_serialize_includes_lineage(self) -> None:
|
|
state = self._make_state()
|
|
state._parent_id = "parent456"
|
|
state._branch = "experiment"
|
|
dumped = json.loads(state.model_dump_json())
|
|
assert dumped["parent_id"] == "parent456"
|
|
assert dumped["branch"] == "experiment"
|
|
assert "checkpoint_id" not in dumped
|
|
|
|
def test_deserialize_restores_lineage(self) -> None:
|
|
state = self._make_state()
|
|
state._parent_id = "parent456"
|
|
state._branch = "experiment"
|
|
raw = state.model_dump_json()
|
|
restored = RuntimeState.model_validate_json(
|
|
raw, context={"from_checkpoint": True}
|
|
)
|
|
assert restored._parent_id == "parent456"
|
|
assert restored._branch == "experiment"
|
|
|
|
def test_deserialize_defaults_missing_lineage(self) -> None:
|
|
state = self._make_state()
|
|
raw = state.model_dump_json()
|
|
data = json.loads(raw)
|
|
data.pop("parent_id", None)
|
|
data.pop("branch", None)
|
|
restored = RuntimeState.model_validate_json(
|
|
json.dumps(data), context={"from_checkpoint": True}
|
|
)
|
|
assert restored._parent_id is None
|
|
assert restored._branch == "main"
|
|
|
|
def test_from_checkpoint_sets_checkpoint_id(self) -> None:
|
|
"""from_checkpoint sets _checkpoint_id from the location, not the blob."""
|
|
state = self._make_state()
|
|
state._provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
loc = state.checkpoint(d)
|
|
written_id = state._checkpoint_id
|
|
|
|
cfg = CheckpointConfig(restore_from=loc)
|
|
restored = RuntimeState.from_checkpoint(
|
|
cfg, context={"from_checkpoint": True}
|
|
)
|
|
assert restored._checkpoint_id == written_id
|
|
assert restored._parent_id == written_id
|
|
|
|
def test_fork_sets_branch(self) -> None:
|
|
state = self._make_state()
|
|
state._checkpoint_id = "abc12345"
|
|
state._parent_id = "abc12345"
|
|
state.fork("my-experiment")
|
|
assert state._branch == "my-experiment"
|
|
assert state._parent_id == "abc12345"
|
|
|
|
def test_fork_auto_branch(self) -> None:
|
|
state = self._make_state()
|
|
state._checkpoint_id = "20260409T120000_abc12345"
|
|
state.fork()
|
|
assert state._branch.startswith("fork/20260409T120000_abc12345_")
|
|
assert len(state._branch) == len("fork/20260409T120000_abc12345_") + 6
|
|
|
|
def test_fork_no_checkpoint_id_unique(self) -> None:
|
|
state = self._make_state()
|
|
state.fork()
|
|
assert state._branch.startswith("fork/")
|
|
assert len(state._branch) == len("fork/") + 8
|
|
# Two forks without checkpoint_id produce different branches
|
|
first = state._branch
|
|
state.fork()
|
|
assert state._branch != first
|
|
|
|
|
|
class TestFlowInitialStateSerialization:
|
|
"""Regression tests for checkpoint serialization of ``Flow.initial_state``."""
|
|
|
|
def test_class_ref_serializes_as_schema(self) -> None:
|
|
class MyState(BaseModel):
|
|
id: str = "x"
|
|
foo: str = "bar"
|
|
|
|
flow = Flow(initial_state=MyState)
|
|
state = RuntimeState(root=[flow])
|
|
dumped = json.loads(state.model_dump_json())
|
|
entity = dumped["entities"][0]
|
|
wrapped = entity["initial_state"]
|
|
assert isinstance(wrapped, dict)
|
|
assert _INITIAL_STATE_CLASS_MARKER in wrapped
|
|
assert wrapped[_INITIAL_STATE_CLASS_MARKER].get("title") == "MyState"
|
|
|
|
def test_class_ref_round_trips_to_basemodel_subclass(self) -> None:
|
|
class MyState(BaseModel):
|
|
id: str = "x"
|
|
foo: str = "bar"
|
|
|
|
flow = Flow(initial_state=MyState)
|
|
raw = RuntimeState(root=[flow]).model_dump_json()
|
|
restored = RuntimeState.model_validate_json(
|
|
raw, context={"from_checkpoint": True}
|
|
)
|
|
rehydrated = restored.root[0].initial_state
|
|
assert isinstance(rehydrated, type)
|
|
assert issubclass(rehydrated, BaseModel)
|
|
assert set(rehydrated.model_fields.keys()) == {"id", "foo"}
|
|
|
|
def test_instance_serializes_as_values(self) -> None:
|
|
class MyState(BaseModel):
|
|
id: str = "x"
|
|
foo: str = "bar"
|
|
|
|
flow = Flow(initial_state=MyState(foo="baz"))
|
|
state = RuntimeState(root=[flow])
|
|
dumped = json.loads(state.model_dump_json())
|
|
entity = dumped["entities"][0]
|
|
assert entity["initial_state"] == {"id": "x", "foo": "baz"}
|
|
|
|
def test_dict_passthrough(self) -> None:
|
|
flow = Flow(initial_state={"id": "x", "foo": "bar"})
|
|
state = RuntimeState(root=[flow])
|
|
dumped = json.loads(state.model_dump_json())
|
|
entity = dumped["entities"][0]
|
|
assert entity["initial_state"] == {"id": "x", "foo": "bar"}
|
|
|
|
def test_dict_round_trips_as_dict(self) -> None:
|
|
flow = Flow(initial_state={"id": "x", "foo": "bar"})
|
|
raw = RuntimeState(root=[flow]).model_dump_json()
|
|
restored = RuntimeState.model_validate_json(
|
|
raw, context={"from_checkpoint": True}
|
|
)
|
|
assert restored.root[0].initial_state == {"id": "x", "foo": "bar"}
|
|
|
|
|
|
|
|
|
|
class TestJsonProviderFork:
|
|
def test_checkpoint_writes_to_branch_subdir(self) -> None:
|
|
provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
path = provider.checkpoint("{}", d, branch="main")
|
|
assert "/main/" in path
|
|
assert path.endswith(".json")
|
|
assert os.path.isfile(path)
|
|
|
|
def test_checkpoint_fork_branch_subdir(self) -> None:
|
|
provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
path = provider.checkpoint("{}", d, branch="fork/exp1")
|
|
assert "/fork/exp1/" in path
|
|
assert os.path.isfile(path)
|
|
|
|
def test_prune_branch_aware(self) -> None:
|
|
provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
for _ in range(3):
|
|
provider.checkpoint("{}", d, branch="main")
|
|
time.sleep(0.01)
|
|
for _ in range(2):
|
|
provider.checkpoint("{}", d, branch="fork/a")
|
|
time.sleep(0.01)
|
|
|
|
# Prune main to 1
|
|
provider.prune(d, max_keep=1, branch="main")
|
|
|
|
main_dir = os.path.join(d, "main")
|
|
fork_dir = os.path.join(d, "fork", "a")
|
|
assert len(os.listdir(main_dir)) == 1
|
|
assert len(os.listdir(fork_dir)) == 2
|
|
|
|
def test_extract_id(self) -> None:
|
|
provider = JsonProvider()
|
|
assert provider.extract_id("/dir/main/20260409T120000_abc12345_p-none.json") == "20260409T120000_abc12345"
|
|
assert provider.extract_id("/dir/main/20260409T120000_abc12345_p-20260409T115900_def67890.json") == "20260409T120000_abc12345"
|
|
|
|
def test_branch_traversal_rejected(self) -> None:
|
|
provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
with pytest.raises(ValueError, match="escapes checkpoint directory"):
|
|
provider.checkpoint("{}", d, branch="../../etc")
|
|
with pytest.raises(ValueError, match="escapes checkpoint directory"):
|
|
provider.prune(d, max_keep=1, branch="../../etc")
|
|
|
|
def test_filename_encodes_parent_id(self) -> None:
|
|
provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
# First checkpoint — no parent
|
|
path1 = provider.checkpoint("{}", d, branch="main")
|
|
assert "_p-none.json" in path1
|
|
|
|
# Second checkpoint — with parent
|
|
id1 = provider.extract_id(path1)
|
|
path2 = provider.checkpoint("{}", d, parent_id=id1, branch="main")
|
|
assert f"_p-{id1}.json" in path2
|
|
|
|
def test_checkpoint_chaining(self) -> None:
|
|
"""RuntimeState.checkpoint() chains parent_id after each write."""
|
|
state = self._make_state()
|
|
state._provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
state.checkpoint(d)
|
|
id1 = state._checkpoint_id
|
|
assert id1 is not None
|
|
assert state._parent_id == id1
|
|
|
|
loc2 = state.checkpoint(d)
|
|
id2 = state._checkpoint_id
|
|
assert id2 is not None
|
|
assert id2 != id1
|
|
assert state._parent_id == id2
|
|
|
|
with open(loc2) as f:
|
|
data2 = json.loads(f.read())
|
|
assert data2["parent_id"] == id1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acheckpoint_chaining(self) -> None:
|
|
"""Async checkpoint path chains lineage identically to sync."""
|
|
state = self._make_state()
|
|
state._provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
await state.acheckpoint(d)
|
|
id1 = state._checkpoint_id
|
|
assert id1 is not None
|
|
|
|
loc2 = await state.acheckpoint(d)
|
|
id2 = state._checkpoint_id
|
|
assert id2 != id1
|
|
assert state._parent_id == id2
|
|
|
|
with open(loc2) as f:
|
|
data2 = json.loads(f.read())
|
|
assert data2["parent_id"] == id1
|
|
|
|
def _make_state(self) -> RuntimeState:
|
|
from crewai import Agent, Crew
|
|
|
|
agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini")
|
|
crew = Crew(agents=[agent], tasks=[], verbose=False)
|
|
return RuntimeState(root=[crew])
|
|
|
|
|
|
|
|
|
|
class TestSqliteProviderFork:
|
|
def test_checkpoint_stores_branch_and_parent(self) -> None:
|
|
provider = SqliteProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
db = os.path.join(d, "cp.db")
|
|
loc = provider.checkpoint("{}", db, parent_id="p1", branch="exp")
|
|
cid = provider.extract_id(loc)
|
|
|
|
with sqlite3.connect(db) as conn:
|
|
row = conn.execute(
|
|
"SELECT parent_id, branch FROM checkpoints WHERE id = ?",
|
|
(cid,),
|
|
).fetchone()
|
|
assert row == ("p1", "exp")
|
|
|
|
def test_prune_branch_aware(self) -> None:
|
|
provider = SqliteProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
db = os.path.join(d, "cp.db")
|
|
for _ in range(3):
|
|
provider.checkpoint("{}", db, branch="main")
|
|
for _ in range(2):
|
|
provider.checkpoint("{}", db, branch="fork/a")
|
|
|
|
provider.prune(db, max_keep=1, branch="main")
|
|
|
|
with sqlite3.connect(db) as conn:
|
|
main_count = conn.execute(
|
|
"SELECT COUNT(*) FROM checkpoints WHERE branch = 'main'"
|
|
).fetchone()[0]
|
|
fork_count = conn.execute(
|
|
"SELECT COUNT(*) FROM checkpoints WHERE branch = 'fork/a'"
|
|
).fetchone()[0]
|
|
assert main_count == 1
|
|
assert fork_count == 2
|
|
|
|
def test_extract_id(self) -> None:
|
|
provider = SqliteProvider()
|
|
assert provider.extract_id("/path/to/db#abc123") == "abc123"
|
|
|
|
def test_checkpoint_chaining_sqlite(self) -> None:
|
|
state = self._make_state()
|
|
state._provider = SqliteProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
db = os.path.join(d, "cp.db")
|
|
state.checkpoint(db)
|
|
id1 = state._checkpoint_id
|
|
|
|
state.checkpoint(db)
|
|
id2 = state._checkpoint_id
|
|
assert id2 != id1
|
|
|
|
with sqlite3.connect(db) as conn:
|
|
row = conn.execute(
|
|
"SELECT parent_id FROM checkpoints WHERE id = ?", (id2,)
|
|
).fetchone()
|
|
assert row[0] == id1
|
|
|
|
def _make_state(self) -> RuntimeState:
|
|
from crewai import Agent, Crew
|
|
|
|
agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini")
|
|
crew = Crew(agents=[agent], tasks=[], verbose=False)
|
|
return RuntimeState(root=[crew])
|
|
|
|
|
|
|
|
|
|
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_agent_kickoff_delegates_to_from_checkpoint(self) -> None:
|
|
mock_restored = MagicMock(spec=Agent)
|
|
mock_restored.kickoff.return_value = "agent_result"
|
|
|
|
cfg = CheckpointConfig(restore_from="/path/to/agent_cp.json")
|
|
with patch.object(Agent, "from_checkpoint", return_value=mock_restored):
|
|
agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini")
|
|
result = agent.kickoff(messages="hello", from_checkpoint=cfg)
|
|
|
|
mock_restored.kickoff.assert_called_once_with(
|
|
messages="hello", response_format=None, input_files=None
|
|
)
|
|
assert mock_restored.checkpoint.restore_from is None
|
|
assert result == "agent_result"
|
|
|
|
def test_agent_kickoff_config_only_sets_checkpoint(self) -> None:
|
|
cfg = CheckpointConfig(on_events=["lite_agent_execution_completed"])
|
|
agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini")
|
|
assert agent.checkpoint is None
|
|
with patch.object(Agent, "_prepare_kickoff", side_effect=RuntimeError("stop")):
|
|
with pytest.raises(RuntimeError, match="stop"):
|
|
agent.kickoff(messages="hello", from_checkpoint=cfg)
|
|
assert isinstance(agent.checkpoint, CheckpointConfig)
|
|
assert agent.checkpoint.on_events == ["lite_agent_execution_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"
|
|
|
|
|
|
|
|
|
|
class TestLegacyMethodOutputsRestore:
|
|
def test_restore_wraps_legacy_plain_value_outputs(self) -> None:
|
|
flow = Flow()
|
|
flow._method_outputs = ["first", "second"]
|
|
state = RuntimeState(root=[flow])
|
|
state._provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
loc = state.checkpoint(d)
|
|
cfg = CheckpointConfig(restore_from=loc)
|
|
restored = Flow.from_checkpoint(cfg)
|
|
|
|
assert restored.method_outputs == ["first", "second"]
|
|
|
|
def test_restore_legacy_outputs_evaluates_expressions(self) -> None:
|
|
from crewai.flow.runtime._expressions import _expression_context
|
|
|
|
flow = Flow()
|
|
flow._method_outputs = ["legacy"]
|
|
state = RuntimeState(root=[flow])
|
|
state._provider = JsonProvider()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
loc = state.checkpoint(d)
|
|
cfg = CheckpointConfig(restore_from=loc)
|
|
restored = Flow.from_checkpoint(cfg)
|
|
|
|
context = _expression_context(restored)
|
|
assert context["outputs"] == {"": "legacy"}
|
|
|
|
def test_raw_legacy_outputs_remain_readable(self) -> None:
|
|
from crewai.flow.runtime._expressions import _expression_context
|
|
|
|
flow = Flow()
|
|
flow._method_outputs = ["legacy"]
|
|
|
|
assert flow.method_outputs == ["legacy"]
|
|
assert _expression_context(flow)["outputs"] == {"": "legacy"}
|
|
|
|
|
|
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"
|
|
|
|
|
|
class _FinalAnswerLLM(BaseLLM):
|
|
"""Stub LLM that always returns a final answer without any API calls."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__(model="stub")
|
|
|
|
def call(
|
|
self,
|
|
messages,
|
|
tools=None,
|
|
callbacks=None,
|
|
available_functions=None,
|
|
from_task=None,
|
|
from_agent=None,
|
|
response_model=None,
|
|
):
|
|
return "Final Answer: done."
|
|
|
|
def supports_function_calling(self) -> bool:
|
|
return False
|
|
|
|
def supports_stop_words(self) -> bool:
|
|
return False
|
|
|
|
def get_context_window_size(self) -> int:
|
|
return 4096
|
|
|
|
async def acall(self, *args, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
|
|
class TestCheckpointReusedExecutor:
|
|
"""Checkpoint serialization stamps every live Flow's completed methods.
|
|
|
|
The agent executor is a Flow reused across a crew's tasks, so the stamp
|
|
must not be read back as a restore signal on the next task — otherwise the
|
|
second task replays as a resume and never reaches a final answer.
|
|
"""
|
|
|
|
def test_second_task_runs_with_checkpointing_enabled(self) -> None:
|
|
agent = Agent(role="r", goal="g", backstory="b", llm=_FinalAnswerLLM())
|
|
task1 = Task(description="first", expected_output="x", agent=agent)
|
|
task2 = Task(description="second", expected_output="y", agent=agent)
|
|
with tempfile.TemporaryDirectory() as d:
|
|
crew = Crew(
|
|
agents=[agent],
|
|
tasks=[task1, task2],
|
|
verbose=False,
|
|
checkpoint=CheckpointConfig(
|
|
provider=JsonProvider(location=d),
|
|
on_events=["task_started", "task_completed"],
|
|
),
|
|
)
|
|
result = crew.kickoff()
|
|
|
|
assert len(result.tasks_output) == 2
|
|
assert result.tasks_output[1].raw
|
|
|
|
|
|
class TestCustomLLMCheckpointRestore:
|
|
"""A custom BaseLLM subclass serializes with the inherited llm_type "base".
|
|
|
|
Restoring it must not try to instantiate the abstract BaseLLM; it is rebuilt
|
|
as a concrete LLM from the saved config instead.
|
|
"""
|
|
|
|
def test_restore_does_not_instantiate_abstract_base_llm(self) -> None:
|
|
agent = Agent(role="r", goal="g", backstory="b", llm=_FinalAnswerLLM())
|
|
task = Task(description="d", expected_output="e", agent=agent)
|
|
crew = Crew(agents=[agent], tasks=[task], verbose=False)
|
|
|
|
raw = RuntimeState(root=[crew]).model_dump_json()
|
|
restored = RuntimeState.model_validate_json(
|
|
raw, context={"from_checkpoint": True}
|
|
)
|
|
|
|
llm = restored.root[0].agents[0].llm
|
|
assert isinstance(llm, BaseLLM)
|
|
assert not inspect.isabstract(type(llm))
|
|
assert llm.model == "stub"
|