simbarag/blueprints/rag/retrieval.py

"""Validated retrieval filters and stable, authenticated source references."""

from datetime import date, datetime, time, timedelta, timezone
import re
from uuid import UUID


def parse_date(value: str | None) -> date | None:
    if not value:
        return None
    try:
        if not isinstance(value, str) or not re.fullmatch(r"\d{4}-\d{2}-\d{2}", value):
            raise ValueError
        return date.fromisoformat(value)
    except (TypeError, ValueError) as exc:
        raise ValueError("Dates must be YYYY-MM-DD.") from exc


def build_filter(
    source: str | None, start_date: str | None = None, end_date: str | None = None
) -> dict | None:
    if source not in {None, "paperless", "obsidian"}:
        raise ValueError("source must be paperless or obsidian")
    start, end = parse_date(start_date), parse_date(end_date)
    if start and end and start > end:
        raise ValueError("start_date must be on or before end_date")
    parts = [{"source": source}] if source else []
    if start or end:
        strings, epochs = {}, {}
        if start:
            strings["$gte"] = start.isoformat()
            epochs["$gte"] = datetime.combine(start, time.min, timezone.utc).timestamp()
        if end:
            try:
                next_day = end + timedelta(days=1)
            except OverflowError as exc:
                raise ValueError("end_date is beyond the supported date range") from exc
            strings["$lt"] = next_day.isoformat()
            epochs["$lt"] = datetime.combine(
                next_day, time.min, timezone.utc
            ).timestamp()

        # Both newly normalized fields and existing indexed metadata remain usable.
        def date_bounds(field, bounds):
            clauses = [{field: {operator: value}} for operator, value in bounds.items()]
            return {"$and": clauses} if len(clauses) > 1 else clauses[0]

        parts.append(
            {
                "$or": [
                    date_bounds("document_date", strings),
                    {
                        "$and": [
                            {"source": "paperless"},
                            date_bounds("created_date", epochs),
                        ]
                    },
                    {"$and": [{"source": "obsidian"}, date_bounds("date", strings)]},
                    {
                        "$and": [
                            {"source": "obsidian"},
                            date_bounds("created_at", strings),
                        ]
                    },
                ]
            }
        )
    return {"$and": parts} if len(parts) > 1 else parts[0] if parts else None


def document_dates(metadata: dict) -> list[str]:
    """Every date the search filter can match on, as YYYY-MM-DD."""
    dates = []
    for key in ("document_date", "date", "created_at"):
        value = metadata.get(key)
        if isinstance(value, (date, datetime)):
            dates.append(value.isoformat()[:10])
        elif isinstance(value, str):
            try:
                dates.append(date.fromisoformat(value[:10]).isoformat())
            except ValueError:
                pass
    try:
        if metadata.get("created_date") is not None:
            dates.append(
                datetime.fromtimestamp(float(metadata["created_date"]), timezone.utc)
                .date()
                .isoformat()
            )
    except (ValueError, TypeError, OverflowError, OSError):
        pass
    return dates


def document_date(metadata: dict) -> str | None:
    """The date shown for a citation: the most specific one available."""
    dates = document_dates(metadata)
    return dates[0] if dates else None


def citation_url(document) -> str | None:
    try:
        return "/sources/" + str(UUID(str(document.id)))
    except (ValueError, TypeError, AttributeError):
        return None