feat: add scoped collection and document management APIs - #2378
feat: add scoped collection and document management APIs#2378KyleZheng1284 wants to merge 33 commits into
Conversation
KyleZheng1284
changed the title
Greptile SummaryThis PR introduces a complete scope-isolated collection and document lifecycle API for NeMo Retriever, layering logical collection management on top of LanceDB without exposing physical table names or storage paths to clients. The implementation spans a new
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py | New 1216-line file implementing the LanceDBCollectionStore with retryable lifecycle, cursor-based pagination, table-user leasing, and crash-recovery via reconcile_collections. Vector write I/O and retrieval are correctly placed outside _write_lock, but the catalog scan for uncommitted documents inside retrieve_collection still runs under the lock. |
| nemo_retriever/src/nemo_retriever/service/vectordb_app.py | Heavily refactored to delegate persistence to the injected VDB contract. VectorDBState now wraps an operator pair instead of direct LanceDB handles. Adds internal credential middleware, collection lifecycle routes, async reconciliation loop, and the query_semaphore. Previously flagged concurrency issues (search serialising under _write_lock) are resolved. |
| nemo_retriever/src/nemo_retriever/service/routers/ingest.py | Extended with collection-aware job creation, stable document IDs, manifest validation, idempotent re-upload detection, and scope-scoped job authorization. _resolve_stable_document_id preserves legacy poll-by-document-id behaviour for non-collection jobs. Multiple _require_job calls per request path are redundant but not harmful. |
| nemo_retriever/src/nemo_retriever/service/services/job_tracker.py | Adds register_document_idempotent, scope-keyed idempotency map, idempotency_fingerprint, and manifest-entry replay. All mutations are guarded by self._lock; the idempotency and manifest checks fold check-and-register into one critical section, correctly eliminating the TOCTOU race flagged in prior review. |
| nemo_retriever/src/nemo_retriever/service/auth.py | Replaces single-token auth with multi-scope ScopeAuthorizer, separates internal-credential validation, and unifies 401 for both invalid-token and wrong-scope cases (fixing the prior status-code oracle). authorized_scope reads from middleware-attached request state, preventing raw-header injection. |
| nemo_retriever/src/nemo_retriever/service/routers/collections.py | New transparent proxy router: forwards collection and document lifecycle requests to the internal VectorDB service, injecting the middleware-authorised scope and internal credential. Correctly rejects non-gateway modes. |
| nemo_retriever/src/nemo_retriever/common/vdb/adt_vdb.py | New abstract VDB base class defining the collection lifecycle contract. Well-structured but missing the required SPDX license header present on every other new file in this PR. |
| nemo_retriever/src/nemo_retriever/common/schemas/collections.py | New public wire models for collection/document lifecycle. DocumentId uses a validated Annotated[str, StringConstraints(..., pattern=...)] type that restricts characters, satisfying the no-sql-injection rule for identifiers interpolated into DataFusion predicates. |
| nemo_retriever/src/nemo_retriever/service/config.py | Adds VectorDB config fields (internal_api_token, vectordb_url, collection management toggle). allow_unscoped_dev=True default (flagged in a prior review) is still present. |
Sequence Diagram
sequenceDiagram
participant C as Client
participant GW as Gateway (ingest router)
participant VDB as VectorDB service
participant LC as LanceDBCollectionStore
participant LDB as LanceDB
C->>GW: POST /v1/collections (X-NRL-Scope: tenant-a)
GW->>VDB: forward + internal token + scope
VDB->>LC: create_collection(scope, name)
LC->>LDB: catalog write (_nrl_collections)
VDB-->>C: 201 CollectionInfo
C->>GW: POST /v1/ingest/job (collection_name, operation)
GW->>VDB: "GET /v1/collections/{name} (validate active)"
VDB-->>GW: 200 CollectionInfo
GW-->>C: 201 JobCreatedResponse (job_id)
C->>GW: "POST /v1/ingest/job/{id}/document (file)"
GW->>GW: idempotent register (stable_document_id)
GW->>VDB: internal write callback (records + context)
VDB->>LC: write_collection(records, context)
LC->>LDB: persist recovery marker (catalog)
LC->>LDB: table.merge_insert or table.add (vectors)
LC->>LDB: persist completed row (catalog)
VDB-->>GW: WriteResponse
C->>GW: "GET /v1/collections/{name}/documents"
GW->>VDB: forward (scope)
VDB->>LC: list_documents(scope, name)
LC->>LDB: scan _nrl_documents (catalog)
VDB-->>C: 200 DocumentPage
C->>GW: POST /v1/query (collection_name)
GW->>VDB: forward (scope)
VDB->>LC: retrieve_collection(vectors, scope, name)
Note over LC: hold _write_lock: resolve table, build visibility filter
LC->>LDB: scan _nrl_documents (uncommitted filter)
Note over LC: release _write_lock
LC->>LDB: backend.retrieval(vectors)
VDB-->>C: 200 QueryResponse
Reviews (19): Last reviewed commit: "Merge branch 'main' into collection-mana..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
d2a7419 to
3833c68
Compare
jperez999
left a comment
jperez999
left a comment
There was a problem hiding this comment.
Lets do this in multiple passes
| "pdf_basename": pdf_basename, | ||
| "pdf_page": f"{pdf_basename}_{page_number}" if pdf_basename and page_number is not None else "", | ||
| } | ||
| chunk_id = entity.get("chunk_id") or hit.get("chunk_id") |
There was a problem hiding this comment.
This is about normalizing hits... shouldnt you just make sure the hit returns the parts of the schema you want? Why do you have dive into the entity? shouldn't the field you want be in a specific place and not multiple places?
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
d05de5d to
6958168
Compare
| from nemo_retriever.common.vdb.lancedb_capabilities import LanceRetrievalMode | ||
|
|
||
|
|
||
| class RetrievalContractError(RuntimeError): |
There was a problem hiding this comment.
Under what circumstances is this not a value error? What is this actually supposed to bubble up? Why do we need a new error type for this, what do we gain from this?
There was a problem hiding this comment.
would like to keep this would be useful for other applications if we can distinguish between scenarios A: client makes a bad request vs B: the backend returned an invalid result
There was a problem hiding this comment.
for instance if invalid top_k from the client → HTTP 422, calling on dense retrieval result missing _distance → HTTP 500.
| return parsed if isinstance(parsed, dict) else {} | ||
|
|
||
|
|
||
| def _flatten_legacy_entity_hit(hit: dict[str, Any]) -> dict[str, Any]: |
There was a problem hiding this comment.
Your taking away nesting here? Like no more json dicts?
There was a problem hiding this comment.
can you provide and example of a legacy and a "new" hit? Also why they are needed?
There was a problem hiding this comment.
there is still nesting, the flattening is only occuring so that the collection hit has a better way of handling both the legacy and new hit shape.
currently it doesn't seem like consumers can support both the legacy and new hit shapes without writing some custom impl of it so flattening it should resolve that portion.
There was a problem hiding this comment.
from what I could gather this seems to be the old backend hit:
{
"entity": {
"text": "retrieved chunk",
"source": {
"source_id": "document.pdf"
},
"content_metadata": {
"page_number": 2
}
},
"_distance": 0.12
}
this is the new internal nrl hit
{
"text": "retrieved chunk",
"source": {
"source_id": "document.pdf"
},
"metadata": {
"page_number": 2
},
"_distance": 0.12
}
and for the collection hit this is how it exists right now:
{
"text": "retrieved chunk",
"source": "document.pdf",
"metadata": {
"page_number": 2
},
"chunk_id": "chunk-123",
"document_id": "document-456",
"filename": "document.pdf",
"ranking": {
"value": 0.12,
"kind": "vector_distance",
"higher_is_better": false
}
}
| kind = "hybrid_relevance" | ||
| higher_is_better = True | ||
| else: | ||
| raise RetrievalContractError(f"unsupported collection retrieval mode: {retrieval_mode}") |
There was a problem hiding this comment.
All of these are just value errors?
There was a problem hiding this comment.
I mean yes essentially it is, but mapping it to be the new RetrievalContractError seems to be perhaps a better UI/UX as we can clearly show if its client vs nrl backend error better. Not sure if breakages on nrl backend is possible but still having these guards set up in case it does I think would be beneficial.
| from nemo_retriever.service.auth import authorized_scope, internal_auth_headers | ||
|
|
||
| scope = authorized_scope(request) | ||
| if body.collection_name: |
There was a problem hiding this comment.
This is now a first class feature, the collection management api, so it should be apart of all deployments? not just an option right?
charlesbluca
left a comment
charlesbluca
left a comment
There was a problem hiding this comment.
Requesting changes for the remaining deployment blockers: authentication must fail closed by default, the published service image needs the TXT/HTML tokenizer dependency, and the published Helm path needs supported internal VectorDB credential wiring. The wrong-scope response documentation should also be aligned with the implemented 401 contract.
The focused NRL suite passed 140 tests. The relevant full run passed 2,791 tests with 4 skips, 10 deselections, and one environment-dependent shared-temp-path failure.
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
| chunk_id: str | ||
| document_id: DocumentId | ||
| text: str | ||
| distance: float = Field( |
There was a problem hiding this comment.
P1: The current QueryHit.distance contract is incompatible with AIQ PR #371
On pinned NRL head fe33d24b, QueryHit requires native distance. The live head of AIQ PR #371 is still 1c73ba49, where QueryHitWire instead requires score. Fresh validation of a representative exact-head NRL payload containing distance: 0.125 fails with score: missing, so collection-query integration fails at wire validation before the AIQ adapter can normalize the result.
For clarity, an unpublished local AIQ successor is not the state of PR #371: it targets the earlier NRL ranking contract and rejects the current distance payload with ranking: missing. It therefore does not resolve the published cross-PR incompatibility.
Please land one coordinated wire contract across NRL and AIQ. If NRL keeps distance, update AIQ PR #371 and its compatibility tests to consume native distance and derive any required bounded Chunk.score at the adapter boundary. If NRL changes the wire contract again, coordinate that change with the AIQ model, adapter, and tests in the same integration round.
There was a problem hiding this comment.
Would make more sense to have this normalized on AIQ side instead of requiring retriever to do this
There was a problem hiding this comment.
looking into adding this fix into the existing 371 PR for AIQ
|
|
||
| if config.auth.api_token: | ||
| from nemo_retriever.service.auth import BearerAuthMiddleware | ||
| from nemo_retriever.service.auth import BearerAuthMiddleware |
There was a problem hiding this comment.
This only happened if there was a config auth api token, now it happens all the time, why the change?
| ] | ||
|
|
||
|
|
||
| def _normalize_expires_at(value: str | None) -> str | None: |
There was a problem hiding this comment.
why do you need this wrapper then?
| job_id: str, | ||
| filename: str | None = None, | ||
| ) -> None: | ||
| content_sha256: str | None = None, |
There was a problem hiding this comment.
Proposal: centralize collection ingestion context resolution
The collection-aware ingestion logic is currently repeated across the gateway, worker, and standalone branches of both submit_document_to_job() and submit_whole_document_to_job().
Each branch independently determines:
- The authorized scope and collection
- The associated job
- Append versus replacement behavior
- Manifest validation and replay detection
- The processing attempt ID
- The stable collection document ID
- The effective pipeline specification
- The fields forwarded to workers and VectorDB
I suggest introducing a typed CollectionIngestContext:
@dataclass(frozen=True, slots=True)
class CollectionIngestContext:
job_id: str
job: JobAggregate | None
scope: str
collection_name: str | None
operation: Literal["append", "replace"]
attempt_id: str
stable_document_id: str
content_sha256: str
manifest_entry_id: str | None
pipeline_spec: PipelineSpec | None
is_replay: bool = False
The identity fields have distinct meanings:
- attempt_id identifies this pipeline execution. It is used for polling, callbacks, and worker tracking.
- stable_document_id is the persistent logical document identity inside the collection.
- manifest_entry_id deterministically identifies one file in an idempotent submission.
- content_sha256 verifies that a replay contains the same file.
- operation identifies whether the attempt appends or replaces a document.
The resolver should account for the service topology:
- In gateway mode, the job comes from the gateway JobTracker. Scope and collection come from the authorized job. The attempt ID is either a new UUID or the ID from an existing replay record.
- In standalone mode, the job comes from the local JobTracker, with the same identity and authorization rules as the gateway.
- In worker mode, the job context, scope, collection, stable ID, and attempt ID come from trusted gateway-injected internal headers. The worker should not generate a new stable ID or trust a public scope header.
Conceptually:
def resolve_collection_ingest_context(
request: Request,
*,
job_id: str,
filename: str,
content_sha256: str,
manifest_entry_id: str | None,
requested_pipeline: PipelineSpec | None,
) -> CollectionIngestContext:
if _is_worker(request):
return _context_from_gateway_headers(
request,
job_id=job_id,
filename=filename,
content_sha256=content_sha256,
)
job = _require_job(job_id, request)
_validate_manifest_entry(
job,
manifest_entry_id,
filename,
content_sha256,
)
_validate_collection_pipeline_spec(job, requested_pipeline)
return _context_from_local_job(
job,
content_sha256=content_sha256,
manifest_entry_id=manifest_entry_id,
pipeline_spec=requested_pipeline,
)
Replay registration should remain an explicit second step because it mutates JobTracker:
context = resolve_collection_ingest_context(...)
acceptance = register_ingest_attempt(context)
That operation could return:
@dataclass(frozen=True, slots=True)
class IngestAcceptance:
context: CollectionIngestContext
record: DocumentRecord
created: bool
When created=False, the route returns the original acceptance without enqueueing duplicate processing.
The context can then construct a WorkItem consistently:
def make_work_item(
context: CollectionIngestContext,
*,
payload: bytes,
filename: str,
callback_url: str | None,
) -> WorkItem:
return WorkItem(
id=context.attempt_id,
payload=payload,
filename=filename,
job_id=context.job_id,
scope=context.scope,
collection_name=context.collection_name,
operation=context.operation,
content_sha256=context.content_sha256,
storage_document_id=context.stable_document_id,
pipeline_spec=(
context.pipeline_spec.model_dump(mode="json")
if context.pipeline_spec
else None
),
callback_url=callback_url,
)
Both ingestion routes could then follow the same high-level flow:
file_bytes = await file.read()
content_sha256 = hashlib.sha256(file_bytes).hexdigest()
context = resolve_collection_ingest_context(...)
acceptance = register_ingest_attempt(context)
if not acceptance.created:
return response_from_existing(acceptance.record)
item = make_work_item(
acceptance.context,
payload=file_bytes,
filename=file.filename,
callback_url=...,
)
await enqueue(item)
return accepted_response(acceptance)
Endpoint-specific behavior would remain limited to:
- File classification
- Realtime versus batch routing
- Response model
- Page-count routing
- Metrics
This would provide:
- One stable-document identity policy
- One manifest and replay policy
- One trusted gateway-to-worker boundary
- Typed propagation instead of loosely related arguments and headers
- Smaller ingestion routes
- Focused unit tests for identity and topology resolution
- Less risk that /document, /whole, gateway, and standalone behavior diverge
This should help reduce 150-200 Lines of code.
| summary="Create a new ingestion job aggregate", | ||
| ) | ||
| async def create_job(request: Request, response: Response, body: JobCreateRequest) -> JobCreatedResponse: | ||
| async def create_job(request: Request, response: Response, body: JobCreateRequest) -> JobCreatedResponse | Response: |
There was a problem hiding this comment.
Proposal: extract collection validation and idempotency fingerprint construction
create_job() currently combines several responsibilities:
- Resolving the authorized scope
- Verifying that VectorDB is enabled
- Looking up the target collection
- Confirming that the collection is active
- Looking up the target document for replacement
- Translating VectorDB errors into gateway responses
- Constructing the idempotency fingerprint
- Creating tracing state
- Registering the job
The collection validation and fingerprint logic occupy a substantial part of create_job() (
NeMo-Retriever/nemo_retriever/src/nemo_retriever/service/routers/ingest.py
Lines 702 to 831 in fe33d24
to read in isolation.
I suggest extracting two separate operations:
await validate_collection_job_target(
request,
scope=scope,
job_request=body,
)
fingerprint = build_job_idempotency_fingerprint(body)
1. Extract collection target validation
The validation helper should own all remote VectorDB resource checks:
async def validate_collection_job_target(
request: Request,
*,
scope: str,
job_request: JobCreateRequest,
) -> None:
if job_request.collection_name is None:
return
config = request.app.state.config
if not config.vectordb.enabled:
raise HTTPException(
status_code=404,
detail="VectorDB is not enabled in the service configuration",
)
headers = {
"X-NRL-Scope": scope,
**internal_auth_headers(config.vectordb.internal_api_token),
}
base_url = config.vectordb.vectordb_url.rstrip("/")
collection_url = (
f"{base_url}/v1/collections/{job_request.collection_name}"
)
async with httpx.AsyncClient(timeout=30.0) as client:
collection = await _get_vectordb_resource(
client,
collection_url,
headers=headers,
unavailable_detail="Failed to validate collection",
)
if collection.get("status") != "active":
raise HTTPException(
status_code=409,
detail="Collection is not active",
)
if (
job_request.operation == "replace"
and job_request.target_document_id is not None
):
await _get_vectordb_resource(
client,
(
f"{collection_url}/documents/"
f"{job_request.target_document_id}"
),
headers=headers,
unavailable_detail="Failed to validate replacement document",
)
The repeated HTTP request and response translation can also be centralized:
async def _get_vectordb_resource(
client: httpx.AsyncClient,
url: str,
*,
headers: dict[str, str],
unavailable_detail: str,
) -> dict[str, Any]:
try:
response = await client.get(url, headers=headers)
except httpx.HTTPError as exc:
raise HTTPException(
status_code=502,
detail=unavailable_detail,
) from exc
if response.status_code >= 400:
detail = response.text[:1000] or "VectorDB request failed"
raise HTTPException(
status_code=response.status_code,
detail=detail,
)
try:
return response.json()
except ValueError as exc:
raise HTTPException(
status_code=502,
detail="VectorDB returned malformed JSON",
) from exc
This has several benefits:
- Collection and replacement-document validation use one AsyncClient.
- Scope and internal authentication headers are constructed once.
- Transport errors are handled consistently.
- Malformed VectorDB responses are handled explicitly.
- create_job() no longer contains internal VectorDB URL construction.
- Additional job operations can reuse the same validation boundary.
The helper name should make clear that this is validation, not authorization. Authorization has already happened when the scope is attached to the request.
2. Extract fingerprint construction
The current fingerprint is assembled inline inside create_job():
fingerprint = hashlib.sha256(
json.dumps(
{
...
},
sort_keys=True,
separators=(",", ":"),
).encode()
).hexdigest()
This is security- and correctness-sensitive logic because it defines when two uses of the same idempotency key are considered the same request.
It should have a dedicated function:
_JOB_FINGERPRINT_VERSION = 1
def build_job_idempotency_fingerprint(
request: JobCreateRequest,
) -> str:
payload = {
"version": _JOB_FINGERPRINT_VERSION,
"request": request.model_dump(
mode="json",
exclude={"idempotency_key"},
),
}
canonical = json.dumps(
payload,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
)
return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
Using model_dump() instead of manually selecting fields has an important advantage: new semantically relevant request fields are less likely to be accidentally omitted from the fingerprint.
The idempotency key itself should be excluded because it identifies the stored fingerprint; it is not part of the request equivalence comparison.
A fingerprint version is useful because changing the canonical representation later would otherwise silently change idempotency behavior.
Current fingerprint omission
The current inline fingerprint includes:
- expected_documents
- collection_name
- operation
- target_document_id
- metadata
- document_manifest
It does not include:
- label
- retain_results
retain_results affects job behavior, so reusing an idempotency key with a different value should probably conflict rather than silently return the original job.
Whether label participates is a product-contract decision. If it is intentionally excluded because it is only descriptive, that should be explicit:
_FINGERPRINT_EXCLUDED_FIELDS = {
"idempotency_key",
"label", # Descriptive only; does not change processing semantics.
}
An explicit exclusion is preferable to maintaining a manual allowlist that can silently omit future fields.
Resulting create_job() flow
After extraction, the core route becomes much easier to follow:
async def create_job(
request: Request,
response: Response,
body: JobCreateRequest,
) -> JobCreatedResponse:
tracker = get_job_tracker()
if tracker is None:
raise HTTPException(
status_code=503,
detail="Job tracker not available",
)
scope = authorized_scope(request)
await validate_collection_job_target(
request,
scope=scope,
job_request=body,
)
fingerprint = build_job_idempotency_fingerprint(body)
job_id = uuid.uuid4().hex
trace_context, trace_id = create_job_trace_context(
request,
expected_documents=body.expected_documents,
)
job = tracker.register_job(
job_id,
expected_documents=body.expected_documents,
label=body.label,
metadata=body.metadata,
retain_results=body.retain_results,
trace_id=trace_id,
trace_context=trace_context,
collection_name=body.collection_name,
scope=scope,
operation=body.operation,
target_document_id=body.target_document_id,
idempotency_key=body.idempotency_key,
idempotency_fingerprint=fingerprint,
document_manifest=[
entry.model_dump(mode="json")
for entry in body.document_manifest
],
)
return job_created_response(job, requested_job_id=job_id)
This separates three concepts that are currently interleaved:
- The target collection and document exist and are usable.
- The request has a deterministic idempotency identity.
- The job is registered and traced.
this extraction could remove approximately 30–60 lines directly
| status_code=202, | ||
| summary="Submit a document to an existing job (auto-routed by page count)", | ||
| ) | ||
| async def submit_document_to_job( |
There was a problem hiding this comment.
Proposal: consolidate the shared document-submission flow
submit_document_to_job() and submit_whole_document_to_job() currently repeat much of the same collection-aware submission sequence across their gateway and standalone/worker branches.
The repeated sequence includes:
- Reading the uploaded file
- Calculating its size and SHA-256 digest
- Validating the file against the job manifest
- Resolving or allocating the processing attempt ID
- Resolving the stable collection document ID
- Registering the attempt with JobTracker
- Detecting idempotent replays
- Constructing a WorkItem
- Enqueueing the work
- Constructing the accepted response
I suggest extracting this into a shared preparation, registration, and work-item construction flow.
1. Prepare the upload once
File reading, size calculation, and hashing should happen in one helper:
@dataclass(frozen=True, slots=True)
class PreparedUpload:
payload: bytes
filename: str
content_sha256: str
size_bytes: int
async def prepare_upload(file: UploadFile) -> PreparedUpload:
payload = await file.read()
return PreparedUpload(
payload=payload,
filename=file.filename or "",
content_sha256=hashlib.sha256(payload).hexdigest(),
size_bytes=len(payload),
)
This gives the remaining ingestion flow one canonical representation of the uploaded file.
It also avoids independently calculating:
- file_bytes
- file_size
- content_sha256
- The effective filename
in each gateway and worker branch.
Upload-size enforcement should still happen before fully buffering the file when the framework provides the size. The helper is intended to centralize the post-validation read and hash, not bypass the existing upload-size guard.
2. Resolve the ingestion context
After preparing the file, resolve the job, scope, collection, operation, pipeline, and identity information:
context = resolve_collection_ingest_context(
request,
job_id=job_id,
filename=upload.filename,
content_sha256=upload.content_sha256,
manifest_entry_id=manifest_entry_id,
requested_pipeline=validated_spec,
)
This context should establish:
- The authoritative job
- The authorized scope
- The target collection
- Append versus replacement
- The attempt ID
- The stable document ID
- The effective pipeline specification
- Whether the request came from a gateway, worker, or standalone service
The topology rules should remain explicit:
- A gateway or standalone service resolves identity from its JobTracker.
- A worker uses identity and scope values supplied through the authenticated internal gateway headers.
- A worker must not generate a different stable document ID.
- A public scope header must not override the scope already authorized by the gateway.
3. Validate the manifest in one place
The manifest check should be part of context resolution for gateway and standalone requests:
_validate_manifest_entry(
job,
manifest_entry_id,
upload.filename,
upload.content_sha256,
)
That check should ensure:
- A manifest entry is supplied when the job declared a manifest.
- The entry exists in the job manifest.
- The filename matches.
- The content hash matches.
- Reuse of the entry with different content produces a conflict.
Worker requests should not independently reinterpret the public manifest. The gateway has already validated the upload before forwarding it.
If the worker needs the manifest entry for diagnostics, the gateway can forward it as authenticated internal metadata.
4. Allocate attempt and stable document IDs consistently
The shared identity logic should preserve the distinction between an ingestion attempt and a persistent collection document.
For an ordinary collection append:
attempt_id = new processing-attempt UUID
stable_document_id = new collection-document UUID
For replacement:
attempt_id = new processing-attempt UUID
stable_document_id = target document ID
For legacy non-collection ingestion:
attempt_id = new processing-attempt UUID
stable_document_id = attempt ID
For a worker:
attempt_id = gateway-provided attempt ID
stable_document_id = gateway-provided stable document ID
This logic currently appears in multiple branches. Centralizing it ensures /document and /whole cannot assign different identities for the same operation.
5. Register the attempt and detect replay atomically
Registration should be a separate explicit operation because it mutates the job tracker:
acceptance = register_ingest_attempt(
context,
filename=upload.filename,
content_sha256=upload.content_sha256,
manifest_entry_id=manifest_entry_id,
)
The result could be represented as:
@dataclass(frozen=True, slots=True)
class IngestAcceptance:
context: CollectionIngestContext
record: DocumentRecord | None
created: bool
created=True means a new attempt was registered and should be enqueued.
created=False means the manifest entry was already accepted. In that case, the route should return the original acceptance without:
- Consuming additional job capacity
- Generating another stable document ID
- Creating another WorkItem
- Starting duplicate pipeline processing
For worker requests, record can be None because the gateway owns the authoritative job tracker.
6. Construct WorkItem through one helper
The gateway, standalone service, /document, and /whole routes currently propagate the same collection fields separately.
A shared constructor can make this consistent:
def make_ingest_work_item(
acceptance: IngestAcceptance,
upload: PreparedUpload,
*,
pool_type: PoolType,
callback_url: str | None,
callback_headers: dict[str, str],
retain_results: bool,
) -> WorkItem:
context = acceptance.context
return WorkItem(
id=context.attempt_id,
payload=upload.payload,
filename=upload.filename,
callback_url=callback_url,
callback_headers=callback_headers,
job_id=context.job_id,
pipeline_spec=(
context.pipeline_spec.model_dump(mode="json")
if context.pipeline_spec is not None
else None
),
retain_results=retain_results,
scope=context.scope,
collection_name=context.collection_name,
operation=context.operation,
content_sha256=upload.content_sha256,
storage_document_id=context.stable_document_id,
)
This prevents errors where one branch forgets to propagate fields such as:
- scope
- collection_name
- operation
- content_sha256
- storage_document_id
- pipeline_spec
- Internal callback authentication
It also gives tests one place to verify the gateway-to-worker contract.
7. Share accepted-response construction
/document and /whole return different Pydantic response models, but the identity and acceptance data used to construct those responses are largely the same.
The shared flow can produce an endpoint-neutral result:
@dataclass(frozen=True, slots=True)
class AcceptedDocument:
job_id: str
attempt_id: str
stable_document_id: str
content_sha256: str
filename: str
size_bytes: int
created_at: str
replayed: bool
Each endpoint can then perform a thin response conversion:
def document_accepted_response(
accepted: AcceptedDocument,
) -> IngestAccepted:
return IngestAccepted(
document_id=accepted.stable_document_id,
attempt_id=accepted.attempt_id,
job_id=accepted.job_id,
content_sha256=accepted.content_sha256,
status="accepted",
created_at=accepted.created_at,
)
def whole_document_accepted_response(
accepted: AcceptedDocument,
) -> DocumentIngestAccepted:
return DocumentIngestAccepted(
document_id=accepted.stable_document_id,
attempt_id=accepted.attempt_id,
filename=accepted.filename,
file_size_bytes=accepted.size_bytes,
content_sha256=accepted.content_sha256,
status="accepted",
created_at=accepted.created_at,
)
This preserves the existing public response models while centralizing the source data.
Resulting route structure
After extraction, submit_document_to_job() could be reduced conceptually to:
async def submit_document_to_job(...):
meta = parse_ingest_metadata(metadata)
_check_upload_size(file, request)
classification = classify_and_validate(file, meta)
validated_spec = _resolve_pipeline_spec(request, meta)
upload = await prepare_upload(file)
route = _route_by_page_count(
upload.payload,
meta,
file_category=classification.category,
)
context = resolve_collection_ingest_context(
request,
job_id=job_id,
filename=upload.filename,
content_sha256=upload.content_sha256,
manifest_entry_id=manifest_entry_id,
requested_pipeline=validated_spec,
)
acceptance = register_ingest_attempt(
context,
filename=upload.filename,
content_sha256=upload.content_sha256,
manifest_entry_id=manifest_entry_id,
)
if not acceptance.created:
return document_accepted_response(
accepted_document_from_replay(
acceptance,
upload,
)
)
item = make_ingest_work_item(
acceptance,
upload,
pool_type=route,
callback_url=...,
callback_headers=...,
retain_results=...,
)
await enqueue_ingest_work(request, route, item)
accepted = accepted_document_from_new_attempt(
acceptance,
upload,
)
record_ingest_metrics(request, classification, accepted)
return document_accepted_response(accepted)
submit_whole_document_to_job() would follow the same flow but would:
- Always use PoolType.BATCH
- Return DocumentIngestAccepted
- Retain its endpoint-specific classification and metrics labels
What should remain endpoint-specific
The extraction should not force all ingestion routes into one oversized generic function.
The following behavior should remain in the endpoint:
- Parsing the endpoint-specific form data
- File classification
- Media dependency enforcement
- Page-count routing for /document
- Fixed batch routing for /whole
- The public response model
- Endpoint-specific metrics
- Endpoint-specific tracing span names
The shared helper should own identity, manifest, job registration, work-item propagation, and replay behavior—not every detail of the endpoint.
Expected benefit
This change would provide:
- One file hashing implementation
- One manifest-validation policy
- One stable-ID allocation policy
- One replay-registration path
- One WorkItem propagation contract
- One source of accepted-response identity data
- Less duplication across /document, /whole, gateway, worker, and standalone modes
- Less risk that one branch omits a collection or authentication field
this portion of the ingestion consolidation could remove approximately 80–140 production lines on its own. Combined with the proposed CollectionIngestContext resolver and shared pipeline forwarding, the total ingestion-related reduction could reach approximately 150–250 lines.
3 participants
Description
This PR adds collection and document lifecycle management to the NeMo Retriever service. It provides a logical, scope-isolated API for creating collections, ingesting and replacing documents, tracking asynchronous jobs, querying collection content, and cleaning up resources without exposing physical LanceDB table names or storage locations to clients.
What changed
Why
Applications that want to integrate NeMo Retriever (like aiq) need a stable service boundary for managing durable knowledge collections. Previously, consumers had to depend on ingestion-specific behavior or physical VectorDB details.
This change establishes a logical API boundary where:
Validation
Focused regression testing was run from commit
edfed55da:The focused coverage includes:
The collection workflow was also exercised end to end with collection creation, asynchronous ingestion, job polling, document listing, retrieval, deletion, cleanup, and persistence across service recreation.
Out of scope
The separate TXT/HTML tokenizer fix is intentionally not included in this PR. The integrated validation environment contained that independent patch, but this PR does not modify the tokenizer implementation.
Checklist