haiku.rag/tests/test_client.py

import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, patch

import httpx
import pytest
from datasets import Dataset

from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document


@pytest.mark.asyncio
async def test_client_document_crud(qa_corpus: Dataset, temp_db_path):
    """Test HaikuRAG CRUD operations for documents."""
    async with HaikuRAG(temp_db_path) as client:
        # Get test data
        first_doc = qa_corpus[0]
        document_text = first_doc["document_extracted"]
        test_uri = "file:///path/to/test.txt"
        test_metadata = {"source": "test", "topic": "testing"}

        # Test create_document
        created_doc = await client.create_document(
            content=document_text, uri=test_uri, metadata=test_metadata
        )

        assert created_doc.id is not None
        assert created_doc.content == document_text
        assert created_doc.uri == test_uri
        assert created_doc.metadata == test_metadata

        # Test get_document_by_id
        retrieved_doc = await client.get_document_by_id(created_doc.id)
        assert retrieved_doc is not None
        assert retrieved_doc.id == created_doc.id
        assert retrieved_doc.content == document_text
        assert retrieved_doc.uri == test_uri

        # Test get_document_by_uri
        retrieved_by_uri = await client.get_document_by_uri(test_uri)
        assert retrieved_by_uri is not None
        assert retrieved_by_uri.id == created_doc.id
        assert retrieved_by_uri.content == document_text

        # Test get_document_by_uri with non-existent URI
        non_existent = await client.get_document_by_uri("file:///non/existent.txt")
        assert non_existent is None

        # Test update_document
        retrieved_doc.content = "Updated content"
        retrieved_doc.uri = "file:///updated/path.txt"
        updated_doc = await client.update_document(retrieved_doc)
        assert updated_doc.content == "Updated content"
        assert updated_doc.uri == "file:///updated/path.txt"

        # Test list_documents
        all_docs = await client.list_documents()
        assert len(all_docs) == 1
        assert all_docs[0].id == created_doc.id

        # Test list_documents with pagination
        limited_docs = await client.list_documents(limit=10, offset=0)
        assert len(limited_docs) == 1

        # Test delete_document
        deleted = await client.delete_document(created_doc.id)
        assert deleted is True

        # Verify document is gone
        retrieved_doc = await client.get_document_by_id(created_doc.id)
        assert retrieved_doc is None

        # Test delete non-existent document
        deleted_again = await client.delete_document(created_doc.id)
        assert deleted_again is False


@pytest.mark.asyncio
async def test_client_update_document_fields(qa_corpus: Dataset, temp_db_path):
    """Test updating document with individual parameters."""
    async with HaikuRAG(temp_db_path) as client:
        # Get test data
        first_doc = qa_corpus[0]
        document_text = first_doc["document_extracted"]
        test_uri = "file:///path/to/test.txt"
        test_metadata = {"source": "test", "topic": "testing"}

        # Create a document
        created_doc = await client.create_document(
            content=document_text,
            uri=test_uri,
            title="Original Title",
            metadata=test_metadata,
        )
        assert created_doc.id is not None
        original_id = created_doc.id

        # Test updating only content
        updated_doc = await client.update_document_fields(
            document_id=original_id, content="Updated content only"
        )
        assert updated_doc.id == original_id
        assert updated_doc.content == "Updated content only"
        assert updated_doc.title == "Original Title"
        assert updated_doc.uri == test_uri

        # Test updating only metadata
        new_metadata = {"source": "updated", "version": "2.0"}
        updated_doc = await client.update_document_fields(
            document_id=original_id, metadata=new_metadata
        )
        assert updated_doc.metadata == new_metadata
        assert (
            updated_doc.content == "Updated content only"
        )  # Should keep previous update

        # Test updating only title
        updated_doc = await client.update_document_fields(
            document_id=original_id, title="New Title"
        )
        assert updated_doc.title == "New Title"
        assert updated_doc.content == "Updated content only"
        assert updated_doc.metadata == new_metadata

        # Test updating multiple fields at once
        custom_chunks = [
            Chunk(content="Custom chunk 1", order=0),
            Chunk(content="Custom chunk 2", order=1),
        ]
        updated_doc = await client.update_document_fields(
            document_id=original_id,
            content="Content with custom chunks",
            title="Final Title",
            metadata={"final": "true"},
            chunks=custom_chunks,
        )
        assert updated_doc.id == original_id
        assert updated_doc.content == "Content with custom chunks"
        assert updated_doc.title == "Final Title"
        assert updated_doc.metadata == {"final": "true"}

        # Verify the custom chunks were created
        doc_chunks = await client.chunk_repository.get_by_document_id(original_id)
        assert len(doc_chunks) == 2
        assert doc_chunks[0].content == "Custom chunk 1"
        assert doc_chunks[1].content == "Custom chunk 2"


