haiku.rag/tests/test_utils.py

import asyncio
import importlib.util
from unittest.mock import AsyncMock

import pytest
from pydantic_ai.models.openai import OpenAIChatModel

from haiku.rag.config import get_config
from haiku.rag.config.models import ModelConfig
from haiku.rag.converters import get_converter
from haiku.rag.store.exceptions import ReadOnlyError
from haiku.rag.utils import gather_all, get_model

# Check for optional dependencies
HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None
HAS_GOOGLE = importlib.util.find_spec("google.genai") is not None
HAS_GROQ = importlib.util.find_spec("groq") is not None
HAS_BEDROCK = importlib.util.find_spec("botocore") is not None


@pytest.mark.asyncio
async def test_text_to_docling_document():
    """Test text to DoclingDocument conversion."""
    # Test basic text conversion
    simple_text = "This is a simple text document."
    converter = get_converter(get_config())
    doc = await converter.convert_text(simple_text)

    # Verify it returns a DoclingDocument
    from docling_core.types.doc.document import DoclingDocument

    assert isinstance(doc, DoclingDocument)

    # Verify the content can be exported back to markdown
    markdown = doc.export_to_markdown()
    assert "This is a simple text document." in markdown


@pytest.mark.asyncio
async def test_text_to_docling_document_with_custom_name():
    """Test text to DoclingDocument conversion with custom name parameter."""
    code_text = """# Python Code

```python
def hello():
    print("Hello, World!")
    return True
```"""

    converter = get_converter(get_config())
    doc = await converter.convert_text(code_text, name="hello.md")

    # Verify it's a valid DoclingDocument
    from docling_core.types.doc.document import DoclingDocument

    assert isinstance(doc, DoclingDocument)

    # Verify the content is preserved
    markdown = doc.export_to_markdown()
    assert "def hello():" in markdown
    assert "Hello, World!" in markdown


@pytest.mark.asyncio
async def test_text_to_docling_document_markdown_content():
    """Test text to DoclingDocument conversion with markdown content."""
    markdown_text = """# Test Document

This is a test document with:

- List item 1
- List item 2

## Code Example

```python
def test():
    return "Hello"
```

**Bold text** and *italic text*."""

    converter = get_converter(get_config())
    doc = await converter.convert_text(markdown_text, name="test.md")

    # Verify it's a DoclingDocument
    from docling_core.types.doc.document import DoclingDocument

    assert isinstance(doc, DoclingDocument)

    # Verify the markdown structure is preserved
    result_markdown = doc.export_to_markdown()
    assert "# Test Document" in result_markdown
    assert "List item 1" in result_markdown
    assert "def test():" in result_markdown


@pytest.mark.asyncio
async def test_text_to_docling_document_empty_content():
    """Test text to DoclingDocument conversion with empty content."""
    converter = get_converter(get_config())
    doc = await converter.convert_text("")

    # Should still create a valid DoclingDocument
    from docling_core.types.doc.document import DoclingDocument

    assert isinstance(doc, DoclingDocument)

    # Export should work even with empty content
    markdown = doc.export_to_markdown()
    assert isinstance(markdown, str)


@pytest.mark.asyncio
async def test_text_to_docling_document_unicode_content():
    """Test text to DoclingDocument conversion with unicode content."""
    unicode_text = """# 测试文档

这是一个包含中文的测试文档。

## Código en Español
```javascript
function saludar() {
    return "¡Hola mundo!";
}
```

Emoji test: 🚀 ✅ 📝"""

    converter = get_converter(get_config())
    doc = await converter.convert_text(unicode_text, name="unicode.md")

    # Verify it's a DoclingDocument
    from docling_core.types.doc.document import DoclingDocument

    assert isinstance(doc, DoclingDocument)

    # Verify unicode content is preserved
    result_markdown = doc.export_to_markdown()
    assert "测试文档" in result_markdown
    assert "¡Hola mundo!" in result_markdown
    assert "🚀" in result_markdown


