simbarag/tests/unit/test_retrieval.py

from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock
from uuid import uuid4

import pytest
from langchain_core.documents import Document
from langchain_postgres import PGVector
from quart import Quart
from quart_jwt_extended import JWTManager, create_access_token
from sqlalchemy import Column
from sqlalchemy.dialects.postgresql import JSONB

from blueprints.rag.retrieval import build_filter, citation_url, document_date
from blueprints.rag import logic, source_api
from blueprints.conversation.artifacts_util import search_documents_artifact


@pytest.mark.parametrize("source", [None, "paperless", "obsidian"])
@pytest.mark.parametrize(
    "start,end",
    [("2026-09-01", "2026-09-30"), ("2026-09-01", None), (None, "2026-09-30")],
)
def test_filters_compile_with_actual_pgvector_operator_rules(source, start, end):
    store = PGVector.__new__(PGVector)
    store.EmbeddingStore = SimpleNamespace(cmetadata=Column("cmetadata", JSONB))
    expression = store._create_filter_clause(build_filter(source, start, end))
    assert expression is not None


@pytest.mark.parametrize(
    "start,end",
    [
        ("20260901", None),
        ("bad", None),
        ("2026-10-01", "2026-09-01"),
        (None, "9999-12-31"),
    ],
)
def test_date_filters_validate_before_vector_lookup(start, end):
    with pytest.raises(ValueError):
        build_filter("obsidian", start, end)


async def test_filtered_query_preserves_citations_and_excludes_unknown_dates(
    monkeypatch,
):
    identifier = str(uuid4())
    docs = [
        Document(
            id=identifier,
            page_content="September groceries plan",
            metadata={
                "source": "obsidian",
                "date": "2026-09-15",
                "filepath": "notes/budget.md",
            },
        ),
        Document(
            id=str(uuid4()),
            page_content="Undated plan",
            metadata={"source": "obsidian"},
        ),
    ]
    store = SimpleNamespace(amax_marginal_relevance_search=AsyncMock(return_value=docs))
    monkeypatch.setattr(logic, "_get_vector_store", lambda: store)
    text, results = await logic.query_vector_store(
        "budget", source="obsidian", start_date="2026-09-01", end_date="2026-09-30"
    )
    assert results == docs[:1]
    assert f"Citation: /sources/{identifier}" in text
    assert store.amax_marginal_relevance_search.call_args.kwargs[
        "filter"
    ] == build_filter("obsidian", "2026-09-01", "2026-09-30")


def test_artifact_citations_identify_exact_retrieved_chunk(monkeypatch):
    monkeypatch.setenv("OBSIDIAN_VAULT_PATH", "/vault")
    identifier = str(uuid4())
    doc = Document(
        id=identifier,
        page_content="Savings target",
        metadata={
            "filepath": "/vault/notes/budget.md",
            "h1": "Budget",
            "h2": "Savings",
        },
    )
    artifact = search_documents_artifact("target", "obsidian", [doc])
    result = artifact["data"]["results"][0]
    assert result["url"] == f"/sources/{identifier}"
    assert result["source_id"] == identifier
    assert result["path"] == "notes/budget.md"
    assert result["headings"] == ["Budget", "Savings"]


def test_document_date_and_unknown_citations():
    assert (
        document_date(
            {"created_date": datetime(2026, 9, 1, tzinfo=timezone.utc).timestamp()}
        )
        == "2026-09-01"
    )
    assert document_date({"date": "bad"}) is None
    assert citation_url(Document(page_content="no stable ID")) is None


async def test_source_read_requires_auth_and_handles_missing_chunk(monkeypatch):
    app = Quart(__name__)
    app.config["JWT_SECRET_KEY"] = "test-secret"
    JWTManager(app)
    app.register_blueprint(source_api.sources_blueprint)
    calls = []

    def load(identifier):
        calls.append(identifier)
        return {
            "id": identifier,
            "source": "paperless",
            "title": "Vet report",
            "content": "Retrieved section",
            "headings": [],
        }

    monkeypatch.setattr(source_api, "_load_chunk", load)
    identifier = str(uuid4())
    assert (
        await app.test_client().get(f"/api/sources/{identifier}")
    ).status_code == 401
    async with app.app_context():
        token = create_access_token(identity=str(uuid4()))
    headers = {"Authorization": f"Bearer {token}"}
    response = await app.test_client().get(
        f"/api/sources/{identifier}", headers=headers
    )
    assert (await response.get_json())["content"] == "Retrieved section"
    assert response.headers["Cache-Control"] == "no-store"
    assert (
        await app.test_client().get("/api/sources/not-an-id", headers=headers)
    ).status_code == 404
    assert calls == [identifier]


async def test_post_filter_keeps_note_matched_by_any_date_field(monkeypatch):
    # The database filter matches if any date field is in range; the post-filter
    # must agree or matching notes silently vanish from results.
    doc = Document(
        id=str(uuid4()),
        page_content="Lease renewal notes",
        metadata={
            "source": "obsidian",
            "date": "2023-05-01",
            "created_at": "2024-06-12T09:00:00",
            "filepath": "notes/lease.md",
        },
    )
    store = SimpleNamespace(
        amax_marginal_relevance_search=AsyncMock(return_value=[doc])
    )
    monkeypatch.setattr(logic, "_get_vector_store", lambda: store)
    _, results = await logic.query_vector_store(
        "lease", source="obsidian", start_date="2024-06-01", end_date="2024-06-30"
    )
    assert results == [doc]