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."""