@pytest.mark.parametrize(
    "kwargs,expected_settings",
    [
        ({"provider": "ollama", "name": "llama3"}, None),
        (
            {"provider": "ollama", "name": "gpt-oss", "enable_thinking": False},
            {"openai_reasoning_effort": "low"},
        ),
        (
            {"provider": "ollama", "name": "gpt-oss", "enable_thinking": True},
            {"openai_reasoning_effort": "high"},
        ),
        (
            {
                "provider": "ollama",
                "name": "llama3",
                "temperature": 0.5,
                "max_tokens": 100,
            },
            {"temperature": 0.5, "max_tokens": 100},
        ),
        ({"provider": "openai", "name": "gpt-4o"}, None),
        (
            {"provider": "openai", "name": "o1", "enable_thinking": True},
            {"openai_reasoning_effort": "high"},
        ),
        (
            {"provider": "openai", "name": "o1", "enable_thinking": False},
            {"openai_reasoning_effort": "low"},
        ),
        (
            {
                "provider": "openai",
                "name": "gpt-4o",
                "enable_thinking": False,
                "temperature": 0.7,
                "max_tokens": 500,
            },
            # gpt-4o is not a reasoning model, so only the common settings land.
            {"temperature": 0.7, "max_tokens": 500},
        ),
    ],
    ids=[
        "ollama",
        "ollama_thinking_off",
        "ollama_thinking_on",
        "ollama_with_settings",
        "openai",
        "openai_reasoning_thinking_on",
        "openai_reasoning_thinking_off",
        "openai_all_settings",
    ],
)
def test_get_model_openai_chat_settings(kwargs, expected_settings):
    """Each ollama/openai configuration maps onto the expected model settings."""
    result = get_model(ModelConfig(**kwargs))

    assert isinstance(result, OpenAIChatModel)
    if expected_settings is None:
        assert result.settings is None
        return
    assert result.settings is not None
    for key, value in expected_settings.items():
        assert result.settings.get(key) == value


def test_get_model_ollama_appends_v1_to_per_model_base_url():
    """Per-model base_url without /v1 should get it appended."""
    model_config = ModelConfig(
        provider="ollama", name="qwen3.6", base_url="http://my-ollama:11434"
    )
    result = get_model(model_config)
    assert isinstance(result, OpenAIChatModel)
    assert str(result.client.base_url).rstrip("/").endswith("/v1")


def test_get_model_ollama_does_not_double_append_v1():
    """If the per-model base_url already ends with /v1, leave it alone."""
    model_config = ModelConfig(
        provider="ollama", name="qwen3.6", base_url="http://my-ollama:11434/v1"
    )
    result = get_model(model_config)
    url = str(result.client.base_url).rstrip("/")
    assert url.endswith("/v1")
    assert not url.endswith("/v1/v1")


def test_get_model_openai_non_reasoning_model_ignores_thinking():
    """Test that non-reasoning OpenAI models don't get reasoning_effort setting."""
    model_config = ModelConfig(
        provider="openai", name="gpt-4o-mini", enable_thinking=False
    )
    result = get_model(model_config)
    assert isinstance(result, OpenAIChatModel)
    # Non-reasoning models should not have reasoning_effort set
    assert result._settings is None


def test_get_model_vllm_model_without_reasoning_profile_sends_no_thinking():
    """A vLLM-served model with no reasoning profile carries no thinking settings.

    Its chat template reads the switch from `chat_template_kwargs`, which only
    `extra_body` can reach, and the endpoint rejects `reasoning_effort`.
    """
    model_config = ModelConfig(
        provider="openai",
        name="Qwen/Qwen3-32B",
        base_url="http://vllm:8000/v1",
        enable_thinking=True,
        temperature=0.2,
    )
    result = get_model(model_config)

    assert isinstance(result, OpenAIChatModel)
    assert result._settings is not None
    assert "thinking" not in result._settings
    assert "openai_reasoning_effort" not in result._settings


def test_get_model_openai_extra_body_forwarded():
    """`extra_body` on ModelConfig is forwarded to ModelSettings.extra_body.

    pydantic-ai's OpenAI model branch reads `model_settings["extra_body"]`
    and passes it verbatim to the OpenAI SDK. Enables vLLM-specific keys
    like `chat_template_kwargs.enable_thinking` without coupling them to
    the high-level `enable_thinking` flag.
    """
    extra = {"chat_template_kwargs": {"enable_thinking": False}}
    model_config = ModelConfig(
        provider="openai",
        name="qwen3.6-35b",
        base_url="http://localhost:11430/v1",
        extra_body=extra,
    )
    result = get_model(model_config)
    assert isinstance(result, OpenAIChatModel)
    assert result._settings is not None
    assert result._settings.get("extra_body") == extra


