core/tests/components/cloud/test_entity.py

"""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.const import AI_TASK_ENTITY_UNIQUE_ID, DOMAIN
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 entity_registry as er, llm, selector
from homeassistant.setup import async_setup_component

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"


async def test_async_handle_chat_log_service_sets_structured_output_non_strict(
    hass: HomeAssistant,
    cloud: MagicMock,
    entity_registry: er.EntityRegistry,
    mock_cloud_login: None,
) -> None:
    """Ensure structured output requests always disable strict validation via service."""
    assert await async_setup_component(hass, DOMAIN, {})
    await hass.async_block_till_done()

    on_start_callback = cloud.register_on_start.call_args[0][0]
    await on_start_callback()
    await hass.async_block_till_done()

    entity_id = entity_registry.async_get_entity_id(
        "ai_task", DOMAIN, AI_TASK_ENTITY_UNIQUE_ID
    )
    assert entity_id is not None

    async def _empty_stream():
        return

    async def _fake_delta_stream(
        self: conversation.ChatLog,
        agent_id: str,
        stream,
    ):
        content = conversation.AssistantContent(
            agent_id=agent_id, content='{"value": "ok"}'
        )
        self.async_add_assistant_content_without_tools(content)
        yield content

    cloud.llm.async_generate_data = AsyncMock(return_value=_empty_stream())

    with patch(
        "homeassistant.components.conversation.chat_log.ChatLog.async_add_delta_content_stream",
        _fake_delta_stream,
    ):
        await hass.services.async_call(
            "ai_task",
            "generate_data",
            {
                "entity_id": entity_id,
                "task_name": "Device Report",
                "instructions": "Provide value.",
                "structure": {
                    "value": {
                        "selector": {"text": None},
                        "required": True,
                    }
                },
            },
            blocking=True,
            return_response=True,
        )

    cloud.llm.async_generate_data.assert_awaited_once()
    _, kwargs = cloud.llm.async_generate_data.call_args

    assert kwargs["response_format"]["json_schema"]["strict"] is False