@pytest.mark.asyncio
async def test_client_create_document_from_source(temp_db_path):
    """Test creating a document from a file source."""
    async with HaikuRAG(temp_db_path) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            test_content = "This is test content from a file."
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text(test_content)

            # Test create_document_from_source with Path
            doc = await client.create_document_from_source(source=temp_path)
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert doc.content == test_content
            assert doc.uri == temp_path.as_uri()
            assert "contentType" in doc.metadata
            assert "md5" in doc.metadata
            assert doc.metadata["contentType"] == "text/plain"

            # Test create_document_from_source with string path
            doc2 = await client.create_document_from_source(source=str(temp_path))
            assert isinstance(doc2, Document)

            assert doc2.id is not None
            assert doc2.content == test_content
            assert doc2.uri == temp_path.as_uri()
            assert "contentType" in doc2.metadata
            assert "md5" in doc2.metadata


@pytest.mark.asyncio
async def test_client_create_document_from_source_with_title(temp_db_path):
    """Test creating a document from a file source with a title."""
    async with HaikuRAG(temp_db_path) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            test_content = "This is test content from a file."
            temp_path = Path(temp_dir) / "test_title.txt"
            temp_path.write_text(test_content)

            doc = await client.create_document_from_source(
                source=temp_path, title="My Doc"
            )
            assert isinstance(doc, Document)
            assert doc.id is not None
            assert doc.title == "My Doc"


@pytest.mark.asyncio
async def test_client_update_title_noop_behavior(temp_db_path):
    """When content is unchanged, updating title should update document without re-chunking."""
    async with HaikuRAG(temp_db_path) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test_update_title.txt"
            temp_path.write_text("Original content")

            doc1 = await client.create_document_from_source(temp_path, title="Title A")
            assert isinstance(doc1, Document)
            assert doc1.id is not None

            # Re-add with same content but new title
            doc2 = await client.create_document_from_source(temp_path, title="Title B")
            assert isinstance(doc2, Document)
            assert doc2.id == doc1.id
            # Fetch and verify title updated
            got = await client.get_document_by_id(doc1.id)
            assert got is not None
            assert got.title == "Title B"


@pytest.mark.asyncio
async def test_client_create_document_from_source_unsupported(temp_db_path):
    """Test creating a document from an unsupported file type."""
    async with HaikuRAG(temp_db_path) as client:
        # Create a temporary file with unsupported extension
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".unsupported", delete=False
        ) as f:
            f.write("content")
            temp_path = Path(f.name)

            # Should raise ValueError for unsupported extension
            with pytest.raises(ValueError, match="Unsupported file extension"):
                await client.create_document_from_source(temp_path)


@pytest.mark.asyncio
async def test_client_create_document_from_source_nonexistent(temp_db_path):
    """Test creating a document from a non-existent file."""
    async with HaikuRAG(temp_db_path) as client:
        non_existent_path = Path("/non/existent/file.txt")

        # Should raise ValueError when file doesn't exist
        with pytest.raises(ValueError, match="File does not exist"):
            await client.create_document_from_source(non_existent_path)