@pytest.mark.asyncio
async def test_get_model_merges_system_messages_for_openai_compatible():
    """Instruction parts from multiple sources (agent preamble, capability
    instructions, dynamic notices) each map to their own system message.
    Strict OpenAI-compatible templates (e.g. Qwen on vLLM) reject more than
    one leading system message, so OpenAI-compatible endpoints merge them;
    plain OpenAI keeps them separate."""
    from pydantic_ai.messages import InstructionPart, ModelRequest, UserPromptPart
    from pydantic_ai.models import ModelRequestParameters

    parts = [
        InstructionPart(content="Base instructions."),
        InstructionPart(content="Limit notice.", dynamic=True),
    ]
    messages = [ModelRequest(parts=[UserPromptPart("hi")])]

    async def mapped_for(model):
        return await model._map_messages(
            messages, ModelRequestParameters(instruction_parts=parts)
        )

    vllm = get_model(
        ModelConfig(provider="openai", name="qwen3.6", base_url="http://vllm:1/v1")
    )
    mapped = await mapped_for(vllm)
    assert [m["role"] for m in mapped] == ["system", "user"]
    assert mapped[0]["content"] == "Base instructions.\n\nLimit notice."

    ollama = get_model(ModelConfig(provider="ollama", name="llama3"))
    assert [m["role"] for m in await mapped_for(ollama)] == ["system", "user"]

    openai_native = get_model(ModelConfig(provider="openai", name="gpt-4o"))
    assert [m["role"] for m in await mapped_for(openai_native)] == [
        "system",
        "system",
        "user",
    ]


def test_get_model_ollama_extra_body_forwarded():
    """`extra_body` is forwarded through the Ollama (openai-compatible) branch too."""
    extra = {"chat_template_kwargs": {"enable_thinking": False}}
    model_config = ModelConfig(provider="ollama", name="qwen3", extra_body=extra)
    result = get_model(model_config)
    assert isinstance(result, OpenAIChatModel)
    assert result._settings is not None
    assert result._settings.get("extra_body") == extra


def test_get_model_extra_body_absent_when_unset():
    """No `extra_body` key appears on the settings when the config omits it."""
    model_config = ModelConfig(provider="openai", name="gpt-4o-mini", temperature=0.3)
    result = get_model(model_config)
    assert isinstance(result, OpenAIChatModel)
    # temperature triggers settings construction; extra_body should not be there.
    assert result._settings is not None
    assert "extra_body" not in result._settings


@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
def test_get_model_anthropic():
    """Test get_model returns AnthropicModel for Anthropic."""
    from pydantic_ai.models.anthropic import AnthropicModel

    model_config = ModelConfig(provider="anthropic", name="claude-3-5-sonnet-20241022")
    result = get_model(model_config)
    assert isinstance(result, AnthropicModel)


@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
def test_get_model_anthropic_with_thinking():
    """Test get_model configures thinking for Anthropic."""
    from pydantic_ai.models.anthropic import AnthropicModel

    model_config = ModelConfig(
        provider="anthropic",
        name="claude-3-5-sonnet-20241022",
        enable_thinking=True,
    )
    result = get_model(model_config)

    assert isinstance(result, AnthropicModel)
    assert result.settings is not None
    assert result.settings.get("thinking") is True


@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
def test_get_model_anthropic_thinking_off_disables_adaptive_models():
    """Adaptive-thinking models think by default, so off must be explicit.

    The unified `thinking=False` omits the request field, which leaves Sonnet
    4.6+ and Opus 4.6+ thinking.
    """
    from pydantic_ai.models.anthropic import AnthropicModel

    model_config = ModelConfig(
        provider="anthropic", name="claude-sonnet-4-6", enable_thinking=False
    )
    result = get_model(model_config)

    assert isinstance(result, AnthropicModel)
    assert result.settings is not None
    assert result.settings.get("anthropic_thinking") == {"type": "disabled"}
    # The explicit disable replaces the unified key rather than joining it.
    assert "thinking" not in result.settings


@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed")
def test_get_model_google():
    """Test get_model returns GoogleModel for Google."""
    from pydantic_ai.models.google import GoogleModel

    model_config = ModelConfig(provider="google", name="gemini-2.0-flash-exp")
    result = get_model(model_config)
    assert isinstance(result, GoogleModel)


@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed")
@pytest.mark.parametrize("enable_thinking", [True, False])
def test_get_model_google_with_thinking(enable_thinking):
    """Test get_model configures thinking for Google."""
    from pydantic_ai.models.google import GoogleModel

    model_config = ModelConfig(
        provider="google",
        name="gemini-2.0-flash-thinking-exp",
        enable_thinking=enable_thinking,
    )
    result = get_model(model_config)

    assert isinstance(result, GoogleModel)
    assert result.settings is not None
    assert result.settings.get("thinking") == enable_thinking


