Resolve worker source by source_id, not just supports(uri)

This commit is contained in:
Yiorgis Gozadinos 2026-05-27 13:36:03 +03:00
parent 922d1d567d
commit d7fdd61fed
No known key found for this signature in database
6 changed files with 56 additions and 11 deletions

View file

@ -216,6 +216,7 @@ class HaikuRAG:
uri: str | None = None,
storage_options: dict[str, str] | None = None,
sources: "list[Source] | None" = None,
source_id: str | None = None,
) -> Document | list[Document]:
from haiku.rag.client.documents import create_document_from_source
@ -227,6 +228,7 @@ class HaikuRAG:
uri=uri,
storage_options=storage_options,
sources=sources,
source_id=source_id,
)
async def update_document(

View file

@ -308,6 +308,7 @@ async def create_document_from_source(
uri: str | None = None,
storage_options: dict[str, str] | None = None,
sources: "list[Source] | None" = None,
source_id: str | None = None,
) -> Document | list[Document]:
"""Create or update document(s) from a file path, directory, or URL.
@ -375,7 +376,10 @@ async def create_document_from_source(
# `sources` (configured, in-order) wins over scheme-based adhoc adapters
# so worker fetches reuse the authenticated source the poller used.
fetcher = resolve_fetcher(
source_str, sources=sources, storage_options=storage_options
source_str,
sources=sources,
source_id=source_id,
storage_options=storage_options,
)
# The stored URI is what we look up + persist by. For an explicit uri

View file

@ -13,19 +13,31 @@ def resolve_fetcher(
uri: str,
sources: Iterable[Source] | None = None,
*,
source_id: str | None = None,
storage_options: dict[str, str] | None = None,
) -> Source:
"""Pick a Source adapter for ``uri``.
Configured ``sources`` win the first whose ``supports(uri)`` returns True
is returned. Without a configured match, an ad-hoc adapter is built from
the URI scheme so one-shot calls (``add-src <uri>``) work without any
configuration.
When ``source_id`` is given (worker path), the source with that id is
used so credentials/headers of the configured source are reused rather
than picking whichever source happens to match ``supports(uri)`` first.
Without ``source_id`` (ad-hoc ``add-src <uri>``), the first configured
source whose ``supports(uri)`` returns True wins, then a scheme-based
ad-hoc adapter.
"""
if sources:
for src in sources:
if src.supports(uri):
return src
if source_id is not None:
for src in sources:
if src.source_id == source_id:
if not src.supports(uri):
raise UnsupportedSourceError(
f"Source {source_id!r} doesn't support URI {uri!r}"
)
return src
else:
for src in sources:
if src.supports(uri):
return src
scheme = urlparse(uri).scheme
if scheme in ("", "file"):

View file

@ -103,6 +103,7 @@ async def run_job(
result = await client.create_document_from_source(
job.uri,
sources=sources,
source_id=job.source_id,
)
# Directory ingestion returns list[Document] — workers ingest single
# resources, so a list here is a programming error in the caller.

View file

@ -58,14 +58,14 @@ async def test_upsert_calls_create_document_from_source_and_returns_metadata():
assert result.content_hash == "abcd"
assert result.deleted is False
client.create_document_from_source.assert_awaited_once_with(
"https://example.com/a.pdf", sources=None
"https://example.com/a.pdf", sources=None, source_id="src"
)
@pytest.mark.asyncio
async def test_upsert_threads_configured_sources_to_client():
"""The list of configured Source adapters reaches the client so
resolve_fetcher can pick the authenticated one over an adhoc adapter."""
resolve_fetcher can pick the right one by source_id."""
from haiku.rag.ingester.sources.http import HTTPSource
client = _mock_client()
@ -76,7 +76,7 @@ async def test_upsert_threads_configured_sources_to_client():
await run_job(client, _job(), sources=[configured])
client.create_document_from_source.assert_awaited_once_with(
"https://example.com/a.pdf", sources=[configured]
"https://example.com/a.pdf", sources=[configured], source_id="src"
)

View file

@ -57,3 +57,29 @@ def test_configured_source_falls_through_when_no_match():
chosen = resolve_fetcher("https://example.com/x", sources=[fs])
assert isinstance(chosen, HTTPSource)
assert chosen is not fs
def test_source_id_prefers_matching_source_over_first_supports():
"""Two HTTPSource configs with different auth headers: passing
source_id picks the right one. Without source_id, the first one to
return True from supports(uri) wins regardless of which config
actually owns the job."""
arxiv = HTTPSource(source_id="arxiv", headers={"Authorization": "Bearer A"})
intranet = HTTPSource(source_id="intranet", headers={"Authorization": "Bearer B"})
by_id = resolve_fetcher(
"https://example.com/x", sources=[arxiv, intranet], source_id="intranet"
)
assert by_id is intranet
without_id = resolve_fetcher("https://example.com/x", sources=[arxiv, intranet])
assert without_id is arxiv
def test_source_id_raises_when_matched_source_rejects_uri():
"""source_id selects the source by identity, but that source still
has to support the URI. Mismatch means a stale enqueue from a
reconfigured source surface it instead of silently picking another."""
fs = FSSource(root=Path("/tmp"), source_id="local")
with pytest.raises(ValueError, match="doesn't support URI"):
resolve_fetcher("https://example.com/x", sources=[fs], source_id="local")