mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-22 14:48:13 +00:00
Fix output_json with custom OpenAI APIs by using PARALLEL_TOOLS mode
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -29,7 +29,14 @@ class InternalInstructor:
|
|||||||
import instructor
|
import instructor
|
||||||
from litellm import completion
|
from litellm import completion
|
||||||
|
|
||||||
self._client = instructor.from_litellm(completion)
|
is_custom_openai = getattr(self.llm, 'model', '').startswith('custom_openai/')
|
||||||
|
|
||||||
|
mode = instructor.Mode.PARALLEL_TOOLS if is_custom_openai else instructor.Mode.TOOLS
|
||||||
|
|
||||||
|
self._client = instructor.from_litellm(
|
||||||
|
completion,
|
||||||
|
mode=mode,
|
||||||
|
)
|
||||||
|
|
||||||
def to_json(self):
|
def to_json(self):
|
||||||
model = self.to_pydantic()
|
model = self.to_pydantic()
|
||||||
@@ -40,4 +47,8 @@ class InternalInstructor:
|
|||||||
model = self._client.chat.completions.create(
|
model = self._client.chat.completions.create(
|
||||||
model=self.llm.model, response_model=self.model, messages=messages
|
model=self.llm.model, response_model=self.model, messages=messages
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if isinstance(model, list) and len(model) > 0:
|
||||||
|
return model[0] # Return the first model from the list
|
||||||
|
|
||||||
return model
|
return model
|
||||||
|
|||||||
72
tests/utilities/test_internal_instructor.py
Normal file
72
tests/utilities/test_internal_instructor.py
Normal file
@@ -0,0 +1,72 @@
|
|||||||
|
import unittest
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from crewai.utilities.internal_instructor import InternalInstructor
|
||||||
|
|
||||||
|
|
||||||
|
class TestOutput(BaseModel):
|
||||||
|
value: str
|
||||||
|
|
||||||
|
|
||||||
|
class TestInternalInstructor(unittest.TestCase):
|
||||||
|
@patch("instructor.from_litellm")
|
||||||
|
def test_tools_mode_for_regular_models(self, mock_from_litellm):
|
||||||
|
mock_llm = MagicMock()
|
||||||
|
mock_llm.model = "gpt-4o"
|
||||||
|
mock_instructor = MagicMock()
|
||||||
|
mock_from_litellm.return_value = mock_instructor
|
||||||
|
|
||||||
|
instructor = InternalInstructor(
|
||||||
|
content="Test content",
|
||||||
|
model=TestOutput,
|
||||||
|
llm=mock_llm
|
||||||
|
)
|
||||||
|
|
||||||
|
import instructor
|
||||||
|
mock_from_litellm.assert_called_once_with(
|
||||||
|
unittest.mock.ANY,
|
||||||
|
mode=instructor.Mode.TOOLS
|
||||||
|
)
|
||||||
|
|
||||||
|
@patch("instructor.from_litellm")
|
||||||
|
def test_parallel_tools_mode_for_custom_openai(self, mock_from_litellm):
|
||||||
|
mock_llm = MagicMock()
|
||||||
|
mock_llm.model = "custom_openai/some-model"
|
||||||
|
mock_instructor = MagicMock()
|
||||||
|
mock_from_litellm.return_value = mock_instructor
|
||||||
|
|
||||||
|
instructor = InternalInstructor(
|
||||||
|
content="Test content",
|
||||||
|
model=TestOutput,
|
||||||
|
llm=mock_llm
|
||||||
|
)
|
||||||
|
|
||||||
|
import instructor
|
||||||
|
mock_from_litellm.assert_called_once_with(
|
||||||
|
unittest.mock.ANY,
|
||||||
|
mode=instructor.Mode.PARALLEL_TOOLS
|
||||||
|
)
|
||||||
|
|
||||||
|
@patch("instructor.from_litellm")
|
||||||
|
def test_handling_list_response_in_to_pydantic(self, mock_from_litellm):
|
||||||
|
mock_llm = MagicMock()
|
||||||
|
mock_llm.model = "custom_openai/some-model"
|
||||||
|
mock_instructor = MagicMock()
|
||||||
|
mock_chat = MagicMock()
|
||||||
|
mock_instructor.chat.completions.create.return_value = [
|
||||||
|
TestOutput(value="test value")
|
||||||
|
]
|
||||||
|
mock_from_litellm.return_value = mock_instructor
|
||||||
|
|
||||||
|
instructor = InternalInstructor(
|
||||||
|
content="Test content",
|
||||||
|
model=TestOutput,
|
||||||
|
llm=mock_llm
|
||||||
|
)
|
||||||
|
result = instructor.to_pydantic()
|
||||||
|
|
||||||
|
assert isinstance(result, TestOutput)
|
||||||
|
assert result.value == "test value"
|
||||||
Reference in New Issue
Block a user