@pytest.mark.asyncio
async def test_client_create_document_from_directory(temp_db_path):
    """Test creating documents from a directory recursively."""
    async with HaikuRAG(temp_db_path) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            test_dir = Path(temp_dir) / "test_docs"
            test_dir.mkdir()

            (test_dir / "doc1.txt").write_text("Content of doc1")
            (test_dir / "doc2.md").write_text("# Content of doc2")

            subdir = test_dir / "subdir"
            subdir.mkdir()
            (subdir / "doc3.py").write_text("print('hello')")

            (test_dir / "unsupported.xyz").write_text("unsupported file")

            result = await client.create_document_from_source(test_dir)

            assert isinstance(result, list)
            assert len(result) == 3

            for doc in result:
                assert doc.id is not None
                assert doc.uri is not None
                assert "md5" in doc.metadata
                assert "contentType" in doc.metadata

            uris = [doc.uri for doc in result if doc.uri]
            assert any("doc1.txt" in uri for uri in uris)
            assert any("doc2.md" in uri for uri in uris)
            assert any("doc3.py" in uri for uri in uris)
            assert not any("unsupported.xyz" in uri for uri in uris)


@pytest.mark.asyncio
async def test_client_create_document_from_directory_with_filters(
    monkeypatch, temp_db_path
):
    """Test creating documents from a directory with ignore and include patterns."""
    # Mock config to have ignore and include patterns
    monkeypatch.setattr(
        "haiku.rag.client.Config.monitor.ignore_patterns", ["**/ignore_me/**", "*.log"]
    )
    monkeypatch.setattr(
        "haiku.rag.client.Config.monitor.include_patterns", ["**/include/**/*.txt"]
    )

    async with HaikuRAG(temp_db_path) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            test_dir = Path(temp_dir) / "test_docs"
            test_dir.mkdir()

            # Create files in include directory - should be included
            include_dir = test_dir / "include"
            include_dir.mkdir()
            (include_dir / "doc1.txt").write_text("Content of doc1")
            (include_dir / "doc2.txt").write_text("Content of doc2")

            # Create files outside include directory - should be excluded by include pattern
            (test_dir / "doc3.txt").write_text("Content of doc3")

            # Create files in ignore directory - should be excluded by ignore pattern
            ignore_dir = test_dir / "ignore_me"
            ignore_dir.mkdir()
            (ignore_dir / "doc4.txt").write_text("Content of doc4")

            # Create log file - should be excluded by ignore pattern
            (test_dir / "debug.log").write_text("log content")

            result = await client.create_document_from_source(test_dir)

            assert isinstance(result, list)
            # Should only include doc1.txt and doc2.txt from include directory
            assert len(result) == 2

            uris = [doc.uri for doc in result if doc.uri]
            assert any("doc1.txt" in uri for uri in uris)
            assert any("doc2.txt" in uri for uri in uris)
            assert not any("doc3.txt" in uri for uri in uris)
            assert not any("doc4.txt" in uri for uri in uris)
            assert not any("debug.log" in uri for uri in uris)


@pytest.mark.asyncio
async def test_client_create_document_from_url(temp_db_path):
    """Test creating a document from a URL."""
    async with HaikuRAG(temp_db_path) as client:
        # Mock the HTTP response
        mock_response = AsyncMock()
        mock_response.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
        mock_response.headers = {"content-type": "text/html"}
        mock_response.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_response):
            doc = await client.create_document_from_source(
                source="https://example.com/test.html", metadata={"source_type": "web"}
            )
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert "Test Page" in doc.content
            assert "test content" in doc.content
            assert doc.uri == "https://example.com/test.html"
            assert doc.metadata["source_type"] == "web"
            assert "contentType" in doc.metadata
            assert "md5" in doc.metadata
            assert doc.metadata["contentType"] == "text/html"


@pytest.mark.asyncio
async def test_client_create_document_from_url_with_different_content_types(
    temp_db_path,
):
    """Test creating documents from URLs with different content types."""
    async with HaikuRAG(temp_db_path) as client:
        # Test JSON content
        mock_json_response = AsyncMock()
        mock_json_response.content = (
            b'{"title": "Test JSON", "content": "This is JSON content"}'
        )
        mock_json_response.headers = {"content-type": "application/json"}
        mock_json_response.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_json_response):
            doc = await client.create_document_from_source(
                "https://api.example.com/data.json"
            )
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert "Test JSON" in doc.content
            assert doc.uri == "https://api.example.com/data.json"
            assert "contentType" in doc.metadata
            assert "md5" in doc.metadata
            assert doc.metadata["contentType"] == "application/json"

        # Test plain text content
        mock_text_response = AsyncMock()
        mock_text_response.content = b"This is plain text content from a URL."
        mock_text_response.headers = {"content-type": "text/plain"}
        mock_text_response.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_text_response):
            doc = await client.create_document_from_source(
                "https://example.com/readme.txt"
            )
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert doc.content == "This is plain text content from a URL."
            assert doc.uri == "https://example.com/readme.txt"
            assert "contentType" in doc.metadata
            assert "md5" in doc.metadata
            assert doc.metadata["contentType"] == "text/plain"