@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed")
def test_get_model_groq():
    """Test get_model returns GroqModel for Groq."""
    from pydantic_ai.models.groq import GroqModel

    model_config = ModelConfig(provider="groq", name="llama-3.3-70b-versatile")
    result = get_model(model_config)
    assert isinstance(result, GroqModel)


@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed")
@pytest.mark.parametrize("enable_thinking", [True, False])
def test_get_model_groq_with_thinking(enable_thinking):
    """Test get_model configures thinking for Groq."""
    from pydantic_ai.models.groq import GroqModel

    model_config = ModelConfig(
        provider="groq",
        name="llama-3.3-70b-versatile",
        enable_thinking=enable_thinking,
    )
    result = get_model(model_config)

    assert isinstance(result, GroqModel)
    assert result.settings is not None
    assert result.settings.get("thinking") == enable_thinking


@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
def test_get_model_bedrock():
    """Test get_model returns BedrockConverseModel for Bedrock."""
    from pydantic_ai.models.bedrock import BedrockConverseModel

    model_config = ModelConfig(
        provider="bedrock", name="anthropic.claude-3-5-sonnet-20241022-v2:0"
    )
    result = get_model(model_config)
    assert isinstance(result, BedrockConverseModel)


@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
@pytest.mark.parametrize(
    "name",
    [
        "anthropic.claude-3-5-sonnet-20241022-v2:0",
        "openai.gpt-oss-120b-1:0",
        "qwen.qwen3-32b-v1:0",
        "meta.llama3-70b-instruct-v1:0",
    ],
    ids=["claude", "gpt_oss", "qwen", "unmapped"],
)
def test_get_model_bedrock_with_thinking(name):
    """Every Bedrock family carries the unified thinking setting."""
    from pydantic_ai.models.bedrock import BedrockConverseModel

    model_config = ModelConfig(
        provider="bedrock",
        name=name,
        enable_thinking=True,
    )
    result = get_model(model_config)

    assert isinstance(result, BedrockConverseModel)
    assert result.settings is not None
    assert result.settings.get("thinking") is True


@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
@pytest.mark.parametrize(
    "name",
    [
        "anthropic.claude-sonnet-4-6-20260514-v1:0",
        "us.anthropic.claude-sonnet-4-6-20260514-v1:0",
    ],
    ids=["plain", "cross_region"],
)
def test_get_model_bedrock_thinking_off_disables_adaptive_claude(name):
    """Bedrock omits the field for adaptive Claude, which leaves it thinking."""
    from pydantic_ai.models.bedrock import BedrockConverseModel

    model_config = ModelConfig(provider="bedrock", name=name, enable_thinking=False)
    result = get_model(model_config)

    assert isinstance(result, BedrockConverseModel)
    assert result.settings is not None
    assert result.settings.get("bedrock_additional_model_requests_fields") == {
        "thinking": {"type": "disabled"}
    }
    # The explicit disable replaces the unified key rather than joining it.
    assert "thinking" not in result.settings


@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
def test_get_model_bedrock_thinking_off_leaves_non_claude_families_alone():
    """Only the Anthropic variant takes a `thinking` request field."""
    from pydantic_ai.models.bedrock import BedrockConverseModel

    model_config = ModelConfig(
        provider="bedrock", name="qwen.qwen3-32b-v1:0", enable_thinking=False
    )
    result = get_model(model_config)

    assert isinstance(result, BedrockConverseModel)
    assert result.settings is not None
    assert "bedrock_additional_model_requests_fields" not in result.settings
    assert result.settings.get("thinking") is False


@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
def test_get_model_bedrock_rejects_mantle_only_model():
    """Proprietary OpenAI models are Bedrock Mantle-only, not served by Converse."""
    from pydantic_ai.exceptions import UserError

    model_config = ModelConfig(provider="bedrock", name="openai.o3-mini-v1:0")

    with pytest.raises(UserError):
        get_model(model_config)


def test_get_model_passthrough_for_unbranched_provider():
    """A provider pydantic-ai knows but we do not branch on passes through as a
    string, so a new pydantic-ai provider needs no haiku.rag release."""
    model_config = ModelConfig(provider="mistral", name="mistral-large-latest")
    result = get_model(model_config)
    assert isinstance(result, str)
    assert result == "mistral:mistral-large-latest"


