"""Tests for helpers in the Home Assistant Cloud conversation entity.""" from __future__ import annotations import base64 from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch from PIL import Image import pytest import voluptuous as vol from homeassistant.components import conversation from homeassistant.components.cloud.entity import ( BaseCloudLLMEntity, _convert_content_to_param, _format_structured_output, ) from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import llm, selector from tests.common import MockConfigEntry @pytest.fixture def cloud_entity(hass: HomeAssistant) -> BaseCloudLLMEntity: """Return a CloudLLMTaskEntity attached to hass.""" cloud = MagicMock() cloud.llm = MagicMock() cloud.is_logged_in = True cloud.valid_subscription = True entry = MockConfigEntry(domain="cloud") entry.add_to_hass(hass) entity = BaseCloudLLMEntity(cloud, entry) entity.entity_id = "ai_task.cloud_ai_task" entity.hass = hass return entity @pytest.fixture def mock_prepare_files_for_prompt( cloud_entity: BaseCloudLLMEntity, ) -> AsyncMock: """Patch file preparation helper on the entity.""" with patch.object( cloud_entity, "_async_prepare_files_for_prompt", AsyncMock(), ) as mock: yield mock class DummyTool(llm.Tool): """Simple tool used for schema conversion tests.""" name = "do_something" description = "Test tool" parameters = vol.Schema({vol.Required("value"): str}) async def async_call(self, hass: HomeAssistant, tool_input, llm_context): """No-op implementation.""" return {"value": "done"} async def test_format_structured_output() -> None: """Test that structured output schemas are normalized.""" schema = vol.Schema( { vol.Required("name"): selector.TextSelector(), vol.Optional("age"): selector.NumberSelector( config=selector.NumberSelectorConfig(min=0, max=120), ), vol.Required("stuff"): selector.ObjectSelector( { "multiple": True, "fields": { "item_name": {"selector": {"text": None}}, "item_value": {"selector": {"text": None}}, }, } ), } ) assert _format_structured_output(schema, None) == { "type": "object", "properties": { "name": {"type": "string"}, "age": {"type": "number", "minimum": 0.0, "maximum": 120.0}, "stuff": { "type": "array", "items": { "type": "object", "properties": { "item_name": {"type": "string"}, "item_value": {"type": "string"}, }, "additionalProperties": False, }, }, }, "required": ["name", "stuff"], "additionalProperties": False, } async def test_prepare_files_for_prompt( cloud_entity: BaseCloudLLMEntity, tmp_path: Path ) -> None: """Test that media attachments are converted to the expected payload.""" image_path = tmp_path / "doorbell.jpg" Image.new("RGB", (2, 2), "blue").save(image_path, "JPEG") pdf_path = tmp_path / "context.pdf" pdf_path.write_bytes(b"%PDF-1.3\nmock\n") attachments = [ conversation.Attachment( media_content_id="media-source://media/doorbell.jpg", mime_type="image/jpeg", path=image_path, ), conversation.Attachment( media_content_id="media-source://media/context.pdf", mime_type="application/pdf", path=pdf_path, ), ] files = await cloud_entity._async_prepare_files_for_prompt(attachments) assert files[0] == { "type": "input_image", "image_url": "data:image/jpeg;base64," + base64.b64encode(image_path.read_bytes()).decode(), "detail": "auto", } assert files[1] == { "type": "input_file", "filename": "context.pdf", "file_data": "data:application/pdf;base64," + base64.b64encode(pdf_path.read_bytes()).decode(), } async def test_prepare_files_for_prompt_invalid_type( cloud_entity: BaseCloudLLMEntity, tmp_path: Path ) -> None: """Test that unsupported attachments raise an error.""" text_path = tmp_path / "notes.txt" text_path.write_text("notes") attachments = [ conversation.Attachment( media_content_id="media-source://media/notes.txt", mime_type="text/plain", path=text_path, ) ] with pytest.raises( HomeAssistantError, match="Only images and PDF are currently supported as attachments", ): await cloud_entity._async_prepare_files_for_prompt(attachments) async def test_prepare_chat_for_generation_appends_attachments( hass: HomeAssistant, cloud_entity: BaseCloudLLMEntity, mock_prepare_files_for_prompt: AsyncMock, ) -> None: """Test chat preparation adds LLM tools, attachments, and metadata.""" chat_log = conversation.ChatLog(hass, "conversation-id") attachment = conversation.Attachment( media_content_id="media-source://media/doorbell.jpg", mime_type="image/jpeg", path=Path(hass.config.path("doorbell.jpg")), ) chat_log.async_add_user_content( conversation.UserContent(content="Describe the door", attachments=[attachment]) ) chat_log.llm_api = MagicMock( tools=[DummyTool()], custom_serializer=None, ) files = [{"type": "input_image", "image_url": "data://img", "detail": "auto"}] mock_prepare_files_for_prompt.return_value = files messages = _convert_content_to_param(chat_log.content) response = await cloud_entity._prepare_chat_for_generation( chat_log, messages, response_format={"type": "json"}, ) assert response["conversation_id"] == "conversation-id" assert response["response_format"] == {"type": "json"} assert response["tool_choice"] == "auto" assert len(response["tools"]) == 2 assert response["tools"][0]["name"] == "do_something" assert response["tools"][1]["type"] == "web_search" assert response["messages"] is messages mock_prepare_files_for_prompt.assert_awaited_once_with([attachment]) async def test_prepare_chat_for_generation_passes_messages_through( hass: HomeAssistant, cloud_entity: BaseCloudLLMEntity ) -> None: """Test that prepared messages are forwarded unchanged.""" chat_log = conversation.ChatLog(hass, "conversation-id") chat_log.async_add_assistant_content_without_tools( conversation.AssistantContent(agent_id="agent", content="Ready") ) messages = _convert_content_to_param(chat_log.content) response = await cloud_entity._prepare_chat_for_generation(chat_log, messages) assert response["messages"] == messages assert response["conversation_id"] == "conversation-id"