From f68fd9e8503868dc6a0bbbcd08ab2404a82b3c5f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Moura?= Date: Mon, 24 Aug 2026 13:39:46 -0300 Subject: [PATCH] feat(flow): let a chat flow declare its own state shape (#7061) * feat(flow): let a chat flow declare its own state shape A conversational declaration could only use `state: {type: pydantic, ref: ...}` pointing at a `ConversationState` subclass. Every other shape loaded clean and then died on the first turn -- inline `json_schema` and a non-subclass ref with `AttributeError: 'StateWithId' object has no attribute 'messages'`, and `type: dict` with `AttributeError: 'dict' object has no attribute 'id'`. `Flow._compose_extension_state_model` is a new runtime extension seam -- the seventh alongside the existing six -- applied to the model built from `state:` before the engine wraps it for `id`. The conversational mixin uses it to add the chat fields to whatever the declaration asked for, so declared fields and defaults survive; a model that already extends `ConversationState` is returned untouched, so today's supported shape is a no-op. `dict` and `unknown` state cannot carry those fields at all, so the default extension state supplies the real shape (seeded from the declared defaults where they fit) rather than forbidding it. Raising instead would break construction, and `Flow[dict]` with `conversational = True` constructs today -- `crewai flow plot` and definition-only consumers would stop working. Co-Authored-By: Claude Opus 5 (1M context) * fix(flow): keep declared defaults and cover an unbuildable state model Two review follow-ups on the declared-state work. A `dict` state's defaults are arbitrary keys, and the fallback kept only the ones matching `ConversationState`, so `{"type": "dict", "default": {"topic": "ai"}}` lost `topic` before the first turn and an action reading `state.topic` would fail. They are carried as extras now. And a declared `pydantic`/`json_schema` state whose model cannot be built -- a bad ref, an invalid schema -- fell through to a plain dict with none of the chat fields, so the turn died on `state.id` instead. The engine now re-asks the extension in that case, as if nothing had been declared. Found by Cursor and CodeRabbit on #7061. * refactor(flows): drop the unreachable extension-state fallback The _initial_state_t branch sat after an unconditional return. The state_definition is None case and _conversation_state_with_defaults now cover every path that used to reach it. Found by Cursor on #7061. --------- Co-authored-by: Claude Opus 5 (1M context) --- .../experimental/conversational_mixin.py | 89 ++++++++++-- .../src/crewai/flow/runtime/__init__.py | 35 ++++- lib/crewai/tests/test_flow_conversation.py | 135 +++++++++++++++++- 3 files changed, 241 insertions(+), 18 deletions(-) diff --git a/lib/crewai/src/crewai/experimental/conversational_mixin.py b/lib/crewai/src/crewai/experimental/conversational_mixin.py index 35422701c..5cac956d7 100644 --- a/lib/crewai/src/crewai/experimental/conversational_mixin.py +++ b/lib/crewai/src/crewai/experimental/conversational_mixin.py @@ -23,9 +23,9 @@ from contextlib import contextmanager from enum import Enum import json import logging -from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast -from pydantic import BaseModel, Field, create_model +from pydantic import BaseModel, ConfigDict, Field, create_model from crewai.events.event_bus import crewai_event_bus from crewai.events.types.flow_events import ( @@ -140,6 +140,37 @@ def _config_from_definition( ) +def _conversation_state_with_defaults( + defaults: dict[str, Any] | None, +) -> ConversationState: + """Build the chat state, keeping declared defaults the shape cannot hold. + + A ``dict`` state's defaults are arbitrary keys, and dropping the ones that + are not conversational fields would break an action reading + ``state.topic``. They are carried as extras instead. + """ + if not defaults: + return ConversationState() + + known = { + name: value + for name, value in defaults.items() + if name in ConversationState.model_fields + } + extra = { + name: value + for name, value in defaults.items() + if name not in ConversationState.model_fields + } + if not extra: + return ConversationState(**known) + + class ConversationStateWithDeclaredDefaults(ConversationState): + model_config = ConfigDict(extra="allow") + + return ConversationStateWithDeclaredDefaults(**known, **extra) + + def _builtin_method(handler: Callable[..., Any], **roles: Any) -> FlowMethodDefinition: """One built-in conversational method, as the code ref the DSL already emits.""" return FlowMethodDefinition(do=_method_action(handler), **roles) @@ -997,23 +1028,53 @@ class _ConversationalMixin: update={"methods": {**definition.methods, **missing}} ) - def _create_default_extension_state(self) -> ConversationState | None: - """Supply ``ConversationState`` only when nothing else declares state. + def _compose_extension_state_model( + self, model_class: type[BaseModel] + ) -> type[BaseModel]: + """Add the conversational fields to a declared state model. - A declared ``state:`` block always wins. This hook is consulted before - ``_create_definition_state``, so returning a state here would discard - every field the declaration asked for. + A turn reads ``messages``, ``current_user_message``, ``last_intent``, + ``ended``, ``events`` and ``agent_threads`` off state, so a declared + model that lacks them fails on the first turn. Composing rather than + replacing keeps every field the declaration asked for. A model that + already extends ``ConversationState`` is returned untouched. + """ + if not self._is_conversational_enabled(): + return model_class + if issubclass(model_class, ConversationState): + return model_class + + class ConversationalState(ConversationState, model_class): # type: ignore[misc, valid-type] + pass + + return ConversationalState + + def _create_default_extension_state( + self, *, ignore_declared_state: bool = False + ) -> ConversationState | None: + """Supply ``ConversationState`` when the declaration cannot carry it. + + A declared ``pydantic`` or ``json_schema`` state is composed with the + conversational fields instead (see + ``_compose_extension_state_model``), so it keeps everything it asked + for. ``dict`` and ``unknown`` state cannot carry them at all, so this + supplies the real shape rather than letting the turn die on + ``state.id`` -- seeded from the declared defaults where they fit. """ if not self._is_conversational_enabled(): return None - if self._conversation_flow_definition().state is not None: - return None - initial_state_t = getattr(self, "_initial_state_t", None) - if not hasattr(self, "_initial_state_t") or isinstance( - initial_state_t, TypeVar - ): + + state_definition = self._conversation_flow_definition().state + if state_definition is None: return ConversationState() - return None + if not ignore_declared_state and state_definition.type in ( + "pydantic", + "json_schema", + ): + # Composed with the declared model instead; see + # ``_compose_extension_state_model``. + return None + return _conversation_state_with_defaults(state_definition.default) def _should_apply_pending_kickoff_context(self) -> bool: return ( diff --git a/lib/crewai/src/crewai/flow/runtime/__init__.py b/lib/crewai/src/crewai/flow/runtime/__init__.py index 0e95d971d..07e53de19 100644 --- a/lib/crewai/src/crewai/flow/runtime/__init__.py +++ b/lib/crewai/src/crewai/flow/runtime/__init__.py @@ -185,6 +185,7 @@ def _condition_satisfied(condition: FlowDefinitionCondition, events: set[str]) - def _build_definition_state_model( state_definition: FlowStateDefinition, + compose: Callable[[type[BaseModel]], type[BaseModel]] | None = None, ) -> BaseModel | None: kwargs = dict(state_definition.default or {}) @@ -216,6 +217,9 @@ def _build_definition_state_model( if model_class is None: return None + if compose is not None: + model_class = compose(model_class) + if not issubclass(model_class, FlowState): class StateWithId(FlowState, model_class): # type: ignore[misc, valid-type] @@ -458,10 +462,28 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): """ return definition - def _create_default_extension_state(self) -> Any | None: - """Return a default state supplied by an optional runtime extension.""" + def _create_default_extension_state( + self, *, ignore_declared_state: bool = False + ) -> Any | None: + """Return a default state supplied by an optional runtime extension. + + ``ignore_declared_state`` is set when a declared ``state:`` block named + a model that could not be built, so the extension is asked again as if + nothing had been declared. + """ return None + def _compose_extension_state_model( + self, model_class: type[BaseModel] + ) -> type[BaseModel]: + """Let an optional runtime extension add bases to a declared state model. + + Applied to the model built from ``state:`` before the engine wraps it + for its ``id`` field, so an extension can require its own fields + alongside whatever the declaration asked for. + """ + return model_class + def _should_apply_pending_kickoff_context(self) -> bool: """Whether an optional runtime extension has pending kickoff context.""" return False @@ -1721,7 +1743,9 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): if state_definition is None: return {"id": str(uuid4())} if state_definition.type in ("pydantic", "json_schema"): - state = _build_definition_state_model(state_definition) + state = _build_definition_state_model( + state_definition, compose=self._compose_extension_state_model + ) if state is not None: return state logger.error( @@ -1730,6 +1754,11 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): self._definition.name, state_definition.type, ) + extension_state: dict[str, Any] | BaseModel | None = ( + self._create_default_extension_state(ignore_declared_state=True) + ) + if extension_state is not None: + return extension_state elif state_definition.type == "unknown": logger.warning( "Flow %r declares state of unknown type; falling back to dict state", diff --git a/lib/crewai/tests/test_flow_conversation.py b/lib/crewai/tests/test_flow_conversation.py index 87d0836ac..6841ecade 100644 --- a/lib/crewai/tests/test_flow_conversation.py +++ b/lib/crewai/tests/test_flow_conversation.py @@ -2420,6 +2420,140 @@ class TestDeclarativeConversationalFlow: assert "cannot carry a model class" in caplog.text +class DeclaredSchemaChatState(ConversationState): + """A ConversationState subclass, i.e. the already-supported state shape.""" + + ticket_id: str | None = None + + +class TestDeclaredConversationalState: + """Any declared state shape works for a chat flow, keeping its own fields.""" + + @staticmethod + def _flow(state: dict[str, Any] | None) -> Flow[Any]: + declaration = _conversational_declaration(conversational={}) + if state is None: + declaration.pop("state", None) + else: + declaration["state"] = state + return Flow.from_declaration(contents=declaration) + + CONVERSATIONAL_FIELDS = ( + "messages", + "current_user_message", + "last_intent", + "ended", + "events", + "agent_threads", + ) + + def test_inline_json_schema_state_gains_the_conversational_fields(self) -> None: + flow = self._flow( + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "ticket_id": {"type": "string"}, + "turns": {"type": "integer"}, + }, + }, + "default": {"ticket_id": "T-1"}, + } + ) + fields = type(flow.state).model_fields + + for name in self.CONVERSATIONAL_FIELDS: + assert name in fields, name + assert "ticket_id" in fields and "turns" in fields + assert flow.state.ticket_id == "T-1" + assert flow.state.id + + def test_a_pydantic_ref_that_is_not_a_conversation_state_is_composed(self) -> None: + flow = self._flow({"type": "pydantic", "ref": "crewai.flow:ChatState"}) + fields = type(flow.state).model_fields + + for name in self.CONVERSATIONAL_FIELDS: + assert name in fields, name + assert "session_ready" in fields + + def test_a_conversation_state_subclass_is_not_composed_twice(self) -> None: + flow = self._flow( + {"type": "pydantic", "ref": f"{__name__}:DeclaredSchemaChatState"} + ) + mro = [cls.__name__ for cls in type(flow.state).__mro__] + + assert mro.count("ConversationState") == 1 + assert "ticket_id" in type(flow.state).model_fields + + def test_dict_state_is_given_the_real_shape(self) -> None: + """A dict cannot carry the fields, so supply them rather than die.""" + flow = self._flow({"type": "dict", "default": {"last_intent": "order"}}) + + assert isinstance(flow.state, ConversationState) + assert flow.state.last_intent == "order" + + def test_dict_state_keeps_declared_defaults_it_does_not_own(self) -> None: + """Dropping them would break an action reading `state.topic`.""" + flow = self._flow({"type": "dict", "default": {"topic": "ai", "limit": 3}}) + + assert flow.state.topic == "ai" + assert flow.state.limit == 3 + assert isinstance(flow.state, ConversationState) + + def test_unknown_state_is_treated_like_dict(self) -> None: + flow = self._flow({"type": "unknown", "ref": "x:Y", "default": {"topic": "ai"}}) + + assert flow.state.topic == "ai" + assert isinstance(flow.state, ConversationState) + + def test_an_unbuildable_declared_model_still_gets_the_chat_shape(self) -> None: + """A bad ref fell back to a plain dict, so the turn died on `state.id`.""" + flow = self._flow({"type": "pydantic", "ref": "no.such.module:Nope"}) + + assert isinstance(flow.state, ConversationState) + assert flow.state.id + + def test_declared_state_still_runs_a_turn(self) -> None: + flow = self._flow( + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": {"ticket_id": {"type": "string"}}, + }, + } + ) + flow._conversation_config.llm = _ScriptedLLM(["Hello."]) + + assert flow.handle_turn("hi") == "Hello." + assert [m.role for m in flow.state.messages] == ["user", "assistant"] + + def test_non_conversational_declaration_keeps_its_plain_state(self) -> None: + flow = Flow.from_declaration( + contents={ + "schema": "crewai.flow/v1", + "name": "Plain", + "state": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": {"topic": {"type": "string"}}, + }, + "default": {"topic": "ai"}, + }, + "methods": { + "begin": { + "do": {"call": "expression", "expr": "state.topic"}, + "start": True, + } + }, + } + ) + + assert "messages" not in type(flow.state).model_fields + assert flow.state.topic == "ai" + class TestDeclaredConversationalLLM: """A conversational block accepts the shapes a crew agent's `llm` accepts.""" @@ -2525,7 +2659,6 @@ class TestDeclaredConversationalLLM: assert declared in seen assert [m.role for m in flow.state.messages][0] == "user" - class TestRoutingArtefactLabels: """A route label echoed by a handler is a routing artefact, not a reply."""