haiku.rag/haiku_rag_slim/haiku/rag/client/scope.py

from dataclasses import dataclass
from pathlib import Path

from haiku.rag.config import AppConfig
from haiku.rag.store.exceptions import (
    AmbiguousDatabaseError,
    UnknownDatabaseError,
)
from haiku.rag.utils import locate_database


@dataclass(frozen=True)
class DatabaseRef:
    """A resolved database location, and the configured name it answers to.

    Exactly one of ``uri`` and ``db_path`` is set. ``name`` is the key from
    ``lancedb.databases``, and the only identity that leaves the configuration:
    it travels in results, citations and errors, where a location must not.
    None where nothing names the database.
    """

    name: str | None
    uri: str
    db_path: Path | None

    def __post_init__(self) -> None:
        if bool(self.uri) == (self.db_path is not None):
            raise ValueError(
                "a database is either a URI or a local path: "
                f"got uri={self.uri!r} and db_path={self.db_path!r}"
            )

    @classmethod
    def at(cls, path: Path | str, *, name: str | None = None) -> "DatabaseRef":
        """A database at a path the caller named, taken as given."""
        return cls(name=name, uri="", db_path=Path(path))

    @classmethod
    def configured(cls, name: str | None, location: str) -> "DatabaseRef":
        """A database the configuration placed, by ``lancedb.uri`` or an entry in
        ``lancedb.databases``. A location carrying a scheme is a URI, anything
        else a local path."""
        uri, db_path = locate_database(location)
        return cls(name=name, uri=uri, db_path=db_path)

    def connection(self, config: AppConfig) -> tuple[AppConfig, Path | None]:
        """The configuration and path to open this one database with.

        A copy: the caller's configuration still names whatever set it named.
        """
        one = config.model_copy(deep=True)
        one.lancedb.databases = {}
        one.lancedb.uri = self.uri
        return one, self.db_path


@dataclass(frozen=True)
class DatabaseScope:
    """The databases an operation covers.

    Resolved once, from configuration plus at most one selector, then passed
    down. Never empty.

    Nothing here reads the environment: ``HAIKU_RAG_DB`` is the capability entry
    point's to honour.
    """

    databases: tuple[DatabaseRef, ...]

    def __post_init__(self) -> None:
        if not self.databases:
            raise ValueError("a scope covers at least one database")

    @classmethod
    def resolve(
        cls,
        config: AppConfig,
        *,
        database_name: str | None = None,
        database_path: Path | str | None = None,
    ) -> "DatabaseScope":
        """The databases named by `config` and at most one selector.

        A path names one database that nothing calls anything; a name selects one
        of the configured set and keeps its name. With no selector the configured
        set is covered in configuration order, a set of one included.
        """
        if database_name is not None and database_path is not None:
            raise AmbiguousDatabaseError(
                "a database name and a database path both name one database; "
                "pass one of them"
            )

        declared = config.lancedb.databases

        if database_path is not None:
            return cls((DatabaseRef.at(database_path),))

        if database_name is not None:
            if database_name not in declared:
                raise UnknownDatabaseError(
                    f"unknown database {database_name!r}; lancedb.databases names "
                    f"{', '.join(sorted(declared)) or 'nothing'}"
                )
            return cls(
                (DatabaseRef.configured(database_name, declared[database_name]),)
            )

        if declared:
            return cls(
                tuple(
                    DatabaseRef.configured(name, location)
                    for name, location in declared.items()
                )
            )

        if config.lancedb.uri:
            return cls((DatabaseRef.configured(None, config.lancedb.uri),))

        return cls((DatabaseRef.at(config.storage.data_dir / "haiku.rag.lancedb"),))

    def select(self, names: list[str]) -> "DatabaseScope":
        """The databases in this scope named by `names`, in the order given.

        Repeats collapse: a name selects its database once.
        """
        if not names:
            raise ValueError(
                "sources=[] selects no database; pass None for all of them"
            )
        by_name = {ref.name: ref for ref in self.databases if ref.name is not None}
        missing = [name for name in names if name not in by_name]
        if missing:
            raise UnknownDatabaseError(
                f"unknown database(s) {', '.join(sorted(missing))}; "
                f"configured: {', '.join(sorted(by_name))}"
            )
        return DatabaseScope(tuple(by_name[name] for name in dict.fromkeys(names)))

    @property
    def covers_multiple(self) -> bool:
        """Whether this scope covers more than one database."""
        return len(self.databases) > 1

    @property
    def names(self) -> tuple[str, ...]:
        """The configured names covered, in order. Empty where none is named."""
        return tuple(ref.name for ref in self.databases if ref.name is not None)