@pytest.mark.asyncio
async def test_client_create_document_from_url_unsupported_content(temp_db_path):
    """Test creating a document from URL with unsupported content type."""
    async with HaikuRAG(temp_db_path) as client:
        # Mock response with unsupported content type
        mock_response = AsyncMock()
        mock_response.content = b"binary content"
        mock_response.headers = {"content-type": "application/octet-stream"}
        mock_response.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_response):
            with pytest.raises(ValueError, match="Unsupported content type"):
                await client.create_document_from_source(
                    "https://example.com/binary.bin"
                )


@pytest.mark.asyncio
async def test_client_create_document_from_url_http_error(temp_db_path):
    """Test handling HTTP errors when creating document from URL."""
    async with HaikuRAG(temp_db_path) as client:
        with patch("httpx.AsyncClient.get") as mock_get:
            mock_get.side_effect = httpx.HTTPStatusError(
                "404 Not Found",
                request=httpx.Request("GET", "https://example.com/notfound.html"),
                response=httpx.Response(404),
            )

            with pytest.raises(httpx.HTTPStatusError):
                await client.create_document_from_source(
                    "https://example.com/notfound.html"
                )


@pytest.mark.asyncio
async def test_get_extension_from_content_type_or_url(temp_db_path):
    """Test the helper method for determining file extensions."""
    async with HaikuRAG(temp_db_path) as client:
        # Test content type mappings
        assert (
            client._get_extension_from_content_type_or_url("", "text/html") == ".html"
        )
        assert (
            client._get_extension_from_content_type_or_url("", "application/pdf")
            == ".pdf"
        )
        assert (
            client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
        )

        # Test URL extension detection
        assert (
            client._get_extension_from_content_type_or_url(
                "https://example.com/doc.pdf", ""
            )
            == ".pdf"
        )
        assert (
            client._get_extension_from_content_type_or_url(
                "https://example.com/data.json", ""
            )
            == ".json"
        )

        # Test default fallback
        assert (
            client._get_extension_from_content_type_or_url("https://example.com/", "")
            == ".html"
        )

        # Test content type priority over URL extension
        assert (
            client._get_extension_from_content_type_or_url(
                "https://example.com/file.txt", "application/pdf"
            )
            == ".pdf"
        )


@pytest.mark.asyncio
async def test_client_metadata_content_type_and_md5(temp_db_path):
    """Test that contentType and md5 metadata are correctly set."""
    import hashlib

    async with HaikuRAG(temp_db_path) as client:
        # Create a temporary file with known content
        test_content = "Test content for MD5 calculation."
        expected_md5 = hashlib.md5(test_content.encode()).hexdigest()

        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text(test_content)

            doc = await client.create_document_from_source(temp_path)
            assert isinstance(doc, Document)

            assert doc.metadata["contentType"] == "text/plain"
            assert doc.metadata["md5"] == expected_md5

            mock_response = AsyncMock()
            mock_response.content = test_content.encode()
            mock_response.headers = {"content-type": "text/plain"}
            mock_response.raise_for_status = AsyncMock()

            with patch("httpx.AsyncClient.get", return_value=mock_response):
                url_doc = await client.create_document_from_source(
                    "https://example.com/test.txt"
                )
                assert isinstance(url_doc, Document)

                assert url_doc.metadata["contentType"] == "text/plain"
                assert url_doc.metadata["md5"] == expected_md5