def test_get_model_accepts_provider_whose_sdk_is_missing(monkeypatch):
    """A missing vendor SDK is not an unknown provider: that ImportError names
    the extra to install, so it must reach the caller unchanged.

    Uses a provider whose SDK *is* installed, so the patch is what produces the
    ImportError rather than the environment.
    """
    import pydantic_ai.providers

    def _missing_sdk(provider: str):
        raise ImportError("Please install the `cohere` package")

    monkeypatch.setattr(pydantic_ai.providers, "infer_provider_class", _missing_sdk)
    result = get_model(ModelConfig(provider="cohere", name="command-r"))

    assert result == "cohere:command-r"


@pytest.mark.parametrize("provider", ["nonsense", "vllm", "gemini"])
def test_get_model_rejects_unknown_provider(provider):
    """An unknown provider is named here rather than passed through to fail
    inside pydantic-ai, where nothing identifies the config it came from.

    `vllm` and `gemini` get no special case: both were haiku.rag's own
    vocabulary, and both fail the same way as a typo.
    """
    with pytest.raises(ValueError, match=provider):
        get_model(ModelConfig(provider=provider, name="whatever"))


def test_get_package_versions():
    """Test get_package_versions returns expected keys."""
    from haiku.rag.utils import get_package_versions

    versions = get_package_versions()

    assert "haiku_rag" in versions
    assert "lancedb" in versions
    assert "docling" in versions
    assert "pydantic_ai" in versions
    assert "docling_document_schema" in versions

    # All should be non-empty strings
    for value in versions.values():
        assert isinstance(value, str)
        assert len(value) > 0


# --- apply_common_settings tests ---


def test_apply_common_settings_no_settings():
    from haiku.rag.config.models import ModelConfig
    from haiku.rag.utils import apply_common_settings

    mc = ModelConfig(provider="openai", name="gpt-4o")
    result = apply_common_settings(None, mc)
    assert result is None


def test_apply_common_settings_temperature():
    from haiku.rag.config.models import ModelConfig
    from haiku.rag.utils import apply_common_settings

    mc = ModelConfig(provider="openai", name="gpt-4o", temperature=0.7)
    result = apply_common_settings(None, mc)
    assert result is not None
    assert result["temperature"] == 0.7


def test_apply_common_settings_max_tokens():
    from haiku.rag.config.models import ModelConfig
    from haiku.rag.utils import apply_common_settings

    mc = ModelConfig(provider="openai", name="gpt-4o", max_tokens=500)
    result = apply_common_settings(None, mc)
    assert result is not None
    assert result["max_tokens"] == 500


def test_apply_common_settings_existing():
    from haiku.rag.config.models import ModelConfig
    from haiku.rag.utils import apply_common_settings

    mc = ModelConfig(provider="openai", name="gpt-4o", temperature=0.5)
    existing = {"some_key": "value"}
    result = apply_common_settings(existing, mc)
    assert result is not None
    assert result["temperature"] == 0.5
    assert result["some_key"] == "value"


# --- format_bytes tests ---


def test_format_bytes():
    from haiku.rag.utils import format_bytes

    assert format_bytes(0) == "0.0 B"
    assert format_bytes(512) == "512.0 B"
    assert format_bytes(1024) == "1.0 KB"
    assert format_bytes(1048576) == "1.0 MB"
    assert format_bytes(1073741824) == "1.0 GB"
    assert format_bytes(1099511627776) == "1.0 TB"
    assert format_bytes(1125899906842624) == "1.0 PB"


# --- format_citations tests ---


def test_format_citations_empty():
    from haiku.rag.utils import format_citations

    assert format_citations([]) == ""


def test_format_citations_with_citation():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations

    citation = Citation(
        document_id="doc1",
        chunk_id="chunk1",
        document_uri="test://doc",
        document_title="Test Doc",
        content="Some content",
        page_numbers=[1],
        headings=["Intro"],
    )
    result = format_citations([citation])
    assert "[1] Test Doc" in result
    assert "doc1" not in result
    assert "chunk1" not in result
    assert "test://doc" in result
    assert "p. 1" in result
    assert "Section: Intro" in result
    assert "Some content" in result


def test_format_citations_multiple_pages():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations

    citation = Citation(
        document_id="doc1",
        chunk_id="chunk1",
        document_uri="test://doc",
        content="Content",
        page_numbers=[1, 2, 3],
    )
    result = format_citations([citation])
    assert "[1] test://doc" in result
    assert "pp. 1-3" in result
    # No title: the URI stands in, and the document id never leaks.
    assert "doc1" not in result


