diff --git a/lib/crewai/src/crewai/llm_overlay.py b/lib/crewai/src/crewai/llm_overlay.py index 1397052c7..8cb520d6a 100644 --- a/lib/crewai/src/crewai/llm_overlay.py +++ b/lib/crewai/src/crewai/llm_overlay.py @@ -6,7 +6,9 @@ an ``Agent`` or ``LiteAgent`` whose ``role`` is a key of the mapping is built with the mapped model instead of its declared ``llm``. Roles that are not keys, and every agent built outside the block, keep their own model. -Roles are matched exactly, by the text they have when the overlay is read. +Roles are matched exactly, by the text they have when the overlay is read; +only whitespace around a role or a key is ignored, so a role a YAML file +leaves with a trailing newline matches a key written for the clean text. An ``Agent`` is read when it is built, with its declared role, and read again when ``crew.kickoff(inputs=...)`` interpolates the inputs into that role and the text changes. The second read is what lets a role declared as @@ -71,22 +73,53 @@ def llm_overlay(mapping: dict[str, str] | None) -> Iterator[None]: Args: mapping: ``{role: model}`` for the block; ``None`` clears any active - overlay for the block. The previous value is always restored on - exit, including when the block raises. + overlay for the block. Whitespace around each role is dropped from + the copy the block uses, so a key matches a role that differs from + it only by its surrounding whitespace; the mapping passed in is + left as it is. The previous value is always restored on exit, + including when the block raises. """ - token = active.set(mapping) + token = active.set(_stripped(mapping)) try: yield finally: active.reset(token) -def overlay_model_for(role: str) -> str | None: +def overlay_model_for(role: str | None) -> str | None: """The model the active overlay assigns to ``role``. + Whitespace around ``role`` is ignored, as it was around the keys when the + overlay was set. An empty or ``None`` role matches nothing. + Returns: The mapped model string, or ``None`` when no overlay is active or ``role`` is not one of its keys. """ mapping = active.get() - return mapping.get(role) if mapping else None + if not mapping or role is None: + return None + role = role.strip() + return mapping.get(role) if role else None + + +def _stripped(mapping: dict[str, str] | None) -> dict[str, str] | None: + """A copy of ``mapping`` with the whitespace around each role dropped. + + Two keys that differ only by whitespace name one role. When they name the + same model the copy holds it once; when they name different models the + mapping is ambiguous and is refused, so the model an agent runs on never + depends on dictionary order. + """ + if mapping is None: + return None + stripped: dict[str, str] = {} + for role, model in mapping.items(): + key = role.strip() + if key in stripped and stripped[key] != model: + raise ValueError( + f"llm_overlay: role {key!r} is mapped twice with different models " + f"({stripped[key]!r} and {model!r}); give each role one model" + ) + stripped[key] = model + return stripped diff --git a/lib/crewai/tests/test_llm_overlay.py b/lib/crewai/tests/test_llm_overlay.py index 9a91a15a7..0e329bd3d 100644 --- a/lib/crewai/tests/test_llm_overlay.py +++ b/lib/crewai/tests/test_llm_overlay.py @@ -85,6 +85,65 @@ def test_nested_overlay_restores_the_outer_one() -> None: assert overlay_model_for("Researcher") == "openai/gpt-4o" +def test_whitespace_around_a_role_or_a_key_is_ignored() -> None: + """A role read from a YAML file often ends in a newline (``role: >`` folds + to one) while the caller writes the key for the clean text; the newline can + just as well land on the key. Either side is stripped, on the direct lookup + and on the read an agent does when it is built.""" + with llm_overlay(OVERLAY): + assert overlay_model_for("Researcher\n") == "openai/gpt-4o" + assert overlay_model_for(" Researcher ") == "openai/gpt-4o" + assert _agent("Researcher\n").llm.model == "gpt-4o" + + with llm_overlay({"Researcher\n": "openai/gpt-4o"}): + assert overlay_model_for("Researcher") == "openai/gpt-4o" + assert _agent("Researcher").llm.model == "gpt-4o" + + +def test_only_the_whitespace_around_the_text_is_forgiven() -> None: + """The text in between is still matched exactly, and nothing matches an + empty role.""" + with llm_overlay(OVERLAY): + assert overlay_model_for("researcher") is None + assert overlay_model_for("Re searcher") is None + assert overlay_model_for("Writer\n") is None + assert overlay_model_for("") is None + assert overlay_model_for(" \n") is None + assert overlay_model_for(None) is None + + +def test_two_keys_that_are_one_role_with_different_models_are_refused() -> None: + """``"Researcher"`` and ``" Researcher "`` name one role. With one model + they collapse to it; with two the mapping is ambiguous and refused, so the + model never depends on dictionary order.""" + with llm_overlay({"Researcher": "openai/gpt-4o", " Researcher ": "openai/gpt-4o"}): + assert overlay_model_for("Researcher") == "openai/gpt-4o" + + with pytest.raises(ValueError, match="'Researcher' is mapped twice with different models"): + with llm_overlay({"Researcher": "openai/gpt-4o", " Researcher ": "openai/gpt-4o-mini"}): + pass # never entered + assert active.get() is None # nothing was set + + +def test_the_mapping_the_caller_passed_is_not_mutated() -> None: + mapping = {"Researcher\n": "openai/gpt-4o"} + with llm_overlay(mapping): + assert active.get() == {"Researcher": "openai/gpt-4o"} + assert mapping == {"Researcher\n": "openai/gpt-4o"} + + +def test_a_yaml_folded_template_role_matches_after_interpolation() -> None: + """The CrewBase shape: a templated role from YAML keeps its trailing newline + through interpolation, and the key is written for the clean text.""" + with llm_overlay(TEMPLATE_OVERLAY): + agent = _agent("Researcher for {repo}\n") + assert agent.llm.model == "gpt-4o-mini" + agent.interpolate_inputs({"repo": "crewAIInc/x"}) + + assert agent.role == "Researcher for crewAIInc/x\n" + assert agent.llm.model == "gpt-4o" + + @pytest.mark.filterwarnings("ignore:LiteAgent is deprecated") def test_lite_agent_gets_the_overlay_model() -> None: with llm_overlay(OVERLAY):