@pytest.mark.asyncio
async def test_client_create_update_no_op_behavior(temp_db_path):
    """Test create/update/no-op behavior based on MD5 changes."""
    async with HaikuRAG(temp_db_path) as client:
        # Create a temporary file
        test_content = "Original content for testing."
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text(test_content)

            # First call - should create new document
            doc1 = await client.create_document_from_source(temp_path)
            assert isinstance(doc1, Document)
            assert doc1.id is not None
            assert doc1.content == test_content
            original_id = doc1.id

            # Second call with same content - should return existing document (no-op)
            doc2 = await client.create_document_from_source(temp_path)
            assert isinstance(doc2, Document)
            assert doc2.id == original_id  # Same document
            assert doc2.content == test_content

            # Modify file content
            updated_content = "Updated content for testing."
            temp_path.write_text(updated_content)

            # Third call with changed content - should update existing document
            doc3 = await client.create_document_from_source(temp_path)
            assert isinstance(doc3, Document)
            assert doc3.id == original_id  # Same document ID
            assert doc3.content == updated_content  # Updated content

            # Verify the document was actually updated in database
            retrieved_doc = await client.get_document_by_id(original_id)
            assert retrieved_doc is not None
            assert retrieved_doc.content == updated_content


@pytest.mark.asyncio
async def test_client_unchanged_file_keeps_timestamp(temp_db_path):
    """Test that unchanged files don't update the updated_at timestamp."""
    async with HaikuRAG(temp_db_path) as client:
        # Create a temporary file
        test_content = "Test content for timestamp check."
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text(test_content)

            # First call - create document
            doc1 = await client.create_document_from_source(temp_path)
            assert isinstance(doc1, Document)
            original_updated_at = doc1.updated_at

            # Second call with same content - should not update timestamp
            doc2 = await client.create_document_from_source(temp_path)
            assert isinstance(doc2, Document)
            assert doc2.id == doc1.id
            assert doc2.updated_at == original_updated_at  # Timestamp should not change


@pytest.mark.asyncio
async def test_client_url_create_update_no_op_behavior(temp_db_path):
    """Test create/update/no-op behavior for URLs based on MD5 changes."""
    async with HaikuRAG(temp_db_path) as client:
        url = "https://example.com/test.txt"
        original_content = b"Original URL content"
        updated_content = b"Updated URL content"

        # Mock first response
        mock_response1 = AsyncMock()
        mock_response1.content = original_content
        mock_response1.headers = {"content-type": "text/plain"}
        mock_response1.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_response1):
            # First call - should create new document
            doc1 = await client.create_document_from_source(url)
            assert isinstance(doc1, Document)
            assert doc1.id is not None
            original_id = doc1.id

            # Second call with same content - should return existing document (no-op)
            doc2 = await client.create_document_from_source(url)
            assert isinstance(doc2, Document)
            assert doc2.id == original_id  # Same document

        mock_response2 = AsyncMock()
        mock_response2.content = updated_content
        mock_response2.headers = {"content-type": "text/plain"}
        mock_response2.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_response2):
            # Third call with changed content - should update existing document
            doc3 = await client.create_document_from_source(url)
            assert isinstance(doc3, Document)
            assert doc3.id == original_id  # Same document ID
            assert doc3.content == updated_content.decode()  # Updated content


@pytest.mark.asyncio
async def test_client_search(temp_db_path):
    """Test HaikuRAG search functionality."""
    async with HaikuRAG(temp_db_path) as client:
        # Add multiple documents to search from
        doc1_text = "Python is a high-level programming language known for its simplicity and readability."
        doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming."
        doc3_text = "Data science combines statistics, programming, and domain expertise to extract insights."

        # Create documents
        doc1 = await client.create_document(
            content=doc1_text, uri="doc1.txt", metadata={"topic": "python"}
        )
        doc2 = await client.create_document(
            content=doc2_text, uri="doc2.txt", metadata={"topic": "ml"}
        )
        await client.create_document(
            content=doc3_text, uri="doc3.txt", metadata={"topic": "data_science"}
        )

        # Test search with keyword that should match doc1
        results = await client.search("Python programming", limit=3)

        assert len(results) > 0
        assert all(len(result) == 2 for result in results)

        # Verify first result is from the Python document (doc1)
        first_chunk, _ = results[0]
        assert first_chunk.document_id == doc1.id

        # Test search with different query
        ml_results = await client.search("machine learning data", limit=2)
        assert len(ml_results) > 0

        # Verify first result is from the machine learning document (doc2)
        first_ml_chunk, _ = ml_results[0]
        assert first_ml_chunk.document_id == doc2.id

        # Test search with limit parameter
        limited_results = await client.search("programming", limit=1)
        assert len(limited_results) <= 1