def test_format_citations_with_index():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations

    citation = Citation(
        index=5,
        document_id="doc1",
        chunk_id="chunk1",
        document_uri="test://doc",
        document_title="Test Doc",
        content="Content",
    )
    result = format_citations([citation])
    assert "[5] Test Doc" in result


def test_format_citations_sequential_indices():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations

    citations = [
        Citation(
            document_id="doc1",
            chunk_id="chunk1",
            document_uri="test://doc1",
            document_title="First",
            content="Content 1",
        ),
        Citation(
            document_id="doc2",
            chunk_id="chunk2",
            document_uri="test://doc2",
            document_title="Second",
            content="Content 2",
        ),
    ]
    result = format_citations(citations)
    assert "[1] First" in result
    assert "[2] Second" in result


# --- format_citations tests (pictures) ---


def test_format_citations_picture_refs_render_as_markers():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations

    citation = Citation(
        document_id="doc1",
        chunk_id="chunk1",
        document_uri="test://doc",
        document_title="Test Doc",
        content="text body",
        picture_refs=["#/pictures/0", "#/pictures/3"],
    )
    result = format_citations([citation])
    assert "[Figure: #/pictures/0]" in result
    assert "[Figure: #/pictures/3]" in result


# --- format_citations_rich tests ---


def _render_rich(renderables: list) -> str:
    from rich.console import Console

    console = Console(record=True, width=200)
    for r in renderables:
        console.print(r)
    return console.export_text()


async def test_format_citations_rich_empty():
    from haiku.rag.utils import format_citations_rich

    assert await format_citations_rich([]) == []


async def test_format_citations_rich_header_and_footer():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations_rich

    citation = Citation(
        document_id="doc-uuid-1",
        chunk_id="chunk-uuid-1",
        document_uri="test://doc",
        document_title="Test Doc",
        content="Body",
        page_numbers=[1, 2, 3],
        headings=["Intro", "Background"],
    )
    output = _render_rich(await format_citations_rich([citation]))
    assert "Citations" in output
    assert "[1] Test Doc (test://doc)" in output
    assert "pp. 1-3" in output
    assert "§Background" in output
    assert "doc: doc-uuid-1" in output
    assert "chunk: chunk-uuid-1" in output


async def test_format_citations_rich_names_the_database_when_federating():
    """Across databases, a citation has to say which one it came from."""
    from unittest.mock import AsyncMock

    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations_rich

    citation = Citation(
        document_id="doc-uuid-1",
        chunk_id="chunk-uuid-1",
        document_uri="test://doc",
        document_title="Test Doc",
        content="Body",
        source="papers",
    )
    client = AsyncMock()
    client.covers_multiple = True
    client.source_names = ("papers", "notes")

    output = _render_rich(await format_citations_rich([citation], client))

    assert "papers" in output


async def test_an_unattributable_picture_renders_its_marker(tmp_path):
    """A citation without a source has no picture owner across databases. One
    unrenderable figure must not cost the answer."""
    from rich.console import Console

    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations_rich

    covering = AsyncMock()
    covering.covers_multiple = True
    citation = Citation(
        document_id="d1",
        chunk_id="c1",
        content="body",
        document_uri="test://doc",
        picture_refs=["#/pictures/0"],
    )

    renderables = await format_citations_rich([citation], covering)

    console = Console(record=True, width=200)
    for renderable in renderables:
        console.print(renderable)
    assert "[Figure: #/pictures/0]" in console.export_text()
    covering.get_picture_bytes.assert_not_awaited()


def test_truncated_marks_what_it_dropped():
    """An unmarked cut reads as the value: a sentence ending "in 1991" becomes
    one ending "in 1"."""
    from haiku.rag.utils import truncated

    sentence = "Station Kestrel sits at 980 metres and was commissioned in 1991."

    assert truncated(sentence, 60) == sentence[:60] + "…"
    assert truncated(sentence, len(sentence)) == sentence
    assert truncated("short", 60) == "short"
    # Trailing space before the mark reads as a gap in the text.
    assert truncated("a bc", 2) == "a…"


async def test_format_citations_rich_omits_the_database_for_one_database():
    """A single database is not worth naming on every citation."""
    from unittest.mock import AsyncMock

    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations_rich

    citation = Citation(
        document_id="doc-uuid-1",
        chunk_id="chunk-uuid-1",
        document_uri="test://doc",
        document_title="Test Doc",
        content="Body",
        source="papers",
    )
    client = AsyncMock()
    client.covers_multiple = False
    client.source_names = ()

    output = _render_rich(await format_citations_rich([citation], client))

    assert "papers" not in output


