Resolve worker source by source_id, not just supports(uri)
This commit is contained in:
parent
922d1d567d
commit
d7fdd61fed
6 changed files with 56 additions and 11 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Reference in a new issue