# from __future__ import annotations # import asyncio # import base64 # from enum import Enum # from pathlib import Path # from typing import Annotated, Callable # import pytest # from livekit.agents import APIConnectionError, llm # from livekit.agents.llm import ChatContext, ToolContext, TypeInfo, ai_callable # from livekit.plugins import anthropic, aws, google, mistralai, openai # from livekit.rtc import VideoBufferType, VideoFrame # class Unit(Enum): # FAHRENHEIT = "fahrenheit" # CELSIUS = "celsius" # class FncCtx(ToolContext): # @ai_callable(description="Get the current weather in a given location", auto_retry=True) # def get_weather( # self, # location: Annotated[ # str, TypeInfo(description="The city and state, e.g. San Francisco, CA") # ], # unit: Annotated[Unit, TypeInfo(description="The temperature unit to use.")] = Unit.CELSIUS, # ) -> None: ... # @ai_callable(description="Play a music") # def play_music( # self, # name: Annotated[str, TypeInfo(description="the name of the Artist")], # ) -> None: ... # # test for cancelled calls # @ai_callable(description="Turn on/off the lights in a room") # async def toggle_light( # self, # room: Annotated[str, TypeInfo(description="The room to control")], # on: bool = True, # ) -> None: # await asyncio.sleep(60) # # used to test arrays as arguments # @ai_callable(description="Select currencies of a specific area") # def select_currencies( # self, # currencies: Annotated[ # list[str], # TypeInfo( # description="The currencies to select", # choices=["usd", "eur", "gbp", "jpy", "sek"], # ), # ], # ) -> None: ... # @ai_callable(description="Update user info") # def update_user_info( # self, # email: Annotated[str | None, TypeInfo(description="The user address email")] = None, # name: Annotated[str | None, TypeInfo(description="The user name")] = None, # address: Annotated[str, TypeInfo(description="The user address")] | None = None, # ) -> None: ... # def test_hashable_typeinfo(): # typeinfo = TypeInfo(description="testing", choices=[1, 2, 3]) # # TypeInfo must be hashable when used in combination of typing.Annotated # hash(typeinfo) # LLMS: list[Callable[[], llm.LLM]] = [ # pytest.param(lambda: openai.LLM(), id="openai"), # # lambda: openai.beta.AssistantLLM( # # assistant_opts=openai.beta.AssistantOptions( # # create_options=openai.beta.AssistantCreateOptions( # # name=f"test-{uuid.uuid4()}", # # instructions="You are a basic assistant", # # model="gpt-4o", # # ) # # ) # # ), # pytest.param(lambda: anthropic.LLM(), id="anthropic"), # pytest.param(lambda: google.LLM(), id="google"), # pytest.param(lambda: google.LLM(vertexai=True), id="google-vertexai"), # pytest.param(lambda: aws.LLM(), id="aws"), # pytest.param(lambda: mistralai.LLM(), id="mistralai"), # ] # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_chat(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # chat_ctx = ChatContext().append( # text='You are an assistant at a drive-thru restaurant "Live-Burger". Ask the customer what they would like to order.', # noqa: E501 # ) # # Anthropic and vertex requires at least one message (system messages don't count) # chat_ctx.append( # text="Hello", # role="user", # ) # stream = input_llm.chat(chat_ctx=chat_ctx) # text = "" # async for chunk in stream: # if not chunk.choices: # continue # content = chunk.choices[0].delta.content # if content: # text += content # assert len(text) > 0 # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_llm_chat_with_consecutive_messages( # llm_factory: callable, # ): # input_llm = llm_factory() # chat_ctx = ChatContext() # chat_ctx.append( # text="Hello, How can I help you today?", # role="assistant", # ) # chat_ctx.append(text="I see that you have a busy day ahead.", role="assistant") # chat_ctx.append(text="Actually, I need some help with my recent order.", role="user") # chat_ctx.append(text="I want to cancel my order.", role="user") # stream = input_llm.chat(chat_ctx=chat_ctx) # collected_text = "" # async for chunk in stream: # if not chunk.choices: # continue # content = chunk.choices[0].delta.content # if content: # collected_text += content # assert len(collected_text) > 0, "Expected a non-empty response from the LLM chat" # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_basic_fnc_calls(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # fnc_ctx = FncCtx() # stream = await _request_fnc_call( # input_llm, # "What's the weather in San Francisco and what's the weather Paris?", # fnc_ctx, # ) # calls = stream.execute_functions() # await asyncio.gather(*[f.task for f in calls]) # await stream.aclose() # assert len(calls) == 2, "get_weather should be called twice" # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_function_call_exception_handling(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # fnc_ctx = FncCtx() # @fnc_ctx.ai_callable(description="Simulate a failure") # async def failing_function(): # raise RuntimeError("Simulated failure") # stream = await _request_fnc_call(input_llm, "Call the failing function", fnc_ctx) # calls = stream.execute_functions() # await asyncio.gather(*[f.task for f in calls], return_exceptions=True) # await stream.aclose() # assert len(calls) == 1 # assert isinstance(calls[0].exception, RuntimeError) # assert str(calls[0].exception) == "Simulated failure" # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_runtime_addition(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # fnc_ctx = FncCtx() # called_msg = "" # @fnc_ctx.ai_callable(description="Show a message on the screen") # async def show_message( # message: Annotated[str, TypeInfo(description="The message to show")], # ): # nonlocal called_msg # called_msg = message # stream = await _request_fnc_call( # input_llm, "Can you show 'Hello LiveKit!' on the screen?", fnc_ctx # ) # fns = stream.execute_functions() # await asyncio.gather(*[f.task for f in fns]) # await stream.aclose() # assert called_msg == "Hello LiveKit!", "send_message should be called" # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_cancelled_calls(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # fnc_ctx = FncCtx() # stream = await _request_fnc_call(input_llm, "Turn off the lights in the bedroom", fnc_ctx) # calls = stream.execute_functions() # await asyncio.sleep(0.2) # wait for the loop executor to start the task # # don't wait for gather_function_results and directly close (this should cancel the # # ongoing calls) # await stream.aclose() # assert len(calls) == 1 # assert isinstance(calls[0].exception, asyncio.CancelledError), ( # "toggle_light should have been cancelled" # ) # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_calls_arrays(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # fnc_ctx = FncCtx() # stream = await _request_fnc_call( # input_llm, # "Can you select all currencies in Europe at once from given choices using function call `select_currencies`?", # noqa: E501 # fnc_ctx, # temperature=0.2, # ) # calls = stream.execute_functions() # await asyncio.gather(*[f.task for f in calls]) # await stream.aclose() # assert len(calls) == 1, "select_currencies should have been called only once" # call = calls[0] # currencies = call.call_info.arguments["currencies"] # assert len(currencies) == 3, "select_currencies should have 3 currencies" # assert "eur" in currencies and "gbp" in currencies and "sek" in currencies, ( # "select_currencies should have eur, gbp, sek" # ) # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_calls_choices(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # fnc_ctx = FncCtx() # # test choices on int # @fnc_ctx.ai_callable(description="Change the volume") # def change_volume( # volume: Annotated[ # int, TypeInfo(description="The volume level", choices=[0, 11, 30, 83, 99]) # ], # ) -> None: ... # if not input_llm.capabilities.supports_choices_on_int: # with pytest.raises(APIConnectionError): # stream = await _request_fnc_call(input_llm, "Set the volume to 30", fnc_ctx) # else: # stream = await _request_fnc_call(input_llm, "Set the volume to 30", fnc_ctx) # calls = stream.execute_functions() # await asyncio.gather(*[f.task for f in calls]) # await stream.aclose() # assert len(calls) == 1, "change_volume should have been called only once" # call = calls[0] # volume = call.call_info.arguments["volume"] # assert volume == 30, "change_volume should have been called with volume 30" # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_optional_args(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # fnc_ctx = FncCtx() # stream = await _request_fnc_call( # input_llm, "Using a tool call update the user info to name Theo", fnc_ctx # ) # calls = stream.execute_functions() # await asyncio.gather(*[f.task for f in calls]) # await stream.aclose() # assert len(calls) == 1, "update_user_info should have been called only once" # call = calls[0] # name = call.call_info.arguments.get("name", None) # email = call.call_info.arguments.get("email", None) # address = call.call_info.arguments.get("address", None) # assert name == "Theo", "update_user_info should have been called with name 'Theo'" # assert email is None, "update_user_info should have been called with email None" # assert address is None, "update_user_info should have been called with address None" # test_tool_choice_cases = [ # pytest.param( # "Default tool_choice (auto)", # "Get the weather for New York and play some music from the artist 'The Beatles'.", # None, # {"get_weather", "play_music"}, # id="Default tool_choice (auto)", # ), # pytest.param( # "Tool_choice set to 'required'", # "Get the weather for Chicago and play some music from the artist 'Eminem'.", # "required", # {"get_weather", "play_music"}, # id="Tool_choice set to 'required'", # ), # pytest.param( # "Tool_choice set to a specific tool ('get_weather')", # "Get the weather for Miami.", # llm.ToolChoice(type="function", name="get_weather"), # {"get_weather"}, # id="Tool_choice set to a specific tool ('get_weather')", # ), # pytest.param( # "Tool_choice set to 'none'", # "Get the weather for Seattle and play some music from the artist 'Frank Sinatra'.", # "none", # set(), # No tool calls expected # id="Tool_choice set to 'none'", # ), # ] # @pytest.mark.parametrize( # "description, user_request, tool_choice, expected_calls", test_tool_choice_cases # ) # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_tool_choice_options( # description: str, # user_request: str, # tool_choice: dict | str | None, # expected_calls: set, # llm_factory: Callable[[], llm.LLM], # ): # input_llm = llm_factory() # fnc_ctx = FncCtx() # stream = await _request_fnc_call( # input_llm, # user_request, # fnc_ctx, # tool_choice=tool_choice, # parallel_tool_calls=True, # ) # calls = stream.execute_functions() # await asyncio.gather(*[f.task for f in calls], return_exceptions=True) # await stream.aclose() # print(calls) # call_names = {call.call_info.function_info.name for call in calls} # if tool_choice == "none": # assert call_names == expected_calls, ( # f"Test '{description}' failed: Expected calls {expected_calls}, but got {call_names}" # ) # async def _request_fnc_call( # model: llm.LLM, # request: str, # fnc_ctx: FncCtx, # temperature: float | None = None, # parallel_tool_calls: bool | None = None, # tool_choice: llm.ToolChoice | None = None, # ) -> llm.LLMStream: # stream = model.chat( # chat_ctx=ChatContext() # .append( # text="You are an helpful assistant. Follow the instructions provided by the user. You can use multiple tool calls at once.", # noqa: E501 # role="system", # ) # .append(text=request, role="user"), # fnc_ctx=fnc_ctx, # temperature=temperature, # tool_choice=tool_choice, # parallel_tool_calls=parallel_tool_calls, # ) # async for _ in stream: # pass # return stream # _HEARTS_RGBA_PATH = Path(__file__).parent / "hearts.rgba" # with open(_HEARTS_RGBA_PATH, "rb") as f: # image_data = f.read() # _HEARTS_IMAGE_VIDEO_FRAME = VideoFrame( # width=512, height=512, type=VideoBufferType.RGBA, data=image_data # ) # _HEARTS_JPEG_PATH = Path(__file__).parent / "hearts.jpg" # with open(_HEARTS_JPEG_PATH, "rb") as f: # _HEARTS_IMAGE_DATA_URL = f"data:image/jpeg;base64,{base64.b64encode(f.read()).decode()}" # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_chat_with_image_data_url(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # chat_ctx = ( # ChatContext() # .append( # text="You are an AI assistant that describes images in detail upon request.", # role="system", # ) # .append( # text="Describe this image", # images=[llm.ChatImage(image=_HEARTS_IMAGE_DATA_URL, inference_detail="low")], # role="user", # ) # ) # stream = input_llm.chat(chat_ctx=chat_ctx) # text = "" # async for chunk in stream: # if not chunk.choices: # continue # content = chunk.choices[0].delta.content # if content: # text += content # assert "heart" in text.lower() # @pytest.mark.parametrize("llm_factory", LLMS) # async def test_chat_with_image_frame(llm_factory: Callable[[], llm.LLM]): # input_llm = llm_factory() # chat_ctx = ( # ChatContext() # .append( # text="You are an AI assistant that describes images in detail upon request.", # role="system", # ) # .append( # text="Describe this image", # images=[llm.ChatImage(image=_HEARTS_IMAGE_VIDEO_FRAME, inference_detail="low")], # role="user", # ) # ) # stream = input_llm.chat(chat_ctx=chat_ctx) # text = "" # async for chunk in stream: # if not chunk.choices: # continue # content = chunk.choices[0].delta.content # if content: # text += content # assert "heart" in text.lower()