async def test_format_citations_rich_truncates_long_content():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import CITATION_PREVIEW_CHARS, format_citations_rich

    citation = Citation(
        document_id="doc1",
        chunk_id="chunk1",
        document_uri="test://doc",
        content="A" * (CITATION_PREVIEW_CHARS + 200),
    )
    output = _render_rich(await format_citations_rich([citation]))
    assert "…" in output
    assert "A" * (CITATION_PREVIEW_CHARS + 1) not in output


async def test_format_citations_rich_full_keeps_the_whole_content():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import CITATION_PREVIEW_CHARS, format_citations_rich

    content = "A" * (CITATION_PREVIEW_CHARS + 200)
    citation = Citation(
        document_id="doc1",
        chunk_id="chunk1",
        document_uri="test://doc",
        content=content,
    )
    output = _render_rich(await format_citations_rich([citation], full=True))
    assert "…" not in output
    # Rich wraps the body across panel lines, so count the content instead.
    assert output.count("A") == len(content)


async def test_format_citations_rich_picture_marker_without_client():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations_rich

    citation = Citation(
        document_id="doc1",
        chunk_id="chunk1",
        document_uri="test://doc",
        content="body",
        picture_refs=["#/pictures/0"],
    )
    output = _render_rich(await format_citations_rich([citation]))
    assert "[Figure: #/pictures/0]" in output


# --- get_default_data_dir tests ---


def test_get_default_data_dir():
    from pathlib import Path

    from haiku.rag.utils import get_default_data_dir

    result = get_default_data_dir()
    assert isinstance(result, Path)
    assert "haiku.rag" in str(result)


# --- build_prompt tests ---


def test_build_prompt_without_preamble():
    from haiku.rag.config.models import AppConfig
    from haiku.rag.utils import build_prompt

    config = AppConfig()
    result = build_prompt("Base prompt", config)
    assert result == "Base prompt"


def test_build_prompt_with_preamble():
    from haiku.rag.config.models import AppConfig, PromptsConfig
    from haiku.rag.utils import build_prompt

    config = AppConfig(prompts=PromptsConfig(domain_preamble="You are a legal expert."))
    result = build_prompt("Base prompt", config)
    assert result == "You are a legal expert.\n\nBase prompt"


# --- is_up_to_date tests ---


def test_cosine_similarity_zero_norm():
    from haiku.rag.utils import cosine_similarity

    assert cosine_similarity([0, 0, 0], [1, 2, 3]) == 0.0
    assert cosine_similarity([1, 2, 3], [0, 0, 0]) == 0.0
    assert cosine_similarity([0, 0], [0, 0]) == 0.0


@pytest.mark.asyncio
async def test_is_up_to_date(monkeypatch):
    from unittest.mock import AsyncMock, MagicMock

    import httpx

    from haiku.rag.utils import is_up_to_date

    mock_response = MagicMock()
    mock_response.json.return_value = {"info": {"version": "0.0.1"}}

    mock_client = AsyncMock()
    mock_client.get = AsyncMock(return_value=mock_response)
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=None)

    monkeypatch.setattr(httpx, "AsyncClient", lambda: mock_client)

    is_current, running, latest = await is_up_to_date()
    assert is_current is True
    assert running >= latest


# --- parse_model_option tests ---


def test_parse_model_option():
    from haiku.rag.utils import parse_model_option

    result = parse_model_option("anthropic:claude-sonnet-4-20250514")
    assert result.provider == "anthropic"
    assert result.name == "claude-sonnet-4-20250514"

    # Colons in name are preserved
    assert parse_model_option("openai:gpt-4o:latest").name == "gpt-4o:latest"

    for bad in ["just-a-name", ":model", "provider:"]:
        with pytest.raises(ValueError, match="Invalid model format"):
            parse_model_option(bad)


def test_cosine_similarity_identical_vectors():
    from haiku.rag.utils import cosine_similarity

    assert cosine_similarity([1.0, 0.0], [1.0, 0.0]) == pytest.approx(1.0)
    assert cosine_similarity([1.0, 0.0], [0.0, 1.0]) == pytest.approx(0.0)


async def test_format_citations_rich_separates_multiple_citations():
    from haiku.rag.store.models.citation import Citation
    from haiku.rag.utils import format_citations_rich

    citations = [
        Citation(
            document_id=f"doc{i}",
            chunk_id=f"chunk{i}",
            document_uri=f"test://doc{i}",
            document_title=f"Doc {i}",
            content=f"Body {i}",
        )
        for i in (1, 2)
    ]

    output = _render_rich(await format_citations_rich(citations))

    assert "[1] Doc 1 (test://doc1)" in output
    assert "[2] Doc 2 (test://doc2)" in output