@pytest.mark.asyncio
async def test_client_async_context_manager(temp_db_path):
    """Test HaikuRAG as async context manager."""

    # Test that context manager works and auto-closes
    async with HaikuRAG(temp_db_path) as client:
        # Create a document to ensure the client works
        doc = await client.create_document(
            content="Test content for context manager",
            uri="test://context",
            metadata={"test": "context_manager"},
        )

        assert doc.id is not None
        assert doc.content == "Test content for context manager"

        # Test search works within context
        results = await client.search("Test content", limit=1)
        assert len(results) > 0

    # Context manager should have automatically closed the connection
    # We can't easily test that the connection is closed without accessing internals,
    # but the test passing means the context manager methods work correctly


@pytest.mark.asyncio
async def test_client_create_document_with_custom_chunks(temp_db_path):
    """Test creating a document with pre-created chunks."""
    async with HaikuRAG(temp_db_path) as client:
        # Create some custom chunks with and without embeddings
        chunks = [
            Chunk(
                content="This is the first chunk",
                metadata={"custom": "metadata1"},
                order=0,
            ),
            Chunk(
                content="This is the second chunk",
                metadata={"custom": "metadata2"},
                embedding=[0.1] * Config.embeddings.vector_dim,
                order=1,
            ),  # With embedding
            Chunk(
                content="This is the third chunk",
                metadata={"custom": "metadata3"},
                order=2,
            ),
        ]

        # Create document with custom chunks
        document = await client.create_document(
            content="Full document content", chunks=chunks
        )

        assert document.id is not None
        assert document.content == "Full document content"

        # Verify the chunks were created correctly
        doc_chunks = await client.chunk_repository.get_by_document_id(document.id)
        assert len(doc_chunks) == 3

        # Check chunks have correct content, document_id, and order from list position
        for i, chunk in enumerate(doc_chunks):
            assert chunk.document_id == document.id
            assert chunk.content == chunks[i].content
            assert chunk.order == i  # Order should be set from list position
            assert (
                chunk.metadata["custom"] == f"metadata{i + 1}"
            )  # Original metadata preserved


@pytest.mark.asyncio
async def test_client_ask_without_cite(monkeypatch, temp_db_path):
    """Test asking questions without citations."""
    from pydantic_ai.models.test import TestModel

    # Mock OpenAIChatModel to return TestModel
    monkeypatch.setattr(
        "haiku.rag.qa.agent.OpenAIChatModel", lambda **kwargs: TestModel()
    )

    async with HaikuRAG(temp_db_path) as client:
        # Create a test document for the agent to search
        await client.create_document(
            content="Python is a high-level programming language.", uri="test.txt"
        )

        # Use real QA agent with TestModel
        answer = await client.ask("What is Python?")

        # TestModel will generate a valid string response
        assert answer is not None
        assert isinstance(answer, str)


@pytest.mark.asyncio
async def test_client_ask_with_cite(monkeypatch, temp_db_path):
    """Test asking questions with citations."""
    from pydantic_ai.models.test import TestModel

    # Mock OpenAIChatModel to return TestModel
    monkeypatch.setattr(
        "haiku.rag.qa.agent.OpenAIChatModel", lambda **kwargs: TestModel()
    )

    async with HaikuRAG(temp_db_path) as client:
        # Create a test document
        await client.create_document(
            content="Python is a high-level programming language.", uri="test.txt"
        )

        # Use real QA agent with TestModel
        answer = await client.ask("What is Python?", cite=True)

        # TestModel will generate a valid string response
        assert answer is not None
        assert isinstance(answer, str)


