Repository: crewAIInc/crewAI
Pull request #2029 (2025-02-04 17:41:27)

Brandon/improve llm structured output


Patch:
diff --git a/docs/concepts/llms.mdx b/docs/concepts/llms.mdx
index 309ae892..b6a3ac32 100644
--- a/docs/concepts/llms.mdx
+++ b/docs/concepts/llms.mdx
@@ -720,6 +720,30 @@ Learn how to get the most out of your LLM configuration:
   </Accordion>
 </AccordionGroup>
 
+## Structured LLM Calls
+
+CrewAI supports structured responses from LLM calls by allowing you to define a `response_format` using a Pydantic model. This enables the framework to automatically parse and validate the output, making it easier to integrate the response into your application without manual post-processing.
+
+For example, you can define a Pydantic model to represent the expected response structure and pass it as the `response_format` when instantiating the LLM. The model will then be used to convert the LLM output into a structured Python object.
+
+```python Code
+from crewai import LLM
+
+class Dog(BaseModel):
+    name: str
+    age: int
+    breed: str
+
+
+llm = LLM(model="gpt-4o", response_format=Dog)
+
+response = llm.call(
+    "Analyze the following messages and return the name, age, and breed. "
+    "Meet Kona! She is 3 years old and is a black german shepherd."
+)
+print(response)
+```
+
 ## Common Issues and Solutions
 
 <Tabs>
diff --git a/src/crewai/llm.py b/src/crewai/llm.py
index 1e52773a..d8f2be23 100644
--- a/src/crewai/llm.py
+++ b/src/crewai/llm.py
@@ -5,15 +5,17 @@ import sys
 import threading
 import warnings
 from contextlib import contextmanager
-from typing import Any, Dict, List, Literal, Optional, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Type, Union, cast
 
 from dotenv import load_dotenv
+from pydantic import BaseModel
 
 with warnings.catch_warnings():
     warnings.simplefilter("ignore", UserWarning)
     import litellm
     from litellm import Choices, get_supported_openai_params
     from litellm.types.utils import ModelResponse
+    from litellm.utils import supports_response_schema
 
 
 from crewai.utilities.exceptions.context_window_exceeding_exception import (
@@ -128,7 +130,7 @@ class LLM:
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         logit_bias: Optional[Dict[int, float]] = None,
-        response_format: Optional[Dict[str, Any]] = None,
+        response_format: Optional[Type[BaseModel]] = None,
         seed: Optional[int] = None,
         logprobs: Optional[int] = None,
         top_logprobs: Optional[int] = None,
@@ -213,6 +215,9 @@ class LLM:
         response = llm.call(messages)
         print(response)
         """
+        # Validate parameters before proceeding with the call.
+        self._validate_call_params()
+
         if isinstance(messages, str):
             messages = [{"role": "user", "content": messages}]
 
@@ -312,6 +317,36 @@ class LLM:
                     logging.error(f"LiteLLM call failed: {str(e)}")
                 raise
 
+    def _get_custom_llm_provider(self) -> str:
+        """
+        Derives the custom_llm_provider from the model string.
+        - For example, if the model is "openrouter/deepseek/deepseek-chat", returns "openrouter".
+        - If the model is "gemini/gemini-1.5-pro", returns "gemini".
+        - If there is no '/', defaults to "openai".
+        """
+        if "/" in self.model:
+            return self.model.split("/")[0]
+        return "openai"
+
+    def _validate_call_params(self) -> None:
+        """
+        Validate parameters before making a call. Currently this only checks if
+        a response_format is provided and whether the model supports it.
+        The custom_llm_provider is dynamically determined from the model:
+          - E.g., "openrouter/deepseek/deepseek-chat" yields "openrouter"
+          - "gemini/gemini-1.5-pro" yields "gemini"
+          - If no slash is present, "openai" is assumed.
+        """
+        provider = self._get_custom_llm_provider()
+        if self.response_format is not None and not supports_response_schema(
+            model=self.model,
+            custom_llm_provider=provider,
+        ):
+            raise ValueError(
+                f"The model {self.model} does not support response_format for provider '{provider}'. "
+                "Please remove response_format or use a supported model."
+            )
+
     def supports_function_calling(self) -> bool:
         try:
             params = get_supported_openai_params(model=self.model)


Test patch:
diff --git a/tests/llm_test.py b/tests/llm_test.py
index 723e28a1..d64639dc 100644
--- a/tests/llm_test.py
+++ b/tests/llm_test.py
@@ -3,6 +3,7 @@ from time import sleep
 from unittest.mock import MagicMock, patch
 
 import pytest
+from pydantic import BaseModel
 
 from crewai.agents.agent_builder.utilities.base_token_process import TokenProcess
 from crewai.llm import LLM
@@ -205,6 +206,52 @@ def test_llm_passes_additional_params():
         assert result == "Test response"
 
 
+def test_get_custom_llm_provider_openrouter():
+    llm = LLM(model="openrouter/deepseek/deepseek-chat")
+    assert llm._get_custom_llm_provider() == "openrouter"
+
+
+def test_get_custom_llm_provider_gemini():
+    llm = LLM(model="gemini/gemini-1.5-pro")
+    assert llm._get_custom_llm_provider() == "gemini"
+
+
+def test_get_custom_llm_provider_openai():
+    llm = LLM(model="gpt-4")
+    assert llm._get_custom_llm_provider() == "openai"
+
+
+def test_validate_call_params_supported():
+    class DummyResponse(BaseModel):
+        a: int
+
+    # Patch supports_response_schema to simulate a supported model.
+    with patch("crewai.llm.supports_response_schema", return_value=True):
+        llm = LLM(
+            model="openrouter/deepseek/deepseek-chat", response_format=DummyResponse
+        )
+        # Should not raise any error.
+        llm._validate_call_params()
+
+
+def test_validate_call_params_not_supported():
+    class DummyResponse(BaseModel):
+        a: int
+
+    # Patch supports_response_schema to simulate an unsupported model.
+    with patch("crewai.llm.supports_response_schema", return_value=False):
+        llm = LLM(model="gemini/gemini-1.5-pro", response_format=DummyResponse)
+        with pytest.raises(ValueError) as excinfo:
+            llm._validate_call_params()
+        assert "does not support response_format" in str(excinfo.value)
+
+
+def test_validate_call_params_no_response_format():
+    # When no response_format is provided, no validation error should occur.
+    llm = LLM(model="gemini/gemini-1.5-pro", response_format=None)
+    llm._validate_call_params()
+
+
 @pytest.mark.vcr(filter_headers=["authorization"])
 def test_o3_mini_reasoning_effort_high():
     llm = LLM(