from typing import ClassVar from unittest.mock import Mock import litellm import numpy as np import pytest from aviary.env import DummyEnv from aviary.message import Message from aviary.tools import Tool, ToolRequestMessage from pydantic import BaseModel, Field from ldp.llms import ( JSONSchemaValidationError, LLMModel, LLMResult, MultipleCompletionLLMModel, validate_json_completion, ) from . import CILLMModelNames def test_json_schema_validation() -> None: # Invalid JSON mock_completion1 = Mock() mock_completion1.choices = [Mock()] mock_completion1.choices[0].message.content = "not a json" # Invalid schema mock_completion2 = Mock() mock_completion2.choices = [Mock()] mock_completion2.choices[0].message.content = '{"name": "John", "age": "nan"}' # Valid schema mock_completion3 = Mock() mock_completion3.choices = [Mock()] mock_completion3.choices[0].message.content = '{"name": "John", "age": 30}' class DummyModel(BaseModel): name: str age: int with pytest.raises(JSONSchemaValidationError): validate_json_completion(mock_completion1, DummyModel) with pytest.raises(JSONSchemaValidationError): validate_json_completion(mock_completion2, DummyModel) validate_json_completion(mock_completion3, DummyModel) @pytest.mark.parametrize( "model_name", ["gpt-3.5-turbo", CILLMModelNames.ANTHROPIC.value] ) @pytest.mark.asyncio async def test_achat(model_name: str) -> None: model = LLMModel(name=model_name) response = await model.achat( messages=[ Message(content="What are three things I should do today?"), ] ) assert len(response.choices) == 1 # Check we can iterate through the response async for chunk in await model.achat_iter( messages=[ Message(content="What are three things I should do today?"), ] ): assert len(chunk.choices) == 1 @pytest.mark.parametrize( "model_name", [CILLMModelNames.OPENAI.value, CILLMModelNames.ANTHROPIC.value] ) @pytest.mark.flaky(reruns=3, only_on=[litellm.exceptions.APIConnectionError]) @pytest.mark.asyncio async def test_tools(dummy_env: DummyEnv, model_name: str) -> None: model = LLMModel(name=model_name) messages = [ Message(content="What are three things I should do today?"), ] def get_todo_list(n: int) -> str: """Get todo list for today. Args: n: number of items to return """ return "\n".join(["Go for a walk", "Read a book", "Call a friend"][:n]) tool = Tool.from_function(get_todo_list) dummy_env.tools = [tool] result = await model.call(messages, tools=dummy_env.tools) assert result.completion_count > 0 # try specifying tool choice result = await model.call(messages, tools=dummy_env.tools, tool_choice=tool) assert result.completion_count > 0, "Tool choice failed to execute tool" assert result.messages (tool_request_message,) = result.messages assert isinstance(tool_request_message, ToolRequestMessage) new_messages = await dummy_env.exec_tool_calls(tool_request_message) (new_message,) = new_messages assert new_message.content == "Go for a walk\nRead a book\nCall a friend" assert new_message.tool_call_id == tool_request_message.tool_calls[0].id def get_todo_list_no_args() -> str: """Get todo list for today.""" return "\n".join(["Go for a walk", "Read a book", "Call a friend"]) tool = Tool.from_function(get_todo_list_no_args) dummy_env.tools = [tool] result = await model.call(messages, tools=dummy_env.tools) assert result.completion_count > 0 assert result.messages (tool_request_message,) = result.messages assert isinstance(tool_request_message, ToolRequestMessage) new_messages = await dummy_env.exec_tool_calls(tool_request_message) (new_message,) = new_messages assert new_message.content == "Go for a walk\nRead a book\nCall a friend" assert new_message.tool_call_id == tool_request_message.tool_calls[0].id # ok now try with multiple functions messages = [ Message( content=( "What items will I have time to accomplish on my todo list today based" " on my calendar?" ) ), ] def get_calendar() -> str: """Get text version of calendar for today.""" return "9:00am Wake-up\n10:00pm Go to bed\n" tool2 = Tool.from_function(get_calendar) dummy_env.tools = [tool, tool2] result = await model.call(messages, tools=dummy_env.tools) assert result.messages (tool_request_message,) = result.messages assert isinstance(tool_request_message, ToolRequestMessage) new_messages = await dummy_env.exec_tool_calls(tool_request_message) if model_name.startswith("claude"): # Anthropic not always so smart assert 1 <= len(new_messages) <= 2 else: assert len(new_messages) == 2 # ok now try continuation - I AM NOT SURE IF THIS IS VALID? # TODO: - supported on openai, but not litellm # messages = messages + result.messages + new_messages # result = await model.call(messages) class DummyOutputSchema(BaseModel): name: str age: int class TestMultipleCompletionLLMModel: NUM_COMPLETIONS: ClassVar[int] = 2 DEFAULT_CONFIG: ClassVar[dict] = {"n": NUM_COMPLETIONS} MODEL_CLS: ClassVar[type[MultipleCompletionLLMModel]] = MultipleCompletionLLMModel async def call_model( self, model: MultipleCompletionLLMModel, *args, **kwargs ) -> list[LLMResult]: return await model.call(*args, **kwargs) @pytest.mark.vcr @pytest.mark.parametrize("model_name", ["gpt-3.5-turbo"]) @pytest.mark.asyncio async def test_model(self, model_name: str) -> None: # Make model_name an arg so that TestLLMModel can parametrize it # only testing OpenAI, as other APIs don't support n>1 model = self.MODEL_CLS(name=model_name, config=self.DEFAULT_CONFIG) messages = [ Message(role="system", content="Respond with single words."), Message(content="Hello, how are you?"), ] results = await self.call_model(model, messages) assert len(results) == self.NUM_COMPLETIONS for result in results: assert result.prompt_count > 0 assert result.completion_count > 0 prompt_cost, completion_cost = result.prompt_and_completion_costs assert prompt_cost > 0 assert completion_cost > 0 assert result.logprob is None or result.logprob <= 0 @pytest.mark.parametrize( "model_name", [CILLMModelNames.ANTHROPIC.value, "gpt-3.5-turbo"] ) @pytest.mark.asyncio async def test_streaming(self, model_name: str) -> None: model = self.MODEL_CLS(name=model_name, config=self.DEFAULT_CONFIG) messages = [ Message(role="system", content="Respond with single words."), Message(content="Hello, how are you?"), ] def callback(_) -> None: return with pytest.raises( NotImplementedError, match="Multiple completions with callbacks is not supported", ): await self.call_model(model, messages, [callback]) @pytest.mark.vcr @pytest.mark.asyncio async def test_parameterizing_tool_from_arg_union(self) -> None: def play(move: int | None) -> None: """Play one turn by choosing a move. Args: move: Choose an integer to lose, choose None to win. """ results = await self.call_model( self.MODEL_CLS(name="gpt-3.5-turbo", config=self.DEFAULT_CONFIG), messages=[Message(content="Please win.")], tools=[Tool.from_function(play)], ) assert len(results) == self.NUM_COMPLETIONS for result in results: assert result.messages assert len(result.messages) == 1 assert isinstance(result.messages[0], ToolRequestMessage) assert result.messages[0].tool_calls assert result.messages[0].tool_calls[0].function.arguments["move"] is None @pytest.mark.asyncio @pytest.mark.vcr async def test_output_schema(self) -> None: model = self.MODEL_CLS(name="gpt-3.5-turbo", config=self.DEFAULT_CONFIG) messages = [ Message( content=( "My name is Claude and I am 1 year old. What is my name and age?" ) ), ] results = await self.call_model(model, messages, output_type=DummyOutputSchema) assert len(results) == self.NUM_COMPLETIONS for result in results: assert result.messages assert len(result.messages) == 1 assert result.messages[0].content DummyOutputSchema.model_validate_json(result.messages[0].content) @pytest.mark.parametrize("model_name", [CILLMModelNames.OPENAI.value]) @pytest.mark.asyncio @pytest.mark.vcr async def test_text_image_message(self, model_name: str) -> None: model = self.MODEL_CLS(name=model_name, config=self.DEFAULT_CONFIG) # An RGB image of a red square image = np.zeros((32, 32, 3), dtype=np.uint8) image[:] = [255, 0, 0] # (255 red, 0 green, 0 blue) is maximum red in RGB results = await self.call_model( model, messages=[ Message.create_message( text="What color is this square? Respond only with the color name.", image=image, ) ], ) assert len(results) == self.NUM_COMPLETIONS for result in results: assert ( result.messages is not None ), "Expected messages in result, but got None" assert ( result.messages[-1].content is not None ), "Expected content in message, but got None" assert "red" in result.messages[-1].content.lower() class TestLLMModel(TestMultipleCompletionLLMModel): NUM_COMPLETIONS: ClassVar[int] = 1 DEFAULT_CONFIG: ClassVar[dict] = {} MODEL_CLS: ClassVar[type[MultipleCompletionLLMModel]] = LLMModel async def call_model(self, model: LLMModel, *args, **kwargs) -> list[LLMResult]: # type: ignore[override] return [await model.call(*args, **kwargs)] @pytest.mark.parametrize( "model_name", [CILLMModelNames.ANTHROPIC.value, "gpt-3.5-turbo"] ) @pytest.mark.asyncio @pytest.mark.vcr async def test_model(self, model_name: str) -> None: await super().test_model(model_name) @pytest.mark.vcr @pytest.mark.parametrize( "model_name", [CILLMModelNames.ANTHROPIC.value, "gpt-3.5-turbo"] ) @pytest.mark.asyncio async def test_streaming(self, model_name: str) -> None: model = LLMModel(name=model_name) messages = [ Message(role="system", content="Respond with single words."), Message(content="Hello, how are you?"), ] content = [] def callback(s): content.append(s) result = await model.call(messages, [callback]) assert result.completion_count > 0 assert content @pytest.mark.vcr @pytest.mark.asyncio async def test_parameterizing_tool_from_arg_union(self) -> None: await super().test_parameterizing_tool_from_arg_union() @pytest.mark.asyncio async def test_output_type_rejected_validation(self) -> None: class InstructionList(BaseModel): instructions: list[str] = Field(description="list of instructions") model = self.MODEL_CLS(name=CILLMModelNames.ANTHROPIC.value) with pytest.raises( litellm.UnsupportedParamsError, match="anthropic does not support parameters", ): await model.call( [Message(content="What are three things I should do today?")], output_type=InstructionList, ) @pytest.mark.parametrize( "model_name", [CILLMModelNames.ANTHROPIC.value, "gpt-4-turbo", CILLMModelNames.OPENAI.value], ) @pytest.mark.asyncio @pytest.mark.vcr async def test_text_image_message(self, model_name: str) -> None: await super().test_text_image_message(model_name)