@pytest.mark.asyncio
async def test_client_expand_context(temp_db_path):
    """Test expanding search results with adjacent chunks."""
    # Mock Config to have CONTEXT_CHUNK_RADIUS = 2
    with patch("haiku.rag.client.Config.processing.context_chunk_radius", 2):
        async with HaikuRAG(temp_db_path) as client:
            # Create chunks manually with precomputed embeddings to avoid network
            dim = client.chunk_repository.embedder._vector_dim
            z = [0.0] * dim
            manual_chunks = [
                Chunk(content="Chunk 0 content", order=0, embedding=z),
                Chunk(content="Chunk 1 content", order=1, embedding=z),
                Chunk(content="Chunk 2 content", order=2, embedding=z),
                Chunk(content="Chunk 3 content", order=3, embedding=z),
                Chunk(content="Chunk 4 content", order=4, embedding=z),
            ]

        doc = await client.create_document(
            content="Full document content",
            uri="test_doc.txt",
            title="test_doc_title",
            chunks=manual_chunks,
        )

        # Get all chunks for the document
        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(chunks) == 5

        # Find the middle chunk (order=2)
        middle_chunk = next(c for c in chunks if c.order == 2)
        search_results = [(middle_chunk, 0.8)]

        # Test expand_context with radius=2 and document title preserved
        expanded_results = await client.expand_context(search_results, radius=2)

        assert len(expanded_results) == 1
        expanded_chunk, score = expanded_results[0]

        # Check that the expanded chunk has combined content and preserves title/uri
        assert expanded_chunk.id == middle_chunk.id
        assert score == 0.8
        assert "Chunk 2 content" in expanded_chunk.content
        assert expanded_chunk.document_title == "test_doc_title"
        assert expanded_chunk.document_uri == "test_doc.txt"

        # Should include all chunks (radius=2 from chunk 2 = chunks 0,1,2,3,4)
        assert "Chunk 0 content" in expanded_chunk.content
        assert "Chunk 1 content" in expanded_chunk.content
        assert "Chunk 2 content" in expanded_chunk.content
        assert "Chunk 3 content" in expanded_chunk.content
        assert "Chunk 4 content" in expanded_chunk.content


@pytest.mark.asyncio
async def test_client_expand_context_radius_zero(temp_db_path):
    """Test expand_context with radius 0 returns original results."""
    async with HaikuRAG(temp_db_path) as client:
        # Create a simple document
        doc = await client.create_document(content="Simple test content")
        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)

        search_results = [(chunks[0], 0.9)]
        expanded_results = await client.expand_context(search_results, radius=0)

        # Should return exactly the same results
        assert expanded_results == search_results


@pytest.mark.asyncio
async def test_client_expand_context_multiple_chunks(temp_db_path):
    """Test expand_context with multiple search results."""
    with patch("haiku.rag.client.Config.processing.context_chunk_radius", 1):
        async with HaikuRAG(temp_db_path) as client:
            # Create first document with manual chunks
            doc1_chunks = [
                Chunk(content="Doc1 Part A", order=0),
                Chunk(content="Doc1 Part B", order=1),
                Chunk(content="Doc1 Part C", order=2),
            ]
            doc1 = await client.create_document(
                content="Doc1 content", uri="doc1.txt", chunks=doc1_chunks
            )

            # Create second document with manual chunks
            doc2_chunks = [
                Chunk(content="Doc2 Section X", order=0),
                Chunk(content="Doc2 Section Y", order=1),
            ]
            doc2 = await client.create_document(
                content="Doc2 content", uri="doc2.txt", chunks=doc2_chunks
            )

        assert doc1.id is not None
        assert doc2.id is not None
        chunks1 = await client.chunk_repository.get_by_document_id(doc1.id)
        chunks2 = await client.chunk_repository.get_by_document_id(doc2.id)

        # Get middle chunk from doc1 (order=1) and first chunk from doc2 (order=0)
        chunk1 = next(c for c in chunks1 if c.order == 1)
        chunk2 = next(c for c in chunks2 if c.order == 0)

        search_results = [(chunk1, 0.8), (chunk2, 0.7)]
        expanded_results = await client.expand_context(search_results, radius=1)

        assert len(expanded_results) == 2

        # Check first expanded result (should include chunks 0,1,2 from doc1)
        expanded1, score1 = expanded_results[0]
        assert expanded1.id == chunk1.id
        assert score1 == 0.8
        assert "Doc1 Part A" in expanded1.content
        assert "Doc1 Part B" in expanded1.content
        assert "Doc1 Part C" in expanded1.content

        # Check second expanded result (should include chunks 0,1 from doc2)
        expanded2, score2 = expanded_results[1]
        assert expanded2.id == chunk2.id
        assert score2 == 0.7
        assert "Doc2 Section X" in expanded2.content
        assert "Doc2 Section Y" in expanded2.content


