Merge branch 'main' into danielfsbarreto/mongodb-persist

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

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):