@pytest.mark.parametrize(
    "stored,renders",
    [
        (None, False),
        (b"not a real image", False),
        ("png", True),
    ],
    ids=["no_bytes", "undecodable_bytes", "valid_png"],
)
async def test_render_picture_handles_stored_bytes(stored, renders):
    from unittest.mock import AsyncMock

    from haiku.rag.utils import _render_picture

    if stored == "png":
        from io import BytesIO

        from PIL import Image as PILImage

        buf = BytesIO()
        PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG")
        stored = buf.getvalue()

    client = AsyncMock()
    client.covers_multiple = False
    client.get_picture_bytes = AsyncMock(return_value=stored)

    result = await _render_picture(client, "doc1", "#/pictures/0")

    if renders:
        from textual_image.renderable import Image as RichImage

        assert isinstance(result, RichImage)
    else:
        assert result is None


async def test_render_picture_without_client_returns_none():
    from haiku.rag.utils import _render_picture

    assert await _render_picture(None, "doc1", "#/pictures/0") is None


def test_get_package_versions_reports_missing_docling(monkeypatch):
    from importlib import metadata as importlib_metadata

    from haiku.rag.utils import get_package_versions

    real_version = importlib_metadata.version

    def fake_version(name):
        if name == "docling":
            raise importlib_metadata.PackageNotFoundError(name)
        return real_version(name)

    monkeypatch.setattr(importlib_metadata, "version", fake_version)

    assert get_package_versions()["docling"] == "not installed"


def test_get_model_openai_api_key_from_config():
    """A config-supplied api_key reaches the client, so several
    openai-compatible endpoints can each carry their own key."""
    result = get_model(
        ModelConfig(
            provider="openai",
            name="qwen3.6",
            base_url="http://vllm:8000/v1",
            api_key="sk-vendor-a",
        )
    )
    assert result.client.api_key == "sk-vendor-a"


def test_get_model_openai_api_key_without_base_url_overrides_env():
    result = get_model(
        ModelConfig(provider="openai", name="gpt-4o", api_key="sk-vendor-b")
    )
    assert result.client.api_key == "sk-vendor-b"


def test_get_model_ollama_api_key_from_config():
    result = get_model(
        ModelConfig(
            provider="ollama",
            name="gpt-oss",
            base_url="http://remote-ollama:11434/v1",
            api_key="sk-proxy",
        )
    )
    assert result.client.api_key == "sk-proxy"


def test_get_model_api_key_rejected_on_unplumbed_provider():
    """Providers whose client we never build read their own vendor variable;
    an api_key there would be silently dropped."""
    with pytest.raises(ValueError, match="api_key is not supported"):
        get_model(
            ModelConfig(provider="anthropic", name="claude-sonnet-4-5", api_key="sk-x")
        )


class TestGatherAll:
    """A fan-out leaves nothing running: a caller unwinding from a failure closes
    the sessions its siblings are still reading through."""

    @pytest.mark.asyncio
    async def test_results_arrive_in_the_order_asked_for(self):
        async def slow(value):
            await asyncio.sleep(0.01)
            return value

        async def fast(value):
            return value

        assert await gather_all(slow("a"), fast("b"), slow("c")) == ["a", "b", "c"]

    @pytest.mark.asyncio
    async def test_a_failure_leaves_no_sibling_running(self):
        started = asyncio.Event()
        unwound = asyncio.Event()

        async def sibling():
            started.set()
            try:
                await asyncio.sleep(60)
            finally:
                unwound.set()

        async def failing():
            await started.wait()
            raise ReadOnlyError("beta is read-only")

        before = asyncio.all_tasks()
        with pytest.raises(ReadOnlyError, match="beta"):
            await gather_all(sibling(), failing())

        assert unwound.is_set()
        assert asyncio.all_tasks() - before == set()

    @pytest.mark.asyncio
    async def test_the_failure_arrives_as_itself(self):
        """A `TaskGroup` drains the siblings too, but raises an `ExceptionGroup`."""

        async def failing():
            raise ReadOnlyError("beta is read-only")

        async def sibling():
            await asyncio.sleep(60)

        with pytest.raises(ReadOnlyError) as raised:
            await gather_all(sibling(), failing())

        assert not isinstance(raised.value, BaseExceptionGroup)