"""Tests for the LiteLLM conversation entity.""" import datetime from unittest.mock import AsyncMock, patch from freezegun import freeze_time import httpx import openai from openai.types import CompletionUsage from openai.types.chat import ( ChatCompletion, ChatCompletionMessage, ChatCompletionMessageFunctionToolCall, ) from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_message_function_tool_call_param import Function import pytest from syrupy.assertion import SnapshotAssertion from homeassistant.components import conversation from homeassistant.const import STATE_UNAVAILABLE, Platform from homeassistant.core import Context, HomeAssistant from homeassistant.helpers import entity_registry as er, intent from homeassistant.helpers.llm import ToolInput from . import setup_integration from tests.common import MockConfigEntry, snapshot_platform from tests.components.conversation import MockChatLog, mock_chat_log # noqa: F401 AGENT_ID = "conversation.gpt_3_5_turbo" @pytest.fixture(autouse=True) def freeze_the_time(): """Freeze the time.""" with freeze_time("2024-05-24 12:00:00", tz_offset=0): yield @pytest.mark.parametrize("enable_assist", [True, False], ids=["assist", "no_assist"]) async def test_all_entities( hass: HomeAssistant, snapshot: SnapshotAssertion, mock_openai_client: AsyncMock, mock_config_entry: MockConfigEntry, entity_registry: er.EntityRegistry, ) -> None: """Test all entities.""" with patch( "homeassistant.components.litellm.PLATFORMS", [Platform.CONVERSATION], ): await setup_integration(hass, mock_config_entry) await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id) async def test_default_prompt( hass: HomeAssistant, mock_config_entry: MockConfigEntry, snapshot: SnapshotAssertion, mock_openai_client: AsyncMock, mock_chat_log: MockChatLog, # noqa: F811 ) -> None: """Test that the default prompt works.""" await setup_integration(hass, mock_config_entry) result = await conversation.async_converse( hass, "hello", mock_chat_log.conversation_id, Context(), agent_id=AGENT_ID, ) assert result.response.response_type is intent.IntentResponseType.ACTION_DONE assert mock_chat_log.content[1:] == snapshot call = mock_openai_client.chat.completions.create.call_args_list[0][1] assert call["model"] == "gpt-3.5-turbo" assert "extra_headers" not in call async def test_empty_api_response( hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_openai_client: AsyncMock, mock_chat_log: MockChatLog, # noqa: F811 ) -> None: """Test that an empty choices response raises an error.""" await setup_integration(hass, mock_config_entry) mock_openai_client.chat.completions.create = AsyncMock( return_value=ChatCompletion( id="chatcmpl-1234567890ABCDEFGHIJKLMNOPQRS", choices=[], created=1700000000, model="gpt-3.5-turbo", object="chat.completion", system_fingerprint=None, usage=CompletionUsage(completion_tokens=0, prompt_tokens=8, total_tokens=8), ) ) result = await conversation.async_converse( hass, "hello", mock_chat_log.conversation_id, Context(), agent_id=AGENT_ID, ) assert result.response.response_type is intent.IntentResponseType.ERROR async def test_api_error( hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_openai_client: AsyncMock, mock_chat_log: MockChatLog, # noqa: F811 ) -> None: """Test that an error talking to the API is handled gracefully.""" await setup_integration(hass, mock_config_entry) mock_openai_client.chat.completions.create = AsyncMock( side_effect=openai.OpenAIError("boom") ) result = await conversation.async_converse( hass, "hello", mock_chat_log.conversation_id, Context(), agent_id=AGENT_ID, ) assert result.response.response_type is intent.IntentResponseType.ERROR async def test_connection_error_availability( hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_openai_client: AsyncMock, mock_chat_log: MockChatLog, # noqa: F811 ) -> None: """Test a connection error marks the entity unavailable until it recovers.""" await setup_integration(hass, mock_config_entry) assert hass.states.get(AGENT_ID).state != STATE_UNAVAILABLE mock_openai_client.chat.completions.create = AsyncMock( side_effect=openai.APIConnectionError( request=httpx.Request("POST", "http://localhost") ) ) result = await conversation.async_converse( hass, "hello", mock_chat_log.conversation_id, Context(), agent_id=AGENT_ID, ) assert result.response.response_type is intent.IntentResponseType.ERROR await hass.async_block_till_done() assert hass.states.get(AGENT_ID).state == STATE_UNAVAILABLE # A successful availability ping restores the entity. await mock_config_entry.runtime_data.async_request_refresh() await hass.async_block_till_done() assert hass.states.get(AGENT_ID).state != STATE_UNAVAILABLE @pytest.mark.parametrize("enable_assist", [True]) async def test_function_call( hass: HomeAssistant, mock_chat_log: MockChatLog, # noqa: F811 mock_config_entry: MockConfigEntry, snapshot: SnapshotAssertion, mock_openai_client: AsyncMock, ) -> None: """Test function call from the assistant.""" await setup_integration(hass, mock_config_entry) mock_chat_log.async_add_user_content( conversation.UserContent(content="What time is it?") ) mock_chat_log.async_add_assistant_content_without_tools( conversation.AssistantContent( agent_id=AGENT_ID, tool_calls=[ ToolInput( tool_name="HassGetCurrentTime", tool_args={}, id="mock_tool_call_id", external=True, ) ], ) ) mock_chat_log.async_add_assistant_content_without_tools( conversation.ToolResultContent( agent_id=AGENT_ID, tool_call_id="mock_tool_call_id", tool_name="HassGetCurrentTime", tool_result={ "speech": {"plain": {"speech": "12:00 PM", "extra_data": None}}, "response_type": "action_done", "speech_slots": {"time": datetime.time(12, 0)}, "data": {"success": [], "failed": []}, }, ) ) mock_chat_log.async_add_assistant_content_without_tools( conversation.AssistantContent( agent_id=AGENT_ID, content="12:00 PM", ) ) mock_chat_log.mock_tool_results( { "call_call_1": "value1", "call_call_2": "value2", } ) mock_openai_client.chat.completions.create.side_effect = ( ChatCompletion( id="chatcmpl-1234567890ABCDEFGHIJKLMNOPQRS", choices=[ Choice( finish_reason="tool_calls", index=0, message=ChatCompletionMessage( content=None, role="assistant", function_call=None, tool_calls=[ ChatCompletionMessageFunctionToolCall( id="call_call_1", function=Function( arguments='{"param1":"call1"}', name="test_tool", ), type="function", ) ], ), ) ], created=1700000000, model="gpt-4", object="chat.completion", system_fingerprint=None, usage=CompletionUsage( completion_tokens=9, prompt_tokens=8, total_tokens=17 ), ), ChatCompletion( id="chatcmpl-1234567890ZYXWVUTSRQPONMLKJIH", choices=[ Choice( finish_reason="stop", index=0, message=ChatCompletionMessage( content="I have successfully called the function", role="assistant", function_call=None, tool_calls=None, ), ) ], created=1700000000, model="gpt-4", object="chat.completion", system_fingerprint=None, usage=CompletionUsage( completion_tokens=9, prompt_tokens=8, total_tokens=17 ), ), ) result = await conversation.async_converse( hass, "Please call the test function", mock_chat_log.conversation_id, Context(), agent_id=AGENT_ID, ) assert result.response.response_type is intent.IntentResponseType.ACTION_DONE # Don't test the prompt, as it's not deterministic assert mock_chat_log.content[1:] == snapshot assert mock_openai_client.chat.completions.create.call_count == 2 assert ( mock_openai_client.chat.completions.create.call_args.kwargs["messages"] == snapshot )