diff --git a/CHANGELOG.md b/CHANGELOG.md index 4855c0e9..ece58266 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Added + +- Custom ingester sources: a source config with `type: plugin` names a source factory registered under the `haiku.rag.sources` entry-point group via its `plugin` field, with an opaque `options` mapping the plugin validates itself. Only the referenced plugin is imported. + ### Fixed - SQLite ingester queue runs with a multi-connection pool (`pool_size=5, max_overflow=5`) instead of a single connection. API reads (`/stats`, `/jobs`) no longer time out with `QueuePool limit of size 1 reached` while workers hold the connection. diff --git a/docs/ingester.md b/docs/ingester.md index b25dc09c..f48b883e 100644 --- a/docs/ingester.md +++ b/docs/ingester.md @@ -228,6 +228,96 @@ override them. A `metadata_provider` name with no installed entry point fails at startup. A provider exception is classified like any other ingestion error (network and timeout errors retry; others go to the DLQ). +### Custom sources + +The four built-in source types (`fs`, `http`, `s3`, `webdav`) cover the +common cases. To ingest from something else (a git host, a ticketing +system, a bespoke API), an external package registers a source factory +under the `haiku.rag.sources` entry-point group and a config references it +with `type: plugin`. + +```yaml + - type: plugin + id: api-docs + plugin: git + options: + owner: acme + repo: api + branch: main + token: ${SCM_TOKEN} +``` + +`plugin` is the entry-point name. `options` is an opaque mapping passed +straight to the factory, which validates it however it likes (for example +with its own Pydantic model). The base fields on every source +(`id`, `poll_interval_s`, `delete_orphans`, `max_file_size`, `retry`, +`circuit_breaker`, `metadata_provider`) are handled by the ingester and +are not part of `options`. + +The factory is called with the source id, the validated `options`, and the +ambient extension and size limits, and returns a `Source`: + +```python +def __call__( + self, + *, + source_id: str, + options: dict, + supported_extensions: list[str] | None, + max_file_size: int | None, +) -> Source: ... +``` + +A `Source` implements this protocol: + +```python +class Source(Protocol): + source_id: str + + def supports(self, uri: str) -> bool: ... + + # Current revision for `uri`, cheaply, or None if there is no cheap + # lookup. Lets the pipeline skip re-ingest when the revision is unchanged. + async def head(self, uri: str) -> str | None: ... + + # Release resources (connection pools, etc.). Called once at shutdown. + async def aclose(self) -> None: ... + + async def fetch(self, uri: str) -> FetchResult: ... + + # Yield UPSERT / UNCHANGED / DELETE events. `since` is the uri -> revision + # snapshot from the previous sweep so the source can emit only deltas. + def discover( + self, + since: RevisionSnapshot | None = None, + *, + known_uris: set[str] | None = None, + ) -> AsyncIterator[SourceEvent]: ... +``` + +`FetchResult`, `SourceEvent`, `SourceEventKind`, and `RevisionSnapshot` +live in `haiku.rag.ingester.sources`. + +```toml +# in the source package's pyproject.toml +[project.entry-points."haiku.rag.sources"] +git = "example_pkg:build_git_source" +``` + +Only the plugin a source references is imported, so an unused plugin with a +missing optional dependency does not break startup. A `plugin` name with no +installed entry point fails at startup, as does a factory that returns +something that is not a `Source`. + +Two limits to know: + +- Custom sources are reached through configured discovery and the job + queue, not through one-shot `haiku-rag add-src `, which only knows + the built-in URI schemes. +- Change detection is per `(source, uri)`. A source that needs a single + per-source cursor (for example a git last-commit SHA) tracks it itself, + by encoding it in each URI's revision or stashing it under a sentinel URI. + ## Workers and retry ```yaml diff --git a/haiku_rag_slim/haiku/rag/config/__init__.py b/haiku_rag_slim/haiku/rag/config/__init__.py index 678ff4bb..55645163 100644 --- a/haiku_rag_slim/haiku/rag/config/__init__.py +++ b/haiku_rag_slim/haiku/rag/config/__init__.py @@ -19,6 +19,7 @@ from haiku.rag.config.models import ( LanceDBConfig, ModelConfig, OllamaConfig, + PluginSourceConfig, ProcessingConfig, PromptsConfig, ProvidersConfig, @@ -47,6 +48,7 @@ __all__ = [ "LanceDBConfig", "ModelConfig", "OllamaConfig", + "PluginSourceConfig", "ProcessingConfig", "PromptsConfig", "ProvidersConfig", diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index ae5a7a92..f6caa659 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Annotated, Literal +from typing import Annotated, Any, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator @@ -447,8 +447,27 @@ class WebDAVSourceConfig(_SourceBase): include_patterns: list[str] = [] +class PluginSourceConfig(_SourceBase): + """A source provided by an external package registered under the + `haiku.rag.sources` entry-point group. `plugin` names the entry point; + `options` is passed through to the plugin's factory, which validates it.""" + + type: Literal["plugin"] + # No natural key to derive an id from; require an explicit one. + id: str + plugin: str = Field( + description="Name of a source factory registered under the " + "'haiku.rag.sources' entry-point group." + ) + options: dict[str, Any] = Field(default_factory=dict) + + SourceConfig = Annotated[ - FSSourceConfig | HTTPSourceConfig | S3SourceConfig | WebDAVSourceConfig, + FSSourceConfig + | HTTPSourceConfig + | S3SourceConfig + | WebDAVSourceConfig + | PluginSourceConfig, Field(discriminator="type"), ] diff --git a/haiku_rag_slim/haiku/rag/ingester/api/schemas.py b/haiku_rag_slim/haiku/rag/ingester/api/schemas.py index fcbccefc..68888880 100644 --- a/haiku_rag_slim/haiku/rag/ingester/api/schemas.py +++ b/haiku_rag_slim/haiku/rag/ingester/api/schemas.py @@ -40,7 +40,7 @@ class ProvidersResponse(BaseModel): class SourceSummary(BaseModel): source_id: str - type: Literal["fs", "http", "s3", "webdav"] + type: Literal["fs", "http", "s3", "webdav", "plugin"] last_polled_at: datetime | None circuit_breaker_open: bool # Reason the most recent sweep attempt was skipped (e.g. "pending_work"), diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/factory.py b/haiku_rag_slim/haiku/rag/ingester/pollers/factory.py index 09a9e721..693905fc 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/factory.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/factory.py @@ -1,6 +1,7 @@ from haiku.rag.config import ( FSSourceConfig, HTTPSourceConfig, + PluginSourceConfig, S3SourceConfig, SourceConfig, WebDAVSourceConfig, @@ -12,6 +13,10 @@ from haiku.rag.ingester.sources import ( Source, WebDAVSource, ) +from haiku.rag.ingester.sources.plugins import ( + ENTRY_POINT_GROUP, + load_source_factories, +) def build_source( @@ -63,6 +68,29 @@ def build_source( supported_extensions=supported_extensions, max_file_size=cfg.max_file_size, ) + if isinstance(cfg, PluginSourceConfig): + factories = load_source_factories() + try: + entry_point = factories[cfg.plugin] + except KeyError: + raise ValueError( + f"Source {cfg.id!r} references unknown source plugin " + f"{cfg.plugin!r}; no entry point registered under " + f"{ENTRY_POINT_GROUP!r}." + ) from None + source = entry_point.load()( + source_id=cfg.id, + options=cfg.options, + supported_extensions=supported_extensions, + max_file_size=cfg.max_file_size, + ) + if not isinstance(source, Source): + raise TypeError( + f"Source plugin {cfg.plugin!r} returned " + f"{type(source).__name__}, which does not satisfy the " + f"Source protocol." + ) + return source raise TypeError( # pragma: no cover - discriminator union exhausts all cases f"Unsupported source config: {type(cfg).__name__}" ) diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py index 0e390693..c05dda5f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -1,5 +1,6 @@ import asyncio import logging +from collections.abc import Sequence from typing import TYPE_CHECKING from haiku.rag.config import FSSourceConfig, SourceConfig @@ -24,7 +25,7 @@ class PollerManager: def __init__( self, *, - configs: list[SourceConfig], + configs: Sequence[SourceConfig], job_repo: "JobRepo", sync_repo: "SyncStateRepo", supported_extensions: list[str] | None = None, diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/plugins.py b/haiku_rag_slim/haiku/rag/ingester/sources/plugins.py new file mode 100644 index 00000000..65bb2f86 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/sources/plugins.py @@ -0,0 +1,48 @@ +from importlib.metadata import entry_points +from typing import Any, Protocol, runtime_checkable + +from haiku.rag.ingester.sources.base import Source + + +@runtime_checkable +class SourceFactory(Protocol): + """Builds a Source from a plugin source's config. A package registers a + factory under the ``haiku.rag.sources`` entry-point group; the ingester + calls it with the source id, the config's opaque ``options`` (which the + plugin validates itself), and the ambient extension/size limits.""" + + def __call__( + self, + *, + source_id: str, + options: dict[str, Any], + supported_extensions: list[str] | None, + max_file_size: int | None, + ) -> Source: ... + + +@runtime_checkable +class LoadableEntryPoint(Protocol): + """The slice of ``importlib.metadata.EntryPoint`` the factory loader needs: + a deferred ``load()`` returning the source factory.""" + + def load(self) -> SourceFactory: ... + + +ENTRY_POINT_GROUP = "haiku.rag.sources" + + +def load_source_factories() -> dict[str, LoadableEntryPoint]: + """Discover registered source-factory entry points, keyed by name. The + entry points are not imported here; ``build_source`` loads only the one a + source references, so an unused plugin with a broken import does not fail + the ingester at startup.""" + return {ep.name: ep for ep in entry_points(group=ENTRY_POINT_GROUP)} + + +__all__ = [ + "ENTRY_POINT_GROUP", + "LoadableEntryPoint", + "SourceFactory", + "load_source_factories", +] diff --git a/tests/ingester/test_config.py b/tests/ingester/test_config.py index b30a1b8e..36bcbf5b 100644 --- a/tests/ingester/test_config.py +++ b/tests/ingester/test_config.py @@ -110,6 +110,26 @@ def test_discriminator_picks_webdav_source(): assert cfg.sources[0].username == "alice" +def test_discriminator_picks_plugin_source(): + cfg = IngesterConfig.model_validate( + { + "sources": [ + { + "type": "plugin", + "id": "git-docs", + "plugin": "git", + "options": {"owner": "acme", "repo": "api"}, + } + ] + } + ) + from haiku.rag.config import PluginSourceConfig + + assert isinstance(cfg.sources[0], PluginSourceConfig) + assert cfg.sources[0].plugin == "git" + assert cfg.sources[0].options == {"owner": "acme", "repo": "api"} + + def test_discriminator_rejects_unknown_type(): with pytest.raises(ValidationError): IngesterConfig.model_validate({"sources": [{"type": "ftp", "uri": "x"}]}) diff --git a/tests/ingester/test_source_plugins.py b/tests/ingester/test_source_plugins.py new file mode 100644 index 00000000..1a8ccf50 --- /dev/null +++ b/tests/ingester/test_source_plugins.py @@ -0,0 +1,180 @@ +import hashlib +from datetime import UTC, datetime + +import pytest + +from haiku.rag.config import PluginSourceConfig +from haiku.rag.ingester.pollers.factory import build_source +from haiku.rag.ingester.pollers.periodic import PeriodicPoller +from haiku.rag.ingester.queue.models import JobOp +from haiku.rag.ingester.sources import plugins as plugins_module +from haiku.rag.ingester.sources import resolve_configured_source +from haiku.rag.ingester.sources.base import ( + FetchResult, + Source, + SourceEvent, + SourceEventKind, +) +from haiku.rag.ingester.sources.plugins import ( + ENTRY_POINT_GROUP, + load_source_factories, +) + + +class _MemorySource: + """A real Source over an in-memory {uri: text} dict. Doubles as its own + factory: build_source calls it with the plugin kwargs.""" + + def __init__( + self, + *, + source_id: str, + options: dict, + supported_extensions: list[str] | None, + max_file_size: int | None, + ): + self.source_id = source_id + self.options = options + self.supported_extensions = supported_extensions + self.max_file_size = max_file_size + self._docs: dict[str, str] = options["docs"] + + def supports(self, uri: str) -> bool: + return uri in self._docs + + async def head(self, uri: str) -> str | None: + return "v1" + + async def aclose(self) -> None: + return None + + async def fetch(self, uri: str) -> FetchResult: + body = self._docs[uri].encode() + return FetchResult( + uri=uri, + body=body, + content_type="text/markdown", + content_hash=hashlib.md5(body).hexdigest(), + revision="v1", + ) + + async def discover(self, since=None, *, known_uris=None): + for uri in self._docs: + yield SourceEvent( + source_id=self.source_id, + uri=uri, + kind=SourceEventKind.UPSERT, + revision="v1", + discovered_at=datetime.now(UTC), + ) + + +class _FakeEntryPoint: + def __init__(self, name, factory): + self.name = name + self._factory = factory + self.loaded = False + + def load(self): + self.loaded = True + return self._factory + + +def _register(monkeypatch, *eps): + captured: dict = {} + + def fake_entry_points(*, group): + captured["group"] = group + return list(eps) + + monkeypatch.setattr(plugins_module, "entry_points", fake_entry_points) + return captured + + +def test_load_keys_entry_points_by_name_without_loading(monkeypatch): + ep = _FakeEntryPoint("memory", _MemorySource) + captured = _register(monkeypatch, ep) + + discovered = load_source_factories() + + assert captured["group"] == ENTRY_POINT_GROUP + assert discovered == {"memory": ep} + assert ep.loaded is False + + +def test_load_is_empty_when_none_registered(monkeypatch): + _register(monkeypatch) + assert load_source_factories() == {} + + +def _config(**overrides): + return PluginSourceConfig( + type="plugin", + id="mem", + plugin="memory", + options={"docs": {"mem://a.md": "hello"}}, + **overrides, + ) + + +def test_build_source_loads_referenced_plugin_with_kwargs(monkeypatch): + ep = _FakeEntryPoint("memory", _MemorySource) + _register(monkeypatch, ep) + + source = build_source(_config(max_file_size=1024), supported_extensions=[".md"]) + + assert ep.loaded is True + assert isinstance(source, _MemorySource) + assert isinstance(source, Source) + assert source.source_id == "mem" + assert source.options == {"docs": {"mem://a.md": "hello"}} + assert source.supported_extensions == [".md"] + assert source.max_file_size == 1024 + + +def test_build_source_raises_on_unknown_plugin(monkeypatch): + _register(monkeypatch, _FakeEntryPoint("memory", _MemorySource)) + + cfg = _config().model_copy(update={"plugin": "missing"}) + with pytest.raises(ValueError, match="unknown source plugin 'missing'"): + build_source(cfg) + + +def test_build_source_raises_when_plugin_returns_non_source(monkeypatch): + _register(monkeypatch, _FakeEntryPoint("memory", lambda **kw: object())) + + with pytest.raises(TypeError, match="does not satisfy the Source protocol"): + build_source(_config()) + + +def test_build_source_does_not_load_unreferenced_plugins(monkeypatch): + def _explode(**kw): + raise ImportError("optional dependency missing") + + used = _FakeEntryPoint("memory", _MemorySource) + unused = _FakeEntryPoint("broken", _explode) + _register(monkeypatch, used, unused) + + build_source(_config()) + + assert unused.loaded is False + + +@pytest.mark.asyncio +async def test_plugin_source_drives_poller_and_fetch(monkeypatch, jobs, sync): + _register(monkeypatch, _FakeEntryPoint("memory", _MemorySource)) + cfg = _config() + source = build_source(cfg) + + poller = PeriodicPoller(source=source, config=cfg, job_repo=jobs, sync_repo=sync) + assert await poller._sweep_once() is True + + queued = await jobs.list_jobs(source_id="mem") + assert len(queued) == 1 + assert queued[0].op is JobOp.UPSERT + assert queued[0].uri == "mem://a.md" + + fetcher = resolve_configured_source("mem://a.md", "mem", [source]) + result = await fetcher.fetch("mem://a.md") + assert result.body == b"hello" + assert result.content_type == "text/markdown"