From e74c4dd5d62f9404b5f731a822ca20329dbcef22 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 11 Mar 2025 08:44:54 +0000
Subject: [PATCH] Fix KeyError when specifying function_calling_llm in
 agents.yaml

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/project/crew_base.py |  9 ++--
 tests/yaml_config_test.py       | 79 +++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100644 tests/yaml_config_test.py

diff --git a/src/crewai/project/crew_base.py b/src/crewai/project/crew_base.py
index 53d3d5f3c..325905add 100644
--- a/src/crewai/project/crew_base.py
+++ b/src/crewai/project/crew_base.py
@@ -172,9 +172,12 @@ def CrewBase(cls: T) -> T:
                 ]
 
             if function_calling_llm := agent_info.get("function_calling_llm"):
-                self.agents_config[agent_name]["function_calling_llm"] = agents[
-                    function_calling_llm
-                ]()
+                try:
+                    self.agents_config[agent_name]["function_calling_llm"] = agents[
+                        function_calling_llm
+                    ]()
+                except KeyError:
+                    self.agents_config[agent_name]["function_calling_llm"] = function_calling_llm
 
             if step_callback := agent_info.get("step_callback"):
                 self.agents_config[agent_name]["step_callback"] = callbacks[
diff --git a/tests/yaml_config_test.py b/tests/yaml_config_test.py
new file mode 100644
index 000000000..9a71e4609
--- /dev/null
+++ b/tests/yaml_config_test.py
@@ -0,0 +1,79 @@
+import os
+import tempfile
+import yaml
+import pytest
+from crewai.project import CrewBase, agent, crew, task, tool
+from crewai import Agent, Crew, Task, Process
+from crewai.tools import BaseTool
+from pydantic import BaseModel, Field
+from typing import Type
+
+def test_function_calling_llm_in_yaml():
+    """Test that function_calling_llm can be specified in YAML."""
+    # Create temporary YAML files
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Create agents.yaml with function_calling_llm
+        agents_yaml = os.path.join(temp_dir, "agents.yaml")
+        with open(agents_yaml, "w") as f:
+            yaml.dump(
+                {
+                    "test_agent": {
+                        "role": "Test Agent",
+                        "goal": "Test Goal",
+                        "backstory": "Test Backstory",
+                        "function_calling_llm": "gpt-4o-mini"
+                    }
+                },
+                f
+            )
+        
+        # Create tasks.yaml
+        tasks_yaml = os.path.join(temp_dir, "tasks.yaml")
+        with open(tasks_yaml, "w") as f:
+            yaml.dump(
+                {
+                    "test_task": {
+                        "description": "Test Task",
+                        "expected_output": "Test Output",
+                        "agent": "test_agent"
+                    }
+                },
+                f
+            )
+        
+        # Create a CrewBase class that uses the YAML files
+        @CrewBase
+        class TestCrew:
+            """Test crew with function_calling_llm in YAML."""
+            agents_config = agents_yaml
+            tasks_config = tasks_yaml
+            
+            @agent
+            def test_agent(self) -> Agent:
+                return Agent(
+                    config=self.agents_config["test_agent"],
+                    verbose=True
+                )
+            
+            @task
+            def test_task(self) -> Task:
+                return Task(
+                    config=self.tasks_config["test_task"]
+                )
+            
+            @crew
+            def crew(self) -> Crew:
+                return Crew(
+                    agents=self.agents,
+                    tasks=self.tasks,
+                    process=Process.sequential,
+                    verbose=True
+                )
+        
+        # Initialize the crew - this should not raise a KeyError
+        test_crew = TestCrew()
+        crew_instance = test_crew.crew()
+        
+        # Verify that function_calling_llm was properly set
+        assert crew_instance.agents[0].function_calling_llm is not None
+        assert crew_instance.agents[0].function_calling_llm.model == "gpt-4o-mini"