@pytest.mark.asyncio
async def test_client_expand_context_merges_overlapping_chunks(temp_db_path):
    """Test that overlapping expanded chunks are merged into one."""
    async with HaikuRAG(temp_db_path) as client:
        # Create document with 5 chunks
        manual_chunks = [
            Chunk(content="Chunk 0", order=0),
            Chunk(content="Chunk 1", order=1),
            Chunk(content="Chunk 2", order=2),
            Chunk(content="Chunk 3", order=3),
            Chunk(content="Chunk 4", order=4),
        ]

        doc = await client.create_document(
            content="Full document content", chunks=manual_chunks
        )

        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)

        # Get adjacent chunks (orders 1 and 2) - these will overlap when expanded
        chunk1 = next(c for c in chunks if c.order == 1)
        chunk2 = next(c for c in chunks if c.order == 2)

        # With radius=1:
        # chunk1 expanded would be [0,1,2]
        # chunk2 expanded would be [1,2,3]
        # These should merge into one chunk containing [0,1,2,3]
        search_results = [(chunk1, 0.8), (chunk2, 0.7)]
        expanded_results = await client.expand_context(search_results, radius=1)

        # Should have only 1 merged result instead of 2 overlapping ones
        assert len(expanded_results) == 1

        merged_chunk, score = expanded_results[0]

        # Should contain all chunks from 0 to 3
        assert "Chunk 0" in merged_chunk.content
        assert "Chunk 1" in merged_chunk.content
        assert "Chunk 2" in merged_chunk.content
        assert "Chunk 3" in merged_chunk.content
        assert "Chunk 4" not in merged_chunk.content  # Should not include chunk 4

        # Should use the higher score (0.8)
        assert score == 0.8


@pytest.mark.asyncio
async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path):
    """Test that non-overlapping expanded chunks remain separate."""
    async with HaikuRAG(temp_db_path) as client:
        # Create document with chunks far apart
        manual_chunks = [
            Chunk(content="Chunk 0", order=0),
            Chunk(content="Chunk 1", order=1),
            Chunk(content="Chunk 2", order=2),
            Chunk(content="Chunk 5", order=5),  # Gap here
            Chunk(content="Chunk 6", order=6),
            Chunk(content="Chunk 7", order=7),
        ]

        doc = await client.create_document(
            content="Full document content", chunks=manual_chunks
        )

        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)

        # Get chunks by index - they will have sequential orders 0,1,2,3,4,5
        # So get chunk with order=0 and chunk with order=5 (far enough apart)
        chunk0 = next(c for c in chunks if c.order == 0)  # Content: "Chunk 0"
        chunk5 = next(
            c for c in chunks if c.order == 5
        )  # Content: "Chunk 7" but now at order 5

        # chunk0 expanded: [0,1] with radius=1 (orders 0,1)
        # chunk5 expanded: [4,5] with radius=1 (orders 4,5)
        search_results = [(chunk0, 0.8), (chunk5, 0.7)]
        expanded_results = await client.expand_context(search_results, radius=1)

        # Should have 2 separate results
        assert len(expanded_results) == 2

        # Sort by score to ensure predictable order
        expanded_results.sort(key=lambda x: x[1], reverse=True)

        chunk0_expanded, score1 = expanded_results[0]
        chunk5_expanded, score2 = expanded_results[1]

        # First chunk (order=0) expanded should contain orders [0,1]
        # Content should be "Chunk 0" + "Chunk 1"
        assert "Chunk 0" in chunk0_expanded.content
        assert "Chunk 1" in chunk0_expanded.content
        assert (
            "Chunk 5" not in chunk0_expanded.content
        )  # Should not have chunk 7 content
        assert score1 == 0.8

        # Second chunk (order=5) expanded should contain orders [4,5]
        # Content should be "Chunk 6" (order 4) + "Chunk 7" (order 5)
        assert "Chunk 6" in chunk5_expanded.content
        assert "Chunk 7" in chunk5_expanded.content
        assert "Chunk 0" not in chunk5_expanded.content
        assert score2 == 0.7