diff --git a/.github/workflows/ci-pull-request.yml b/.github/workflows/ci-pull-request.yml index ba5bb356e7..05c1be93d8 100644 --- a/.github/workflows/ci-pull-request.yml +++ b/.github/workflows/ci-pull-request.yml @@ -56,8 +56,6 @@ jobs: os-label: windows-x64 - runner: macos-26 os-label: macos-arm64 - - runner: macos-26-intel - os-label: macos-x64 steps: - name: Check out repository code uses: actions/checkout@v4 diff --git a/.github/workflows/docker-build-arm.yml b/.github/workflows/docker-build-arm.yml index 5cf05fd587..f5d4e004dd 100644 --- a/.github/workflows/docker-build-arm.yml +++ b/.github/workflows/docker-build-arm.yml @@ -37,7 +37,7 @@ jobs: SECRET_ARG="--secret id=hf_token,src=./scripts/private_local/hf_token.txt" fi docker buildx create --use - docker buildx build -f Dockerfile --platform linux/arm64 --load --target service --build-arg GIT_COMMIT=${GITHUB_SHA} --build-arg BASE_IMG=ubuntu --build-arg BASE_IMG_TAG=jammy-20250415.1 --build-arg DOWNLOAD_LLAMA_TOKENIZER=True $SECRET_ARG -t nrl-service:latest . + docker buildx build -f Dockerfile --platform linux/arm64 --load --target service --build-arg GIT_COMMIT=${GITHUB_SHA} --build-arg BASE_IMG=ubuntu --build-arg BASE_IMG_TAG=jammy-20250415.1 --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True $SECRET_ARG -t nrl-service:latest . - name: Cleanup HF token file if: always() diff --git a/.github/workflows/docker-nightly-publish.yml b/.github/workflows/docker-nightly-publish.yml index 1088a1acd4..5c1584ca9a 100644 --- a/.github/workflows/docker-nightly-publish.yml +++ b/.github/workflows/docker-nightly-publish.yml @@ -45,7 +45,7 @@ jobs: - name: Build Docker image run: | docker buildx create --use - docker buildx build -f Dockerfile --platform linux/amd64 --push --target service --build-arg GIT_COMMIT=${GITHUB_SHA} --build-arg DOWNLOAD_LLAMA_TOKENIZER=True --secret id=hf_token,src=./scripts/private_local/hf_token.txt -t ${{ secrets.DOCKER_REGISTRY }}/nrl-service:${{ env.CURRENT_DATE }} . + docker buildx build -f Dockerfile --platform linux/amd64 --push --target service --build-arg GIT_COMMIT=${GITHUB_SHA} --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True --secret id=hf_token,src=./scripts/private_local/hf_token.txt -t ${{ secrets.DOCKER_REGISTRY }}/nrl-service:${{ env.CURRENT_DATE }} . - name: Cleanup HF token file if: always() diff --git a/.github/workflows/docker-release-publish.yml b/.github/workflows/docker-release-publish.yml index ab0263ba97..d3e7346f0d 100644 --- a/.github/workflows/docker-release-publish.yml +++ b/.github/workflows/docker-release-publish.yml @@ -41,7 +41,7 @@ jobs: - name: Build Docker image run: | docker buildx create --use - docker buildx build -f Dockerfile --platform linux/amd64 --push --target service --build-arg GIT_COMMIT=${GITHUB_SHA} --build-arg DOWNLOAD_LLAMA_TOKENIZER=True --secret id=hf_token,src=./scripts/private_local/hf_token.txt -t ${{ secrets.DOCKER_REGISTRY }}/nrl-service:${{ env.SHORT_BRANCH_NAME }} . + docker buildx build -f Dockerfile --platform linux/amd64 --push --target service --build-arg GIT_COMMIT=${GITHUB_SHA} --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True --secret id=hf_token,src=./scripts/private_local/hf_token.txt -t ${{ secrets.DOCKER_REGISTRY }}/nrl-service:${{ env.SHORT_BRANCH_NAME }} . - name: Cleanup HF token file if: always() diff --git a/.github/workflows/integration-test-library-mode.yml b/.github/workflows/integration-test-library-mode.yml index a39911c9ac..5de7ed9f72 100644 --- a/.github/workflows/integration-test-library-mode.yml +++ b/.github/workflows/integration-test-library-mode.yml @@ -30,8 +30,6 @@ jobs: os-label: windows-x64 - runner: macos-26 os-label: macos-arm64 - - runner: macos-26-intel - os-label: macos-x64 env: # Hosted NIM / integrate.api auth (ExtractParams / EmbedParams read NVIDIA_API_KEY / NGC_API_KEY). diff --git a/.github/workflows/nrl-docs-nvidia-publish.yml b/.github/workflows/nrl-docs-nvidia-publish.yml index e93e7f69e1..be4a213e29 100644 --- a/.github/workflows/nrl-docs-nvidia-publish.yml +++ b/.github/workflows/nrl-docs-nvidia-publish.yml @@ -1,252 +1,225 @@ -# NeMo Retriever Library (NRL) — publish MkDocs site to docs.nvidia.com (S3 + Akamai). -# -# Replaces the manual brightspot S3 upload and Akamai ECCU flush for NRL guide content. -# MkDocs-only (no Sphinx / mike deploy). GitHub Pages remains nrl-docs-github-pages.yml. -# -# Required GitHub configuration (NVIDIA/NeMo-Retriever org/repo): -# Environment: docs-nvidia-prod (recommended approval gate before first live publish) -# Secrets: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_ASSUME_ROLE_ARN, -# S3_BUCKET_NAME, AKAMAI_HOST, AKAMAI_CLIENT_TOKEN, AKAMAI_CLIENT_SECRET, -# AKAMAI_ACCESS_TOKEN -# Variables (optional): DOCS_AWS_REGION, NRL_DOCS_PUBLISH_VERSION, DOCS_RELEASE_EMAILS -# -# S3 layout written under developer/docs/nemo/retriever/: -# index.html, versions.json, latest/, / (e.g. 26.5.0/) -# -# Per-release operator steps (before merge or tag): -# 1. Edit docs/publish/versions.json — add the new version, move the "latest" alias, -# keep older entries unless retiring them. -# 2. Set NRL_DOCS_PUBLISH_VERSION to the release version, or push a version tag -# (e.g. v26.5.0 or release-v26.5.0) so publish-docs and SITE_URL stay aligned. -# 3. Merge doc changes to main (or push the tag). CI copies versions.json into the -# artifact, builds with SITE_URL=.../retriever//, and runs publish-docs. -# Backports to old branches: workflow_dispatch with publish-as-latest: false (or /not-latest -# in the commit message); do not move the "latest" alias in versions.json. -name: NRL documentation — docs.nvidia.com publish - -on: - push: - branches: - - main - tags: - - "v[0-9]*.[0-9]*" - - "*-v[0-9]*.[0-9]*" - paths: - - "docs/**" - - "nemo_retriever/**" - - ".github/workflows/nrl-docs-nvidia-publish.yml" - workflow_dispatch: - inputs: - dry-run: - description: Skip S3 sync and Akamai flush (download/build only) - required: true - type: boolean - default: true - docs-version-override: - description: Version folder to publish (e.g. 26.5.0). Empty uses tag, NRL_DOCS_PUBLISH_VERSION, or 26.5.0. - required: false - type: string - default: "" - publish-as-latest: - description: Also sync the same build to latest/ (set false for backports) - required: false - type: boolean - default: true - notify-emails: - description: Extra Akamai notification emails (comma-separated) - required: false - type: string - default: "" - -permissions: - contents: read - -concurrency: - group: nrl-docs-nvidia-publish - cancel-in-progress: false - -env: - # Immutable pin for publish-docs (release v0.80.2). Do not use the mutable v0.80.2 tag. - FW_CI_TEMPLATES_REF: 563a769458cf1428194be49cb4de575b72d739e8 - DOCS_SITE_URL_BASE: https://docs.nvidia.com/nemo/retriever - S3_TARGET_PATH: developer/docs/nemo/retriever - ARTIFACT_NAME: nrl-docs-nvidia-site - -jobs: - resolve: - name: Resolve publish version and flags - runs-on: ubuntu-latest - outputs: - docs_version: ${{ steps.publish.outputs.docs_version }} - dry_run: ${{ steps.publish.outputs.dry_run }} - publish_latest: ${{ steps.publish.outputs.publish_latest }} - emails_csv: ${{ steps.publish.outputs.emails_csv }} - site_url: ${{ steps.publish.outputs.site_url }} - steps: - - name: Resolve publish inputs - id: publish - env: - DISPATCH_DRY_RUN: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run }} - VERSION_INPUT: ${{ github.event_name == 'workflow_dispatch' && inputs.docs-version-override || '' }} - PUBLISH_LATEST_INPUT: ${{ github.event_name == 'workflow_dispatch' && inputs.publish-as-latest }} - NOTIFY_INPUT: ${{ github.event_name == 'workflow_dispatch' && inputs.notify-emails || '' }} - NRL_DOCS_PUBLISH_VERSION_VAR: ${{ vars.NRL_DOCS_PUBLISH_VERSION }} - DOCS_RELEASE_EMAILS_VAR: ${{ vars.DOCS_RELEASE_EMAILS }} - HEAD_COMMIT_MSG: ${{ github.event.head_commit.message || '' }} - GITHUB_REF_NAME: ${{ github.ref }} - run: | - if [[ "${{ github.event_name }}" == "workflow_dispatch" ]]; then - echo "dry_run=${DISPATCH_DRY_RUN}" >> "$GITHUB_OUTPUT" - echo "publish_latest=${PUBLISH_LATEST_INPUT}" >> "$GITHUB_OUTPUT" - else - echo "dry_run=false" >> "$GITHUB_OUTPUT" - echo "publish_latest=true" >> "$GITHUB_OUTPUT" - fi - - VERSION="${VERSION_INPUT}" - if [[ -z "${VERSION}" ]]; then - if [[ "${GITHUB_REF_NAME}" =~ -v([0-9]+\.[0-9]+(\.[0-9]+)?)$ ]]; then - VERSION="${BASH_REMATCH[1]}" - elif [[ "${GITHUB_REF_NAME}" =~ ^refs/tags/v([0-9]+\.[0-9]+(\.[0-9]+)?)$ ]]; then - VERSION="${BASH_REMATCH[1]}" - fi - fi - if [[ -z "${VERSION}" ]]; then - VERSION="${NRL_DOCS_PUBLISH_VERSION_VAR}" - fi - if [[ -z "${VERSION}" ]]; then - VERSION="26.5.0" - fi - echo "docs_version=${VERSION}" >> "$GITHUB_OUTPUT" - echo "site_url=${{ env.DOCS_SITE_URL_BASE }}/${VERSION}/" >> "$GITHUB_OUTPUT" - - if [[ "${{ github.event_name }}" != "workflow_dispatch" ]]; then - if [[ "${HEAD_COMMIT_MSG}" =~ /not-latest ]] || [[ "${GITHUB_REF_NAME}" =~ not-latest ]]; then - echo "publish_latest=false" >> "$GITHUB_OUTPUT" - fi - fi - - EMAILS="${DOCS_RELEASE_EMAILS_VAR}" - if [[ -n "${NOTIFY_INPUT}" ]]; then - if [[ -n "${EMAILS}" ]]; then - EMAILS="${EMAILS},${NOTIFY_INPUT}" - else - EMAILS="${NOTIFY_INPUT}" - fi - fi - if [[ -z "${EMAILS}" ]]; then - EMAILS="kheiss@nvidia.com" - fi - echo "emails_csv=${EMAILS}" >> "$GITHUB_OUTPUT" - - build: - name: Build NRL docs for docs.nvidia.com - needs: resolve - runs-on: ubuntu-latest - steps: - # Publish the doc pages from the 26.05 release branch, not main. main keeps - # moving toward the next release, so its docs must not be published under the - # 26.5.0 label. mkdocs.yml, requirements.txt, and docs/docs/** come from 26.05. - - name: Checkout 26.05 docs content - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - ref: "26.05" - - # versions.json (the version picker + "latest" alias) is canonical on main and - # does not exist on 26.05, so fetch just that file from main into a side path. - - name: Checkout version-picker metadata from main - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - ref: main - path: .docs-main-meta - sparse-checkout: docs/publish/versions.json - sparse-checkout-cone-mode: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5 - with: - python-version: "3.12" - cache: pip - cache-dependency-path: docs/requirements.txt - - - name: Install Python dependencies - run: | - python -m pip install --upgrade pip - pip install -r docs/requirements.txt - pip install -e ./nemo_retriever - - - name: Build MkDocs (production site_url) - working-directory: docs - env: - SITE_URL: ${{ needs.resolve.outputs.site_url }} - DISABLE_MKDOCS_2_WARNING: "true" - run: mkdocs build -f mkdocs.yml --strict - - # MkDocs does not emit versions.json. publish-docs uploads the artifact copy to S3 - # (it does not merge with the live S3 file). versions.json is canonical on main. - - name: Stage version picker metadata for publish-docs - run: | - cp .docs-main-meta/docs/publish/versions.json docs/site/versions.json - - - name: Upload docs.nvidia.com site artifact - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 - with: - name: ${{ env.ARTIFACT_NAME }} - path: docs/site - if-no-files-found: error - - publish: - name: Publish to S3 and flush Akamai - needs: [resolve, build] - if: github.repository == 'NVIDIA/NeMo-Retriever' - runs-on: ubuntu-latest - environment: docs-nvidia-prod - steps: - - name: Checkout FW-CI-templates (publish-docs action) - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - repository: NVIDIA-NeMo/FW-CI-templates - ref: ${{ env.FW_CI_TEMPLATES_REF }} - path: FW-CI-templates - sparse-checkout: .github/actions/publish-docs - - - name: Publish versioned and latest docs (S3 + Akamai) - uses: ./FW-CI-templates/.github/actions/publish-docs - with: - dry-run: ${{ needs.resolve.outputs.dry_run }} - artifacts-name: ${{ env.ARTIFACT_NAME }} - artifacts-path: docs/site - project-type: single-docset - run-on-version-tag-only: "false" - docs-version-override: ${{ needs.resolve.outputs.docs_version }} - overwrite-latest-on-tag: ${{ needs.resolve.outputs.publish_latest }} - update-version-picker: "true" - request-name: nemo-retriever-docs-${{ github.run_id }} - emails-csv: ${{ needs.resolve.outputs.emails_csv }} - aws-region: ${{ vars.DOCS_AWS_REGION }} - aws-role-to-assume: ${{ secrets.AWS_ASSUME_ROLE_ARN }} - aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - akamai-host: ${{ secrets.AKAMAI_HOST }} - akamai-client-token: ${{ secrets.AKAMAI_CLIENT_TOKEN }} - akamai-client-secret: ${{ secrets.AKAMAI_CLIENT_SECRET }} - akamai-access-token: ${{ secrets.AKAMAI_ACCESS_TOKEN }} - s3-target-root: ${{ secrets.S3_BUCKET_NAME }} - s3-target-path: ${{ env.S3_TARGET_PATH }} - - # publish-docs leaves AWS credentials on the runner; only fetch the redirect file here. - - name: Checkout repository (root redirect) - if: needs.resolve.outputs.dry_run == 'false' - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - path: nemo-retriever-src - sparse-checkout: docs/publish/index.html - - - name: Upload retriever root index.html - if: needs.resolve.outputs.dry_run == 'false' - env: - S3_TARGET_ROOT: ${{ secrets.S3_BUCKET_NAME }} - run: | - S3_ROOT="${S3_TARGET_ROOT%/}" - DEST="${S3_ROOT}/${{ env.S3_TARGET_PATH }}/index.html" - aws s3 cp --quiet nemo-retriever-src/docs/publish/index.html "${DEST}" +# NeMo Retriever Library (NRL) — publish MkDocs site to docs.nvidia.com (S3 + Akamai). +# +# Replaces the manual brightspot S3 upload and Akamai ECCU flush for NRL guide content. +# MkDocs-only (no Sphinx / mike deploy). GitHub Pages remains nrl-docs-github-pages.yml. +# +# Required GitHub configuration (NVIDIA/NeMo-Retriever org/repo): +# Environment: docs-nvidia-prod (recommended approval gate before first live publish) +# Secrets: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_ASSUME_ROLE_ARN, +# S3_BUCKET_NAME, AKAMAI_HOST, AKAMAI_CLIENT_TOKEN, AKAMAI_CLIENT_SECRET, +# AKAMAI_ACCESS_TOKEN +# Variables (optional): DOCS_AWS_REGION, NRL_DOCS_PUBLISH_VERSION, DOCS_RELEASE_EMAILS +# +# S3 layout written under developer/docs/nemo/retriever/: +# index.html, versions.json, latest/, / (e.g. 26.5.0/) +# +# Manual publish only (workflow_dispatch). Do not re-enable push/tag auto-publish +# while docs.nvidia.com 26.5.0 / latest are frozen to the 26.05 release docs. +# +# Operator steps when ready to publish: +# 1. Merge doc changes to the 26.05 branch (content source for the build). +# 2. Keep docs/publish/versions.json on main accurate (version picker + latest alias). +# 3. Actions → "NRL documentation — docs.nvidia.com publish" → Run workflow: +# dry-run: false +# docs-version-override: 26.5.0 (or leave empty to use NRL_DOCS_PUBLISH_VERSION / default) +# publish-as-latest: true only when intentionally refreshing latest/ +# Backports: publish-as-latest: false; do not move the "latest" alias in versions.json. +name: NRL documentation — docs.nvidia.com publish + +on: + workflow_dispatch: + inputs: + dry-run: + description: Skip S3 sync and Akamai flush (download/build only) + required: true + type: boolean + default: true + docs-version-override: + description: Version folder to publish (e.g. 26.5.0). Empty uses NRL_DOCS_PUBLISH_VERSION, or 26.5.0. + required: false + type: string + default: "" + publish-as-latest: + description: Also sync the same build to latest/ (set false for backports) + required: false + type: boolean + default: true + notify-emails: + description: Extra Akamai notification emails (comma-separated) + required: false + type: string + default: "" + +permissions: + contents: read + actions: write + +concurrency: + group: nrl-docs-nvidia-publish + cancel-in-progress: false + +env: + # Immutable pin for publish-docs (release v0.80.2). Do not use the mutable v0.80.2 tag. + FW_CI_TEMPLATES_REF: 563a769458cf1428194be49cb4de575b72d739e8 + DOCS_SITE_URL_BASE: https://docs.nvidia.com/nemo/retriever + S3_TARGET_PATH: developer/docs/nemo/retriever + ARTIFACT_NAME: nrl-docs-nvidia-site + +jobs: + resolve: + name: Resolve publish version and flags + runs-on: ubuntu-latest + outputs: + docs_version: ${{ steps.publish.outputs.docs_version }} + dry_run: ${{ steps.publish.outputs.dry_run }} + publish_latest: ${{ steps.publish.outputs.publish_latest }} + emails_csv: ${{ steps.publish.outputs.emails_csv }} + site_url: ${{ steps.publish.outputs.site_url }} + steps: + - name: Resolve publish inputs + id: publish + env: + DISPATCH_DRY_RUN: ${{ inputs.dry-run }} + VERSION_INPUT: ${{ inputs.docs-version-override }} + PUBLISH_LATEST_INPUT: ${{ inputs.publish-as-latest }} + NOTIFY_INPUT: ${{ inputs.notify-emails }} + NRL_DOCS_PUBLISH_VERSION_VAR: ${{ vars.NRL_DOCS_PUBLISH_VERSION }} + DOCS_RELEASE_EMAILS_VAR: ${{ vars.DOCS_RELEASE_EMAILS }} + run: | + echo "dry_run=${DISPATCH_DRY_RUN}" >> "$GITHUB_OUTPUT" + echo "publish_latest=${PUBLISH_LATEST_INPUT}" >> "$GITHUB_OUTPUT" + + VERSION="${VERSION_INPUT}" + if [[ -z "${VERSION}" ]]; then + VERSION="${NRL_DOCS_PUBLISH_VERSION_VAR}" + fi + if [[ -z "${VERSION}" ]]; then + VERSION="26.5.0" + fi + echo "docs_version=${VERSION}" >> "$GITHUB_OUTPUT" + echo "site_url=${{ env.DOCS_SITE_URL_BASE }}/${VERSION}/" >> "$GITHUB_OUTPUT" + + EMAILS="${DOCS_RELEASE_EMAILS_VAR}" + if [[ -n "${NOTIFY_INPUT}" ]]; then + if [[ -n "${EMAILS}" ]]; then + EMAILS="${EMAILS},${NOTIFY_INPUT}" + else + EMAILS="${NOTIFY_INPUT}" + fi + fi + if [[ -z "${EMAILS}" ]]; then + EMAILS="kheiss@nvidia.com" + fi + echo "emails_csv=${EMAILS}" >> "$GITHUB_OUTPUT" + + build: + name: Build NRL docs for docs.nvidia.com + needs: resolve + runs-on: ubuntu-latest + steps: + # Publish the doc pages from the 26.05 release branch, not main. main keeps + # moving toward the next release, so its docs must not be published under the + # 26.5.0 label. mkdocs.yml, requirements.txt, and docs/docs/** come from 26.05. + - name: Checkout 26.05 docs content + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + ref: "26.05" + + # versions.json (the version picker + "latest" alias) is canonical on main and + # does not exist on 26.05, so fetch just that file from main into a side path. + - name: Checkout version-picker metadata from main + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + ref: main + path: .docs-main-meta + sparse-checkout: docs/publish/versions.json + sparse-checkout-cone-mode: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: docs/requirements.txt + + - name: Install Python dependencies + run: | + python -m pip install --upgrade pip + pip install -r docs/requirements.txt + pip install -e ./nemo_retriever + + - name: Build MkDocs (production site_url) + working-directory: docs + env: + SITE_URL: ${{ needs.resolve.outputs.site_url }} + DISABLE_MKDOCS_2_WARNING: "true" + run: mkdocs build -f mkdocs.yml --strict + + # MkDocs does not emit versions.json. publish-docs uploads the artifact copy to S3 + # (it does not merge with the live S3 file). versions.json is canonical on main. + - name: Stage version picker metadata for publish-docs + run: | + cp .docs-main-meta/docs/publish/versions.json docs/site/versions.json + + - name: Upload docs.nvidia.com site artifact + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: ${{ env.ARTIFACT_NAME }} + path: docs/site + if-no-files-found: error + + publish: + name: Publish to S3 and flush Akamai + needs: [resolve, build] + if: github.repository == 'NVIDIA/NeMo-Retriever' + runs-on: ubuntu-latest + environment: docs-nvidia-prod + steps: + - name: Checkout FW-CI-templates (publish-docs action) + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + repository: NVIDIA-NeMo/FW-CI-templates + ref: ${{ env.FW_CI_TEMPLATES_REF }} + path: FW-CI-templates + sparse-checkout: .github/actions/publish-docs + + - name: Publish versioned and latest docs (S3 + Akamai) + uses: ./FW-CI-templates/.github/actions/publish-docs + with: + dry-run: ${{ needs.resolve.outputs.dry_run }} + artifacts-name: ${{ env.ARTIFACT_NAME }} + artifacts-path: docs/site + project-type: single-docset + run-on-version-tag-only: "false" + docs-version-override: ${{ needs.resolve.outputs.docs_version }} + overwrite-latest-on-tag: ${{ needs.resolve.outputs.publish_latest }} + update-version-picker: "true" + request-name: nemo-retriever-docs-${{ github.run_id }} + emails-csv: ${{ needs.resolve.outputs.emails_csv }} + aws-region: ${{ vars.DOCS_AWS_REGION }} + aws-role-to-assume: ${{ secrets.AWS_ASSUME_ROLE_ARN }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + akamai-host: ${{ secrets.AKAMAI_HOST }} + akamai-client-token: ${{ secrets.AKAMAI_CLIENT_TOKEN }} + akamai-client-secret: ${{ secrets.AKAMAI_CLIENT_SECRET }} + akamai-access-token: ${{ secrets.AKAMAI_ACCESS_TOKEN }} + s3-target-root: ${{ secrets.S3_BUCKET_NAME }} + s3-target-path: ${{ env.S3_TARGET_PATH }} + + # publish-docs leaves AWS credentials on the runner; only fetch the redirect file here. + - name: Checkout repository (root redirect) + if: needs.resolve.outputs.dry_run == 'false' + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + path: nemo-retriever-src + sparse-checkout: docs/publish/index.html + + - name: Upload retriever root index.html + if: needs.resolve.outputs.dry_run == 'false' + env: + S3_TARGET_ROOT: ${{ secrets.S3_BUCKET_NAME }} + run: | + S3_ROOT="${S3_TARGET_ROOT%/}" + DEST="${S3_ROOT}/${{ env.S3_TARGET_PATH }}/index.html" + aws s3 cp --quiet nemo-retriever-src/docs/publish/index.html "${DEST}" diff --git a/.github/workflows/perform-release.yml b/.github/workflows/perform-release.yml index 30641e8a17..d97d3a63df 100644 --- a/.github/workflows/perform-release.yml +++ b/.github/workflows/perform-release.yml @@ -151,7 +151,7 @@ jobs: target: service platforms: linux/amd64 build-args: | - DOWNLOAD_LLAMA_TOKENIZER=True + DOWNLOAD_DEFAULT_TOKENIZER=True GIT_COMMIT=${{ github.sha }} tags: ${{ steps.meta.outputs.image }} secret-files: hf_token=./scripts/private_local/hf_token.txt diff --git a/.github/workflows/release-docker.yml b/.github/workflows/release-docker.yml index adb45bc4f3..32afd7ed7e 100644 --- a/.github/workflows/release-docker.yml +++ b/.github/workflows/release-docker.yml @@ -63,7 +63,7 @@ jobs: --push \ --target service \ --build-arg HF_ACCESS_TOKEN=${{ secrets.HF_ACCESS_TOKEN }} \ - --build-arg DOWNLOAD_LLAMA_TOKENIZER=True \ + --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True \ --build-arg GIT_COMMIT=${GITHUB_SHA} \ -t ${{ secrets.DOCKER_REGISTRY }}/nrl-service:${{ needs.determine-version.outputs.version }} \ . diff --git a/.github/workflows/reusable-docker-build-and-test.yml b/.github/workflows/reusable-docker-build-and-test.yml index c56f980a49..f9925512a5 100644 --- a/.github/workflows/reusable-docker-build-and-test.yml +++ b/.github/workflows/reusable-docker-build-and-test.yml @@ -36,8 +36,8 @@ on: required: false type: boolean default: false - download-llama-tokenizer: - description: 'Set Docker build arg DOWNLOAD_LLAMA_TOKENIZER' + download-default-tokenizer: + description: 'Set Docker build arg DOWNLOAD_DEFAULT_TOKENIZER' required: false type: boolean default: true @@ -141,9 +141,9 @@ jobs: SECRET_ARG="--secret id=hf_token,src=./scripts/private_local/hf_token.txt" fi - DOWNLOAD_LLAMA="True" - if [ "${{ inputs.download-llama-tokenizer }}" != "true" ]; then - DOWNLOAD_LLAMA="False" + DOWNLOAD_DEFAULT_TOKENIZER="True" + if [ "${{ inputs.download-default-tokenizer }}" != "true" ]; then + DOWNLOAD_DEFAULT_TOKENIZER="False" fi docker buildx build \ @@ -154,7 +154,7 @@ jobs: --build-arg GIT_COMMIT=${GITHUB_SHA} \ --build-arg BASE_IMG="${{ inputs.base-image }}" \ --build-arg BASE_IMG_TAG="${{ inputs.base-image-tag }}" \ - --build-arg DOWNLOAD_LLAMA_TOKENIZER=$DOWNLOAD_LLAMA \ + --build-arg DOWNLOAD_DEFAULT_TOKENIZER=$DOWNLOAD_DEFAULT_TOKENIZER \ $SECRET_ARG \ $TAG_ARGS \ . @@ -204,6 +204,47 @@ jobs: exit 1 fi + if [ "${{ inputs.download-default-tokenizer }}" = "true" ]; then + # Verify the production image before test-only dependencies are + # installed. Disabling networking proves both service variants + # can use the tokenizer cache populated by the build argument. + docker run --rm -i --network none --platform "$PLATFORM" \ + -e NEMO_RETRIEVER_IMAGE_TARGET="${{ inputs.target }}" \ + --entrypoint bash "$IMAGE" \ + -lc 'source /opt/retriever_runtime/bin/activate && python -' <<'PY' + import importlib.util + import os + + target = os.environ["NEMO_RETRIEVER_IMAGE_TARGET"] + assert importlib.util.find_spec("tokenizers") is not None + assert importlib.util.find_spec("huggingface_hub") is not None + if target == "service": + assert importlib.util.find_spec("transformers") is None + assert os.environ.get("HF_HUB_OFFLINE") == "1" + elif target == "service-gpu": + assert os.environ.get("HF_HUB_OFFLINE") != "1" + + from nemo_retriever.common.modality.html.convert import html_bytes_to_chunks_df + from nemo_retriever.common.modality.txt.split import ( + DEFAULT_TOKENIZER_MODEL_ID, + txt_bytes_to_chunks_df, + ) + from nemo_retriever.common.modality.txt.tokenizer_provider import load_chunk_tokenizer + + tokenizer = load_chunk_tokenizer(DEFAULT_TOKENIZER_MODEL_ID) + assert tokenizer.encode("offline tokenizer smoke test") + + txt = txt_bytes_to_chunks_df(b"offline tokenizer smoke test", "smoke.txt") + html = html_bytes_to_chunks_df( + b"

offline tokenizer smoke test

", + "smoke.html", + ) + assert not txt.empty + assert not html.empty + print(f"{target} offline TXT/HTML chunking smoke test passed") + PY + fi + if [ "${{ inputs.test-selection }}" != "full" ] && [ "${{ inputs.test-selection }}" != "random" ]; then echo "Error: test-selection must be 'full' or 'random'" exit 1 diff --git a/.github/workflows/reusable-docker-build.yml b/.github/workflows/reusable-docker-build.yml index d7e6af4406..c24705ae8b 100644 --- a/.github/workflows/reusable-docker-build.yml +++ b/.github/workflows/reusable-docker-build.yml @@ -41,8 +41,8 @@ on: required: false type: boolean default: false - download-llama-tokenizer: - description: 'Set Docker build arg DOWNLOAD_LLAMA_TOKENIZER' + download-default-tokenizer: + description: 'Set Docker build arg DOWNLOAD_DEFAULT_TOKENIZER' required: false type: boolean default: true @@ -141,9 +141,9 @@ jobs: OUTPUT_FLAG="--push" fi - DOWNLOAD_LLAMA="True" - if [ "${{ inputs.download-llama-tokenizer }}" != "true" ]; then - DOWNLOAD_LLAMA="False" + DOWNLOAD_DEFAULT_TOKENIZER="True" + if [ "${{ inputs.download-default-tokenizer }}" != "true" ]; then + DOWNLOAD_DEFAULT_TOKENIZER="False" fi docker buildx build \ @@ -154,7 +154,7 @@ jobs: --build-arg GIT_COMMIT=${GITHUB_SHA} \ --build-arg BASE_IMG="${{ inputs.base-image }}" \ --build-arg BASE_IMG_TAG="${{ inputs.base-image-tag }}" \ - --build-arg DOWNLOAD_LLAMA_TOKENIZER=$DOWNLOAD_LLAMA \ + --build-arg DOWNLOAD_DEFAULT_TOKENIZER=$DOWNLOAD_DEFAULT_TOKENIZER \ $SECRET_ARG \ $TAG_ARGS \ . diff --git a/.github/workflows/scheduled-nightly.yml b/.github/workflows/scheduled-nightly.yml index 6f2bed65e4..0c363d7ee3 100644 --- a/.github/workflows/scheduled-nightly.yml +++ b/.github/workflows/scheduled-nightly.yml @@ -66,7 +66,7 @@ jobs: --push \ --target service \ --build-arg HF_ACCESS_TOKEN=${{ secrets.HF_ACCESS_TOKEN }} \ - --build-arg DOWNLOAD_LLAMA_TOKENIZER=True \ + --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True \ --build-arg GIT_COMMIT=${GITHUB_SHA} \ -t ${{ secrets.DOCKER_REGISTRY }}/nrl-service:${{ needs.determine-version.outputs.docker-tag }} \ . diff --git a/.gitignore b/.gitignore index 402eb807d6..568c81edb0 100644 --- a/.gitignore +++ b/.gitignore @@ -189,6 +189,7 @@ tags /build*/ dask-worker-space data/* +!data/jp20_query_gt.csv docs/source/_lib docs/source/_modules mlruns/* diff --git a/.greptile/config.json b/.greptile/config.json index 47000e57c3..d2f1ebd3cb 100644 --- a/.greptile/config.json +++ b/.greptile/config.json @@ -179,6 +179,20 @@ "rule": "Non-obvious design decisions, performance trade-offs, and environmental assumptions must be documented in code comments or docstrings.", "scope": ["**/*.py"], "severity": "low" + }, + { + "id": "documentation-stays-current", + "rule": "When this PR changes user-facing code, configuration, examples, or documentation, verify that the affected documentation remains accurate. Compare documented APIs, CLI commands and flags, configuration keys and defaults, supported formats, error behavior, and code examples against the PR's checked-in source and tests. Flag concrete stale, incorrect, contradictory, or missing documentation, and identify the affected page. Do not request documentation for internal-only changes or speculate without a specific mismatch.", + "scope": [ + "nemo_retriever/src/**", + "nemo_retriever/helm/**", + "nemo_retriever/README.md", + "nemo_retriever/docs/**", + "docs/docs/**", + "examples/**", + "README.md" + ], + "severity": "medium" } ], diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..a2962673fb --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,69 @@ +# Agent Instructions — NeMo Retriever + +These standards apply to AI coding agents working in this repository (Cursor, Claude Code, and compatible hosts). + +## Project Overview + +NVIDIA NeMo Retriever Library (NRL) is the multimodal extraction and retrieval library formerly known in docs as NV-Ingest. Published customer documentation lives primarily under `docs/docs/extraction/` and builds with MkDocs (`docs/mkdocs.yml`). Library and Helm sources live under `nemo_retriever/`. + +## Repository Map + +| Path | Purpose | +|------|---------| +| `docs/docs/extraction/` | Published NRL extraction documentation (MkDocs) | +| `docs/mkdocs.yml` | Nav, redirects, and MkDocs config | +| `nemo_retriever/` | Library source, CLI, Helm chart, and in-repo README examples | +| `nemo_retriever/tests/` | Library tests | +| `.github/workflows/` | CI, including NRL docs publish workflows | + +## Documentation + +- Treat `docs/docs/extraction/` and related MkDocs pages as the source of truth for user-facing NRL documentation. Follow [`docs/AGENTS.md`](docs/AGENTS.md). +- Before completing a code change, determine whether it changes a **user-visible** surface. This includes a public API, CLI, configuration, Helm values or defaults, workflow, error message or error contract, supported file type, or other supported product behavior. +- When it does and the host supports subagents, start a documentation authoring subagent while the primary agent continues the implementation. Direct it to read `docs/AGENTS.md`, update the affected docs, and run validation. Give it the changed sources and user-visible impact. +- Reconcile the authoring subagent's documentation changes and validation evidence before completing the implementation. Include the required documentation in the same change when the repository workflow allows a combined PR. If policy requires a **docs-only** follow-up PR, open that PR in the same task and link it from the code PR. +- If the host cannot run subagents, read `docs/AGENTS.md` in the primary task, complete the documentation work, and run its documented validation. Do not omit required documentation because parallel execution is unavailable. +- Do not document defaults or behavior that `main` does not have yet. +- Documentation PRs that change published NRL prose must stay docs-scoped. Do not change `nemo_retriever/src/**`, tests, Helm chart behavior, lockfiles, or runtime CI env on a docs PR unless the user explicitly requests eng work. +- Verified product surfaces that often need docs updates: Python `create_ingestor` / `GraphIngestor` APIs, `retriever` CLI, Helm chart README and values, support matrix and NIM defaults, authentication and environment variables, error and troubleshoot guidance, and release notes. + +### NVIDIA DORI Routing + +Select the documentation path from current host capabilities. +Do not ask the user to classify themselves or store repository-scoped identity +state during a normal documentation task. + +1. Check whether the current agent exposes `dori_handle` or `dori_route` and + `dori_collections`. + If the user explicitly asks not to use DORI, use the + [Writing Style Guide](docs/AGENTS.md#writing-style-guide) instead. +2. When those tools are available, list the installed collections. + - If a collection source contains `tech-docs/skill-library`, use DORI for + task routing. + - If the collection is missing, inaccessible, or cannot be verified, + continue with the + [Writing Style Guide](docs/AGENTS.md#writing-style-guide). +3. When the DORI tools are unavailable, continue with the Writing Style Guide. + Do not inspect a shell-visible CLI, install software, or configure the host + during a normal documentation task. +4. Use [NVIDIA DORI Setup](docs/DORI_SETUP.md) only when the user explicitly + asks to install or configure DORI. + +Capability detection does not approve installation or host configuration. +DORI unavailability must not block documentation work. + +## Engineering Guardrails + +- Prefer small, focused diffs that match existing style. +- Do not invent APIs, CLI flags, Helm keys, or defaults. Verify against checked-in source or tests. +- Never commit secrets, API keys, or credentials. +- Do not add lint, hooks, or CI from agent guidance alone. Those require a separately reviewed repository change. +- Do not create or modify `CLAUDE.md` as part of documentation-agent setup. + +## Validation Shortcuts + +| Change type | Validation | +|-------------|------------| +| Docs under `docs/` | From `docs/`: `python -m mkdocs build --strict --config-file mkdocs.yml` when the environment supports it | +| Library code | Run the targeted tests that cover the changed modules | +| Docs-only PR scope | `git diff --name-only upstream/main...HEAD` (or `origin/main...HEAD`) and confirm no runtime/out-of-scope paths | diff --git a/Dockerfile b/Dockerfile index 1f09a20e1c..4c8785f1b5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,8 +4,8 @@ # syntax=docker/dockerfile:1.3 # # Build from repo root: docker build -f Dockerfile -t nemo-retriever . -# Service (NIM-forwarding): docker build -f Dockerfile --target service -t nemo-retriever-service . -# Service + in-pod HF: docker build -f Dockerfile --target service-gpu -t nemo-retriever-service-gpu . +# Service (NIM-forwarding): docker build -f Dockerfile --target service --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True -t nemo-retriever-service . +# Service + in-pod HF: docker build -f Dockerfile --target service-gpu --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True -t nemo-retriever-service-gpu . # Runtime ffmpeg/ffprobe install for service image: docker run -e INSTALL_FFMPEG=true nemo-retriever-service # Run: docker run nemo-retriever (shell with venv active) # Run with dev mount: docker run -v $(pwd):/workspace -it nemo-retriever (code changes reflect without rebuild) @@ -16,6 +16,10 @@ ARG BASE_IMG_TAG=jammy-20250619 FROM $BASE_IMG:$BASE_IMG_TAG AS base +ARG DOWNLOAD_DEFAULT_TOKENIZER="False" + +ENV HF_HOME=/opt/nemo-retriever/huggingface + RUN apt-get update && apt-get install -y --no-install-recommends \ bzip2 \ ca-certificates \ @@ -116,7 +120,10 @@ SHELL ["/bin/bash", "-c"] RUN --mount=type=cache,target=/root/.cache/pip \ --mount=type=cache,target=/root/.cache/uv \ . /opt/retriever_runtime/bin/activate \ - && uv pip install -e "./nemo_retriever[service]" + && uv pip install -e "./nemo_retriever[service,multimedia]" \ + && if [ "${DOWNLOAD_DEFAULT_TOKENIZER}" = "True" ]; then \ + python -c "from nemo_retriever.common.modality.txt.split import DEFAULT_TOKENIZER_MODEL_ID; from nemo_retriever.common.modality.txt.tokenizer_provider import load_chunk_tokenizer; load_chunk_tokenizer(DEFAULT_TOKENIZER_MODEL_ID)"; \ + fi # GPU service install: in-pod Hugging Face models + multimedia (ASR, SVG). # Build target: service-gpu @@ -132,7 +139,10 @@ SHELL ["/bin/bash", "-c"] RUN --mount=type=cache,target=/root/.cache/pip \ --mount=type=cache,target=/root/.cache/uv \ . /opt/retriever_runtime/bin/activate \ - && uv pip install -e "./nemo_retriever[service,local,multimedia]" + && uv pip install -e "./nemo_retriever[service,local,multimedia]" \ + && if [ "${DOWNLOAD_DEFAULT_TOKENIZER}" = "True" ]; then \ + python -c "from nemo_retriever.common.modality.txt.split import DEFAULT_TOKENIZER_MODEL_ID; from nemo_retriever.common.modality.txt.tokenizer_provider import load_chunk_tokenizer; load_chunk_tokenizer(DEFAULT_TOKENIZER_MODEL_ID)"; \ + fi # Default: run in-process pipeline (help if no args) CMD ["/bin/bash"] @@ -141,6 +151,7 @@ CMD ["/bin/bash"] # Service profile: run the FastAPI ingest service. # # Build: docker build -f Dockerfile --target service \ +# --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True \ # -t nemo-retriever-service . # # Run with the bundled default config: @@ -158,6 +169,7 @@ CMD ["/bin/bash"] FROM install AS service ENV NEMO_RETRIEVER_SERVICE_CONFIG=/etc/nemo-retriever/retriever-service.yaml +ENV HF_HUB_OFFLINE=1 ENV PATH=/opt/retriever_runtime/bin:$PATH @@ -173,10 +185,10 @@ RUN chmod a+rx /usr/local/bin/uv /usr/local/bin/uvx /usr/local/bin/retriever-ser > /etc/sudoers.d/nemo-ffmpeg \ && chmod 0440 /etc/sudoers.d/nemo-ffmpeg \ && visudo -cf /etc/sudoers.d/nemo-ffmpeg \ - && mkdir -p /etc/nemo-retriever /var/lib/nemo-retriever /data/vectordb \ + && mkdir -p /etc/nemo-retriever /var/lib/nemo-retriever /data/vectordb "${HF_HOME}" \ && cp /workspace/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml \ "${NEMO_RETRIEVER_SERVICE_CONFIG}" \ - && chown -R nemo:nemo /workspace /etc/nemo-retriever /var/lib/nemo-retriever /data/vectordb /opt/retriever_runtime + && chown -R nemo:nemo /workspace /etc/nemo-retriever /var/lib/nemo-retriever /data/vectordb "${HF_HOME}" /opt/retriever_runtime EXPOSE 7670 @@ -190,6 +202,7 @@ CMD ["retriever", "service", "start", "--config", "/etc/nemo-retriever/retriever # GPU service profile: FastAPI ingest service with in-pod Hugging Face models. # # Build: docker build -f Dockerfile --target service-gpu \ +# --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True \ # -t nemo-retriever-service-gpu . # # Run (requires --gpus all and local_models.enabled in config): @@ -215,10 +228,10 @@ RUN chmod a+rx /usr/local/bin/uv /usr/local/bin/uvx /usr/local/bin/retriever-ser > /etc/sudoers.d/nemo-ffmpeg \ && chmod 0440 /etc/sudoers.d/nemo-ffmpeg \ && visudo -cf /etc/sudoers.d/nemo-ffmpeg \ - && mkdir -p /etc/nemo-retriever /var/lib/nemo-retriever /data/vectordb \ + && mkdir -p /etc/nemo-retriever /var/lib/nemo-retriever /data/vectordb "${HF_HOME}" \ && cp /workspace/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml \ "${NEMO_RETRIEVER_SERVICE_CONFIG}" \ - && chown -R nemo:nemo /workspace /etc/nemo-retriever /var/lib/nemo-retriever /data/vectordb /opt/retriever_runtime + && chown -R nemo:nemo /workspace /etc/nemo-retriever /var/lib/nemo-retriever /data/vectordb "${HF_HOME}" /opt/retriever_runtime EXPOSE 7670 diff --git a/README.md b/README.md index 98dff8ec66..c80093c2aa 100644 --- a/README.md +++ b/README.md @@ -141,7 +141,8 @@ Cat is the animal whose activity (jumping onto a laptop) matches the location of - **[Official Documentation](https://docs.nvidia.com/nemo/retriever/extraction/)** - Complete user guides, API references, and deployment instructions - **[Getting Started Guide](https://docs.nvidia.com/nemo/retriever/extraction/overview/)** - Overview and prerequisites for production deployments -- **[Benchmarking Guide](nemo_retriever/docs/cli/benchmarking.md)** - Performance testing and recall evaluation framework +- **[Retriever Harness](nemo_retriever/harness/README.md)** - Repeatable end-to-end ingest and retrieval benchmarks +- **[Stage Benchmarking](nemo_retriever/docs/cli/benchmarking.md)** - Internal per-stage throughput measurements - **[MIG Deployment](nemo_retriever/helm/README.md)** - Multi-Instance GPU configurations for Kubernetes - **[API Documentation](docs/docs/extraction/nemo-retriever-api-reference.md)** - Python client and API reference @@ -160,13 +161,17 @@ https://pypi.org/project/pdfservices-sdk/ [license agreement](https://github.com/adobe/pdfservices-python-sdk?tab=License-1-ov-file) for the pdfservices-sdk before enabling this option. - **Built With Llama**: - - **Description**: The NeMo Retriever ingestion container comes with the `meta-llama/Llama-3.2-1B` tokenizer pre-downloaded so - that the split task can use it for token-based splitting without making a network request. The [Llama 3.2 Community License Agreement](https://huggingface.co/meta-llama/Llama-3.2-1B/blob/main/LICENSE.txt) governs your use of these Llama materials. - - If you're building the container yourself and want to pre-download this model, you'll first need to set - `DOWNLOAD_LLAMA_TOKENIZER` to `True`. Because this is a gated model, you'll also need to - [request access](https://huggingface.co/meta-llama/Llama-3.2-1B) and set `HF_ACCESS_TOKEN` to your HuggingFace - access token in order to use it. + - **Description**: Published NeMo Retriever service images pre-cache the + revision-pinned tokenizer for the default embedding model, currently + [`nvidia/llama-nemotron-embed-vl-1b-v2`](https://huggingface.co/nvidia/llama-nemotron-embed-vl-1b-v2), + so TXT and HTML splitting does not need network access at runtime. The model + is governed by the [NVIDIA Open Model License](https://huggingface.co/nvidia/llama-nemotron-embed-vl-1b-v2/blob/main/LICENSE) + and identifies additional Llama 3.2 terms in its model card. + + Source builds opt in to the same pinned tokenizer cache with + `--build-arg DOWNLOAD_DEFAULT_TOKENIZER=True`. The default tokenizer + repository is not gated, so this download does not require a Hugging Face + access token. Before contributing to this project, please review our [Contributor Guide](CONTRIBUTING.md). diff --git a/THIRD_PARTY_LICENSES.md b/THIRD_PARTY_LICENSES.md index 479a8a5ef7..01d47bcaeb 100644 --- a/THIRD_PARTY_LICENSES.md +++ b/THIRD_PARTY_LICENSES.md @@ -2,7 +2,10 @@ This project uses the following third-party components: -## Llama 3.2 +## Llama Nemotron Embed VL 1B v2 -- **License**: [Llama 3.2 Community License Agreement](https://huggingface.co/meta-llama/Llama-3.2-1B/blob/main/LICENSE.txt) -- **Copyright**: © Meta Platforms, Inc. All Rights Reserved. +- **Component**: The revision-pinned tokenizer artifact from + [`nvidia/llama-nemotron-embed-vl-1b-v2`](https://huggingface.co/nvidia/llama-nemotron-embed-vl-1b-v2). +- **License**: [NVIDIA Open Model License](https://huggingface.co/nvidia/llama-nemotron-embed-vl-1b-v2/blob/main/LICENSE). + The model card also identifies the Llama 3.2 Community License Agreement as + additional terms and includes the Built with Llama attribution. diff --git a/data/jp20_query_gt.csv b/data/jp20_query_gt.csv new file mode 100644 index 0000000000..8e29b8a82e --- /dev/null +++ b/data/jp20_query_gt.csv @@ -0,0 +1,116 @@ +query,pdf,page,pdf_page +"In FY19, what were the earnings per share of the Saudi Research and Marketing Group?",1243103,12,1243103_13 +What was the gross profit of Saudi Research and Marketing Group in FY20?,1243103,12,1243103_13 +Was the total reserves spent on the City of Cockburn in the 2019 to 2020 year within budget?,1321440,617,1321440_618 +"What is the closing balance of the City of Cockburn as of April 30, 2020?",1321440,617,1321440_618 +"For 2021, how much budget was requested for general fund revenue in total in the city of Morgantown?",1602122,18,1602122_19 +"In the city of Morgantown, what was the total amended budget for fiscal year 2020 for all personnel services in city hall?",1602122,28,1602122_29 +"By 2017, which fastest-growing occupation had the highest-paid individuals?",1015168,60,1015168_61 +Which occupation will probably lead with maximum employment changes? ,1015168,60,1015168_61 +What is the C/D highway range for 2019 Kia Niro EV model?,1067991,15,1067991_16 +How much C/D highway range does 2018 Tesla Model 3 Long Range have?,1067991,15,1067991_16 +During what period did most people begin living in their current residence in San Marino?,1167316,225,1167316_226 +What is the most commonly occurring mortgage payment in the SCAG region?,1167316,236,1167316_237 +what percentage of mortgage payments in san marino are between $500-$1000?,1167316,236,1167316_237 +How much was the CIP expenditure in the March FY 2020-2021?,1167316,370,1167316_371 +Which month has the maximum CIP expenditure in FY 2019-2020?,1167316,370,1167316_371 +A scientific agreement is more promising with or without a pre-message estimate?,1179117,11,1179117_12 +What is the approximate actual level of scientific agreement?,1179117,11,1179117_12 +What is the expected NPL ratio at the end of 2020?,1262031,20,1262031_21 +What was the NLP ratio in August of 2017?,1262031,20,1262031_21 +Name the US regions which maintained material expenditures per capita at par with the State Average in FY 2019?,1303168,40,1303168_41 +Which county had the maximum material expenditures per capita in 2019?,1303168,40,1303168_41 +"In which year, Business profits tax of New Hampshire reach to its peak?",1312679,12,1312679_13 +"In which year, there is a biggest downfall in the Business Enterprise tax of New hampshire?",1312679,12,1312679_13 +"In 2021, what is the decline rate in employment of the white population in United States?",1312679,36,1312679_37 +"Within the Hispanic community, which gender experienced the most change in employment from February 2020 to February 2021?",1312679,36,1312679_37 +"In New Hampshire, for the year 2020, what kind of tax brought the most revenue?",1312679,59,1312679_60 +What was the tax revenue for business profits in New Hampshire for the year 2020? ,1312679,59,1312679_60 +What proves to be an effective marketing platform for Cockburn city's football games?,1321440,176,1321440_177 +Least popular source of advertising for fremantle football club 2019?,1321440,176,1321440_177 +Who were the organizers of the City of Cockburn Christmas Event in 2019?,1321440,179,1321440_180 +Did majority of people think that the Dockers organized the Cockburn Christmas Collective? ,1321440,179,1321440_180 +Which province in Thailand had the least impact on walking and driving during the COVID-19 pandemic in mid-June of 2021?,1333360,11,1333360_12 +"In mid-June of 2021, how many new deaths were accounted for due to covid-19 in Bangkok? ",1333360,2,1333360_3 +"During mid-June of 2021, how many new cases of covid did the province of Kalasin reported?",1333360,2,1333360_3 +What was the reason that caused an outburst of COVID-19 cases at a rapid level within Thailand in June 2018?,1333360,4,1333360_5 +What province of Thailand experienced the most Covid cases in the workplace during mid-June of 2021?,1333360,4,1333360_5 +What percentage of individuals that heavily use alcohol also use Marijuana?,1334401,35,1334401_36 +What percentage of individuals that heavily use opioids also heavily use alcohol?,1334401,35,1334401_36 +"In 2010, what was the greenhouse gas emissions from land use value in Kazakhstan? ",1381956,5,1381956_6 +What was the conditional number in 2020 of greenhouse gas emissions excluded from land use in Chile? ,1381956,6,1381956_7 +What is expected of greenhouse gas emission rate from land use with current policies in Chilie?,1381956,6,1381956_7 +What's the cheapest place on earth to mine nickel?,1416620,9,1416620_10 +Is Koniambo the most expensive place to mine nickel?,1416620,9,1416620_10 +What is the highest percentage of Metallurgical Corporation of China when it comes to pre-loan repayment? ,1416620,3,1416620_4 +Does JV interest increase from the pre-loan repayment phase to the post-loan repayment phase?,1416620,3,1416620_4 +Was DPT vaccinations in India covered more during pre-covid or post?,1507643,11,1507643_12 +what happened to the maternal and child healthcare once covid took place in india?,1507643,11,1507643_12 +"Between Ethiopia and India, which one has a higher child mortality percentage? ",1507643,2,1507643_3 +What government is thought by the most people to have much more to do to address climate change? ,1515108,15,1515108_16 +What country has the highest percentage of citizens participating in a citizen's campaign?,1515108,19,1515108_20 +Where would people be the least likely to participate in actions for climate change? ,1515108,19,1515108_20 +Where do people think climate change is happening the most?,1515108,6,1515108_7 +Where do people think climate change is happening the least? ,1515108,6,1515108_7 +What year had the best satisfaction for mobile AR in games? ,1547364,10,1547364_11 +Which year of mobile AR was at its lowest rating of neither satisfied or dissatisfied?,1547364,10,1547364_11 +What kind of mobile AR experiences are people most excited about? ,1547364,16,1547364_17 +Top 3 most popular mobile AR experiences 2019?,1547364,16,1547364_17 +What is the main reason for someone to not be interested in mobile AR?,1547364,18,1547364_19 +"In the year 2019, how many people were unsure if their phone was compatible for mobile AR?",1547364,18,1547364_19 +In what population is hypertension more prevalent? ,1547441,11,1547441_12 +What portion of men that are 18 and over experience hypertension?,1547441,11,1547441_12 +What type of amphetamines has had the biggest increase?,1697708,27,1697708_28 +What year did amphetamines nfd decrease the most in percent usage? ,1697708,27,1697708_28 +Which age group shows the most use of cannabis as a drug of concern? ,1697708,31,1697708_32 +How many percent of population in South Australia has English as a preferred language?,1697708,82,1697708_83 +Which age groups had the highest amount of drug treatment services needed in South Australia during 2019-2020? ,1697708,82,1697708_83 +What substance has the highest proportion of closed treatment episodes? ,1697708,96,1697708_97 +Around how much revenue did Apple Airpods generate in 2023?,1821485,2,1821485_3 +The US Hearable Hardware market has been dominated by which company since 2018?,1821485,2,1821485_3 +What was the net asset value of Global Media company and its Subsidiaries in UK for FY18?,1243103,17,1243103_18 +"What percentage of shares did Arab Media company purchase in the final quarter of 2017 from the limited liability company Argaam Investment Trading company (""Argaam"")?",1243103,17,1243103_18 +"By comparing the carrying value of property, plant and equipment in Saudi Printing and Packaging Company, what amount did the management acquire after the impairment assessment in 2020?",1243103,6,1243103_7 +3. What is the total revenue of SAUDI RESEARCH AND MARKETING GROUP for the year December 2020?,1243103,5,1243103_6 +Which two countries had the fastest growing rate of urbanisation of major population bases from 1990 to 2015?,1015168,10,1015168_11 +"What was the purpose of the new ""Skills Agenda for Europe"" that the European Union (EU) approved in 2016?",1015168,5,1015168_6 +Name the basic four aspects of globalization that was identified by the International Monetary Fund in 2000.,1015168,9,1015168_10 +What was the rate of decline in heroin initiation in the United States in 2018?,1334401,18,1334401_19 +What was the decrease in opioid use disorder from 2018 to 2019?,1334401,18,1334401_19 +How many individuals are surveyed annually by National Survey on Drug Use and Health?,1334401,1,1334401_2 +"In FY 2019, what was the Georgia's per capita state funding?",1303168,5,1303168_6 +"By the end of fiscal year 2019, what percentage of public library systems in Georgia were above the state average of 1.86 items per capita?",1303168,35,1303168_36 +What place does Georgia hold in terms of state funding per person for public libraries in FY2018?,1303168,5,1303168_6 +What is the actual average risk of breast cancer in women?,1179117,8,1179117_9 +"Who approved the research ""Simple Messages Help Set the Record Straight about Scientific Agreement on Human-Caused Climate Change: The Results of Two Experiments"" ? ",1179117,3,1179117_4 +"n 2013, what percentage of American adults believed that ""most scientists think global warming is happening""?",1179117,1,1179117_2 +When was the Waste Local Law 2020 adopted by the City of Cockburn council?,1321440,6,1321440_7 +What causes higher volatility in the financial markets of Italy?,1262031,4,1262031_5 +How much was the stock of bad loans in Italy at November 2015?,1262031,2,1262031_3 +What was the percentage increase in Italian GDP by the conclusion of 2014?,1262031,0,1262031_1 +Which country has the largest emitter?,1381956,4,1381956_5 +Which African country ranks for having the fourth largest non-G20 emitter of fossil CO2?,1381956,2,1381956_3 +"As of May 2019, what was the capacity of installed renewable energy in Iran?",1381956,4,1381956_5 +How many electric vehicles were registered in the TJPDC region in 2020?,1067991,7,1067991_8 +What type of vehicle model is more expensive if purchased new?,1067991,19,1067991_20 +When is the effective date for the implementation of the $88 highway user fee on fuel-efficient 4 and electric vehicles in Virginia?,1067991,20,1067991_21 +Which large-scale commercial plants were established as a result of Martin Vydra's efforts in accessing various markets?,1416620,7,1416620_8 +What areas of finance does Justin Cochrane specialize in?,1416620,7,1416620_8 +How does Nickel 28 Capital Corp. focus its business strategy in the battery metals industry?,1416620,0,1416620_1 +What are some considerable strategies for making maternal and child health services more accessible?,1507643,4,1507643_5 +What was the primary focus of the research on Maternal & Child Health Services in Ethiopia and India in 2020?,1507643,2,1507643_3 +"As of February 2021, what are the projected increases in child and maternal mortality in Ethiopia due to the covid pandemic?",1507643,2,1507643_3 +"Which countries are the most and least likely to believe that their governments should do ""much more"" or ""more"" to address climate change?",1515108,4,1515108_5 +"Based on public opinion, what country are respondents most likely to say they need more information about climate change?",1515108,3,1515108_4 +What type of augmented reality experiences are most commonly utilized?,1547364,3,1547364_4 +What types of mobile augmented reality experiences have the highest potential for monetization through in-app purchases?,1547364,11,1547364_12 +how does user engagement in mobile augmented reality compare to the typical usage patterns observed in other mobile applications?,1547364,11,1547364_12 +What type of scan should be performed to assess the extent of myocardial damage?,1547441,0,1547441_1 +What are the classic signs and symptoms of chest pain that could be associated with myocardial infarction?,1547441,0,1547441_1 +How can healthcare systems assist with hypertension intervention?,1547441,10,1547441_11 +How are Business and Occupation Taxes on construction primarily allocated?,1602122,10,1602122_11 +What are some of the significant markers of the recent economic performance in the workforce of Central West Virginia?,1602122,38,1602122_39 +What additional funding is being allocated to Board Of Park And Recreation Commissioners and the Morgantown Library?,1602122,7,1602122_8 +What are the projected AirPods unit sales for the year 2023?,1821485,1,1821485_2 +"According to Artillery Intelligence, what is the estimated size of the ""hearables"" market currently?",1821485,1,1821485_2 +What is Artillery Data Briefs?,1821485,5,1821485_6 diff --git a/docs/AGENTS.md b/docs/AGENTS.md new file mode 100644 index 0000000000..1106cec676 --- /dev/null +++ b/docs/AGENTS.md @@ -0,0 +1,107 @@ +# Documentation Agent Guide + +You are a documentation engineer and writer for NeMo Retriever Library (NRL) user-facing docs. +Treat `docs/docs/extraction/` as the primary source of truth for published extraction content. +MkDocs config and redirects live in `docs/mkdocs.yml`. + +## Role + +- Write clear, accurate, task-oriented documentation for developers who install, configure, and run NeMo Retriever Library. +- Preserve the reader's workflow: explain what to do, when to do it, and how to verify it. +- Prefer small, focused edits that match the structure of the current page. +- Verify commands, defaults, API names, and behavior against checked-in source, tests, Helm values, or CLI help. +- Use existing documentation, issues, and PRs to locate claims and rationale, not as behavior authority. +- Keep product naming consistent: **NeMo Retriever Library** (NRL). Avoid reintroducing NV-Ingest as the current product name except in historical release-note context. + +## Writing Style Guide + +Apply these rules to documentation, examples, headings, UI text, and release notes that you create or edit. + +- Write in a professional, active, conversational voice. +- Use active voice whenever possible. Use present tense for product behavior. + Address the reader in second person as "you." +- Keep sentences concise. Prefer sentences with fewer than 30 words. +- End every sentence with a period. +- Use plain English and precise technical terms. Avoid jargon, filler, + colloquialisms, and flowery marketing claims. +- Avoid contractions in technical documentation. Write "do not," "cannot," + and "it is." +- Write "NVIDIA" in all caps and use "an NVIDIA," not "a NVIDIA." +- Spell out uncommon abbreviations on first use. Spell out LLM, RAG, SLM, VLM, + and MoE on first use when the audience may not know them. +- Use NVIDIA spellings such as data center, dataset, open source, pretrained, + startup, webpage, website, and Wi-Fi. +- Replace Latinisms with plain English. Use "for example," "that is," "and so + on," "through," and "compared to." +- Use "refer to" instead of "see," "can" instead of "may" for possibility, + and "after" instead of "once" for time. +- Do not use "please" in technical instructions. +- Use numerals for specific values, parameters, measurements, and values of 10 + or more. Spell out zero through nine in general prose. +- Include a space between a number and its unit. Use a comma in numbers with + four or more digits. +- Prefer statement-style headings. Use question headings only on FAQ pages. +- Use the Oxford comma. Put periods inside quotation marks in U.S. style. +- Use hyphens only for compound modifiers before nouns. Do not hyphenate an + adverb that ends in "ly." +- Format commands, code, filenames, paths, flags, environment variables, API identifiers, and literal values as code. +- Use bold for UI elements and the greater-than sign for UI navigation. +- Avoid rhetorical questions, emoji, em dashes, and unnecessary bold text. +- Introduce lists, tables, code examples, and images with a complete sentence. + Use parallel construction in lists. +- Use descriptive link text. Do not use raw URLs in running text or generic + link text such as "click here" or "read more." +- Write dates as Month DD, YYYY. Omit the year when it matches the publication + year. +- Provide useful alt text and preserve a logical heading hierarchy. +- Verify commands, flags, API names, defaults, and technical claims against + source code or another checked-in source of truth. +- Do not rewrite literal code, identifiers, commands, URLs, or quoted terminal + and API output to satisfy prose rules. +- Apply rules to improve clarity. Do not make mechanical changes that reduce + technical accuracy or readability. + +### NRL documentation patterns + +- Prefer end-to-end examples that include `.ingest()` when showing `create_ingestor` / `GraphIngestor` usage, unless the page intentionally stops before ingest for inspection. +- Keep Helm, CLI, and Python guidance aligned with current `main` defaults. Soften or omit claims until code matches. +- For NIM catalogs and build links, prefer the support matrix and topic pages over inventing new tables. +- Preserve MkDocs redirects in `docs/mkdocs.yml` when renaming or retiring pages. +- On documentation PRs, change what readers are told, not what the library does by default, unless the user explicitly requests eng work. + +## Use Additional NVIDIA Documentation Tools + +Follow [NVIDIA DORI Routing](../AGENTS.md#nvidia-dori-routing). +Use the following DORI workflow only when current host capabilities include the +verified NVIDIA documentation Skill Library. Complete the documentation before +the developer opens the pull or merge request. + +1. Route the documentation task through DORI. Include the changed source files, + the user-visible impact, the documentation that might need updates, and the + required validation. +2. Follow the skill or workflow that DORI returns. Verify product behavior + against checked-in sources before drafting. +3. When the host supports subagents, start a documentation subagent while the + primary developer finishes the implementation. Reconcile the documentation + changes and validation evidence before opening the pull or merge request. +4. When the host does not support subagents, complete the same documentation + work in the primary task. + +If the verified Skill Library is unavailable, inaccessible, or fails, skip DORI. +Do not attempt routing, prompt for setup, or ask for or persist a user +classification. Continue using the Writing Style Guide above. + +## Before Editing + +- Read the full target page before editing it. +- Map code changes to existing pages before proposing a new page. +- Check `docs/mkdocs.yml` navigation and redirects before adding, renaming, or removing pages. +- Prefer updating an existing topic page over creating a parallel page with overlapping guidance. +- For docs-only work, keep the PR scoped to documentation paths. Do not mix runtime code changes. + +## Verification + +- From the `docs/` directory, run `python -m mkdocs build --strict --config-file mkdocs.yml` when the environment supports it. +- Run `git diff --check` on changed Markdown files when available. +- For docs-only PRs, run `git diff --name-only upstream/main...HEAD` (or `origin/main...HEAD`) and confirm the diff stays in allowed documentation paths. +- Leave validation items unchecked unless you actually ran the applicable command. diff --git a/docs/DORI_SETUP.md b/docs/DORI_SETUP.md new file mode 100644 index 0000000000..5756d1814b --- /dev/null +++ b/docs/DORI_SETUP.md @@ -0,0 +1,106 @@ +# NVIDIA DORI Setup + +Use this guide only when the user explicitly asks to install or configure NVIDIA DORI. +Before inspecting or installing private components, ask the user to confirm that they can access `gitlab-master.nvidia.com`. +If the user does not confirm access, stop this setup and use the checked-in [Writing Style Guide](AGENTS.md#writing-style-guide). +Access confirmation does not approve installation or host configuration. + +Use these internal sources for the current installation and registration instructions: + +- [NVIDIA Skill Library](https://gitlab-master.nvidia.com/tech-docs/skill-library) contains documentation-focused Agent Skills and guidance for installing them with DORI and other supported hosts. +- [NVIDIA Template Library](https://gitlab-master.nvidia.com/tech-docs/template-library) contains reusable documentation templates and guidance for installing its template skills with DORI. + +## Inspect the Environment + +1. Check for DORI MCP tools. + - If the current agent exposes `dori_handle` or `dori_route`, do not reconfigure the host. + - When `dori_collections` is available, verify that a collection source contains `tech-docs/skill-library`. + - If the Skill Library is missing, identify it as the only missing component and continue to [Confirm Changes](#confirm-changes). +2. When DORI MCP tools are unavailable, inspect the command-line interface (CLI). + - Run `command -v dori` (or the host equivalent). + - If the CLI exists, run `dori collections list --json`. + - Treat a collection whose source contains `tech-docs/skill-library` as the installed Skill Library. +3. Identify the host from explicit runtime context. + Do not infer the host from the model name or repository files. +4. Run `dori setup auto --dry-run` as a cross-check when the CLI exists. + - If auto-detection conflicts with the explicit host, use the explicit host. + - If no explicit host exists and auto-detection is uncertain, ask which host is running. + +Use the following host commands: + +| Explicit Host | Setup Command | +|---|---| +| Codex CLI or Desktop | `dori setup codex` | +| Cursor | `dori setup cursor --scope user` | +| Claude Code | `dori setup claude-code --scope user` | +| Claude Desktop | `dori setup claude` | +| VS Code with GitHub Copilot | `dori setup vscode --scope user` | +| Kiro | `dori setup kiro` | +| Google Antigravity | `dori setup antigravity` | + +Keep the listed `--scope user` option. +Project or combined scope requires separate repository-owner authorization because it can create a repository MCP configuration file. + +## Confirm Changes + +Report each missing component. +Before an installation or host configuration change, ask: + +> DORI setup is incomplete: ``. +> Do you want me to install or configure these components in your user environment? + +Continue only after explicit approval. +The user's private-source access confirmation does not approve these changes. +If the user declines, use the [Writing Style Guide](AGENTS.md#writing-style-guide). + +## Install Missing Components + +When DORI MCP is unavailable and `dori` is missing, require an existing `uv` command. + +- If `uv` is missing, stop and direct the user to the [internal DORI installation guide](https://gitlab-master.nvidia.com/tech-docs/dori/-/blob/main/docs/get-started/install.md). + Do not download or execute an installer script. +- If `uv` exists, install the pinned DORI tool version from the internal guide for your host. + Prefer the version and index URL documented in the Skill Library or DORI install guide at the time of setup. + +When DORI MCP is unavailable and the Skill Library is missing, run: + +```bash +dori install gitlab:tech-docs/skill-library --all --yes +``` + +When DORI MCP is available but the Skill Library is missing: + +1. Run `dori_collections(action="install", source="gitlab:tech-docs/skill-library")`. +2. Run `dori_refresh`. +3. Verify the source with `dori_collections(action="list")`. + +Do not depend on a shell-visible CLI or reconfigure the host on the DORI MCP path. + +## Configure and Validate the Host + +Complete this section only when DORI MCP is unavailable. +After the CLI becomes available, run `dori setup auto --dry-run` if it did not run during inspection. +If auto-detection conflicts with the explicit host, use the explicit host. +If no explicit host exists and auto-detection is uncertain, ask which host is running. +After approval, run the setup command for the resolved host. +Then perform the following checks: + +1. Run the selected command with `--validate`. +2. Run `dori doctor health --json`. +3. Require a passing host validation and `"ok": true` health. + +Follow the activation action that DORI reports. +The action can require an application restart, a new session, a window reload, or enabling the MCP server. + +Until the current agent exposes DORI tools, continue the original task with the [Writing Style Guide](AGENTS.md#writing-style-guide). + +## Protect Credentials and Repository State + +- Never search for, request, print, copy, export, or embed a token, password, cookie, SSH key, or credential-bearing URL. +- Let `uv`, Git, and DORI use credentials that the user already configured. + If access is denied or authentication is missing, stop and refer to the internal DORI installation guide. +- Do not create repository-scoped identity or authorization files. + Confirm private-source access only for an explicit setup request. +- Do not bypass approval controls for writes outside the repository. +- Do not create or commit project-scoped DORI state or MCP configuration without separate repository-owner authorization. +- Do not retry a failed installation in the same task. diff --git a/docs/docs/extraction/api-keys.md b/docs/docs/extraction/api-keys.md index 1cf6073cd2..e448a8148d 100644 --- a/docs/docs/extraction/api-keys.md +++ b/docs/docs/extraction/api-keys.md @@ -29,7 +29,7 @@ When you call hosted object-detection NIMs (Page Elements, Table Structure, Grap The `NVIDIA_API_KEY` from build.nvidia.com is not the same string as your NGC personal key used for Helm and `nvcr.io` access. Do not substitute one for the other unless your tooling explicitly documents that mapping. -## Credential references in persisted graphs +## Credential references in persisted graphs { #credential-references-in-persisted-graphs } Persisted pipeline graphs never contain literal API keys. Configure a graph with an explicit worker-side environment reference such as: @@ -41,7 +41,9 @@ Use the provider's own variable name, for example `os.environ/OPENAI_API_KEY` fo Literal keys remain available for non-persisted local execution, but attempting to serialize one raises an error. This prevents graph persistence from silently substituting an NVIDIA credential for another provider's key. -## NGC personal key (Helm and `nvcr.io`) +For how persisted graphs store credential references, refer to [Persisted graphs are trusted configuration](nemo-retriever-api-reference.md#persisted-graphs-are-trusted-configuration) in the Python API guide. + +## NGC personal key (Helm and `nvcr.io`) { #ngc-personal-key } Many public assets on NGC can be used without authentication. For a Kubernetes deployment, the cluster must still pull NIM and microservice images from `nvcr.io` and may need NGC API access; the Helm chart expects credentials derived from an NGC personal key. @@ -59,6 +61,6 @@ When you create an NGC key, select the following for **Services Included**. ![Generate Personal Key](images/generate_personal_key.png) -## Using your NGC key with Helm +## Using your NGC key with Helm { #using-your-ngc-key-with-helm } Configure your key through the chart values described in the [NeMo Retriever Helm chart README](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md): for example `imagePullSecret.create` / `imagePullSecret.password` for pulls from `nvcr.io`, `nimApiKey` (inline value or `existingSecret`) for the retriever service, and `nims.ngcApiKey` when `nims.enabled=true`. Exact paths are versioned—use the **Secrets** section in that README and [`values.yaml`](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/values.yaml) as the source of truth. diff --git a/docs/docs/extraction/audio-video.md b/docs/docs/extraction/audio-video.md index 3f2df2cf9c..0d9477bef2 100644 --- a/docs/docs/extraction/audio-video.md +++ b/docs/docs/extraction/audio-video.md @@ -16,7 +16,7 @@ This documentation describes two ways to run [NeMo Retriever Library](overview.m Supported file types for speech extraction today: - `mp3`, `wav` -- `mp4`, `mov`, `mkv`, `avi` — common video containers; the audio track is transcribed (same extensions as in [What is NeMo Retriever Library?](overview.md)) +- `mp4`, `mov`, `mkv`, `avi` — common video containers; the audio track is transcribed (same extensions as in [NeMo Retriever Library Overview](overview.md)) [NeMo Retriever Library](overview.md) supports extracting speech from audio for Retrieval Augmented Generation (RAG). Similar to how the multimodal document pipeline uses detection and OCR microservices, NeMo Retriever Library uses the [parakeet-1-1b-ctc-en-us ASR NIM](https://docs.nvidia.com/nim/speech/latest/asr/deploy-asr-models/parakeet-ctc-en-us.html) to transcribe speech to text, then embeddings through the NeMo Retriever embedding path. @@ -131,7 +131,7 @@ For video assets, NeMo Retriever Library can combine audio or speech processing For OCR-oriented extract methods on scanned or image-heavy content, refer to [OCR and scanned documents](multimodal-extraction.md#ocr-and-scanned-documents), [text and layout extraction](multimodal-extraction.md#text-and-layout-extraction), and [Nemotron Parse](https://build.nvidia.com/nvidia/nemotron-parse) for advanced visual parsing. -Container formats and early-access video types are listed under [supported file types and formats](multimodal-extraction.md#supported-file-types-and-formats) (refer to [What is NeMo Retriever Library?](overview.md) for the full list). +Container formats and early-access video types are listed under [supported file types and formats](multimodal-extraction.md#supported-file-types-and-formats) (refer to [NeMo Retriever Library Overview](overview.md) for the full list). For end-to-end RAG stacks that include multimodal ingestion, refer to the [NVIDIA AI Blueprints catalog](https://build.nvidia.com/explore/discover) and related solution pages on [NVIDIA Build](https://build.nvidia.com/). diff --git a/docs/docs/extraction/concepts.md b/docs/docs/extraction/concepts.md index 407449e277..b57f28769f 100644 --- a/docs/docs/extraction/concepts.md +++ b/docs/docs/extraction/concepts.md @@ -23,14 +23,14 @@ Optionally, the library can compute **embeddings** for extracted content and sto Chunking is built into the `.extract()` task and depends on **content type**: - **PDF, DOCX, and PPTX** — Text is grouped using built-in **page** boundaries (one chunk per page where the format has pages). -- **Plain text (`.txt`) and HTML** — Formats without natural page breaks are split into segments of **1024 tokens** by default, using the [Llama 3.2 1B tokenizer](https://huggingface.co/meta-llama/Llama-3.2-1B) so chunk boundaries stay aligned with the default embedding tokenizer. The NeMo Retriever container image bundles this tokenizer, so default text chunking does not require a Hugging Face access token. Refer to [Token-based splitting](#token-based-splitting) and [Environment variables](environment-config.md) for overrides and other runtimes. +- **Plain text (`.txt`) and HTML** — Formats without natural page breaks are split into segments of **1024 tokens** by default, using the revision-pinned [Llama Nemotron Embed VL 1B v2 tokenizer](https://huggingface.co/nvidia/llama-nemotron-embed-vl-1b-v2) so chunk boundaries stay aligned with the default embedding model. Published service images bundle this tokenizer artifact without model weights, so default text chunking does not require Hugging Face access at runtime. Refer to [Token-based splitting](#token-based-splitting) and [Environment variables](environment-config.md) for overrides and other runtimes. - **Audio and video** — Media is split into **segments** for decoding and ASR using ffmpeg-based rules (configurable **size**, **time**, or **frame** split modes in the media chunking stage). With the Parakeet ASR path, you can optionally emit **sentence-like segments** using `extract_audio_params={"segment_audio": True}`; refer to [Speech and audio extraction](audio-video.md#speech-and-audio-extraction). For PDF parallelism before Ray processing (large files), refer to [PDF pre-splitting for parallel ingest](nemo-retriever-api-reference.md#pdf-pre-splitting-for-parallel-ingest). ### Token-based splitting { #token-based-splitting } -Token-based splitting uses the Llama 3.2 1B tokenizer (default `meta-llama/Llama-3.2-1B`) with configurable `max_tokens` and `overlap_tokens` when you add an explicit `.split(...)` stage or when the pipeline applies the default text segmentation for unstructured text. In the shipped NeMo Retriever container, tokenizer assets are included locally, so you do not need `HF_ACCESS_TOKEN` for this default path. If your runtime loads the tokenizer from the Hugging Face Hub instead (for example, some library-only installs), set `HF_ACCESS_TOKEN` or pass `hf_access_token` in task params when the Hub requires it. For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md). +Token-based splitting uses the revision-pinned tokenizer for the default embedding model (`nvidia/llama-nemotron-embed-vl-1b-v2`) with configurable `max_tokens` and `overlap_tokens` when you add an explicit `.split(...)` stage or when the pipeline applies the default text segmentation for unstructured text. Published service images and the documented source builds include the tokenizer locally; source builds enable this with `--build-arg DOWNLOAD_DEFAULT_TOKENIZER=True`. The `service` image disables runtime Hub access, while `service-gpu` remains online for its other Hugging Face models. The base library install includes the tokenizer Python dependencies; pre-populate the Hugging Face cache before offline use. For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md). ## Deployment modes { #deployment-modes } diff --git a/docs/docs/extraction/customize-extend.md b/docs/docs/extraction/customize-extend.md index 31fb87e7bc..52ca2efb95 100644 --- a/docs/docs/extraction/customize-extend.md +++ b/docs/docs/extraction/customize-extend.md @@ -23,7 +23,7 @@ The following table maps common needs to the right section: ## Start with task configuration { #start-with-task-configuration } -Most customization does not require new code. Chain tasks on `create_ingestor(...)` and pass keyword arguments to control extraction, chunking, embedding, and storage—for example `extract_method`, chunking and splitting options on `.extract()`, `embed_modality` on `.embed()`, and `vdb_op` / `vdb_kwargs` on `.vdb_upload()`. +Most customization does not require new code. Chain tasks on `create_ingestor(...)` and pass keyword arguments to control extraction, chunking, embedding, and storage—for example `method`, chunking and splitting options on `.extract()`, `embed_modality` on `.embed()`, and `vdb_op` / `vdb_kwargs` on `.vdb_upload()`. For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md). For chunking behavior and pipeline concepts, refer to [Concepts](concepts.md). diff --git a/docs/docs/extraction/deployment-options.md b/docs/docs/extraction/deployment-options.md index 7684326dbe..dbae4311b2 100644 --- a/docs/docs/extraction/deployment-options.md +++ b/docs/docs/extraction/deployment-options.md @@ -22,7 +22,7 @@ Build and run the NeMo Retriever service image with the [Docker service image gu 3. **Published Library Helm charts (supported):** cluster install and upgrade procedures are covered in [About getting started](getting-started-about.md) — use alongside the NeMo Retriever chart README for your release 4. [Environment variables](environment-config.md) and [Troubleshoot](troubleshoot.md) as needed -**Core NIMs for the default extraction pipeline:** `page_elements`, `table_structure`, `ocr`, and `vlm_embed` (`llama-nemotron-embed-vl-1b-v2:1.12.0`). These four are auto-wired into the retriever service. **Nemotron Parse**, **Nemotron 3 Nano Omni**, the **VL reranker**, and **Parakeet ASR** are optional and not auto-wired. For a minimal GPU footprint, disable optional keys you do not need (refer to [Recommended minimal install](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md#recommended-minimal-install-2605)). Refer to [Pre-Requisites & Support Matrix — Default NIMs](prerequisites-support-matrix.md#default-helm-nims) and [Default NVCF endpoints](prerequisites-support-matrix.md#default-nvcf-endpoints). +**Core NIMs for the default extraction pipeline:** `page_elements`, `table_structure`, `ocr`, and `vlm_embed` (`llama-nemotron-embed-vl-1b-v2:2.3.0`). These four are auto-wired into the retriever service. **Nemotron Parse**, **Nemotron 3 Nano Omni**, the **VL reranker**, and **Parakeet ASR** are optional and not auto-wired. For a minimal GPU footprint, disable optional keys you do not need (refer to [Recommended minimal install](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md#recommended-minimal-install-2608)). Refer to [Pre-Requisites & Support Matrix — Default NIMs](prerequisites-support-matrix.md#default-helm-nims) and [Default NVCF endpoints](prerequisites-support-matrix.md#default-nvcf-endpoints). For audio and video extraction in Kubernetes, refer to [Audio and video](audio-video.md). diff --git a/docs/docs/extraction/embedding.md b/docs/docs/extraction/embedding.md index b6a139fac3..90a56909a3 100644 --- a/docs/docs/extraction/embedding.md +++ b/docs/docs/extraction/embedding.md @@ -12,14 +12,14 @@ The model can embed documents in the form of an image, text, or a combination of Documents can then be retrieved given a user query in text form. The model supports images that contain text, tables, charts, and infographics. -## Example with Default Text-Based Embedding +## Example with Default Text-Based Embedding { #example-with-default-text-based-embedding } When you use the multimodal model, by default, all extracted content (text, tables, charts) is treated as plain text. The following example provides a strong baseline for retrieval. - The `embed` method is called with no arguments. -For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md). +For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md) (`create_ingestor` and `.embed()`). ```python from nemo_retriever import create_ingestor @@ -34,7 +34,7 @@ results = ingestor.ingest() ``` -## Example with Embedding Structured Elements as Text + Images +## Example with Embedding Structured Elements as Text + Images { #example-with-embedding-structured-elements-as-text-images } It is common to process PDFs by embedding standard text as text and embed visual elements such as tables and charts as images. The following example enables the multimodal model to capture the spatial and structural information of the visual content. @@ -42,7 +42,7 @@ The following example enables the multimodal model to capture the spatial and st - The `embed` method is configured with `embed_modality="text_image"` to embed the extracted tables and charts as images. - This configuration is more accurate than text only, with a performance cost. -For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md). +For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md) (`create_ingestor` and `.embed()`). ```python from nemo_retriever import create_ingestor @@ -59,7 +59,7 @@ results = ingestor.ingest() ``` -## Example with Embedding Entire PDF Pages as Images +## Example with Embedding Entire PDF Pages as Images { #example-with-embedding-entire-pdf-pages-as-images } For documents where the entire page layout is important (such as infographics, complex diagrams, or forms), you can configure NeMo Retriever Library to treat every page as a single image. @@ -67,7 +67,7 @@ The following example extracts and embeds each page as an image. - The `embed` method processes the page images. -For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md). +For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md) (`create_ingestor` and `.embed()`). ```python from nemo_retriever import create_ingestor @@ -84,7 +84,7 @@ ingestor = ( results = ingestor.ingest() ``` -## Related Topics +## Related Topics { #related-topics } - [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md) - [Troubleshoot Nemo Retriever Extraction](troubleshoot.md) diff --git a/docs/docs/extraction/environment-config.md b/docs/docs/extraction/environment-config.md index f5c6e96d5b..43fe00e1ca 100644 --- a/docs/docs/extraction/environment-config.md +++ b/docs/docs/extraction/environment-config.md @@ -4,17 +4,19 @@ The following are the environment variables that you can use to configure [NeMo You can specify these in a .env file in your working directory or directly as shell environment variables. -## General Environment Variables +## General Environment Variables { #general-environment-variables } | Name | Example | Description | |----------------------------------|--------------------------------|-----------------------------------------------------------------------| -| `HF_ACCESS_TOKEN` | - | A token for Hugging Face Hub downloads when your runtime needs it (default Llama 3.2 chunking tokenizer is bundled in the NeMo Retriever container; a token is not required there). Refer to [Token-based splitting](concepts.md#token-based-splitting). | +| `HF_ACCESS_TOKEN` | - | A token for Hugging Face Hub downloads when your runtime needs one. The default chunking tokenizer is public; refer to [Token-based splitting](concepts.md#token-based-splitting) for container caching and offline behavior. | | `INGEST_LOG_LEVEL` | - `DEBUG`
- `INFO`
- `WARNING`
- `ERROR`
- `CRITICAL`
| The log level for the ingest service, which controls the verbosity of the logging output. | | `NVIDIA_API_KEY` | `nvapi-*************`
| An authorized build.nvidia.com API key, used to interact with NVIDIA-hosted NIMs. Create through build.nvidia.com or through [NGC](https://org.ngc.nvidia.com/setup/api-keys). | | `NGC_API_KEY` | — | The key that NIM microservices in the cluster use to access NGC resources. | | `OTEL_EXPORTER_OTLP_ENDPOINT` | `http://otel-collector:4317`
| The endpoint for the OpenTelemetry exporter, used for sending telemetry data. | -## Related Topics +## Related Topics { #related-topics } -- [Configure Ray Logging](https://docs.nvidia.com/nemo/retriever/latest/extraction/ray-logging/) +- [Configure Ray Logging](ray-logging.md) +- [Authentication and API keys](api-keys.md) +- [Python API guide](nemo-retriever-api-reference.md) diff --git a/docs/docs/extraction/evaluate-on-your-data.md b/docs/docs/extraction/evaluate-on-your-data.md index d441cd81e8..e7be1adab3 100644 --- a/docs/docs/extraction/evaluate-on-your-data.md +++ b/docs/docs/extraction/evaluate-on-your-data.md @@ -2,15 +2,15 @@ Retrieval and ingestion performance **depend on your documents**, hardware, and pipeline settings. Use the following when measuring quality and throughput on **your** datasets. -## Benchmarking and baselines +## Benchmarking and baselines { #benchmarking-and-baselines } Use this page as the baseline for methodology and expectations. Use [Operational tuning](#operational-tuning) below to observe production-like runs. -## Throughput and dataset effects +## Throughput and dataset effects { #throughput-and-dataset-effects } Read [Throughput is dataset-dependent](multimodal-extraction.md#extraction-limitations-and-quality) for why raw numbers from generic benchmarks may not match your corpus (layout complexity, file types, image density, and so on). -## Operational tuning +## Operational tuning { #operational-tuning } - [Ray and distributed ingest](ray-logging.md) - [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md) for supported configurations diff --git a/docs/docs/extraction/faq.md b/docs/docs/extraction/faq.md index b39f78a48a..1a025e4de2 100644 --- a/docs/docs/extraction/faq.md +++ b/docs/docs/extraction/faq.md @@ -2,52 +2,52 @@ This documentation contains the Frequently Asked Questions (FAQ) for [NeMo Retriever Library](overview.md). -## Is the NeMo Retriever Library supported under NVIDIA AI Enterprise (NVAIE)? +## Is the NeMo Retriever Library supported under NVIDIA AI Enterprise (NVAIE)? { #nvaie-support } No. The NeMo Retriever Library, including its container image and Helm chart artifacts, is not supported under NVIDIA AI Enterprise (NVAIE). Some NIM microservices and models that the library calls may be individually covered by NVAIE. That coverage does not extend to the NeMo Retriever Library or its end-to-end extraction workflow. For more information, refer to [NVIDIA AI Enterprise (NVAIE) support](overview.md#nvidia-ai-enterprise-nvaie-support). -## What if I already have a retrieval pipeline? Can I just use NeMo Retriever Library? +## What if I already have a retrieval pipeline? Can I just use NeMo Retriever Library? { #use-with-existing-retrieval-pipeline } You can use the CLI or Python APIs to perform extraction only, and then consume the results. Using the Python API, `results` is a list object with one entry. For code examples, refer to the Jupyter notebooks [Multimodal RAG with LlamaIndex](https://github.com/NVIDIA/NeMo-Retriever/blob/main/examples/llama_index_multimodal_rag.ipynb) and [Multimodal RAG with LangChain](https://github.com/NVIDIA/NeMo-Retriever/blob/main/examples/langchain_multimodal_rag.ipynb). -## Where does NeMo Retriever Library ingest to? +## Where does NeMo Retriever Library ingest to? { #where-does-nrl-ingest-to } NeMo Retriever Library supports extracting text representations of various forms of content, and ingesting to a vector database. **[LanceDB](https://lancedb.com/)** stores vectors as local Lance files on disk for the supported ingestion path. You can ingest to other data stores; however, you must configure other data stores yourself. For more information, refer to [Vector databases](vdbs.md). -## How would I process unstructured images? +## How would I process unstructured images? { #process-unstructured-images } For images that `nemoretriever-page-elements-v3` does not classify as tables, charts, or infographics, you can use our VLM caption task to create a dense caption of the detected image. That caption is then embedded along with the rest of your content. For chart-labeled PDF regions and other caption scope limits, refer to [Are PDF chart or figure regions captioned when Omni is enabled?](#are-pdf-chart-or-figure-regions-captioned-when-omni-is-enabled). For more information, refer to [Extract Captions from Images](nemo-retriever-api-reference.md). -## Are PDF chart or figure regions captioned when Omni is enabled? +## Are PDF chart or figure regions captioned when Omni is enabled? { #are-pdf-chart-or-figure-regions-captioned-when-omni-is-enabled } No. Chart-labeled PDF regions are not routed through Omni captioning. Refer to [Charts and infographics](multimodal-extraction.md#charts-and-infographics) and [Image captioning](multimodal-extraction.md#image-captioning) for caption scope and validation. -## When should I consider advanced visual parsing? +## When should I consider advanced visual parsing? { #advanced-visual-parsing } For scanned documents, or documents with complex layouts, -you can use [nemotron-parse](https://build.nvidia.com/nvidia/nemotron-parse) as an alternate PDF extraction method by setting `extract_method="nemotron_parse"`. +you can use [nemotron-parse](https://build.nvidia.com/nvidia/nemotron-parse) as an alternate PDF extraction method by setting `method="nemotron_parse"`. Nemotron Parse does not produce chart modality rows. For chart detection and chart-filtered retrieval, use the default **pdfium** layout path instead (refer to [Charts and infographics](multimodal-extraction.md#charts-and-infographics)). For more information, refer to [Nemotron Parse](https://build.nvidia.com/nvidia/nemotron-parse). -## Why are the environment variables different between library mode and self-hosted mode? +## Why are the environment variables different between library mode and self-hosted mode? { #library-vs-self-hosted-env-vars } -### Self-Hosted Deployments +### Self-Hosted Deployments { #self-hosted-deployments } For [self-hosted deployments](deployment-options.md#when-to-self-host-nims), you should set the environment variables `NGC_API_KEY` and `NIM_NGC_API_KEY`. For more information, refer to [Authentication and API keys](api-keys.md). -### Library Mode +### Library Mode { #library-mode } For production environments, you should use the provided Helm charts. When you run the NeMo Retriever Library from Python without those charts, set `NVIDIA_API_KEY` only when you call [build.nvidia.com](https://build.nvidia.com/) hosted inference—it is not required for locally deployed Hugging Face models or self-hosted NIM endpoints. For more information, refer to [Deployment options](deployment-options.md) and [Authentication and API keys](api-keys.md). @@ -55,9 +55,9 @@ For advanced scenarios, you might want to use library mode with self-hosted NIM You can set custom endpoints for each NIM. For examples of `*_ENDPOINT` variables, refer to [Environment variables](environment-config.md) and the [Helm chart README](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md). -When you explicitly configure remote NIM endpoints in Python library mode, graph ingestion raises a `GraphIngestionError` if a stage reports row-level connection or inference errors. This makes unreachable services visible to callers instead of returning a DataFrame that looks successful. To intentionally keep partial results with row-level error payloads, pass `error_policy="collect"` to `GraphIngestor` or `create_ingestor`. +When you explicitly configure remote NIM endpoints in Python library mode, graph ingestion raises a `GraphIngestionError` if a stage reports row-level connection or inference errors. This makes unreachable services visible to callers instead of returning a DataFrame that looks successful. To intentionally keep partial results with row-level error payloads, pass `error_policy="collect"` to `GraphIngestor` or `create_ingestor`. Refer to the [Python API error contract](nemo-retriever-api-reference.md#error-and-failure-contract) and [Python API error triage](troubleshoot.md#python-api-error-triage) for error signals, extraction-path mappings, and escalation criteria. -## What parameters or settings can I adjust to optimize extraction from my documents or data? +## What parameters or settings can I adjust to optimize extraction from my documents or data? { #optimize-extraction-parameters } Refer to [Evaluate on your data](evaluate-on-your-data.md) for extraction tuning and optimization guidance. diff --git a/docs/docs/extraction/getting-started-about.md b/docs/docs/extraction/getting-started-about.md index 0e6b5ec8db..6c222d3d72 100644 --- a/docs/docs/extraction/getting-started-about.md +++ b/docs/docs/extraction/getting-started-about.md @@ -5,10 +5,10 @@ This section walks you from **access and prerequisites** through **first deploym Typical order: 1. [Get your API key](api-keys.md) (NGC / API access as required by your workflow). -2. Confirm the [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md) for your OS, GPU, and software stack. +2. Confirm the [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md) for your OS, GPU, and software stack. Local GPU inference requires Linux; remote NIM workflows can use the base package on Windows x64 and macOS Apple Silicon (arm64) as well. macOS Intel (x86_64) is not supported. 3. Choose a path in [Deployment options](deployment-options.md) — local library, hosted NIMs, the Helm chart for Kubernetes, or a standalone Docker service. 4. Explore [Jupyter Notebooks](https://github.com/NVIDIA/NeMo-Retriever/blob/main/examples/README.md) for end-to-end examples. The NeMo Retriever Library and its Helm chart are not supported under NVIDIA AI Enterprise (NVAIE). For more information, refer to [NVIDIA AI Enterprise (NVAIE) support](overview.md#nvidia-ai-enterprise-nvaie-support). -If you are new to the product, read [What is NeMo Retriever Library?](overview.md) and [Concepts](concepts.md) under **Introduction** first. +If you are new to the product, read [NeMo Retriever Library Overview](overview.md) and [Concepts](concepts.md) under **Introduction** first. diff --git a/docs/docs/extraction/multimodal-extraction.md b/docs/docs/extraction/multimodal-extraction.md index 345daf21b8..ba935ad682 100644 --- a/docs/docs/extraction/multimodal-extraction.md +++ b/docs/docs/extraction/multimodal-extraction.md @@ -15,7 +15,7 @@ NeMo Retriever Library classifies and extracts text, tables, charts, infographic ## Supported file types and formats { #supported-file-types-and-formats } -NeMo Retriever Library accepts multiple document and media types. A current list (including PDF, Office formats, HTML, images, audio, and video, some early access) appears in [What is NeMo Retriever Library?](overview.md) under **NeMo Retriever Library supports the following file types**. +NeMo Retriever Library accepts multiple document and media types. A current list (including PDF, Office formats, HTML, images, audio, and video, some early access) appears in [NeMo Retriever Library Overview](overview.md) under **NeMo Retriever Library supports the following file types**. **Related** @@ -24,14 +24,14 @@ NeMo Retriever Library accepts multiple document and media types. A current list ## Text and layout extraction { #text-and-layout-extraction } -For PDFs, NeMo Retriever Library typically uses **pdfium**-based extraction with configurable depth and paths. Scanned or mixed pages may use hybrid, OCR-oriented, or Nemotron Parse methods. For `extract_method` options such as `pdfium`, `pdfium_hybrid`, `ocr`, and `nemotron_parse`, refer to the [Python API reference](nemo-retriever-api-reference.md). +For PDFs, NeMo Retriever Library typically uses **pdfium**-based extraction with configurable depth and paths. Scanned or mixed pages may use hybrid, OCR-oriented, or Nemotron Parse methods. For `method` options such as `pdfium`, `pdfium_hybrid`, `ocr`, and `nemotron_parse`, refer to the [Python API reference](nemo-retriever-api-reference.md). !!! note - `extract_method="nemotron_parse"` requires the Nemotron Parse NIM client dependencies. Install them with the `nemotron-parse` extra, for example `pip install "nemo-retriever[nemotron-parse]"`, before running PDF extraction through Nemotron Parse. This path does not produce chart modality rows; for chart detection, refer to [Charts and infographics](#charts-and-infographics). + `method="nemotron_parse"` requires the Nemotron Parse NIM client dependencies. Install them with the `nemotron-parse` extra, for example `pip install "nemo-retriever[nemotron-parse]"`, before running PDF extraction through Nemotron Parse. This path does not produce chart modality rows; for chart detection, refer to [Charts and infographics](#charts-and-infographics). **Related** -- [What is NeMo Retriever Library?](overview.md) +- [NeMo Retriever Library Overview](overview.md) - [OCR and scanned documents](#ocr-and-scanned-documents) - [Chunking](concepts.md#chunking) @@ -41,7 +41,7 @@ NeMo Retriever Library detects tables as structured page elements, processes the **Related** -- [What is NeMo Retriever Library?](overview.md) for artifact classification +- [NeMo Retriever Library Overview](overview.md) for artifact classification - [Nemotron Parse](https://build.nvidia.com/nvidia/nemotron-parse) for advanced visual parsing - [Metadata reference](content-metadata.md) @@ -52,13 +52,13 @@ Charts and infographic regions are classified with other page layout elements (t !!! important "Chart modality requires the default layout path" [Nemotron Parse v1.2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2) semantic classes do not include `Chart` or `Infographic`. The model labels regions as `Text`, `Table`, `Picture`, `Caption`, `List-item`, `Section-header`, and similar types instead. - When you set `extract_method="nemotron_parse"`: + When you set `method="nemotron_parse"`: - The pipeline does not produce `chart` or `infographic` modality rows, even when `extract_charts=True` or `extract_infographics=True`. - Chart- and infographic-filtered retrieval (for example, queries scoped to figure or chart content) returns no hits. - Chart-heavy and infographic-heavy pages are typically emitted as `Picture` or other non-chart modalities. - For chart and infographic detection and modality-specific retrieval, use the default **pdfium** layout path (page-elements detection and OCR), not `extract_method="nemotron_parse"`. + For chart and infographic detection and modality-specific retrieval, use the default **pdfium** layout path (page-elements detection and OCR), not `method="nemotron_parse"`. Chart-labeled PDF regions are **not** routed through the Omni caption stage; they remain on the layout-and-OCR path. For scope and validation guidance, refer to [Image captioning](#image-captioning). @@ -66,7 +66,7 @@ For natural-language infographic descriptions, optionally enable [image captioni **Related** -- [What is NeMo Retriever Library?](overview.md) +- [NeMo Retriever Library Overview](overview.md) - [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md) - [Multimodal embeddings (VLM)](embedding.md) when you treat graphics as images for embedding diff --git a/docs/docs/extraction/nemo-retriever-api-reference.md b/docs/docs/extraction/nemo-retriever-api-reference.md index 72ad608087..01d38682e6 100644 --- a/docs/docs/extraction/nemo-retriever-api-reference.md +++ b/docs/docs/extraction/nemo-retriever-api-reference.md @@ -1,5 +1,143 @@ # NeMo Retriever API Reference +## Error and failure contract { #error-and-failure-contract } + +The Python API does not define a separate set of numeric NeMo Retriever +extraction error codes. Depending on the run mode and failing stage, callers +observe one or more of the following: + +- Python configuration or dependency exceptions, such as `ValueError`, + `ImportError`, or `RuntimeError`. +- `GraphIngestionError` for row-level failures from explicitly configured + remote NIM stages in `run_mode="inprocess"` or `"batch"` when + `error_policy="raise"` (the default). +- HTTP status codes or gRPC errors returned by a remote NIM or by the + Retriever service. These are transport or upstream-service statuses, not + NeMo Retriever-specific error codes. +- Per-document failures in `ServiceIngestResult.failures` when + `run_mode="service"`. + +The generated API signatures and parameter models below are the API contract. +Exception text and upstream response bodies can change between releases; do +not parse them as stable codes. The stable text-generation codes documented in +[One-shot text generation](#one-shot-text-generation) apply to generation +operator output columns, not to document extraction. + +### Choose raise or collect behavior + +For graph run modes, `error_policy="raise"` raises `GraphIngestionError` when +an explicitly configured remote NIM stage reports a row-level error. The +exception retains the underlying records in `exc.records`. When available, its +message identifies the stage, invoke URL, and HTTP status in a form similar to +`[stage=OCR NIM url=https://... http=503]`, followed by a troubleshooting hint. + +Use `error_policy="collect"` when partial results are useful and your +application inspects the error fields in every returned row. Alternatively, +pass `return_failures=True` to `.ingest()` to receive a `(result, failures)` +tuple. When no remote invoke URL is configured, `return_failures=True` scans +all output columns for row-level error fields so local failures are still +visible. In service mode, failures are also available from +`ServiceIngestResult.failures`. + +### What the raise error policy covers + +The strict policy applies only to stages where you explicitly configure a +remote NIM invoke URL. It does not raise for local-only PDFium parsing, +caption, audio or video, or ASR failures, even when those stages populate +row-level error fields. + +| Configured invoke URL | DataFrame column scanned | Stage label in messages | +| --- | --- | --- | +| `page_elements_invoke_url` | `output_column` (default `page_elements_v3`) | Page Elements NIM | +| `ocr_invoke_url` | `ocr` | OCR NIM | +| `table_structure_invoke_url` | `table_structure_ocr_v1` | Table Structure NIM | +| `nemotron_parse_invoke_url` or `invoke_url` | `nemotron_parse_v1_2` | Nemotron Parse NIM | +| `embed_invoke_url` or `embedding_endpoint` | `output_column` (default `text_embeddings_1b_v2`) | Embedding NIM | + +Caption and ASR use remote endpoints but are outside this raise path today. +Remote caption failures can abort the whole ingest instead of returning a +partial DataFrame. ASR failures can omit affected rows while logging a +warning, which can look like an empty transcript unless you inspect logs. + +### Row-level error payloads + +Most extraction stages write errors into the result row instead of raising +immediately. The common nested shape is: + +```json +{ + "error": { + "stage": "ocr_page_elements", + "type": "HTTPError", + "message": "HTTP 503 from https://example/v1/infer: ...", + "traceback": "..." + } +} +``` + +The `stage` string is a semi-stable operator identifier (for example +`remote_inference`, `nemotron_parse_pages`, or `split_pdf`). It is not a +product-wide error-code enum. HTTP status codes usually appear inside +`message` text rather than as a separate `status_code` field; when a +structured status is present, `GraphIngestionError` can include it in the +rendered exception. + +```python +import os + +from nemo_retriever import GraphIngestionError, create_ingestor +from nemo_retriever.common.params import ExtractParams + +pipeline = ( + create_ingestor(run_mode="inprocess", error_policy="raise") + .files(["document.pdf"]) + .extract( + ExtractParams( + method="ocr", + ocr_invoke_url=os.environ["OCR_INVOKE_URL"], + ) + ) +) + +try: + result = pipeline.ingest() +except GraphIngestionError as exc: + # Records can contain source paths, endpoint details, and upstream + # response text. Extract only known-safe diagnostic fields before + # logging or sending them to your support workflow. + for record in exc.records: + payload = record.get("error") if isinstance(record, dict) else record + if isinstance(payload, dict): + print( + { + "column": record.get("column"), + "stage": payload.get("stage"), + "type": payload.get("type"), + "message": payload.get("message"), + } + ) + else: + print( + { + "column": record.get("column") if isinstance(record, dict) else None, + "message": str(payload), + } + ) +``` + +For a support-oriented mapping of extraction paths, error signals, corrective +actions, and escalation criteria, refer to +[Python API error triage](troubleshoot.md#python-api-error-triage). + +!!! note "Version-specific behavior" + + This reference describes the current NeMo Retriever Library. Older + NV-Ingest releases, including `25.4.2`, can use different exception text + and result shapes and might not include enriched `GraphIngestionError` + diagnostics. When troubleshooting an older deployment, use the package and + container versions from that deployment and include them in the support + case. + ## PDF pre-splitting for parallel ingest { #pdf-pre-splitting-for-parallel-ingest } Large PDFs are split into page batches before Ray processing so extraction can run in parallel. This happens on the default ingest path; you do not need extra configuration for typical workloads. @@ -8,7 +146,7 @@ To tune splitter throughput from the CLI, use `--pdf-split-batch-size` (Ray acto **Python client (`pdf_split_config`):** Only `create_ingestor(run_mode="service")` implements `.pdf_split_config(pages_per_chunk=...)`, which records page-chunking settings in the request pipeline spec for the remote gateway. Local graph ingest (`run_mode="inprocess"` or `"batch"`) raises `NotImplementedError` if you call this method; PDFs are split automatically on the default ingest path without client-side configuration. -## One-shot text generation +## One-shot text generation { #one-shot-text-generation } `TextGenerationOperator` is the reusable base for synchronous, one-request-per-row text generation. It is a provisional text-only API: it does not support tool calls, agent loops, streaming, multiple choices, or structured domain results. @@ -54,7 +192,7 @@ To define another one-request/one-text-result task, subclass `TextGenerationTask Generation failures are collected per row using stable error codes: `empty_input`, `request_error`, `transport_error`, `unsupported_response`, `parse_error`, `empty_output`, and the RAG-specific `thinking_truncated`. Raw provider exceptions and credentials are not written to DataFrame outputs. -## Persisted graphs are trusted configuration +## Persisted graphs are trusted configuration { #persisted-graphs-are-trusted-configuration } Graph loading imports operator classes and invokes their constructors. Load graph JSON only from trusted sources; do not expose graph payloads, callable references, or class names as model- or user-controlled agent tools. diff --git a/docs/docs/extraction/overview.md b/docs/docs/extraction/overview.md index d898b3d550..ab5995a099 100644 --- a/docs/docs/extraction/overview.md +++ b/docs/docs/extraction/overview.md @@ -1,9 +1,8 @@ -# What is NeMo Retriever Library? +# NeMo Retriever Library Overview { #what-is-nemo-retriever-library } -NVIDIA NeMo Retriever Library (NRL) is a high retrieval accuracy, performant, and scalable framework for content and metadata extraction from various media types (PDFs, HTML, Word docs, Powerpoint, audio, video, and image files). It supports both NVIDIA NIM microservices and a range of models to find, contextualize, and extract text, tables, charts, infographics, and transcripts for use in downstream generative and retrieval-augmented applications. +NVIDIA NeMo Retriever Library (NRL) extracts text, tables, charts, infographics, and transcripts from PDFs, HTML, Office documents, audio, video, and images. Run it as a Python library or Kubernetes deployment, and route inference through NVIDIA NIM microservices or local Nemotron models for downstream RAG and generative applications. -NeMo Retriever Library enables parallelization of splitting documents into pages where sub-page content is classified (such as text paragraphs, tables, charts, and infographics), extracted, and further contextualized through optical character recognition (OCR) into a standard schema. From there, NeMo Retriever Library manages computation of embeddings for the extracted content, -and can store vectors in [LanceDB](https://lancedb.com/) for the recommended embedded path when you pass `vdb_op="lancedb"` to upload (refer to [Vector databases](vdbs.md)). +NeMo Retriever Library splits documents into pages, classifies sub-page content (text, tables, charts, and infographics), extracts it, and applies optical character recognition (OCR) where needed into a standard schema. It can compute embeddings for extracted content and store vectors in [LanceDB](https://lancedb.com/) when you pass `vdb_op="lancedb"` to upload (refer to [Vector databases](vdbs.md)). ## NVIDIA AI Enterprise (NVAIE) support { #nvidia-ai-enterprise-nvaie-support } @@ -13,7 +12,7 @@ and can store vectors in [LanceDB](https://lancedb.com/) for the recommended emb Some individual NIM microservices and models that the library calls—for example, the default NIMs in the [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md#default-helm-nims)—may be covered by NVAIE on their own. That coverage applies only to those individual NIMs and models. It does **not** extend to the NeMo Retriever Library or its end-to-end extraction workflow. Using NVAIE-supported NIMs or models through the NeMo Retriever Library does not make the library, its container, or its chart NVAIE-supported. -## What NeMo Retriever Library Is ✔️ +## What NeMo Retriever Library Is ✔️ { #what-nemo-retriever-library-is } The following diagram shows the retriever pipeline. @@ -21,13 +20,13 @@ The following diagram shows the retriever pipeline. NeMo Retriever Library does the following: -- Accept directories of input files and a series of configurable ingestion tasks to perform on that input -- Allow the extracted content be retrieved from a VDB containing discrete metadata element -- Support multiple extraction methods per document type—for example, PDFs can use **pdfium** or [Nemotron Parse](https://build.nvidia.com/nvidia/nemotron-parse) as an alternate method (`extract_method="nemotron_parse"`) -- Support various types of pre- and post- processing operations, including text splitting and chunking, transform and filtering, embedding generation, and image offloading to storage. +- Accept directories of input files and configurable ingestion tasks +- Store extracted content in a vector database (VDB) with discrete metadata elements +- Support multiple extraction methods per document type—for example, PDFs can use **pdfium** or [Nemotron Parse](https://build.nvidia.com/nvidia/nemotron-parse) as an alternate method (`method="nemotron_parse"`) +- Apply pre- and post-processing: text splitting and chunking, transforms and filtering, embedding generation, and image offloading to storage !!! note - To use `extract_method="nemotron_parse"` with PDFs, install the Nemotron Parse client dependencies with the `nemotron-parse` extra, for example `pip install "nemo-retriever[nemotron-parse]"`. + To use `method="nemotron_parse"` with PDFs, install the Nemotron Parse client dependencies with the `nemotron-parse` extra, for example `uv pip install "nemo-retriever[nemotron-parse]"`. You can use the equivalent `pip install` command if you do not use UV. NeMo Retriever Library supports the following file types: @@ -51,7 +50,7 @@ NeMo Retriever Library supports the following file types: - `txt` - `wav` -## Related Topics +## Related Topics { #related-topics } - [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md) - [Deployment options](deployment-options.md) — library, Helm, hosted vs self-hosted NIMs in one place diff --git a/docs/docs/extraction/prerequisites-support-matrix.md b/docs/docs/extraction/prerequisites-support-matrix.md index db489af75e..7036727e62 100644 --- a/docs/docs/extraction/prerequisites-support-matrix.md +++ b/docs/docs/extraction/prerequisites-support-matrix.md @@ -1,30 +1,32 @@ # Pre-Requisites & Support Matrix -Before you begin using [NeMo Retriever Library](overview.md), confirm your software stack, deployment hardware, and—if you use them—advanced features (audio and video, Nemotron Parse, VLM image captioning, reranking) against the guidance in this page. +Before you begin using [NeMo Retriever Library](overview.md), confirm your software stack, deployment hardware, and—if you use them—advanced features (audio and video, Nemotron Parse, VLM image captioning, reranking) against the guidance on this page. -!!! note "NVIDIA AI Enterprise (NVAIE) support" +**Platform summary:** Supported **local GPU inference** requires **Linux** and CUDA 13. For **remote NIM inference**, the base Python package also installs on **Windows x64** and **macOS Apple Silicon (arm64)**; local GPU inference is not supported on those platforms. **macOS Intel (x86_64) is not supported** — `pip`/`uv` installs fail because Ray no longer publishes Intel Mac wheels. - The NeMo Retriever Library, including its container image and Helm chart artifacts, is not supported under NVIDIA AI Enterprise (NVAIE), even though some NIM microservices and models it uses may be individually covered by NVAIE. For more information, refer to [NVIDIA AI Enterprise (NVAIE) support](overview.md#nvidia-ai-enterprise-nvaie-support). +> **Note — NVIDIA AI Enterprise (NVAIE) support** +> +> The NeMo Retriever Library, including its container image and Helm chart artifacts, is not supported under NVIDIA AI Enterprise (NVAIE), even though some NIM microservices and models it uses may be individually covered by NVAIE. For more information, refer to [NVIDIA AI Enterprise (NVAIE) support](overview.md#nvidia-ai-enterprise-nvaie-support). ## Software Requirements { #software-requirements } -- Linux operating systems (Ubuntu 22.04 or later recommended) -- [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads) (NVIDIA Driver >= `580`, CUDA >= `13.0`) +- Linux operating systems (Ubuntu 22.04 or later recommended) for supported local GPU inference. For remote NIM inference, the base package can also be installed on Windows x64 and macOS Apple Silicon (arm64); local GPU inference is not supported on those platforms. macOS Intel (x86_64) is not supported: package installation fails because Ray `>=2.56.1` has no Intel Mac wheels (including in-process library mode). +- [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads) (local GPU inference only; NVIDIA Driver >= `580`, CUDA >= `13.0`) - [Python](https://www.python.org/downloads/) `3.12` — required to install and run the NeMo Retriever Library Python API, CLI, and related packages from PyPI (for example `pip` or `uv`). Older Python versions will fail dependency resolution without a clear error. - [UV Python package and environment manager](https://docs.astral.sh/uv/getting-started/installation/) (optional; recommended for creating isolated environments) - For audio and video, `ffmpeg` and `ffprobe` must be on `PATH` (for example `sudo apt-get install -y --no-install-recommends ffmpeg` on Debian/Ubuntu). `ffmpeg-python` and `nemo-retriever[multimedia]` do not install these binaries. For container and Kubernetes guidance, refer to [Audio and video](audio-video.md). -- For PDF extraction with `extract_method="nemotron_parse"`, install the Nemotron Parse - client dependencies with `pip install "nemo-retriever[nemotron-parse]"` (pulls +- For PDF extraction with `method="nemotron_parse"`, install the Nemotron Parse + client dependencies with `uv pip install "nemo-retriever[nemotron-parse]"` (pulls `open-clip-torch`, which provides the `open_clip` module required by the Nemotron Parse NIM client). The base `nemo-retriever` install and `[local]` extra do not include this - package. + package. You can use the equivalent `pip install` command if you do not use UV. -!!! note - - When you use UV, create the environment with Python 3.12 — for example, `uv venv --python 3.12`. This matches the `requires-python` metadata in the library packages. +> **Note** +> +> When you use UV, create the environment with Python 3.12 — for example, `uv venv --python 3.12`. This matches the `requires-python` metadata in the library packages. ## Hardware Requirements { #hardware-requirements } @@ -40,9 +42,9 @@ For per-feature GPU memory, disk, and co-residency rules, refer to [Model hardwa - **CPU Cores**: At least 32 CPU cores - **GPU**: NVIDIA GPU with at least 24 GB VRAM (for example, A100, H100, L40S, or equivalent) -!!! note - - Using less powerful systems or lower resource limits is still viable, but performance will suffer. +> **Note** +> +> Using less powerful systems or lower resource limits is still viable, but performance will suffer. ### Resource Consumption Notes @@ -73,20 +75,20 @@ Optional advanced features—audio and video transcription, Nemotron Parse, Omni ### Default NIMs { #default-helm-nims } -!!! important "NVAIE support applies to individual NIMs only" - - A NIM or model listed in the default and optional NIM rows in the table below might be supported under NVIDIA AI Enterprise (NVAIE) as an individual product. That support does **not** cover its use through NeMo Retriever Library or extend to the library, its container image, its Helm chart, or the end-to-end extraction workflow. +> **Important — NVAIE support applies to individual NIMs only** +> +> A NIM or model listed in the default and optional NIM rows in the table below might be supported under NVIDIA AI Enterprise (NVAIE) as an individual product. That support does **not** cover its use through NeMo Retriever Library or extend to the library, its container image, its Helm chart, or the end-to-end extraction workflow. -The production Helm chart reconciles NIM microservices through `nimOperator..enabled`. Four core NIMs are **enabled by default** and auto-wired into the retriever service; optional NIMs reconcile only when you opt in. For chart keys, image overrides, and enablement, refer to the [NeMo Retriever Helm chart README](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md#nim-operator-sub-stack) and [Recommended minimal install](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md#recommended-minimal-install-2605). +The production Helm chart reconciles NIM microservices through `nimOperator..enabled`. Four core NIMs are **enabled by default** and auto-wired into the retriever service; optional NIMs reconcile only when you opt in. For chart keys, image overrides, and enablement, refer to the [NeMo Retriever Helm chart README](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md#nim-operator-sub-stack) and [Recommended minimal install](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md#recommended-minimal-install-2608). | Helm flag | NIM | Default image (`repository:tag`) | Role | Enabled by default | |-----------|-----|----------------------------------|------|--------------------| | `page_elements` | [nemotron-page-elements-v3](https://build.nvidia.com/nvidia/nemotron-page-elements-v3) | `nvcr.io/nim/nvidia/nemotron-page-elements-v3:1.8.0` | Page layout and element detection | Yes | | `table_structure` | [nemotron-table-structure-v1](https://build.nvidia.com/nvidia/nemotron-table-structure-v1) | `nvcr.io/nim/nvidia/nemotron-table-structure-v1:1.8.0` | Table structure extraction | Yes | | `ocr` | [nemotron-ocr-v2](https://build.nvidia.com/nvidia/nemotron-ocr-v2) | `nvcr.io/nim/nvidia/nemotron-ocr-v2:1.4.0` | Image OCR | Yes | -| `vlm_embed` | [llama-nemotron-embed-vl-1b-v2](https://build.nvidia.com/nvidia/llama-nemotron-embed-vl-1b-v2) | `nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2:1.12.0` | Multimodal (VL) embedding | Yes | -| `rerankqa` | [llama-nemotron-rerank-vl-1b-v2](https://build.nvidia.com/nvidia/llama-nemotron-rerank-vl-1b-v2) | `nvcr.io/nim/nvidia/llama-nemotron-rerank-vl-1b-v2:1.11.0` | Reranking for improved retrieval accuracy | No | -| `nemotron_parse` | [nemotron-parse](https://build.nvidia.com/nvidia/nemotron-parse) | `nvcr.io/nim/nvidia/nemotron-parse-v1.2:1.7.0-variant` | Optional PDF `extract_method="nemotron_parse"` (default PDF extraction uses **pdfium**) | No | +| `vlm_embed` | [llama-nemotron-embed-vl-1b-v2](https://build.nvidia.com/nvidia/llama-nemotron-embed-vl-1b-v2) | `nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2:2.3.0` | Multimodal (VL) embedding | Yes | +| `rerankqa` | [llama-nemotron-rerank-vl-1b-v2](https://build.nvidia.com/nvidia/llama-nemotron-rerank-vl-1b-v2) | `nvcr.io/nim/nvidia/llama-nemotron-rerank-vl-1b-v2:2.3.0` | Reranking for improved retrieval accuracy | No | +| `nemotron_parse` | [nemotron-parse](https://build.nvidia.com/nvidia/nemotron-parse) | `nvcr.io/nim/nvidia/nemotron-parse-v1.2:1.7.0-variant` | Optional PDF `method="nemotron_parse"` (default PDF extraction uses **pdfium**) | No | | `nemotron_3_nano_omni_30b_a3b_reasoning` | [nemotron-3-nano-omni-30b-a3b-reasoning](https://build.nvidia.com/nvidia/nemotron-3-nano-omni-30b-a3b-reasoning) | `nvcr.io/nim/nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:1.7.0-variant` | Image captioning when you enable the caption stage | No | | `audio` | [parakeet-1-1b-ctc-en-us](https://docs.nvidia.com/nim/speech/latest/reference/support-matrix/index.html) | `nvcr.io/nim/nvidia/parakeet-1-1b-ctc-en-us:1.5.0` | [Audio and video](audio-video.md) transcription | No | | `answer_llm` | [llama-3.3-nemotron-super-49b-v1.5](https://build.nvidia.com/nvidia/llama-3.3-nemotron-super-49b-v1.5) | `nvcr.io/nim/nvidia/llama-3.3-nemotron-super-49b-v1.5:2.0.5` | Optional `/v1/answer` generation LLM (not part of the default extraction pipeline) | No | @@ -104,11 +106,24 @@ When you call [NVIDIA-hosted NIMs](deployment-options.md#when-to-use-nvidia-host | nemotron-ocr-v2 | `https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-ocr-v2` | Chart default OCR SKU; library CPU actors default to this URL when no OCR invoke URL is set. **Local OCR language selectors (`--ocr-lang`, API `ocr_lang`) are not sent on remote requests** — hosted OCR v2 uses its own language behavior | | llama-nemotron-embed-vl-1b-v2 | `https://integrate.api.nvidia.com/v1/embeddings` with model ID `nvidia/llama-nemotron-embed-vl-1b-v2` | Core multimodal embedding | | llama-nemotron-rerank-vl-1b-v2 | `https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-nemotron-rerank-vl-1b-v2/reranking` | Optional VL reranker | -| nemotron-parse | `https://integrate.api.nvidia.com/v1/chat/completions` with model ID `nvidia/nemotron-parse` | Optional `extract_method="nemotron_parse"` | +| nemotron-parse | `https://integrate.api.nvidia.com/v1/chat/completions` with model ID `nvidia/nemotron-parse` | Optional `method="nemotron_parse"`. Hosted Build and self-hosted `nemotron-parse-v1.2` use different request contracts; the library selects the matching contract automatically. Refer to [Nemotron Parse: hosted Build endpoint vs self-hosted NIM](#nemotron-parse-hosted-vs-self-hosted) | | nemotron-3-nano-omni-30b-a3b-reasoning | `https://integrate.api.nvidia.com/v1/chat/completions` with model ID `nvidia/nemotron-3-nano-omni-30b-a3b-reasoning` | Optional image captioning | | llama-3.3-nemotron-super-49b-v1.5 | `https://integrate.api.nvidia.com/v1/chat/completions` with model ID `nvidia/llama-3.3-nemotron-super-49b-v1.5` | Optional `/v1/answer` (Helm `answer_llm`) and OpenAI-compatible agentic RAG endpoint mode; not part of the default extraction pipeline. Agentic query/harness runs default to local in-process vLLM instead. Helm auto-wires to the in-cluster NIM when `nimOperator.answer_llm` is enabled | | parakeet-1-1b-ctc-en-us | `grpc.nvcf.nvidia.com:443` (function ID from [build.nvidia.com](https://build.nvidia.com/)) | Optional ASR; refer to [Parakeet hosted inference](audio-video.md#parakeet-hosted-inference-build-nvidia) | + + +!!! note "Nemotron Parse: hosted Build endpoint vs self-hosted NIM" + + Hosted NVIDIA Build and self-hosted Nemotron Parse use **different request contracts**. The library selects the matching contract from the endpoint and model: + + - **Hosted Build** (`https://integrate.api.nvidia.com/v1/chat/completions`) resolves to model ID `nvidia/nemotron-parse` and uses an image-only tool-call contract. + - **Self-hosted chat endpoints** default to model ID `nvidia/nemotron-parse-v1.2` and use the tagged text-prompt contract. + + To use hosted Build, set `nemotron_parse_invoke_url` to the Build chat-completions URL (and set `method="nemotron_parse"`). You can normally omit `nemotron_parse_model` so the library selects the model automatically. If you set `nemotron_parse_model` explicitly, it must match the endpoint contract. Mixed Build and self-hosted endpoint lists require an explicit model. + + For model/endpoint mismatch symptoms, refer to [Nemotron Parse model and endpoint mismatch](troubleshoot.md#nemotron-parse-model-endpoint-mismatch). + For local Hugging Face OCR language mode (`multi` vs `english`), Helm OCR image overrides, and local model install, refer to [OCR and scanned documents](multimodal-extraction.md#ocr-and-scanned-documents), [OCR NIM configuration](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md#ocr-nim-configuration), and [CLI — OCR language mode](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/docs/cli/README.md#ocr-language-mode). ### Image captioning { #image-captioning } diff --git a/docs/docs/extraction/ray-logging.md b/docs/docs/extraction/ray-logging.md index 3db3b07ef6..d6a43da4f1 100644 --- a/docs/docs/extraction/ray-logging.md +++ b/docs/docs/extraction/ray-logging.md @@ -11,7 +11,7 @@ In addition, NeMo Retriever Library provides preset configurations that you can -## Quick Start - Use Preset Configurations +## Quick Start - Use Preset Configurations { #quick-start-use-preset-configurations } To get started quickly, use one of the NeMo Retriever Library package-level preset variables. Run the code below that corresponds to your use case; production, development, or debugging. @@ -33,7 +33,7 @@ export INGEST_RAY_LOG_LEVEL=DEBUG ``` -### PRODUCTION Log Level +### PRODUCTION Log Level { #production-log-level } The `PRODUCTION` log level is optimized for production deployments with minimal logging overhead. @@ -51,7 +51,7 @@ This log level uses the following settings: - **Encoding** – TEXT -### DEVELOPMENT Log Level +### DEVELOPMENT Log Level { #development-log-level } The `DEVELOPMENT` log level is a balanced configuration for development work, and is the default log level. @@ -70,7 +70,7 @@ This log level uses the following settings: - **Encoding** – TEXT -### DEBUG Log Level +### DEBUG Log Level { #debug-log-level } The `DEBUG` log level provides maximum visibility for troubleshooting issues. @@ -89,7 +89,7 @@ This log level uses the following settings: -## Configuration Reference +## Configuration Reference { #configuration-reference } The following are the environment variables that you can set to control Ray logging behavior. If you specify an invalid value, the variable reverts to the default value with a warning message. @@ -109,7 +109,7 @@ If you specify an invalid value, the variable reverts to the default value with -## Configuration Examples +## Configuration Examples { #configuration-examples } ### Use a Preset With A Manual Override @@ -174,7 +174,7 @@ export RAY_LOGGING_ROTATE_BACKUP_COUNT=9 -## Log Output Examples +## Log Output Examples { #log-output-examples } ### INFO level (Default) @@ -208,6 +208,6 @@ export RAY_LOGGING_ROTATE_BACKUP_COUNT=9 ``` -## Related Topics +## Related Topics { #related-topics } -- [Environment Variables](https://docs.nvidia.com/nemo/retriever/latest/extraction/environment-config/) +- [Environment Variables](environment-config.md) diff --git a/docs/docs/extraction/releasenotes.md b/docs/docs/extraction/releasenotes.md index 8f5ef23848..c709c6d245 100644 --- a/docs/docs/extraction/releasenotes.md +++ b/docs/docs/extraction/releasenotes.md @@ -2,88 +2,95 @@ This documentation contains the release notes for [NeMo Retriever Library](overview.md). -## 26.05 Release Notes (26.5.0) +## 26.08 Release Notes (26.8.0) { #release-2608 } -NVIDIA® NeMo Retriever Library version 26.05 builds on the 26.03 foundation with a graph-based ingest architecture, expanded multimodal and tabular capabilities, production-oriented service deployment, and documentation aligned to a Helm-first supported path. +NVIDIA® NeMo Retriever Library version 26.08 builds on the 26.05 foundation with a graph-based ingest architecture, expanded multimodal and tabular capabilities, production-oriented service deployment, and documentation aligned to a Helm-first supported path. -To upgrade the Helm charts for this release, refer to the [NeMo Retriever Library Helm Charts](https://github.com/NVIDIA/NeMo-Retriever/blob/26.05/nemo_retriever/helm/README.md). +To upgrade the Helm charts for this release, refer to the [NeMo Retriever Library Helm Charts](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md). -Highlights for the 26.05 release include: +Highlights for the 26.08 release include: -### Upgrade notes +### Upgrade notes { #upgrade-notes } - Text splitting for graph and library ingest moved into `.extract(split_config=...)` instead of standalone `.split()` on the graph ingest path (the service ingestor API may still expose `.split()` separately) - Direct `Retriever(...)` construction uses `vdb_kwargs`, `embed_kwargs`, and `rerank` instead of flat `lancedb_uri`, `lancedb_table`, `embedder`, `embedding_endpoint`, `local_query_embed_backend`, and `reranker` arguments - For Helm audio and video extraction, set `service.installFfmpeg: true` in `values.yaml` (or pass `--set service.installFfmpeg=true`) when images no longer bundle `ffmpeg` and `ffprobe` by default - `nemo_retriever` requires Python 3.12 -### Pipeline and ingestion +### Pipeline and ingestion { #pipeline-and-ingestion } - Legacy `nv-ingest` and compatibility pipeline CLI code paths removed; `retriever ingest` and the graph stage registry are the canonical ingestion paths - Manifest-based ingest routing replaces input-type routing; `retriever ingest` is input-aware for PDF, image, audio, video, text, HTML, DOCX/PPTX, SVG, and related types - `allow_no_gpu` option to skip GPU requirement during ingest for CPU-only experimentation -### CLI +### CLI { #cli } - Root CLI adds first-class `retriever ingest` and `retriever query` commands with NIM URL flags, batch tuning, and LanceDB overwrite/append controls - `retriever ingest` and `retriever query` replace the retired compatibility pipeline command. Other top-level subcommands—including `eval`, `benchmark`, `harness`, and `skill-eval`—are development and experimental -### Retriever Service and deployment +### Retriever Service and deployment { #retriever-service-and-deployment } - Retriever Service v2 adds a scalable multi-pod architecture with gateway, process isolation, and VectorDB integration - OpenTelemetry basic support for pipeline and service observability - Expanded air-gapped deployment guidance in [deployment options](deployment-options.md) and the Helm chart README -### Models, OCR, and captioning +### Models, OCR, and captioning { #models-ocr-and-captioning } - Nemotron OCR v2 is the default OCR engine for HuggingFace, with CLI language selectors and unified OCR actors. For Helm NIM deployments, Nemotron OCR v1 is the default. - Nemotron Parse is available as an alternate PDF extraction method (v1.2 HTTP interface; optional Helm NIM; local inference via vLLM where configured) - VLM image captioning via vLLM (including Omni caption model profiles) addresses the capability deferred in 26.03 - vLLM-backed text and vision-language embedders, multimodal VL reranker, and torch 2.11 for local GPU installs -### Multimodal extraction +### Multimodal extraction { #multimodal-extraction } - Video retrieval pipeline with frame extraction, OCR, audio-visual fusion, and text deduplication - Long-audio Parakeet chunking with time-aligned segments; punctuation-based audio segmenting; ASR batch/streaming improvements -### Retrieval and RAG +### Retrieval and RAG { #retrieval-and-rag } - Live RAG SDK with `Retriever.retrieve()`, reference answer generation `Retriever.answer()`, and optional batch operator graphs via LiteLLM (`[llm]` extra) -### Vector database +### Vector database { #vector-database } - Vector database operators integrated directly in the pipeline; custom metadata support; LanceDB hybrid search guidance updated - LanceDB is documented as the first-party vector path for new deployments; Milvus/MinIO guidance removed from the primary extraction doc set -### Evaluation +### Evaluation { #evaluation } - BEIR-centric evaluation overhaul and `retriever skill-eval` benchmark CLI for the NeMo Retriever skill (experimental) - Text-to-SQL agent graph and tabular tooling for structured data retrieval, including tabular data ingestion -### Packaging and platform +### Packaging and platform { #packaging-and-platform } - Optional install extras (`[local]`, `[multimedia]`, `[llm]`, `[tabular]`, `[nemotron-parse]`, `[service]`, and others), including slim remote/NIM-only installs on Mac and Windows -### Helm chart +### Helm chart { #helm-chart } -- Helm chart refresh under `nemo_retriever/helm/` with GA VL embedder defaults and optional Nemotron Parse and Omni caption NIMs +- Helm chart refresh under `nemo_retriever/helm/` with VL embedder defaults and optional Nemotron Parse and Omni caption NIMs -### Documentation +### Documentation { #documentation } - Documentation aligned to a Helm-first supported path for NIM and service deployment -- Documentation consolidates extraction concepts, ingest workflow, embeddings, audio/video guides, prerequisites and support matrix, and UDF/custom stages in the [graph README](https://github.com/NVIDIA/NeMo-Retriever/tree/26.05/nemo_retriever/src/nemo_retriever/graph#nemo-retriever-graph) +- Documentation consolidates extraction concepts, ingest workflow, embeddings, audio/video guides, prerequisites and support matrix, and UDF/custom stages in the [graph README](https://github.com/NVIDIA/NeMo-Retriever/tree/main/nemo_retriever/src/nemo_retriever/graph#nemo-retriever-graph) -## Release Notes for Previous Versions +## Release Notes for Previous Versions { #previous-versions } - [26.05](https://docs.nvidia.com/nemo/retriever/26.5.0/extraction/releasenotes-nv-ingest/) - [26.03](https://docs.nvidia.com/nemo/retriever/26.3.0/extraction/releasenotes-nv-ingest/) +- [26.1.2](https://archive.docs.nvidia.com/nemo/retriever/26.1.2/extraction/releasenotes-nv-ingest/) +- [26.1.1](https://archive.docs.nvidia.com/nemo/retriever/26.1.1/extraction/releasenotes-nv-ingest/) +- [25.9.0](https://archive.docs.nvidia.com/nemo/retriever/25.9.0/extraction/releasenotes-nv-ingest/) +- [25.6.3](https://archive.docs.nvidia.com/nemo/retriever/25.6.3/extraction/releasenotes-nv-ingest/) +- [25.6.2](https://archive.docs.nvidia.com/nemo/retriever/25.6.2/extraction/releasenotes-nv-ingest/) +- [25.4.2](https://archive.docs.nvidia.com/nemo/retriever/25.4.2/extraction/releasenotes-nv-ingest/) +- [25.3.0](https://archive.docs.nvidia.com/nemo/retriever/25.3.0/extraction/releasenotes-nv-ingest/) -Release notes for earlier versions (26.1.2, 26.1.1, 25.9.0, 25.6.3, 25.6.2, 25.4.2, 25.3.0, 24.12.1, and 24.12.0) — *These links have been archived.* +Release notes for 24.12.1 and 24.12.0 are on the [25.3.0 archived release notes](https://archive.docs.nvidia.com/nemo/retriever/25.3.0/extraction/releasenotes-nv-ingest/). -## Related Topics +## Related Topics { #related-topics } - [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md) - [Deployment options](deployment-options.md) -- [NeMo Retriever Library Helm Charts](https://github.com/NVIDIA/NeMo-Retriever/blob/26.05/nemo_retriever/helm/README.md) +- [NeMo Retriever Library Helm Charts](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md) diff --git a/docs/docs/extraction/troubleshoot.md b/docs/docs/extraction/troubleshoot.md index c2b1553cc2..9017513b4f 100644 --- a/docs/docs/extraction/troubleshoot.md +++ b/docs/docs/extraction/troubleshoot.md @@ -2,13 +2,136 @@ Use this documentation to troubleshoot issues that arise when you use [NeMo Retriever Library](overview.md). -## Can't process long, non-language text strings +## Python API error triage { #python-api-error-triage } + +NeMo Retriever Library does not assign product-specific numeric error codes to +each extraction method. The Python API surfaces Python exception types, +row-level failure records, and HTTP or gRPC statuses from the service or +upstream NIM. Treat the exception type, failing stage, upstream status, and +response detail together as the error identifier. + +The current graph API enriches `GraphIngestionError` for explicitly configured +Page Elements, OCR, Table Structure, Nemotron Parse, and embedding endpoints. +When the upstream payload contains an HTTP status, the message includes the +stage, configured URL, and status. For example: + +```text +Graph ingestion detected row-level errors from an explicitly configured +remote NIM endpoint. row 0, column ocr +[stage=OCR NIM url=https://example.invalid/v1/infer http=503], path error: ... +Troubleshooting: OCR NIM ... returned a 5xx server error ... +``` + +This enrichment is not a new error-code namespace. The `503` in this example +is the upstream HTTP status. Exception wording and upstream response bodies are +not stable API fields and should not be parsed programmatically. + +In many row-level payloads, the HTTP status appears only inside +`error.message` (for example `HTTP 503 from ...`), not as a separate +`status_code` field. Inspect `exc.records` or the failing DataFrame column +when troubleshooting. The `error` value can be a nested object or a string. + +### Coverage limits of the raise error policy + +`error_policy="raise"` scans only remote NIM stages with an explicitly +configured invoke URL: Page Elements, OCR, Table Structure, Nemotron Parse, +and embedding. It does not automatically raise for: + +- Local-only pipelines (`pdfium` without remote URLs), even when rows contain + `metadata.error` or column-level error payloads. +- Caption or remote VLM stages. Missing credentials fail at actor setup; + inference failures can abort the entire ingest. +- Audio or video ASR over gRPC or HTTP. Failures can drop individual rows and + log warnings instead of raising `GraphIngestionError`. + +For those paths, use `error_policy="collect"`, pass `return_failures=True`, or +inspect row columns and service logs directly. + +### Error signals and first response + +| Signal | Typical meaning | L1/L2 response | +| --- | --- | --- | +| `ValueError` or Pydantic validation error before execution | An unsupported run mode, parameter value, protocol, or parameter combination | Compare the call with the current [Python API reference](nemo-retriever-api-reference.md). Remove unknown parameters and reproduce with the smallest valid pipeline. | +| `ImportError`, `ModuleNotFoundError`, or a missing-dependency `RuntimeError` | The selected local extraction path requires a package or executable that is not installed | Install the documented package extra or system dependency. Confirm that the Python environment running the worker, not only the client shell, contains it. | +| `GraphIngestionError` with no HTTP status | The named remote stage returned a row-level error, its error payload omitted a status, or the endpoint was unreachable | Check DNS, routing, TLS, the endpoint URL, and the NIM readiness endpoint from the worker or service pod. Inspect `exc.records` after removing secrets and document content. | +| HTTP `401` or `403` from a NIM | Missing, expired, or unauthorized credentials | Verify `NVIDIA_API_KEY`, `NGC_API_KEY`, or the stage-specific credential in the environment that makes the request. Do not attach API keys to a support case. | +| HTTP `403` from the Retriever service | Authentication failure or a deployment policy that disallows the requested endpoint, stage, sink, or override | Read the response `detail`. Verify the service token and compare the requested pipeline with `/v1/ingest/pipeline-config`. | +| HTTP `404` or `410` while opening a service ingest job | The Python SDK and Retriever service can be on incompatible API versions | A current client raises `RetrieverServiceCompatibilityError`. Align the Python package, service image, and Helm chart versions. | +| Other HTTP `4xx` | The upstream service rejected the request | Check file type, rendered page or image size, model name, endpoint path, and request schema. For `413` or `422`, reduce the payload or image size and verify the endpoint's input limits. | +| HTTP `429` | The remote service is rate-limiting requests | Reduce concurrency or batch size and retry with backoff. Escalate only if throttling persists within the service quota. | +| HTTP `5xx`, including `503` | The upstream NIM is unavailable, overloaded, not ready, or failed during inference | Check readiness, pod restarts, GPU memory, server logs, and request volume. Retry a minimal input after the NIM is healthy. | +| Timeout, connection reset, DNS, TLS, or gRPC transport error | The client could not complete transport to the service or NIM | Test connectivity from the process or pod that runs the stage. Verify protocol, port, certificate trust, proxy, and network policy. Preserve the gRPC status and details when present. | +| A per-document entry in `ServiceIngestResult.failures` | Upload or pipeline processing failed after a service job was created | Correlate the document ID with the job ID and service logs. Other documents in the same result can still have succeeded. | +| Successful ingest with fewer rows than inputs (caption or ASR enabled) | Caption inference failed before row collection, or ASR dropped failed rows and logged warnings | Re-run with logging enabled. For caption, verify endpoint credentials and payload limits. For ASR, verify gRPC endpoint, `function_id`, and `NVIDIA_API_KEY`. | +| OOM, worker exit, or pod restart | Host or GPU resources were exhausted, or an orchestrator terminated the worker | Reduce batch size or concurrency, use smaller document groups, and inspect host, Ray, Kubernetes, and NIM resource telemetry. | + +The service can retry some transient transport, `429`, and `5xx` failures. +Report the final status returned after retries, not an intermediate warning. + +### Representative extraction paths + +Use the failing stage—not only the top-level `method` value—to select the +troubleshooting path. A single document can pass through several stages. + +| API path | Components that can fail | Representative signals | +| --- | --- | --- | +| `ExtractParams(method="pdfium")` | File loading, PDF splitting, PDFium parsing, page rendering; optionally Page Elements and Table Structure when enabled | Malformed or encrypted input, `pypdfium2` import failure, local Python exception, or remote-stage `GraphIngestionError` when an invoke URL is explicitly configured | +| `ExtractParams(method="pdfium_hybrid")` | PDFium plus Page Elements, OCR, and optionally Table Structure | The local PDF signals above, or a row-level/HTTP failure attributed to Page Elements, OCR, or Table Structure | +| `ExtractParams(method="ocr")` | Page rendering, Page Elements, and the local or remote OCR backend | Missing local model dependencies, invalid image payload, authentication/transport status, or OCR row-level failure | +| `ExtractParams(method="nemotron_parse")` | PDF rendering and local Nemotron Parse model or configured Nemotron Parse NIM | Missing `open_clip`, missing local model configuration, unsupported image input, or Nemotron Parse row-level/HTTP failure | +| `.caption(...)` | Local caption model or remote VLM endpoint | `ValueError` at setup when credentials or endpoint/protocol are invalid; remote inference failures can abort the whole ingest rather than populate a row error column | +| `.embed(...)` | Local embedding model or configured embedding NIM | Model/dependency error, input-size or schema rejection, authentication/transport status, or embedding row-level failure; `GraphIngestionError` when a remote embed URL is configured | +| Audio or video extraction | `ffmpeg`/`ffprobe`, media decoding, frame/chunk creation, and local or remote ASR | Missing executable, malformed media, codec failure, gRPC status, or credential error; ASR failures may omit rows and log warnings instead of raising, so verify logs when output is unexpectedly empty | + +`pdfium` itself is primarily a local parser, so a Page Elements, Table +Structure, OCR, caption, or embedding HTTP status comes from an enabled +downstream stage rather than from PDFium. + +### Collect diagnostics safely + +Before escalating, collect the following: + +1. Package, image or Helm versions, and `run_mode`. +2. Exception class and sanitized message. For `GraphIngestionError`, include + sanitized `exc.records`. For row-level failures, include `stage`, `type`, + and `message` when present. +3. Extraction method, enabled stages, and endpoint hostnames with credentials + and signed query parameters removed. +4. HTTP or gRPC status, response detail, job ID, document ID, trace ID, and + timestamp when available. +5. Whether the endpoint readiness check succeeds from the worker or service + pod. +6. A minimal non-confidential reproducing input, or characteristics such as + format, page count, dimensions, and size. +7. Relevant client, service, Ray, NIM, and Kubernetes logs for the same + timestamp. + +Never include API keys, bearer tokens, document contents, or unredacted signed +URLs in logs or support cases. + +Escalate to NVIDIA L3 when the failure is reproducible on a supported, +version-aligned configuration after L1 and L2 support has verified input +validity, credentials, endpoint readiness, connectivity, and resource +availability. Escalate immediately for repeatable crashes, incorrect +successful output, or a `5xx` from a healthy NVIDIA-owned NIM with a minimal +valid input. Keep configuration, dependency, customer network, quota, and +malformed-input issues with L1 and L2 support unless the documented behavior +is incorrect. + +!!! note "Older NV-Ingest releases" + + Error text and result shapes differ by release. NV-Ingest `25.4.2` + predates some current enriched diagnostics. Do not assume that a field + shown in current NeMo Retriever Library output exists in `25.4.2`; include + the exact old exception and logs when escalating. + +## Can't process long, non-language text strings { #cant-process-long-non-language-text-strings } NeMo Retriever Library is designed to process language and language-length strings. If you submit a document that contains extremely long, or non-language text strings, such as a DNA sequence, errors or unexpected results occur. -## Can't process malformed input files +## Can't process malformed input files { #cant-process-malformed-input-files } When you run a job you might see errors similar to the following: @@ -58,7 +181,7 @@ service: This path fails with `allowPrivilegeEscalation: false` or `readOnlyRootFilesystem: true`. -## Can't start new thread error +## Can't start new thread error { #cant-start-new-thread-error } In rare cases, when you run a job you might an see an error similar to `can't start new thread`. This error occurs when the maximum number of processes available to a single user is too low. @@ -74,7 +197,7 @@ ulimit -u 10000 -## Out-of-Memory (OOM) Error when Processing Large Datasets +## Out-of-Memory (OOM) Error when Processing Large Datasets { #out-of-memory-oom-error-when-processing-large-datasets } When you process a very large dataset with thousands of documents, you might encounter an Out-of-Memory (OOM) error. This happens because NeMo Retriever Library materializes extraction results in system memory (RAM) while the job runs. @@ -89,7 +212,7 @@ To reduce memory pressure, try one or more of the following: -## Embedding service fails to start with an unsupported batch size error +## Embedding service fails to start with an unsupported batch size error { #embedding-service-fails-unsupported-batch-size } On certain hardware, for example RTX 6000, the embedding service might fail to start and you might see an error similar to the following. @@ -106,7 +229,7 @@ You can set the variable in your .env file or directly in your environment. ## ModuleNotFoundError: No module named open_clip when using nemotron_parse { #modulenotfounderror-no-module-named-open-clip-when-using-nemotron-parse } -When you run PDF extraction with `extract_method="nemotron_parse"`, you might see an error similar to the following: +When you run PDF extraction with `method="nemotron_parse"`, you might see an error similar to the following: ```text ModuleNotFoundError: No module named 'open_clip' @@ -126,12 +249,24 @@ For local GPU inference with Nemotron Parse, combine extras: pip install "nemo-retriever[local,nemotron-parse]" ``` -Also refer to [What is NeMo Retriever Library?](overview.md) and [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md#software-requirements). +Also refer to [NeMo Retriever Library Overview](overview.md) and [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md#software-requirements). -## Extract method nemotron-parse doesn't support image files +## Extract method nemotron-parse doesn't support image files { #extract-method-nemotron-parse-doesnt-support-image-files } Currently, extraction with Nemotron parse doesn't support image files, only scanned PDFs. -To work around this issue, convert image files to PDFs before you use `extract_method="nemotron_parse"`. +To work around this issue, convert image files to PDFs before you use `method="nemotron_parse"`. + +## Nemotron Parse model and endpoint mismatch { #nemotron-parse-model-endpoint-mismatch } + +When you run PDF extraction with `method="nemotron_parse"`, a mismatched model and endpoint can fail with an error similar to the following: + +```text +HTTP 400: Content cannot be a plain string. The model does not support text input. +``` + +This can occur when you send a tagged or versioned `v1.2` model (for example `nvidia/nemotron-parse-v1.2`) to the NVIDIA-hosted Build endpoint, which expects the image-only `nvidia/nemotron-parse` contract. The library may replace the raw HTTP error with a targeted model/contract mismatch hint. + +To use hosted Build, omit `nemotron_parse_model` so the library selects `nvidia/nemotron-parse` automatically, or set `nemotron_parse_model="nvidia/nemotron-parse"` explicitly. Send versioned `v1.2` models only to a compatible self-hosted chat endpoint. For more information, refer to [Nemotron Parse: hosted Build endpoint vs self-hosted NIM](prerequisites-support-matrix.md#nemotron-parse-hosted-vs-self-hosted). ## Hosted Page Elements NIM image size limits { #hosted-page-elements-nim-image-size-limits } @@ -227,7 +362,7 @@ For the request schema, refer to the [Object Detection NIM API reference](https: Supported inline formats remain **PNG** and **JPEG**, encoded as `data:image/;base64,` or `data:image/;asset_id,`. OpenAPI specs for Page Elements v2 and v3 are linked from the [Object Detection NIM API reference](https://docs.nvidia.com/nim/ingestion/object-detection/latest/api-reference.html#openapi-reference-for-page-elements). -## Too many open files error +## Too many open files error { #too-many-open-files-error } In rare cases, when you run a job you might an see an error similar to `too many open files` or `max open file descriptor`. This error occurs when the open file descriptor limit for your service user account is too low. @@ -243,7 +378,7 @@ ulimit -n 10000 -## Triton server INFO messages incorrectly logged as errors +## Triton server INFO messages incorrectly logged as errors { #triton-server-info-messages-incorrectly-logged-as-errors } Sometimes messages are incorrectly logged as errors, when they are information. When this happens, you can ignore the errors, and treat the messages as information. @@ -276,7 +411,7 @@ ERROR 2025-04-24 22:49:44.434 nimutils.py:68] } -## Related Topics +## Related Topics { #related-topics } - [Pre-Requisites & Support Matrix](prerequisites-support-matrix.md) - [Deployment options](deployment-options.md) diff --git a/docs/docs/extraction/vdbs.md b/docs/docs/extraction/vdbs.md index ed69f2244c..a0f05fbfef 100644 --- a/docs/docs/extraction/vdbs.md +++ b/docs/docs/extraction/vdbs.md @@ -5,7 +5,7 @@ Use this documentation to learn how [NeMo Retriever Library](overview.md) stores ## On this page { #on-this-page } - [Overview](#overview) -- [Why LanceDB?](#why-lancedb) +- [LanceDB Overview](#why-lancedb) - [Upload to LanceDB](#upload-to-lancedb) - [Semantic retrieval](#semantic-retrieval) - [Metadata and filtering](#metadata-and-filtering) @@ -37,7 +37,7 @@ Currently, data upload is not supported through the [CLI](https://github.com/NVI -## Why LanceDB? { #why-lancedb } +## LanceDB Overview { #why-lancedb } LanceDB is optimized for low-latency retrieval in this stack: diff --git a/docs/docs/extraction/workflow-agentic-retrieval.md b/docs/docs/extraction/workflow-agentic-retrieval.md index 2299d22b7a..0f354790ef 100644 --- a/docs/docs/extraction/workflow-agentic-retrieval.md +++ b/docs/docs/extraction/workflow-agentic-retrieval.md @@ -23,23 +23,86 @@ For custom or already deployed chat models, opt into the endpoint path: ```bash retriever query "find documents about parser behavior" \ --agentic \ - --agentic-llm-backend openai_compatible \ --agentic-llm-model custom-remote-model \ --agentic-invoke-url http://localhost:9000/v1/chat/completions ``` +Providing `--agentic-invoke-url` routes the agent to that remote endpoint; the LLM +client defaults to `callable`, which calls the endpoint over the shared +chat-completions HTTP client and needs no LLM SDK installed. + ## MCP access for agents -`retriever service start` mounts a FastMCP HTTP endpoint at `/mcp` by default. Agents can use that endpoint to call the running service for health checks, pipeline introspection, document ingestion, job status, VectorDB query, and answer generation. If service auth is enabled, the MCP endpoint uses the same bearer-token middleware as the REST API. +`retriever service start` mounts a FastMCP HTTP endpoint at `/mcp` by default. +Agents can use that endpoint to call the running service for health checks, +pipeline introspection, document ingestion, job status, VectorDB query, agentic +retrieval, and answer generation. If service auth is enabled, the MCP endpoint +uses the same bearer-token middleware as the REST API. + +Plain and agentic retrieval share `POST /v1/query` and the same hits response +envelope. They are separate MCP tools so agents can choose explicitly: + +- `query` calls `POST /v1/query` with `agentic=false` for one-pass dense or hybrid retrieval. +- `agentic_query` calls `POST /v1/query` with `agentic=true` and runs the ReAct + retrieval workflow. It is added to MCP when `agentic.enabled` is true. + Agentic results are document-level: `source` is the selected `doc_id`, + `metadata` carries `result_source` and `rank`, and chunk-level fields + (`text`, `page_number`, scores, …) are unset. + +Enable agentic retrieval in `retriever-service.yaml`: + +```yaml +agentic: + enabled: true + llm_model: your-openai-compatible-model + invoke_url: https://your-llm.example/v1/chat/completions + reasoning_effort: high + backend_top_k: 20 + react_max_steps: 50 + request_timeout_s: 1800 +``` + +The VectorDB process owns the LanceDB volume and executes the agentic workflow. +Start it with matching `--agentic`, `--agentic-llm-model`, and +`--agentic-invoke-url` options. The LLM and embedding credentials are resolved +from the service process environment. Service mode requires remote +OpenAI-compatible LLM and embedding endpoints; local in-process models remain +available through the one-shot CLI and harness paths. + +REST clients set the flag on `/v1/query`: + +```bash +curl -X POST http://localhost:7670/v1/query \ + -H 'Content-Type: application/json' \ + -d '{"query": "find documents about parser behavior", "top_k": 5, "agentic": true}' +``` + +Successful responses include ``query_mode``: ``"agentic"`` for this path and +``"classic"`` for dense/hybrid ``/v1/query`` (including ``format=evidence``). + +Requests with `agentic: true` return HTTP `400` when agentic retrieval is not +configured on the service. Agentic runs use a small dedicated worker pool in the +VectorDB process so they cannot exhaust the capacity used by plain queries. A +ReAct run cannot be interrupted once started, so a worker stays occupied until +it finishes even if the caller times out or disconnects. When every worker is +busy the endpoint sheds load with `503` and a `Retry-After` header instead of +queueing behind a multi-minute run. Agentic queries are capped at 4096 +characters to bound prompt size and cost across the multi-step loop. For local stdio-based agents, run the MCP server as a shim that points at an existing retriever service: ```bash retriever service mcp-stdio \ --service-url http://localhost:7670 \ + --query-methods agentic \ --api-token "$NEMO_RETRIEVER_API_TOKEN" ``` +Use `--query-methods classic` (default), `agentic`, or `all` to choose which retrieval +tools the MCP server registers. Mounted `/mcp` uses the same knob via +`mcp.query_methods` in the service config; agentic tools are omitted unless +`agentic.enabled` is also true. + For remote agents, expose the retriever service URL and configure the agent to connect to: ```text diff --git a/docs/docs/index.md b/docs/docs/index.md index fc891e7283..20a2e48af6 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -1,4 +1,4 @@ -# What is NVIDIA NeMo Retriever? +# NVIDIA NeMo Retriever Overview { #what-is-nvidia-nemo-retriever } NVIDIA NeMo Retriever is a collection of microservices for building and scaling multimodal data extraction, embedding, and reranking pipelines diff --git a/docs/docs/reference/collection-management-api.md b/docs/docs/reference/collection-management-api.md new file mode 100644 index 0000000000..2affa9bf89 --- /dev/null +++ b/docs/docs/reference/collection-management-api.md @@ -0,0 +1,201 @@ +# Collection management Python API + +`RetrieverServiceClient` is the supported boundary for long-lived agentic +applications. The client talks to the NeMo Retriever service; applications do +not open LanceDB, choose table names, or reproduce ingestion stages. + +## End-to-end workflow + +```python +import time +from nemo_retriever import RetrieverServiceClient + +client = RetrieverServiceClient( + base_url="http://nemo-retriever:7670", # Published service endpoint + api_token="...", + scope="workspace-123", +) + +collection = client.create_collection("research-session") +job = client.submit_documents( + collection.name, + ["report.pdf"], + idempotency_key="agent-request-42", +) + +# Submission means the job and uploads were accepted. It does not mean that +# extraction, OCR, splitting, captioning, embedding, and indexing are done. +while True: + job = client.get_job(job.job_id) + if job.status in {"completed", "failed", "partial_success"}: + break + time.sleep(2) + +hits = client.query( + "What are the major findings?", collection_name=collection.name, top_k=10, +) +documents = client.list_documents(collection.name) +client.delete_document(collection.name, documents.items[0].document_id) +client.delete_collection(collection.name) +``` + +For local Docker Compose deployments, use the published gateway address, such +as `http://localhost:7670`. For other deployments, use the published gateway +endpoint reachable by the calling application. Authentication, tracing, +retryable upload handling, collection routing, and result normalization remain +server/SDK responsibilities. + +## Sync and async methods + +Every lifecycle method has a native async equivalent prefixed with `a`: +`create_collection`/`acreate_collection`, `submit_documents`/`asubmit_documents`, +`get_job`/`aget_job`, `list_documents`/`alist_documents`, and +`query`/`aquery`. Use async methods inside an event loop. + +Collection methods include create, get, list, update, and delete. Document +methods include get, list, delete, and atomic replace. Job methods expose the +aggregate and paginated per-file status. List operations use bounded `limit` +values and opaque continuation tokens; callers must not interpret tokens. + +## Append, idempotency, and replacement + +Normal submission appends documents without changing existing documents. An +idempotency key replay with the same request returns the original job. The SDK +then safely replays every manifest entry, including after the client loses a +response before, during, or after upload. Each file has a deterministic +`manifest_entry_id` derived from its position, filename, and SHA-256. The +service returns the original acceptance for entries it already accepted, +without consuming capacity or starting duplicate processing. Reusing the key +or an entry ID with different content returns +`RetrieverServiceConflictError` (HTTP 409). + +Before the first physical append, the VectorDB records a pending-version +recovery marker and writes deterministic chunk IDs with an idempotent merge. +After an interrupted write, reconciliation either finalizes committed chunks +or removes an empty marker, so retrying the same document version does not +duplicate chunks. +Pending initial appends remain hidden from document reads and collection +queries until reconciliation commits them. + +Job document status separates `attempt_id` (one processing attempt) from +`document_id` (the stable collection identity). Append creates a new stable +document ID. Replacement creates a new attempt but retains the target document +ID. Collection document APIs show only indexed materializations; pending, +processing, and failed attempts remain visible through job APIs. + +`replace_document()` submits one replacement file. NeMo Retriever records a +pending-version recovery marker, uses a single LanceDB merge transaction to +insert the new chunks and remove obsolete chunks for that document, and then +finalizes the catalog. The VectorDB reconciler inspects stored chunk versions +after a crash and either finalizes the new version or preserves the old one. +Failed processing never removes the prior version, and queries never expose +mixed versions. + +## Errors, scopes, expiration, and compatibility + +The SDK raises `RetrieverServiceNotFoundError`, +`RetrieverServiceConflictError`, `RetrieverServiceValidationError`, or the base +`RetrieverServiceError`. Resources are isolated by `scope`; cross-scope reads +return 404. `expires_at` can be set at collection creation or update time for an +operator cleanup process. Deletion is retryable and `if_exists=True` makes +repeated deletion safe. Delete results report `existed`, `deleted`, `status`, +and `cleanup_pending`; synchronous completion returns HTTP 200 and a retryable +pending cleanup may return HTTP 202. + +Production deployments map bearer tokens to allowed workspace scopes. Missing +or invalid credentials and valid tokens requesting an unauthorized scope +receive the same 401 response, preventing callers from distinguishing token +validity. Once authorized for a scope, looking up a resource owned by another +scope returns 404 so its existence is not disclosed. Configure either a single +token bound to `default_scope`, or mount a Secret-backed JSON file: + +```json +{"tokens":[{"token":"","scopes":["workspace-123"]}]} +``` + +Set `allow_unscoped_dev` only for an explicitly auth-disabled development +deployment. The gateway records the authorized scope on the request. Pod-only +callback routes and VectorDB calls require the separate internal credential; +an external bearer token is never used to authorize those internal routes or +forwarded to VectorDB. + +`expires_at` must be timezone-aware RFC3339 and is normalized to UTC. For an +expiring collection, successful append and replace indexing activity refreshes +the expiration while preserving the configured window between `updated_at` and +`expires_at`. Collection metadata updates do the same when they omit +`expires_at`; supplying `expires_at` establishes a new window and setting it to +null disables expiration. Writes that do not commit vector data, including +empty writes, do not refresh collection activity. During recovery from an +interrupted write, NRL records activity refresh as durable recovery work. If +the collection update fails, reconciliation retains the marker and retries the +refresh. After the refresh succeeds, it clears the marker without refreshing +the collection again, so retries do not extend the expiration more than once. +Expired collections enter the +same retryable deletion state machine as explicit deletion. The local VectorDB +reconciler runs every 60 seconds by default, applies exponential retry capped at +one hour, and resumes replacement, document deletion, collection deletion, and +expiration cleanup after a crash. +Run one VectorDB replica while this reconciler is enabled; durable distributed +coordination remains separate infrastructure work. An interval of zero is +reserved for deployments where an external reconciler owns cleanup. + +`StoreOperator` artifact persistence remains an independent pipeline and storage +concern. Collection deletion removes collection and document catalog entries, +chunk/vector rows, and the backend-owned physical collection table; it does not +delete extracted artifacts from S3, NFS, or the local filesystem. Configure +artifact retention and garbage collection at the storage/operator boundary, +where the corresponding credentials and ownership policy already live. + +Legacy fixed-table ingestion and query remain available when +`collection_name` is omitted, but only against the operator-configured table. +No service request may specify a raw table name, storage URI, or physical +LanceDB location. `/document` is the canonical ingestion route and `/whole` is +supported; collection-aware `/page` returns 422 before work is registered. + +Continuation tokens are versioned keyset cursors rather than offsets. +Collection cursors advance by collection name; document cursors advance by +`(created_at, document_id)`. Tokens are bound to their resource type, scope, +and collection and return 422 when reused in another context. This keeps pages +stable while resources are inserted or deleted. + +VectorDB health and metrics expose only aggregate catalog schema health, +active/deleting/expired counts, pending cleanup count and oldest age, +reconciliation successes/failures, and open-table cache size. Physical table +names and tenant identifiers are never emitted as public values or labels. + +## Docker Compose operations + +The default development stack lives at +`nemo_retriever/dev/compose/service-mode.compose.yaml` and runs the Retriever +and VectorDB as separate services. Set `NRL_API_TOKEN` to opt into a public +bearer credential and `NRL_INTERNAL_VDB_TOKEN` to protect the private service +hop; leaving them unset preserves the existing unauthenticated development +behavior. Runtime tokens must not be committed. Production deployments can +continue to use the service's Secret-backed multi-scope token-file support. +The same SDK workflow targets `http://localhost:7670`. + +## Application integration and query-result contract + +External applications should construct a `RetrieverServiceClient` from the +service URL, token, and workspace scope, then call the SDK directly. Applications +should orchestrate calls and translate their own configuration only; NeMo +Retriever owns processing status, stable chunk/document identity, retrieval +ordering, citation provenance, retries, idempotency, and lifecycle truth. Clients +must not open LanceDB directly or reproduce the ingestion pipeline. + +Collection query hits provide stable `chunk_id` and `document_id`, non-null +`text`, a finite native `distance`, filename, a one-based page number when +known, content type, source/source ID, stored image URI, bounding box, and +metadata. Collection queries use dense vector retrieval in this release; +lower distances are more similar and list order is authoritative. NRL does +not reinterpret distance as a normalized similarity or confidence. Consumers +that require a bounded score must translate the complete result set at their +own adapter boundary. `page_number` is `null` for non-paginated content or +invalid/unknown page provenance. Audio segments, video frames, and timestamps +keep their existing modality-specific metadata rather than being converted +into document pages. This contract is identical regardless of the network +path used to reach the service. + +For `format=evidence`, each evidence item's `score` is the same native dense +vector distance, not a normalized confidence or probability. Lower is better, +and values are not comparable across queries. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index c499080d7b..fd9075a20e 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -74,7 +74,7 @@ extra_css: # bar when there is only one top-level tab. nav: - "1. Introduction": - - "What is NeMo Retriever?": extraction/overview.md + - "NeMo Retriever Library Overview": extraction/overview.md - Key concepts: extraction/concepts.md - Release notes: extraction/releasenotes.md - "2. Get started": @@ -109,6 +109,7 @@ nav: # TODO: after nv-ingest code removal, update this link when CLI docs are relocated. - "CLI reference": https://github.com/NVIDIA/NeMo-Retriever/tree/main/nemo_retriever/docs/cli - "Quickstart: retriever CLI": reference/retriever-cli-quickstart.md + - "Collection management Python API": reference/collection-management-api.md - Environment variables: extraction/environment-config.md - "Metadata reference": extraction/content-metadata.md - "12. Support & community": @@ -196,6 +197,8 @@ plugins: markdown_extensions: - attr_list - md_in_html + - toc: + permalink: true - pymdownx.details - pymdownx.superfences - pymdownx.snippets: diff --git a/evaluation/bo767_recall.md b/evaluation/bo767_recall.md index 2545f6c4b4..ff0f1b2dd2 100644 --- a/evaluation/bo767_recall.md +++ b/evaluation/bo767_recall.md @@ -1,70 +1,27 @@ -# Evaluate BO767 Retrieval with the Retriever Harness +# Evaluate BO767 Retrieval -Use the Retriever harness for the canonical BO767 end-to-end ingest, query, -and BEIR evaluation. The checked-in runfile selects batch execution and carries -the benchmark's worker tuning and required file, page, and query counts. - -Run these commands from the repository root. - -## Configure the Dataset Paths - -Copy the example path map to an untracked location: - -```bash -cp nemo_retriever/harness/dataset_paths.example.yaml /tmp/retriever-dataset-paths.yaml -``` - -Edit the `bo767` entry so `path` points to the directory containing the 767 -documents and `query_file` points to the BO767 query/qrels CSV available on the -machine running the benchmark. - -## Validate the Resolved Run - -Resolve the run without launching ingestion or evaluation: +BO767 is a checked-in batch benchmark. Configure the host's dataset paths, dry +run the exact request, and then execute it. ```bash -uv run --project nemo_retriever retriever harness run-files \ - --session-name bo767_beir_check \ - --output-dir /tmp/retriever-harness-bo767-check \ - --dataset-paths /tmp/retriever-dataset-paths.yaml \ - --dry-run \ - --json \ - nemo_retriever/harness/runfiles/bo767_beir.json -``` - -Inspect `session_summary.json`, `expanded_runs.json`, and the child run's -`resolved_benchmark.json` before starting the full run. - -## Run the Evaluation +cp nemo_retriever/harness/dataset_paths.example.yaml \ + /local/path/to/dataset_paths.yaml +${EDITOR:-vi} /local/path/to/dataset_paths.yaml -Remove `--dry-run` and choose a durable artifact directory: - -```bash uv run --project nemo_retriever retriever harness run-files \ --session-name bo767_beir \ --output-dir /local/path/to/retriever-artifacts/bo767-beir \ - --dataset-paths /tmp/retriever-dataset-paths.yaml \ - --json \ + --dataset-paths /local/path/to/dataset_paths.yaml \ + --dry-run \ nemo_retriever/harness/runfiles/bo767_beir.json ``` -BO767 is a large batch benchmark. Keep the terminal process alive until the -harness reaches a terminal state; model startup and ingestion can be quiet for -extended periods. - -## Read the Results - -The harness artifacts are the evaluation contract: +Confirm that `session_summary.json` succeeds and inspect the child +`resolved_benchmark.json`. Remove `--dry-run` to execute the benchmark. -- `status.json` reports the current phase and concise failure state. -- `results.json` is the authoritative terminal result and summary metrics. -- `session_summary.json` is the terminal result for the runfile session. -- `beir_metrics.json` contains the complete BEIR metric family. -- `query_results.jsonl` contains per-query latency and ranked hits. -- `environment.json` records the commit and runtime context. +Read `session_summary.json` first and the child `results.json` for terminal +metrics. Use `run.log` only when deeper diagnostics are needed. -Refer to -[`nemo_retriever/harness/EXPECTED_RESULTS.md`](../nemo_retriever/harness/EXPECTED_RESULTS.md#bo767) -for the expected BO767 counts and current reference metrics. For the complete -harness contract and troubleshooting guidance, refer to -[`nemo_retriever/harness/README.md`](../nemo_retriever/harness/README.md). +- [Library harness guide](../nemo_retriever/harness/docs/library.md) +- [BO767 dataset facts and observations](../nemo_retriever/harness/docs/expected-results.md#bo767-observations) +- [Shared artifact contract](../nemo_retriever/harness/README.md#results-and-artifacts) diff --git a/nemo_retriever/README.md b/nemo_retriever/README.md index b49bf2bc60..06af01044f 100644 --- a/nemo_retriever/README.md +++ b/nemo_retriever/README.md @@ -2,11 +2,12 @@ NeMo Retriever Library is a retrieval-augmented generation (RAG) ingestion pipeline for documents that can parse text, tables, charts, and infographics. NeMo Retriever Library parses documents, creates embeddings, optionally stores embeddings in LanceDB, and performs recall evaluation. -This quick start guide shows how to run NeMo Retriever Library as a library all within local Python processes without containers. NeMo Retriever Library supports two inference options: -- Pull and run [Nemotron RAG models from Hugging Face](https://huggingface.co/collections/nvidia/nemotron-rag) on your local GPU(s). -- Make over the network inference calls to build.nvidia.com hosted or locally deployed NeMo Retriever NIM endpoints. +This quick start guide shows how to run NeMo Retriever Library as a library in local Python processes without containers. Choose one inference path: -You’ll set up a CUDA 13–compatible environment, install the library and its dependencies, and run GPU‑accelerated ingestion pipelines that convert PDFs, HTML, plain text, audio, or video into vector embeddings stored in LanceDB (on local disk), with Ray‑based scaling and built‑in recall benchmarking. +- **Local GPU (Linux):** Pull and run [Nemotron RAG models from Hugging Face](https://huggingface.co/collections/nvidia/nemotron-rag) on your GPU(s). Requires CUDA 13.x and the `[local]` extra. +- **Remote NIM:** Call build.nvidia.com hosted or self-hosted NeMo Retriever NIM endpoints over the network. The base package installs on Linux, Windows x64, and macOS Apple Silicon (arm64); no local GPU is required. macOS Intel (x86_64) is not supported. + +The steps below cover environment setup, installation, and a first ingestion run. For Kubernetes or container deployments, refer to [Deployment at a glance](#deployment-at-a-glance) and the [Pre-Requisites & Support Matrix](https://docs.nvidia.com/nemo/retriever/latest/extraction/prerequisites-support-matrix/). ## Deployment at a glance @@ -16,12 +17,12 @@ For standalone service-image builds and local container runs, see **[`docker.md` ## Prerequisites -Before starting, make sure your system meets the following requirements: +Before starting, confirm requirements for your inference path (refer to the [Pre-Requisites & Support Matrix](https://docs.nvidia.com/nemo/retriever/latest/extraction/prerequisites-support-matrix/)): + +- **Local GPU inference (Linux):** CUDA 13.x with `libcudart.so.13` available, and GPUs visible to the system. +- **Remote NIM inference:** Python 3.12 and network access to your NIM endpoints; CUDA is not required on the client host. -- The host is running CUDA 13.x so that `libcudart.so.13` is available. -- Your GPUs are visible to the system and compatible with CUDA 13.x. -​ -If optical character recognition (OCR) fails with a `libcudart.so.13` error, install the CUDA 13 runtime for your platform and update `LD_LIBRARY_PATH` to include the CUDA lib64 directory, then rerun the pipeline. +If optical character recognition (OCR) fails with a `libcudart.so.13` error on a local GPU path, install the CUDA 13 runtime for your platform and update `LD_LIBRARY_PATH` to include the CUDA lib64 directory, then rerun the pipeline. For example, the following command can be used to update the `LD_LIBRARY_PATH` value. @@ -31,7 +32,7 @@ export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/lib64 ## Setup your environment -Complete the following steps to setup your environment. You will create and activate isolated Python and project virtual environments, install the NeMo Retriever Library and its dependencies, and then run the provided ingestion snippets to validate your setup. +Complete the following steps to set up your environment. You will create and activate isolated Python and project virtual environments, install the NeMo Retriever Library and its dependencies, and then run the provided ingestion snippets to validate your setup. 1. Create and activate the NeMo Retriever Library environment @@ -39,7 +40,9 @@ Before installing NeMo Retriever Library, create an isolated Python environment In your terminal, run the following commands from any location. -For **local GPU inference** (Nemotron models running on your GPU), install with the `[local]` extra, which includes the model packages, transformers, and GPU tooling: +**Local GPU (Linux)** + +Install with the `[local]` extra, which includes Nemotron model packages, transformers, and GPU tooling: ```bash uv venv retriever --python 3.12 @@ -52,12 +55,12 @@ try prerelease/nightly Nemotron packages from PyPI within the same supported major-version windows, opt in with `--pre`: ```bash -uv pip install --pre "nemo-retriever[local]==26.05-RC1" +uv pip install --pre "nemo-retriever[local]==26.08-RC1" ``` -Install matching **ingestion client** and **ingestion runtime** wheels at the same version when your workflow expects them (refer to the [NeMo Retriever Library prerequisites](https://docs.nvidia.com/nemo/retriever/latest/extraction/overview/) for the exact PyPI coordinates for your release). +**Remote NIM (no local GPU)** -For **remote NIM inference only** (no local GPU required), the base package is sufficient: +The base package is sufficient when all inference runs through hosted or self-hosted NIM endpoints: ```bash uv python install 3.12 @@ -68,9 +71,9 @@ uv pip install nemo-retriever Install matching **ingestion client** and **ingestion runtime** wheels at the same version when your workflow expects them (refer to the [NeMo Retriever Library prerequisites](https://docs.nvidia.com/nemo/retriever/latest/extraction/overview/) for the exact PyPI coordinates for your release). -This creates a dedicated Python environment and installs the `nemo-retriever` PyPI package, the canonical distribution for the NeMo Retriever Library. +This creates a dedicated Python environment and installs the `nemo-retriever` PyPI package, the canonical distribution for the NeMo Retriever Library. The base install includes the lightweight tokenizer dependencies used for TXT/HTML chunking (no Transformers or local model weights). -If your PDF pipeline uses `extract_method="nemotron_parse"`, install the Nemotron Parse client dependencies with the `nemotron-parse` extra: +If your PDF pipeline uses `method="nemotron_parse"`, install the Nemotron Parse client dependencies with the `nemotron-parse` extra: ```bash uv pip install "nemo-retriever[nemotron-parse]" @@ -86,7 +89,7 @@ The `[local]` extra pulls PyTorch from PyPI, which defaults to a CPU build on Li ```bash uv pip uninstall torch torchvision -uv pip install torch==2.10.0 torchvision -i https://download.pytorch.org/whl/cu130 +uv pip install torch==2.11.0 torchvision==0.26.0 -i https://download.pytorch.org/whl/cu130 ``` Skip this step if you are using remote NIM inference only. @@ -124,6 +127,35 @@ ingestor = ( ) ``` +### Ingest inline text + +Python callers can pass raw text documents directly to the same text splitting, +embedding, and vector database graph without creating temporary files. Inline +text is supported in `inprocess`, `batch`, and `service` run modes. + +```python +from nemo_retriever import create_ingestor + +texts = ["some text", "another longer text"] + +chunks = ( + create_ingestor(run_mode="batch") + .texts(texts) + .embed() + .vdb_upload() + .ingest() +) +``` + +Each string is treated as a raw document and split with `TextChunkParams` +defaults. To override chunking for every text source in the ingest, add +`.extract(split_config={"text": {"max_tokens": 512, "overlap_tokens": 64}})`. +Inline text can be combined with +`.files(...)` and, in modes that support them, `.buffers(...)`; each source is +routed through its matching extractor before the results enter the shared +embedding and sink stages. Inline corpora remain resident in client or driver +memory, so prefer file ingestion when the corpus may exceed the available memory. + ### Optional extras - **`multimedia`** — Audio/video extraction and SVG rendering support. Install this extra when using Parakeet ASR through `extract_method="audio"` so audio decoding and resampling dependencies are available: @@ -189,8 +221,9 @@ retriever ingest /your-example-dir \ > v2 selector. Remote OCR NIM endpoints decide their own model and language > behavior, and the local OCR selectors are not added to remote request payloads. -When you use the remote embedder, pair the `Retriever` with matching -`embed_kwargs` overrides shown in [Run a recall query](#run-a-recall-query). +When you use a remote embedder, the endpoint and provider prefix remain runtime +configuration. The query model is read from LanceDB metadata when available; +pass an explicit model only for an override or a legacy table without metadata. ### Inspect extracts You can inspect how recall accuracy optimized text chunks for various content types were extracted into text representations: @@ -360,8 +393,7 @@ CUDA_VISIBLE_DEVICES=0 retriever query "What is RAG?" \ --table-name nemo-retriever \ --embed-model-name nvidia/llama-nemotron-embed-1b-v2 \ --top-k 1 \ - --agentic-react-max-steps 1 \ - --agentic-backend-top-k 1 + --agentic-react-max-steps 1 ``` You can run the same flow from Python. Omit `invoke_url` for the default local @@ -426,7 +458,7 @@ query vectors land in the same embedding space as the stored chunks. ```python from nemo_retriever.graph.retriever import Retriever -from nemo_retriever.llm import LiteLLMClient +from nemo_retriever.models.llm import LiteLLMClient retriever = Retriever( vdb_kwargs={"uri": "lancedb", "table_name": "nemo-retriever"}, @@ -456,7 +488,7 @@ the bundled `VL_EMBED_MODEL`. Live RAG with scoring and an LLM judge (requires a ground-truth `reference`): ```python -from nemo_retriever.llm import LLMJudge +from nemo_retriever.models.llm import LLMJudge judge = LLMJudge.from_kwargs( model="nvidia_nim/nvidia/llama-3.3-nemotron-super-49b-v1.5", @@ -617,9 +649,14 @@ ingestor = ( ) ``` -You can use a different ingestion pipeline based on [Nemotron-Parse](https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2) combined with the default embedder: +You can use a different ingestion pipeline based on [Nemotron Parse](https://build.nvidia.com/nvidia/nemotron-parse) hosted on NVIDIA Build and combined with the default embedder: + ```python -ingestor = ingestor.files(documents).extract(method="nemotron_parse") +ingestor = ingestor.files(documents).extract( + method="nemotron_parse", + nemotron_parse_invoke_url="https://integrate.api.nvidia.com/v1/chat/completions", + nemotron_parse_model="nvidia/nemotron-parse", +) ``` ## Run with remote inference, no local GPU required: @@ -731,9 +768,12 @@ After installing the headers, restart the pipeline. ## Retriever Harness -The developer harness runs code-owned benchmarks through `retriever harness`. -Use `retriever harness list --runsets` to see available benchmark names and -runsets, then run one benchmark with `retriever harness run `. +The developer harness runs registered ingest and retrieval benchmarks through +`retriever harness`. Start with the +[harness guide](harness/README.md), then choose +[library execution](harness/docs/library.md) or +[service execution](harness/docs/service.md). Recurring workstation runs use the +[nightly launcher](../ops/retriever-nightly/README.md). ### Ingest image storage diff --git a/nemo_retriever/dev/compose/README.md b/nemo_retriever/dev/compose/README.md index 8a4938ecee..5854d7c2b5 100644 --- a/nemo_retriever/dev/compose/README.md +++ b/nemo_retriever/dev/compose/README.md @@ -31,6 +31,22 @@ docker compose -f nemo_retriever/dev/compose/service-mode.compose.yaml up --buil curl -fsSL http://localhost:7670/v1/health ``` +The same default stack exposes collection and document lifecycle APIs. Optional +runtime tokens enable a public bearer credential and a separate credential for +the private Retriever-to-VectorDB hop: + +```bash +export NRL_API_TOKEN="" +export NRL_INTERNAL_VDB_TOKEN="" +docker compose -f nemo_retriever/dev/compose/service-mode.compose.yaml up --build -d +``` + +Leave both values unset to preserve the existing unauthenticated development +experience. In this default Compose configuration, `NRL_API_TOKEN` is bound to +the `default` scope, so clients use `X-NRL-Scope: default` (or SDK +`scope="default"`). Multi-scope deployments use the service's Secret-backed +token mapping. Tokens are runtime inputs and must not be committed. + Endpoints, models, ports, worker counts, and the service image can all be overridden explicitly. The most commonly tuned variables are `NIM_PAGE_ELEMENTS_URL`, `NIM_TABLE_STRUCTURE_URL`, `NIM_OCR_URL`, @@ -45,7 +61,7 @@ Use the NVCR authentication described above before pulling a self-hosted NIM. Keep `NGC_API_KEY` exported so it is also available to the NIM containers. The hosted-only stack does not require `NGC_API_KEY` at runtime. -Start the four core extraction/retrieval NIMs with their checked-in internal +Start the four core extraction/retrieval NIM services with their checked-in internal endpoint wiring: ```bash @@ -66,7 +82,7 @@ docker compose \ -f nemo_retriever/dev/compose/service-mode.compose.yaml up --build -d ``` -To start all nine NIMs, layer the four wiring presets and select every NIM +To start all nine NIM services, layer the four wiring presets and select every NIM profile: ```bash @@ -84,16 +100,17 @@ Reranker and Parse need only `--profile nim-reranker` or `--profile nim-parse`. They are lifecycle/API-only and intentionally are not injected into retriever service configuration, matching Helm. -Every NIM has a persistent cache volume, a configurable GPU assignment, and a +Every NIM has a persistent model or cache volume, a configurable GPU assignment, and a configurable host port. Variables follow the service prefix, for example `NIM_OCR_GPU_ID`, `NIM_OCR_HOST_PORT`, `NIM_OCR_CACHE_VOLUME`, `NIM_OCR_CACHE_PATH`, `NIM_OCR_IMAGE`, and `NIM_OCR_TAG`. The defaults form a collision-free assignment for the combined-profile example: core NIMs use GPUs -0-3, reranker uses 4, parse uses 5, caption uses 6, answer uses 7-8, and audio -uses 9. Change the defaults to match the active profiles and host before +0, 1, 2, and 3 (page-elements, table-structure, OCR, and embedding), reranker uses 4, +parse uses 5, caption uses 6, answer uses 7-8, and audio uses 9. Change the +defaults to match the active profiles and host before startup; for example, an answer-only run on a two-GPU host can set `NIM_ANSWER_GPU_ID_0=0` and `NIM_ANSWER_GPU_ID_1=1`. Compose lifecycle support -means image pull, startup, readiness, persistent cache, restart, logs, and +means image pull, startup, readiness, persistent model or cache data, restart, logs, and teardown; NIM Operator reconciliation, NIMCache CRDs, and model-profile selection remain Kubernetes-only. @@ -164,8 +181,8 @@ These checks pull large images/models and require suitable NVIDIA GPUs. Use `docker compose config` for configuration-only validation without launching the stack. -1. Core extraction/retrieval: start `nims-core.env`, wait for all six services - to report healthy, then ingest a representative PDF with the service CLI: +1. Core extraction/retrieval: start `nims-core.env`, wait for the four NIMs plus + Retriever and VectorDB to report healthy, then ingest a representative PDF with the service CLI: ```bash retriever ingest service /path/to/document.pdf \ diff --git a/nemo_retriever/dev/compose/presets/nims-core.env b/nemo_retriever/dev/compose/presets/nims-core.env index 2580d5737e..4afe969a81 100644 --- a/nemo_retriever/dev/compose/presets/nims-core.env +++ b/nemo_retriever/dev/compose/presets/nims-core.env @@ -1,4 +1,6 @@ -NIM_PAGE_ELEMENTS_URL=http://nim-page-elements:8000/v1/infer -NIM_TABLE_STRUCTURE_URL=http://nim-table-structure:8000/v1/infer -NIM_OCR_URL=http://nim-ocr:8000/v1/infer +# Separate Compose services for page-elements and table-structure (matching Helm). +# Both run the combined nemotron-object-detection:2.0.0 image; OCR stays on nim-ocr. +NIM_PAGE_ELEMENTS_URL=http://nim-page-elements:8000/v1/page-elements +NIM_TABLE_STRUCTURE_URL=http://nim-table-structure:8000/v1/table-structure +NIM_OCR_URL=http://nim-ocr:8000/v1/ocr NIM_EMBED_URL=http://nim-embedding:8000/v1/embeddings diff --git a/nemo_retriever/dev/compose/service-mode.compose.yaml b/nemo_retriever/dev/compose/service-mode.compose.yaml index b9e8c1bda4..864ef0e8c1 100644 --- a/nemo_retriever/dev/compose/service-mode.compose.yaml +++ b/nemo_retriever/dev/compose/service-mode.compose.yaml @@ -7,7 +7,11 @@ name: nemo-retriever-service-dev # Development-only service-mode stack. Production deployments use Helm. x-retriever-image: &retriever-image image: ${NEMO_RETRIEVER_IMAGE:-nemo-retriever-service:dev} - build: {context: ../../.., dockerfile: Dockerfile, target: service} + build: + context: ../../.. + dockerfile: Dockerfile + target: service + args: {DOWNLOAD_DEFAULT_TOKENIZER: "True"} x-host-access: &host-access extra_hosts: ["host.docker.internal:host-gateway"] x-nim-environment: &nim-environment @@ -39,6 +43,8 @@ services: ports: ["${RETRIEVER_HTTP_PORT:-7670}:7670"] environment: NVIDIA_API_KEY: ${NVIDIA_API_KEY:-} + NRL_API_TOKEN: ${NRL_API_TOKEN:-} + NRL_INTERNAL_VDB_TOKEN: ${NRL_INTERNAL_VDB_TOKEN:-} NEMO_RETRIEVER_LLM_API_KEY: ${ANSWER_LLM_API_KEY:-${NGC_API_KEY:-}} AUDIO_GRPC_ENDPOINT: ${NIM_AUDIO_GRPC_ENDPOINT:-} INSTALL_FFMPEG: ${INSTALL_FFMPEG:-false} @@ -79,6 +85,7 @@ services: command: [python, -m, nemo_retriever.service.vectordb_app, --lancedb-uri, /data/vectordb, --table-name, nemo_retriever, --embed-endpoint, "${NIM_EMBED_URL:-https://integrate.api.nvidia.com/v1/embeddings}", --embed-model, "${NIM_EMBED_MODEL:-nvidia/llama-nemotron-embed-vl-1b-v2}", --port, "7671"] environment: NVIDIA_API_KEY: ${NVIDIA_API_KEY:-${NGC_API_KEY:-}} + NRL_INTERNAL_VDB_TOKEN: ${NRL_INTERNAL_VDB_TOKEN:-} volumes: [vectordb_data:/data/vectordb] depends_on: nim-embedding: {condition: service_healthy, required: false} @@ -90,53 +97,73 @@ services: start_period: 10s restart: unless-stopped + # Page-elements and table-structure are separate Compose services (matching + # Helm) but both run the combined nemotron-object-detection:2.0.0 image. nim-page-elements: <<: *nim-service profiles: [nims-core] - image: ${NIM_PAGE_ELEMENTS_IMAGE:-nvcr.io/nim/nvidia/nemotron-page-elements-v3}:${NIM_PAGE_ELEMENTS_TAG:-1.8.0} + image: ${NIM_PAGE_ELEMENTS_IMAGE:-nvcr.io/nim/nvidia/nemotron-object-detection}:${NIM_PAGE_ELEMENTS_TAG:-2.0.0} ports: ["${NIM_PAGE_ELEMENTS_HOST_PORT:-8001}:8000"] environment: <<: *nim-environment - NIM_HTTP_API_PORT: "8000" NIM_OTEL_SERVICE_NAME: nemotron-page-elements-v3 - NIM_TRITON_LOG_VERBOSE: "1" - NIM_TRITON_MAX_BATCH_SIZE: "32" - NIM_TRITON_CPU_THREADS_PRE_PROCESSOR: "2" - NIM_TRITON_CPU_THREADS_POST_PROCESSOR: "1" - OMP_NUM_THREADS: "2" - volumes: ["nim_page_elements_cache:${NIM_PAGE_ELEMENTS_CACHE_PATH:-/opt/nim/.cache}"] + NIM_SERVER_BIND_ADDR: 0.0.0.0:8000 + NIM_PERFORMANCE_MODE: "0" + NIM_SERVER_MODE: latency + NIM_SERVER_MAX_WAIT_MS: "0" + NIM_ENGINE_COUNT: "1" + NIM_PIPELINE_MAX_BATCH_SIZE: "1" + NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER: ngc + NIM_ENGINE_MODEL_NAME: nvidia/nemotron-page-elements-v3 + NIM_ENGINE_MODEL_PATH: ${NIM_PAGE_ELEMENTS_CACHE_PATH:-/model-store}/page-elements + volumes: ["nim_page_elements_cache:${NIM_PAGE_ELEMENTS_CACHE_PATH:-/model-store}"] + healthcheck: + test: [CMD, curl, --fail, --silent, http://localhost:8000/v1/health/ready] deploy: {resources: {reservations: {devices: [{<<: *nim-gpu, device_ids: ["${NIM_PAGE_ELEMENTS_GPU_ID:-0}"]}]}}} nim-table-structure: <<: *nim-service profiles: [nims-core] - image: ${NIM_TABLE_STRUCTURE_IMAGE:-nvcr.io/nim/nvidia/nemotron-table-structure-v1}:${NIM_TABLE_STRUCTURE_TAG:-1.8.0} + image: ${NIM_TABLE_STRUCTURE_IMAGE:-nvcr.io/nim/nvidia/nemotron-object-detection}:${NIM_TABLE_STRUCTURE_TAG:-2.0.0} ports: ["${NIM_TABLE_STRUCTURE_HOST_PORT:-8002}:8000"] environment: <<: *nim-environment - NIM_HTTP_API_PORT: "8000" NIM_OTEL_SERVICE_NAME: nemotron-table-structure-v1 - NIM_TRITON_LOG_VERBOSE: "1" - NIM_TRITON_RATE_LIMIT: "3" - NIM_TRITON_MAX_BATCH_SIZE: "32" - NIM_TRITON_CUDA_MEMORY_POOL_MB: "2048" - OMP_NUM_THREADS: "1" - volumes: ["nim_table_structure_cache:${NIM_TABLE_STRUCTURE_CACHE_PATH:-/opt/nim/.cache}"] + NIM_SERVER_BIND_ADDR: 0.0.0.0:8000 + NIM_PERFORMANCE_MODE: "0" + NIM_SERVER_MODE: latency + NIM_SERVER_MAX_WAIT_MS: "0" + NIM_ENGINE_COUNT: "1" + NIM_PIPELINE_MAX_BATCH_SIZE: "1" + NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER: ngc + NIM_ENGINE_MODEL_NAME: nvidia/nemotron-table-structure-v1 + NIM_ENGINE_MODEL_PATH: ${NIM_TABLE_STRUCTURE_CACHE_PATH:-/model-store}/table-structure + volumes: ["nim_table_structure_cache:${NIM_TABLE_STRUCTURE_CACHE_PATH:-/model-store}"] + healthcheck: + test: [CMD, curl, --fail, --silent, http://localhost:8000/v1/health/ready] deploy: {resources: {reservations: {devices: [{<<: *nim-gpu, device_ids: ["${NIM_TABLE_STRUCTURE_GPU_ID:-1}"]}]}}} nim-ocr: <<: *nim-service profiles: [nims-core] - image: ${NIM_OCR_IMAGE:-nvcr.io/nim/nvidia/nemotron-ocr-v2}:${NIM_OCR_TAG:-1.4.0} + image: ${NIM_OCR_IMAGE:-nvcr.io/nim/nvidia/nemotron-ocr-v2}:${NIM_OCR_TAG:-2.0.0} ports: ["${NIM_OCR_HOST_PORT:-8003}:8000"] environment: <<: *nim-environment - NIM_HTTP_API_PORT: "8000" NIM_OTEL_SERVICE_NAME: nemotron-ocr-v2 - NIM_TRITON_LOG_VERBOSE: "1" - NIM_TRITON_MAX_BATCH_SIZE: "32" - OMP_NUM_THREADS: "8" - volumes: ["nim_ocr_cache:${NIM_OCR_CACHE_PATH:-/opt/nim/.cache}"] + NIM_SERVER_BIND_ADDR: 0.0.0.0:8000 + NIM_PERFORMANCE_MODE: "0" + NIM_SERVER_MODE: latency + NIM_SERVER_MAX_WAIT_MS: "0" + NIM_ENGINE_COUNT: "1" + NIM_PIPELINE_MAX_BATCH_SIZE: "1" + NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER: ngc + NIM_ENGINE_MODEL_NAME: nvidia/nemotron-ocr-v2 + NIM_ENGINE_MODEL_PATH: ${NIM_OCR_CACHE_PATH:-/model-store}/ocr + NIM_ENGINE_MODEL_VARIANT: multilingual + volumes: ["nim_ocr_cache:${NIM_OCR_CACHE_PATH:-/model-store}"] + healthcheck: + test: [CMD, curl, --fail, --silent, http://localhost:8000/v1/health/ready] deploy: {resources: {reservations: {devices: [{<<: *nim-gpu, device_ids: ["${NIM_OCR_GPU_ID:-2}"]}]}}} nim-embedding: @@ -252,9 +279,9 @@ configs: file: "/var/lib/nemo-retriever/retriever-service.log" format: "%(asctime)s | %(levelname)s | %(name)s | %(message)s" nim_endpoints: - page_elements_invoke_url: "${NIM_PAGE_ELEMENTS_URL-https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-page-elements-v3}" - table_structure_invoke_url: "${NIM_TABLE_STRUCTURE_URL-https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-table-structure-v1}" - ocr_invoke_url: "${NIM_OCR_URL-https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-ocr-v1}" + page_elements_invoke_url: "${NIM_PAGE_ELEMENTS_URL-http://nim-page-elements:8000/v1/page-elements}" + table_structure_invoke_url: "${NIM_TABLE_STRUCTURE_URL-http://nim-table-structure:8000/v1/table-structure}" + ocr_invoke_url: "${NIM_OCR_URL-http://nim-ocr:8000/v1/ocr}" embed_invoke_url: "${NIM_EMBED_URL-https://integrate.api.nvidia.com/v1/embeddings}" embed_model_name: "${NIM_EMBED_MODEL:-nvidia/llama-nemotron-embed-vl-1b-v2}" caption_invoke_url: ${NIM_CAPTION_URL:-null} @@ -271,7 +298,6 @@ configs: extract: enabled: ${LOCAL_EXTRACT_ENABLED:-true} use_table_structure: ${LOCAL_EXTRACT_USE_TABLE_STRUCTURE:-true} - use_graphic_elements: ${LOCAL_EXTRACT_USE_GRAPHIC_ELEMENTS:-true} ocr_version: "${LOCAL_EXTRACT_OCR_VERSION:-v2}" ocr_lang: ${LOCAL_EXTRACT_OCR_LANG_YAML:-null} embed: @@ -302,6 +328,9 @@ configs: batch_queue_size: ${PIPELINE_BATCH_QUEUE_SIZE:-8192} mcp: enabled: false + auth: + api_token: "${NRL_API_TOKEN:-}" + allow_unscoped_dev: true vectordb: enabled: true lancedb_uri: "/data/vectordb" diff --git a/nemo_retriever/dev/compose/service-mode.local-models.compose.yaml b/nemo_retriever/dev/compose/service-mode.local-models.compose.yaml index 02b07d8357..23661e33f4 100644 --- a/nemo_retriever/dev/compose/service-mode.local-models.compose.yaml +++ b/nemo_retriever/dev/compose/service-mode.local-models.compose.yaml @@ -7,9 +7,8 @@ services: retriever: image: ${NEMO_RETRIEVER_GPU_IMAGE:-nemo-retriever-service-gpu:dev} - build: {target: service-gpu} + build: {target: service-gpu, args: {DOWNLOAD_DEFAULT_TOKENIZER: "True"}} environment: - HF_HOME: /var/cache/huggingface CUDA_VISIBLE_DEVICES: ${LOCAL_MODELS_GPU_ID:-0} volumes: [huggingface_cache:/var/cache/huggingface] deploy: @@ -19,10 +18,9 @@ services: - {driver: nvidia, device_ids: ["${LOCAL_MODELS_GPU_ID:-0}"], capabilities: [gpu]} vectordb: image: ${NEMO_RETRIEVER_GPU_IMAGE:-nemo-retriever-service-gpu:dev} - build: {target: service-gpu} + build: {target: service-gpu, args: {DOWNLOAD_DEFAULT_TOKENIZER: "True"}} command: [python, -m, nemo_retriever.service.vectordb_app, --lancedb-uri, /data/vectordb, --table-name, nemo_retriever, --local-embed, --local-embed-backend, "${LOCAL_EMBED_BACKEND:-hf}", --hf-cache-dir, /var/cache/huggingface, --device, cuda, --gpu-memory-utilization, "${LOCAL_EMBED_GPU_MEMORY_UTILIZATION:-0.45}", --embed-model, "${LOCAL_EMBED_MODEL:-nvidia/llama-nemotron-embed-vl-1b-v2}", --port, "7671"] environment: - HF_HOME: /var/cache/huggingface CUDA_VISIBLE_DEVICES: ${LOCAL_VECTORDB_GPU_ID:-1} volumes: [huggingface_cache:/var/cache/huggingface] deploy: diff --git a/nemo_retriever/developer_docs/README.md b/nemo_retriever/developer_docs/README.md index 3f06086398..e76f1ec990 100644 --- a/nemo_retriever/developer_docs/README.md +++ b/nemo_retriever/developer_docs/README.md @@ -9,6 +9,6 @@ architecture, subsystems, and developer-facing tools. |-------|-------------| | [Graph Pipeline Registry](graph_pipeline_registry.md) | Central registry for managing, inspecting, comparing, and serializing golden pipeline graphs. | | [NimClient and Custom NIM Endpoints](nimclient.md) | Developer guide for custom NIM integrations with `NimClient`, `ModelInterface`, and UDFs. | -| [Retriever Harness README](../harness/README.md) | Operator and agent instructions for the artifact-first Retriever benchmark harness. | -| [Retriever Harness PRD](harness_retriever_ingest_query_prd.md) | Product requirements for the artifact-first Retriever ingest/query benchmark harness revamp. | +| [Retriever Harness](../harness/README.md) | Current user and agent guide for registered Retriever benchmarks. | +| [Retriever Harness Design History](harness_retriever_ingest_query_prd.md) | Implemented design decisions and product boundaries. | | [Root Ingest CLI Design](root_ingest_cli_design.md) | Reviewer guide for the `retriever ingest` local, batch, and service CLI ownership split. | diff --git a/nemo_retriever/developer_docs/harness_retriever_ingest_query_prd.md b/nemo_retriever/developer_docs/harness_retriever_ingest_query_prd.md index 2224eb6b17..4fc69440a4 100644 --- a/nemo_retriever/developer_docs/harness_retriever_ingest_query_prd.md +++ b/nemo_retriever/developer_docs/harness_retriever_ingest_query_prd.md @@ -1,774 +1,45 @@ -# Retriever Harness PRD: End-to-End Ingest/Query Benchmarks +# Retriever Harness Design History -Last updated: 2026-07-21 +> This design has been implemented. The +> [Retriever Harness README](../harness/README.md) is the current user and agent +> contract. -## Implementation Status +The harness was rebuilt as an internal end-to-end benchmark runner for NeMo +Retriever engineers. It ingests registered datasets, queries the resulting +index, evaluates retrieval quality, and writes durable artifacts. -The current implementation includes the core runner described here plus three -orchestration-neutral extensions: `run-files` applies a machine-local dataset -path map to one or more checked-in runfiles and gives each real child a fresh -process; `check-vidore-access` validates remote ViDoRe evaluation data; and -`post-slack` renders or posts completed artifacts. It does not include -recurring scheduling, deployment, locking, retry policy, or secret -distribution. The harness README is the normative user and agent guide; this -PRD records the design rationale. +## Decisions That Still Apply -## Summary +| Decision | Consequence | +| --- | --- | +| Keep benchmark definitions in a typed Python registry. | Recurring benchmarks are reviewed code, not an open-ended YAML system. | +| Treat runfiles as concrete requests. | A runfile selects a registered benchmark; it cannot define a new one. | +| Keep one harness contract across execution targets. | Library and service runs emit the same results, gates, and artifacts. | +| Treat local and batch as library ingest modes. | Batch changes how ingest runs; it is not a separate harness. | +| Treat service as a system-under-test mode. | The harness uses service ingest and query APIs while retaining the same evaluation contract. | +| Keep Helm outside benchmark semantics. | `run-helm` provisions around `run-files`; Helm is not a runfile mode. | +| Make artifacts the API. | Callers use exit codes, `status.json`, `results.json`, and `session_summary.json`, not stdout. | +| Keep execution separate from reporting. | `post-slack` reads completed artifacts and never reruns or mutates them. | +| Keep scheduling outside the harness. | The nightly launcher and its caller own recurrence, locking, Git selection, and secrets. | +| Require explicit gates. | The harness records quality and performance but does not impose a global score. | -Rebuild `retriever harness` as the internal benchmark runner for Retriever -engineers. The harness should run the new library direction end to end: +## Product Boundary -1. Ingest documents through the same shared code path as `retriever ingest`. -2. Query the resulting LanceDB table through the same shared code path as - `retriever query`. -3. Run full BEIR-style evaluation. -4. Emit stable `summary_metrics` and machine-readable artifacts for humans, - agents, and downstream reporting. +Use `retriever ingest` and `retriever query` for direct product workflows. Use +`retriever harness` for registered benchmark and evaluation work. -This is a total revamp, not a compatibility wrapper around the old harness. The -retired pipeline CLI is not part of the design. The old `sweep`, `compare`, and -graph-pipeline command builder can be rewritten or removed if they get in the -way. +The harness calls the same library workflow code or corresponding service APIs +used by the product. It owns benchmark resolution, run lifecycle, evaluation, +metric gates, and artifacts. Retriever owns ingest and query behavior. -The most important design choice: use a small typed benchmark registry in code, -not a sprawling YAML configuration system. YAML/runfiles can exist as an escape -hatch for one-off experiments, but day-to-day benchmark definitions should live -in the repository as reviewed Python objects. +## Non-Goals -The second most important design choice: stdout is not an API. The harness can -print concise human summaries, but agents and orchestrators must rely on stable -files such as `status.json`, `results.json`, `beir_metrics.json`, and -`query_results.jsonl`. +The harness is not: -## Research Notes +- a public supported product API +- a scheduler, retry system, or secret distributor +- a benchmark history database or regression-policy engine +- a general deployment manager +- a compatibility wrapper for retired pipeline or sweep commands -The industry-standard pattern is not "one huge YAML that can do anything." The -useful patterns are: - -- **Separate the benchmark runner from the system under test.** MLPerf - Inference uses a load generator outside the submitted system under test, so - timing, workload generation, logging, and validation are not buried inside the - model/backend implementation. Retriever should mirror that boundary: the - harness owns run lifecycle, timing, query iteration, metrics, and artifacts; - Retriever ingest/query own retrieval behavior. -- **Keep benchmark suites close to code.** ASV benchmarks Python packages over - time and treats benchmark definitions as part of the project rather than as a - pile of external config files. For Retriever engineers, code-owned benchmark - specs are easier to review, type-check, and refactor with the library. -- **Emit durable artifacts, not just console output.** Benchmark systems such as - ASV and Phoronix Test Suite treat saved run artifacts, histories, and - comparisons as first-class outputs. Retriever should make `results.json`, - `summary_metrics`, BEIR metrics, runfiles, and per-query outputs the product - of a run. -- **Use BEIR conventions for retrieval quality.** BEIR evaluates retrieval with - NDCG@k, MAP@k, Recall@k, Precision@k, and custom metrics like MRR. Retriever - harness should emit that full metric family, then choose a small subset for - `summary_metrics`. -- **Avoid adopting Hydra-style composition unless we truly need it.** Hydra is - powerful, but its value comes from config groups, launchers, sweepers, output - directory patterns, and multi-file composition. That is more machinery than - this internal harness needs right now. - -Reference links are at the end of this document. - -## Goals - -- Make `retriever harness run ` the standard internal entry point - for end-to-end Retriever benchmarks. -- Default to local, in-process execution suitable for developer laptops, - workstations, and agents. -- Run full BEIR evaluation for benchmark datasets that have qrels. -- Emit a compact, stable `summary_metrics` object every run. -- Emit detailed artifacts for debugging, reproducibility, and future reporting. -- Make benchmark definitions easy to discover, review, and extend. -- Avoid duplicating the `retriever ingest` and `retriever query` option surface. -- Avoid YAML sprawl and config inheritance chains. -- Make ablations simple enough for an agent to run without guessing paths or - settings. -- Make every run inspectable by agents and orchestrators without parsing stdout. - -## Non-goals - -- Do not preserve the graph-pipeline harness path. -- Do not preserve current `sweep` or `compare` behavior for its own sake. -- Do not introduce a user-facing `--engine` flag. -- Do not adopt Hydra, MLflow Projects, W&B Sweeps, or another orchestration - framework in phase one. -- Do not make this a public supported product API. -- Do not couple benchmark execution to Slack or another reporting transport. -- Do not add recurring scheduling or deployment infrastructure to the harness. -- Do not make CLI text formatting part of the run contract. -- Do not make pytest the sole validation strategy for the harness. Unit tests - protect artifact and reporting contracts; real harness exit codes and - artifacts validate end-to-end behavior. - -## Users - -- Retriever engineers validating ingest/query behavior. -- Performance owners running throughput and quality ablations. -- Agents asked to run a named benchmark or ablation. -- External automation that needs stable result files. - -## Design Principles - -- **Artifact-first:** every durable result is written to a documented file. - Console output is for humans only. -- **Code-owned defaults:** recurring benchmarks and runsets live in reviewed - Python registry entries. -- **Small overrides:** agents and engineers can apply targeted `--set` changes - without creating new config files. -- **One real execution path:** the harness calls shared Retriever workflow code - directly instead of shelling out to the CLI per phase. -- **Phase visibility:** long runs expose current phase, elapsed time, and - partial artifacts as they progress. -- **Typed failures:** failures include phase, reason, retryability, and pointers - to the relevant logs/artifacts. -- **Orchestration-neutral:** the harness has no dependency on MLflow, Airflow, - Argo, Ray Tune, or similar systems, but its inputs and outputs are easy for - those systems to wrap later. - -## Recommended Shape - -### One Harness Execution Path - -The harness should run in the current Python process and call shared Retriever -workflow functions directly. - -There should not be a user-facing `--engine inprocess|subprocess` flag. That -idea came from separating two implementation options: - -- direct Python calls into shared ingest/query workflow code -- shelling out to the `retriever` CLI as subprocesses - -For this harness, subprocess mode is not worth the extra concept. It would make -query benchmarks less honest by repeatedly measuring process startup and model -warmup unless carefully special-cased. It would also complicate artifacts and -error handling. - -Instead: - -- The harness always uses the shared Python workflow path. -- The harness writes replay hints that show equivalent `retriever ingest` and - `retriever query` commands where useful. -- Console-script smoke tests can live elsewhere. - -### Local vs Batch Is an Ingest Setting - -There is still a real distinction between local and Ray-backed ingest: - -- `local`: maps to `ingest.run_mode = "inprocess"`. -- `batch`: maps to `ingest.run_mode = "batch"` and may set Ray/tuning kwargs. - -Expose that as benchmark intent, not as a generic execution engine: - -```bash -retriever harness run jp20_smoke -retriever harness run jp20_beir --mode batch -``` - -Internally, `--mode batch` only changes the ingest/query benchmark spec fields -that need to change. Query execution should still use one constructed Retriever -object per run so measured query latency is not dominated by setup. - -### Code-Owned Benchmark Registry - -Create a canonical registry, for example: - -```text -nemo_retriever/src/nemo_retriever/harness/ - benchmark_specs.py - benchmark_registry.py - resolution.py - execution.py - beir_runner.py - metrics.py - metric_gates.py - artifact_writer.py - json_io.py - runfile.py - runsets.py - diff.py -``` - -Core types: - -```python -@dataclass(frozen=True) -class DatasetSpec: - name: str - path: str - query_csv: str | None = None - input_type: str = "pdf" - beir_loader: str | None = None - beir_doc_id_field: str = "pdf_page" - beir_ks: tuple[int, ...] = (1, 3, 5, 10) - - -@dataclass(frozen=True) -class BenchmarkSpec: - name: str - dataset: str - ingest: Mapping[str, Any] - query: Mapping[str, Any] - evaluation: Mapping[str, Any] - summary_keys: tuple[str, ...] - tags: tuple[str, ...] = () -``` - -The registry should include named benchmarks such as: - -- `jp20_smoke` -- `jp20_beir` -- `bo767_beir` -- `financebench_beir` -- `bo10k_beir_fast_text` -- `earnings_beir` once its query/qrels file is available in the repo or dataset - mount - -Commands: - -```bash -retriever harness list -retriever harness show jp20_beir -retriever harness run jp20_beir -retriever harness run jp20_beir --set query.top_k=20 -retriever harness run jp20_beir --set ingest.profile=fast-text -``` - -This is the default workflow. Engineers add benchmarks by editing reviewed -Python specs, not by dropping new YAML files around the repo. - -### Narrow Runfiles - -Support one optional runfile path for reproducible one-off runs, agent -handoffs, and orchestrator inputs: - -```bash -retriever harness run --runfile nemo_retriever/harness/runfiles/jp20_beir.json -``` - -Runfiles should be intentionally small JSON/YAML objects: - -```json -{ - "schema_version": 1, - "name": "jp20_beir_expected", - "benchmark": "jp20_beir", - "mode": "local", - "require": ["files==20", "pages==1940", "query_count==115"], - "set": { - "query.top_k": 10 - } -} -``` - -Rules: - -- A runfile must extend a named registry benchmark. -- A runfile cannot define a new schema from scratch in phase one. -- The resolved benchmark spec is always written into artifacts. -- The source runfile payload is copied into `runfile.json` in the artifact - directory. -- Runfiles are for reproducible run requests and agent instructions, not the - default source of truth for recurring benchmark definitions. -- During `--dry-run`, gates for unavailable execution metrics are skipped and - reported in `results.json`; static dataset gates are still evaluated. - -Machine-specific document and query paths belong in a separate, untracked -dataset path map. `run-files --dataset-paths ` applies that map after -runfile overrides and before CLI `--set` overrides. Passing one runfile creates -a one-dataset session; passing multiple runfiles creates a suite session. -`run-files` owns dry-run behavior for the session so a session cannot mix -planned and executed children. - -### Ablations - -Prefer explicit runsets in code for recurring ablations. Phase one runsets are -intentionally literal lists of named benchmarks; they do not expand matrices yet. - -```python -RunSet( - name="jp20_profile_x_rerank", - runs=("jp20_beir", "jp20_beir_rerank"), -) -``` - -Commands: - -```bash -retriever harness list --runsets -retriever harness run-set jp20_profile_x_rerank -``` - -For phase one, `run-set` can simply expand to individual benchmark runs and -write `expanded_runs.json`. We do not need a separate legacy `sweep` concept. - -## Agent And Orchestration Contract - -Agents will use this harness for continuous benchmark/performance loops. They -may inspect artifacts, compare metrics, choose the next ablation, and rerun. -That changes the contract: the harness must behave like a protocol, not just a -pretty CLI. - -### Required Inputs - -Every run should support: - -```bash -retriever harness run \ - --run-id \ - --output-dir \ - --set query.top_k=20 \ - --dry-run -``` - -Rules: - -- `--run-id` is optional for humans but required by orchestrated jobs that need - deterministic artifact paths. -- `--output-dir` controls exactly where artifacts are written. -- `--dry-run` writes the resolved benchmark and planned artifacts without - running ingest or query. -- All prompts are forbidden. The command must be non-interactive. - -### Required Live State - -Write `status.json` early and update it at phase transitions: - -```json -{ - "run_id": "jp20_beir_20260623_120000", - "benchmark": "jp20_beir", - "status": "running", - "phase": "query", - "started_at": "2026-06-23T12:00:00Z", - "updated_at": "2026-06-23T12:34:56Z", - "artifact_dir": "/artifacts/jp20_beir_20260623_120000", - "results_path": null, - "failure": null -} -``` - -Allowed statuses: - -- `planned` -- `running` -- `complete` -- `failed` - -Allowed phases: - -- `resolve` -- `ingest_plan` -- `ingest` -- `query_plan` -- `query` -- `evaluate` -- `write_artifacts` - -Also write append-only `events.jsonl` for phase changes and major milestones. -Agents can tail this file or poll `status.json`; they should never scrape -stdout. - -### Failure Shape - -On failure, `status.json` and `results.json` should include: - -```json -{ - "failed_phase": "query", - "failure_reason": "lancedb_table_missing", - "retryable": false, - "message": "LanceDB table nv-ingest was not found", - "debug_artifacts": [ - "ingest_plan.json", - "logs/query.log" - ] -} -``` - -Failure reasons should be stable enough for agents to branch on. Examples: - -- `invalid_benchmark` -- `invalid_override` -- `dataset_missing` -- `ingest_plan_failed` -- `ingest_failed` -- `query_plan_failed` -- `query_failed` -- `evaluation_failed` -- `metric_gate_failed` -- `artifact_write_failed` - -### Exit Codes - -Use coarse but stable exit codes: - -- `0`: success -- `2`: invalid benchmark/config/override -- `3`: dataset or input missing -- `10`: ingest failure -- `11`: query failure -- `12`: evaluation failure -- `20`: metric gate failure -- `30`: artifact write failure -- `70`: unexpected internal error - -Agents should poll `status.json` while a run is active, read `results.json` when -it is terminal, and use exit codes for coarse process control. - -### Metric Gates - -Support gates after metrics are stable: - -```bash -retriever harness run jp20_beir \ - --require recall_5>=0.80 \ - --require ndcg_10>=0.70 \ - --require query_latency_p95_ms<=150 -``` - -Gate failures should still write all artifacts and exit with -`metric_gate_failed`. - -## Reuse Points - -Use the current root CLI's shared implementation functions as the source of -truth: - -- `nemo_retriever.adapters.cli.sdk_workflow.resolve_ingest_plan()` -- `nemo_retriever.adapters.cli.sdk_workflow.ingest_documents()` -- `nemo_retriever.adapters.cli.sdk_workflow.query_documents()` - -Add one query planning helper so the harness can avoid rebuilding a Retriever -for every query: - -```python -@dataclass(frozen=True) -class ResolvedQueryPlan: - top_k: int - lancedb_uri: str - table_name: str - embed_kwargs: dict[str, Any] - rerank: bool - rerank_kwargs: dict[str, Any] - - def create_retriever(self) -> Retriever: ... - - -def resolve_query_plan(...) -> ResolvedQueryPlan: ... -``` - -Then: - -- `retriever query` remains a thin single-query CLI wrapper. -- `retriever harness` creates one Retriever per benchmark run and queries the - full BEIR query set. - -Validate nested `ingest` and `query` keys against the signatures of these -workflow helpers. Unknown keys should fail before execution with suggestions. -Do not copy Typer option definitions into the harness. - -## Run Lifecycle - -For each benchmark run: - -1. Resolve a `BenchmarkSpec` from the registry plus CLI `--set` overrides or a - tiny runfile. -2. Resolve dataset paths and query/qrels files. -3. Create a run artifact directory. -4. Set the run's LanceDB URI under the artifact directory unless explicitly - overridden. -5. Dry-run ingest with `resolve_ingest_plan()` and write a redacted plan. -6. Execute ingest with `ingest_documents()` and measure wall-clock time. -7. Count input files, pages, and LanceDB rows. -8. Build one Retriever from `resolve_query_plan()`. -9. Execute optional warmup queries. -10. Execute the full measured BEIR query set. -11. Write a BEIR/TREC runfile and raw per-query hits. -12. Compute BEIR metrics. -13. Write `results.json` with `summary_metrics`. -14. Write `status.json` with `status = "complete"`. -15. Print a concise terminal summary for humans. - -## BEIR Evaluation - -BEIR-style evaluation is a phase-one requirement, not a later add-on. - -For every BEIR benchmark, write: - -- `beir_run.trec`: runfile suitable for reranking/debugging. -- `beir_metrics.json`: full metric family. -- `query_results.jsonl`: query text, latency, ranked hits, and resolved doc IDs. - -Metric family: - -- `ndcg@k` -- `map@k` -- `recall@k` -- `precision@k` -- `mrr@k` when supported - -Default k values: - -```python -(1, 3, 5, 10) -``` - -`summary_metrics` should include the small set that engineers and agents need -first: - -```json -{ - "files": 20, - "pages": 496, - "rows_processed": 12345, - "ingest_secs": 321.5, - "pages_per_sec_ingest": 1.54, - "query_count": 200, - "query_latency_p50_ms": 42.1, - "query_latency_p95_ms": 87.3, - "ndcg_10": 0.72, - "recall_5": 0.81, - "recall_10": 0.86 -} -``` - -The exact values above are illustrative. The key names should be stable. - -## Artifact Contract - -Per run: - -- `results.json`: authoritative run result. -- `status.json`: current and final phase/status state. -- `events.jsonl`: append-only phase changes and run milestones. -- `runfile.json`: original runfile payload when one was used. -- `resolved_benchmark.json`: fully resolved benchmark spec. -- `ingest_plan.json`: redacted dry-run ingest plan. -- `query_plan.json`: resolved query execution plan. -- `run.log`: captured lower-level output and full exception tracebacks. -- `beir_metrics.json`: full BEIR metrics. -- `beir_run.trec`: BEIR/TREC runfile. -- `query_results.jsonl`: per-query hits and latency. -- `environment.json`: git SHA, package version, Python, host, GPU count, CUDA - driver, Ray version where available. -- `lancedb/`: default vector store. - -Session/runset: - -- `session_summary.json`: one row per run, centered on `summary_metrics`. -- `expanded_runs.json`: resolved run order for runsets. - -Do not make `compare` phase-one critical. Comparing runs can be rebuilt later -from `results.json` and `session_summary.json`. - -`results.json` should include relative pointers to every artifact path so -external systems can ingest one file, discover the rest, and move the complete -run directory without rewriting its manifest. - -## CLI - -Phase-one CLI: - -```bash -retriever harness list -retriever harness show -retriever harness run -retriever harness run --run-id --output-dir -retriever harness run --mode local -retriever harness run --mode batch -retriever harness run --set query.top_k=20 -retriever harness run --set ingest.profile=fast-text -retriever harness run --runfile /tmp/ablation.yaml -retriever harness run-set -retriever harness run-files --dataset-paths /local/dataset_paths.yaml ... -retriever harness post-slack --preview ... -retriever harness post-slack ... -retriever harness diff --json -``` - -Notes: - -- `--mode local` is the default and maps to Retriever ingest/query settings. -- `--mode batch` is benchmark intent, not a separate harness engine. -- `--set` parses values with YAML/JSON scalar semantics. -- `--dry-run` should be available on `run` and `run-set`. -- `--json` on read-only commands writes machine-readable output to stdout. - Human formatting remains non-contractual. -- `post-slack` consumes existing artifacts and never runs ingest or query. -- No CLI command installs or owns recurring scheduling. - -Defer or remove: - -- legacy `sweep` -- legacy `compare` -- legacy `nightly` -- legacy `portal` -- legacy `runner` - -These can return after the new artifact contract is stable. - -`diff` can be a small new command, not the legacy compare implementation. It -should read two artifact directories and emit changed `summary_metrics` plus -selected BEIR deltas. Agents need this primitive more than humans need a full -reporting UI. - -## Functional Validation - -The harness should be validated both as a library contract and as an evaluation -runner. Focused tests protect resolution, artifact, and reporting behavior; -developers should also prove execution behavior by running harness commands and -inspecting stable artifacts. - -Minimum local validation commands: - -```bash -retriever harness list --json -retriever harness show jp20_beir --json -retriever harness run jp20_beir --dry-run --output-dir /tmp/retriever-harness-dry-run -retriever harness run jp20_beir --dry-run --require 'files>=20' -``` - -Functional validation should assert: - -- expected exit code -- parseable `--json` output for read-only commands -- `status.json` exists and has the expected final status -- `events.jsonl` exists and includes phase transitions -- `resolved_benchmark.json` and `results.json` exist -- `run.log` exists for non-dry execution runs and contains suppressed - lower-level stdout/stderr -- invalid overrides exit with code `2` -- missing datasets exit with code `3` -- metric gate failures write artifacts and exit with code `20` - -Longer validation runs should use real benchmark datasets and BEIR evaluation: - -```bash -retriever harness run jp20_smoke --output-dir /tmp/retriever-harness-jp20-smoke -retriever harness run jp20_beir --output-dir /tmp/retriever-harness-jp20-beir -retriever harness run jp20_beir \ - --output-dir /tmp/retriever-harness-jp20-beir-gated \ - --require 'files==20' \ - --require 'pages==1940' \ - --require 'query_count==115' \ - --require 'recall_5>=0.85' \ - --require 'ndcg_10>=0.75' -``` - -`jp20_smoke` is a cheap fast-text ingest check over the JP20 corpus and does not -run BEIR queries. `jp20_beir` runs the full JP20 end-to-end path: ingest, BEIR -query iteration, BEIR runfile output, and summary recall/NDCG metrics. -Known dataset facts, benchmark result ranges, and suggested gates should live in -`nemo_retriever/harness/EXPECTED_RESULTS.md`, not in benchmark Python code. - -The harness may include tiny no-GPU fixtures or fake benchmark specs to make -developer validation cheap, but execution changes still require functional -validation through the CLI and artifact contract. - -## Implementation Plan - -### Phase 1: New Core Runner - -- Add typed `DatasetSpec`, `BenchmarkSpec`, and `RunSet` models. -- Add benchmark registry with a small set of named BEIR benchmarks. -- Add `resolve_query_plan()` to shared CLI workflow code. -- Implement `retriever harness list`, `show`, `run`, and `run-set`. -- Implement full BEIR query execution and metrics output. -- Implement stable `summary_metrics`. -- Implement `status.json`, `events.jsonl`, typed failure payloads, and stable - exit codes. -- Persist suppressed non-dry execution logs in `run.log`. -- Support explicit `--require` gates. -- Add a cheap functional validation path that exercises artifact writing without - requiring a large GPU run. -- Retire legacy pytest harness coverage that assumes the old graph-pipeline - harness design. Functional CLI runs are the validation contract. - -### Phase 2: Real Dataset Validation - -- Run `jp20_smoke` locally. -- Run one full BEIR benchmark on expected hardware. -- Validate that `summary_metrics`, BEIR metrics, and runfiles are sufficient for - debugging failures. -- Tune default benchmark registry entries. - -### Phase 3: Reporting - -- Rebuild `summary` around the new artifact contract. -- Build `diff --json` around `summary_metrics` and BEIR deltas. -- Rebuild richer `compare` only if the team needs it after stable artifacts - exist. -- Keep Slack as an optional post-hoc artifact sink. -- Keep recurring execution and deployment in separately reviewed infrastructure. - -## Open Decisions - -- Which named benchmarks should be phase-one defaults? -- Should `bo10k` be included immediately or wait until smaller BEIR datasets are - stable? -- What is the canonical page/doc ID mapping for each dataset's BEIR qrels? -- Which query latency metric should gate regressions: p50, p95, or both? -- Do we want `--mode batch` to be accepted for all benchmarks or only for specs - that declare batch-safe tuning defaults? - -## Risks - -- If we expose too much arbitrary config, this becomes config hell. Keep the - registry as the source of truth and make runfiles extend named benchmarks. -- If query execution shells out per query, latency metrics become mostly startup - noise. Build one Retriever per run. -- If BEIR ID mapping is inconsistent across datasets, summary quality metrics - will be misleading. Dataset specs need explicit doc ID policy. -- If `summary_metrics` changes frequently, downstream reporting and agents will - become brittle. Treat key names as a contract. -- If agents need to scrape stdout, the artifact contract failed. Add or fix a - machine-readable file instead of documenting text output. - -## Acceptance Criteria - -- `retriever harness list` shows named built-in benchmarks. -- `retriever harness show jp20_beir --json` emits the resolved benchmark as - machine-readable JSON. -- `retriever harness run jp20_beir --dry-run` resolves ingest, query, - evaluation, and artifact paths without touching graph-pipeline code. -- `retriever harness run jp20_beir` runs ingest, queries the full BEIR query - set, writes BEIR outputs, and emits stable `summary_metrics`. -- `retriever harness run-set ` expands a code-owned ablation and writes - `expanded_runs.json`. -- `retriever harness run-files` runs one or more checked-in requests with an - optional machine-local dataset path map. Real children execute sequentially - in fresh processes while retaining one terminal session summary; dry-runs - stay in the parent process. -- `retriever harness check-vidore-access` validates authenticated access to the - ViDoRe queries, qrels, and corpus partitions without downloading them. -- `retriever harness post-slack --preview` reads artifacts without requiring a - webhook or contacting Slack. -- Every run writes `status.json`, `events.jsonl`, and `results.json`. -- Failed runs write typed failure data without requiring stdout inspection. -- Unknown `--set` keys fail before execution. -- The harness has no user-facing `--engine` flag. -- The harness does not duplicate Typer options from `retriever ingest` or - `retriever query`. -- No phase-one code path invokes a retired CLI adapter. -- CLI text formatting is explicitly non-contractual; machine consumers use - artifact files or `--json` read-only commands. -- Validation combines focused contract tests with functional, artifact-driven - benchmark execution. - -## References - -- [MLPerf Inference benchmark suite](https://github.com/mlcommons/inference): - benchmark suite for measuring inference speed across deployment scenarios. -- [MLPerf Inference paper](https://arxiv.org/abs/1911.02549): describes the - LoadGen/SUT split, accuracy/performance modes, and reproducibility goals. -- [BEIR repository](https://github.com/beir-cellar/beir): heterogeneous IR - benchmark and evaluation framework. -- [ASV documentation](https://asv.readthedocs.io/en/latest/): Python benchmark - suites over time with saved results. -- [Phoronix Test Suite](https://www.phoronix-test-suite.com/): benchmark - profiles, suites, batch operation, saved results, and comparisons. -- [Hydra configuration overview](https://hydra.cc/docs/configure_hydra/intro/): - useful reference for why we should avoid importing a full multi-file config - composition framework unless the harness truly needs it. +Git history retains the original implementation PRD and its research notes. diff --git a/nemo_retriever/developer_docs/ocr_cross_page_batching/README.md b/nemo_retriever/developer_docs/ocr_cross_page_batching/README.md new file mode 100644 index 0000000000..9e1f3e7137 --- /dev/null +++ b/nemo_retriever/developer_docs/ocr_cross_page_batching/README.md @@ -0,0 +1,521 @@ +# Local Nemotron OCR v2 cross-page batching + +Issue: [#2323](https://github.com/NVIDIA/NeMo-Retriever/issues/2323) + +## Decision + +Batch compatible local OCR crops across every page row delivered to one +`OCRActor`. Split table (`word`) and paragraph jobs, bound each model list by +`inference_batch_size`, and stitch ordered results back to their source row and +detection. + +The old local path created and invoked jobs inside the page-row loop, so sparse +pages reached the persistent model as singleton calls. Collection now spans the +Ray batch while retaining row identity. + +```mermaid +flowchart LR + R["Ray page-row batch"] --> J["collect crops with row identity"] + J --> W["word queue"] + J --> P["paragraph queue"] + W --> B1["bounded model lists"] + P --> B2["bounded model lists"] + B1 --> S["ordered stitch"] + B2 --> S +``` + +Three differently scoped controls remain independent: + +| Control | Unit | Owner | Controlled actor A/B | +|---|---|---|---:| +| Ray supply batch | page rows per actor call | Ray graph | 32 | +| Outer OCR list | crops per model call | `OCRActor` | 8 | +| Internal detector batch | images per detector forward | `nemotron-ocr` | 8 | + +The patch changes only the middle layer. It does not change defaults or wire +the actor setting into Nemotron's internal detector policy. + +## Evidence + +### Correctness + +The red/green regression uses one pandas batch with two page rows, one chart per +row, `inference_batch_size=2`, and a recording list-input model. + +| Behavior | Upstream | Patch | +|---|---|---| +| Paragraph calls | `[page A]`, `[page B]` | `[page A, page B]` | +| Invocation count | 2 | 1 | +| Crops per invocation | `[1, 1]` | `[2]` | +| Row, bbox, and fake output identity | preserved | preserved | + +The four focused tests also cover merge-level separation, bounded chunking, +empty and malformed pages, native-text preservation, batch exceptions, +wrong-result-count fallback, and per-crop failure isolation. Fallback occurs +only on an exception or wrong result count. + +### Result summary + +![Issue 2323 OCR batching and deployment validation](proof-summary.svg) + +The actor/model, one-GPU BO767, and one-GPU local-HF ViDoRe comparisons are +attributable A/B evidence for this patch. Every GPU measurement in this report +used NVIDIA H100 80GB HBM3 hardware. NIM and service results are deployment +context only: their backend paths are unchanged and their GPU counts are shown +explicitly. + +The controlled GPU A/B traversed Ray Data, the real `OCRActor`, one persistent +local wrapper, and the locked Nemotron OCR v2 model. It used 128 fixed real +crops (64 tables and 64 charts), one warmup, and five measured trials on one +H100 80GB. + +| Measurement | Upstream | Patch | Effect | +|---|---:|---:|---:| +| Model invocations per 128 crops | 128 x scalar | 16 x list-of-8 | 87.5% fewer | +| Median actor throughput | 34.418 crops/s | 55.787 crops/s | 1.621x | +| Median model throughput | 40.654 crops/s | 73.389 crops/s | 1.805x | + +### Whole-ingest speedup + +**On BO767, the patch reduced mean whole-ingest runtime by 7.75% and increased +throughput by 8.41%.** The harness timer starts immediately before +`run_ingest_workflow(...)` and stops after it returns. It therefore includes the +complete batch ingest through extraction, OCR, embedding, and LanceDB indexing; +query evaluation is outside this timer. + +BO767 supplied a sustained OCR-heavy test: 767 PDFs, 54,730 pages, 79,233-79,234 +extracted rows, and 991 scored queries. It used Ray page-row batches of 24, +OCR crop lists capped at 8, one H100, and counterbalanced order (`upstream, +patch, patch, upstream`). + +| Run | Ingest runtime | Throughput | Rows | +|---|---:|---:|---:| +| Upstream 1 | 1,612.811 s | 33.935 pages/s | 79,234 | +| Patch 1 | 1,476.246 s | 37.074 pages/s | 79,233 | +| Patch 2 | 1,490.440 s | 36.721 pages/s | 79,234 | +| Upstream 2 | 1,603.227 s | 34.137 pages/s | 79,234 | + +The configuration means were 1,608.019 seconds upstream and 1,483.343 seconds +patched: **124.676 seconds (7.75%) less ingest time** and **8.41% higher +throughput**. Both matched pairs were faster (`-8.47%` and `-7.03%`). This is +the whole-ingest speedup result for the corrected batching behavior. + +The retained Ray progress trace also identifies why the gain reaches the +whole-ingest timer here. The final OCR/embedding frontier completed at a mean +1,270.069 seconds upstream and 1,141.847 seconds patched, approximately 128.2 +seconds earlier. Mean whole ingest moved 124.7 seconds, while the remaining +indexing tail stayed similar (338.0 seconds upstream and 341.5 seconds +patched). Ray reports these progress snapshots at roughly ten-second cadence, +so they establish the critical-path relationship rather than sub-second stage +timing. + +The shorter `vidore_v3_computer_science_beir` run is retained as a useful +limit: it ran the same complete ingest boundary over 1,360 pages, but was too +short to separate the patch from warm-run variation. + +| Run | Ingest runtime | Throughput | +|---|---:|---:| +| Upstream 1 (coldest) | 181.616 s | 7.488 pages/s | +| Patch 1 | 168.388 s | 8.077 pages/s | +| Patch 2 | 169.506 s | 8.023 pages/s | +| Upstream 2 | 168.330 s | 8.079 pages/s | + +Its first pair favored the patch by 7.3%, while the counterbalanced warm pair +was effectively tied. That does not contradict the BO767 result; it shows that +a two-document run is underpowered for a whole-ingest performance claim. + +
+Separate follow-up evidence: default multi-GPU scheduling + +### Default multi-GPU batch scaling + +A separate current-commit run exposed two, four, or eight H100s to batch mode +and left every worker count on `auto`. The same eight ViDoRe runfiles, dataset +order, cache, lock, package versions, and OCR crop-list cap of 8 were used at +each point. The eight-GPU comparison was counterbalanced as patch, exact +`upstream/main`, patch. + +| Source | Visible GPUs | Ingest runtime | Throughput | Recall@5 | nDCG@10 | Peak memory on one GPU | +|---|---:|---:|---:|---:|---:|---:| +| Patch | 2 | 1,142.447 s | 16.852 pages/s | 0.4615 | 0.5174 | 34,938 MiB | +| Patch | 4 | 951.175 s | **20.240 pages/s** | 0.4650 | 0.5205 | 36,544 MiB | +| Patch A | 8 | 989.582 s | 19.455 pages/s | 0.4633 | 0.5188 | 43,034 MiB | +| Exact upstream | 8 | 965.978 s | 19.930 pages/s | 0.4635 | 0.5189 | 39,532 MiB | +| Patch B | 8 | 991.837 s | 19.410 pages/s | 0.4620 | 0.5170 | 43,208 MiB | + +All five runs passed all eight datasets without an OCR error or OOM. The two +patch eight-GPU runs differed by only 0.23% in ingest time. Their mean was +990.710 seconds (19.433 pages/s), so the exact upstream control was 2.50% +faster and the four-GPU patch run was 4.16% faster than the patch eight-GPU +mean. The supplied nightly reference, 959.58 seconds and 20.06 pages/s, was +close to the exact upstream control (-0.65% throughput). + +This proves the corrected path executes successfully when eight GPUs are +exposed, but it does **not** demonstrate positive eight-GPU scaling for this +suite. Six datasets were fastest on four GPUs, `industrial` was fastest on +eight, and `physics` was fastest on two: + +| Dataset | 2 GPU | 4 GPU | 8 GPU patch mean | Fastest | +|---|---:|---:|---:|---:| +| `computer_science` | 102.007 s | **91.290 s** | 104.433 s | 4 GPU | +| `energy` | 138.966 s | **115.108 s** | 122.040 s | 4 GPU | +| `finance_en` | 199.473 s | **160.173 s** | 163.430 s | 4 GPU | +| `finance_fr` | 150.908 s | **118.884 s** | 127.161 s | 4 GPU | +| `hr` | 95.646 s | **92.709 s** | 103.058 s | 4 GPU | +| `industrial` | 241.673 s | 175.424 s | **152.791 s** | 8 GPU | +| `pharmaceuticals` | 123.499 s | **103.250 s** | 110.217 s | 4 GPU | +| `physics` | **90.275 s** | 94.337 s | 107.579 s | 2 GPU | + +The default resource policy creates three initial OCR actors per visible GPU +and reserves 0.1 GPU per actor. NVML process sampling observed 6, 12, and 24 +OCR actors per dataset at the two-, four-, and eight-GPU points, but those OCR +processes occupied only 2, 2, and 3 physical GPUs respectively. All eight GPUs +did receive work across the complete eight-GPU pipeline; OCR placement itself +was not even. Together with the dataset crossover, this makes actor +startup/placement and crop-list occupancy the leading scale explanation, not +GPU capacity. A falsifiable next investigation is to record outer list-size +histograms and actor startup time at 2/4/8 GPUs before changing Ray defaults. +Worker policy is deliberately outside this correctness patch. + +
+ +### Retrieval non-regression + +The same ViDoRe run evaluated 1,290 queries and 6,294 qrels. Because repeated +full runs changed page text and rankings even within one configuration, every +query was embedded once and those same vectors were applied to all four stored +corpora using deterministic exact top-10 search. + +| Metric | Upstream mean | Patch mean | Delta | +|---|---:|---:|---:| +| nDCG@10 | 0.7093766 | 0.7093501 | -0.0000265 | +| Recall@5 | 0.6000979 | 0.6000979 | 0.0000000 | +| Recall@10 | 0.7306446 | 0.7307415 | +0.0000969 | + +Both matched pairs had identical Recall@5; the warm pair also had identical +Recall@10. Mean top-10 overlap was 99.91%. No retrieval regression was observed +after controlling query and index execution randomness. + +BO767's configuration means were likewise neutral within run variance: + +| Metric | Upstream mean | Patch mean | Delta | +|---|---:|---:|---:| +| nDCG@10 | 0.752780 | 0.753359 | +0.000579 | +| Recall@5 | 0.850151 | 0.849647 | -0.000505 | +| Recall@10 | 0.899092 | 0.900101 | +0.001009 | + +Cross-configuration mean top-10 overlap was 98.58%; upstream-versus-upstream +overlap was lower at 98.22%. Exact stored text also varied substantially across +identical upstream runs, while structural row overlap remained above 99.4%. +The one-row variation appeared in one patch replicate only, so it was not a +stable patch effect. Every run had the same two non-fatal overlength embedding +failures, zero OCR warnings, and zero OOMs. + +### Full ViDoRe v3 one-GPU local-HF A/B + +The exact upstream parent and patched local path each completed all eight +ViDoRe v3 runfiles on one H100: 189 PDFs, 19,252 pages, 19,252 output rows, and +14,514 scored queries. Both runs used the same lock, persistent local models, +dataset order, cache, default batch worker policy, Ray page-row batches of 24, +and OCR crop lists capped at 8. + +| Source | Ingest runtime | Throughput | Recall@5 | nDCG@10 | +|---|---:|---:|---:|---:| +| Exact upstream parent | 2,074.654 s | **9.280 pages/s** | 0.4633 | 0.5184 | +| Patch | 2,077.167 s | **9.268 pages/s** | 0.4645 | 0.5204 | +| Patch effect | +2.513 s | **-0.12%** | +0.0012 | +0.0020 | + +The one-GPU full-suite result is effectively flat. It rules out a claimed +doubling of ViDoRe ingest throughput, while the small measured difference is +too narrow to interpret as a regression from one run per side collected two +days apart. ViDoRe therefore supplies full-suite correctness and quality +evidence; BO767 remains the demonstrated whole-ingest speedup. + +The Ray progress trace explains the difference from BO767. Across the eight +patched ViDoRe runs, OCR completed after an average 66.9 seconds from the +execution-plan start, while page-image embedding completed after an average +242.5 seconds. Depending on the dataset, embedding continued for another +60-381 seconds after OCR had finished. The exact-parent runs showed the same +pipeline shape. Faster OCR is therefore off the ViDoRe completion path under +this page-granularity, vision-language embedding configuration. + +Per-dataset throughput changes ranged from -5.47% (`hr`) to +4.39% +(`finance_fr`) without a consistent direction. No run failed or OOMed. + +### Self-hosted NIM and service compatibility + +The current `service-mode.compose.yaml` pins the relevant core services to: + +| Service | Image | Local digest | +|---|---|---| +| Page elements and table structure | `nemotron-object-detection:2.0.0` | `sha256:de21875223e4cc26b79e44a4f30ff06dcc8fe97c731c6b9f500a48eb54fa99bf` | +| OCR | `nemotron-ocr-v2:2.0.0` | `sha256:3ac2ea60a83d7aab6275e08ea27a959de46fdab0689594f54f8374f590f416b8` | +| Embedding | `llama-nemotron-embed-vl-1b-v2:1.12.0` | `sha256:58c40b920840be6e2f4ad5d77c32c65d61e048070fe45d51fb4bdb6f84a71e21` | + +The first compatibility smoke placed the four containers on separate H100s. +A subsequent full-suite comparison placed page detection, table structure, +OCR, and embedding together on GPU 0. Docker device requests for every +container named only device `0`; sampled memory and utilization on GPUs 1-7 +remained zero. GPU 0 peaked at 16,331 MiB and 100% utilization. The Compose +default `NIM_PIPELINE_MAX_BATCH_SIZE=1` was retained for page detection, table +structure, and OCR; this validation did not tune NIM internals or use an +unbounded request. + +A direct OCR endpoint probe sent one ordered list containing two fixed +paragraph crops. It received two results in the same order and preserved both +crop anchors. This proves the deployed OCR NIM's bounded list-input contract; +it does not prove that the unchanged application remote path forms cross-page +lists. + +#### Full one-GPU batch comparison + +The one-GPU NIM deployment completed the same eight ViDoRe runfiles as the +one-GPU local-HF path: 19,252 pages and 14,514 queries, with no failed run, +container restart, OCR error, or OOM. Both used the batch CLI and LanceDB. The +NIM harness ran with an empty `CUDA_VISIBLE_DEVICES`, ensuring that all model +work went through the four endpoints on GPU 0. + +| Dataset | NIM ingest | NIM pages/s | Local-HF ingest | NIM time effect | Recall@5 delta vs local | nDCG@10 delta vs local | +|---|---:|---:|---:|---:|---:|---:| +| `computer_science` | 149.002 s | 9.127 | 171.200 s | -12.97% | -0.0095 | -0.0100 | +| `energy` | 237.900 s | 9.353 | 257.985 s | -7.79% | -0.0147 | -0.0169 | +| `finance_en` | 353.536 s | 8.322 | 343.808 s | +2.83% | -0.0163 | -0.0217 | +| `finance_fr` | 269.278 s | 8.853 | 279.635 s | -3.70% | -0.0185 | -0.0205 | +| `hr` | 145.796 s | 7.613 | 151.708 s | -3.90% | -0.0097 | -0.0124 | +| `industrial` | 510.831 s | 10.266 | 536.309 s | -4.75% | +0.0009 | +0.0016 | +| `pharmaceuticals` | 200.500 s | 11.536 | 212.344 s | -5.58% | +0.0018 | +0.0036 | +| `physics` | 127.785 s | 13.100 | 124.178 s | +2.90% | -0.0170 | -0.0111 | + +The total NIM ingest was 1,994.628 seconds (9.652 pages/s), versus 2,077.167 +seconds (9.268 pages/s) for local HF: **3.97% less ingest time and 4.14% higher +throughput**. Six datasets favored NIM and two favored local HF. This is a +single full-suite backend comparison, not a controlled attribution to the +local-only batching patch. + +Relevance did not establish parity: + +| Macro average | NIM Recall@5 | Local HF | Nightly | NIM nDCG@10 | Local HF | Nightly | +|---|---:|---:|---:|---:|---:|---:| +| English | 0.4751 | 0.4843 | 0.485 | 0.5350 | 0.5445 | 0.545 | +| All datasets | 0.4541 | 0.4645 | 0.465 | 0.5095 | 0.5204 | 0.521 | + +The NIM gap versus local HF was -0.0092/-0.0095 on English Recall@5/nDCG@10 +and -0.0104/-0.0109 across all datasets. A preceding NIM +`computer_science` pilot scored 0.5975/0.7070, while the suite repeat scored +0.5901/0.6990, so real run-to-run variation exists. The broadly lower suite +macro still requires backend isolation before treating the paths as equivalent. +NIM query p50 was 198.4-201.5 ms across datasets, versus 41.0-47.5 ms for +local HF; query time is outside ingest pages/s. + +#### Full four-GPU service comparison + +The standalone service completed the full eight-dataset suite with page +detection, table structure, OCR, and embedding assigned to GPUs 0-3. GPUs 4-7 +remained unused. The model services stayed persistent across the suite; only +retriever and vector-store state was reset between datasets. + +| Dataset | Service ingest | Pages/s | Query p50 | Recall@5 | Delta vs nightly | nDCG@10 | Delta vs nightly | +|---|---:|---:|---:|---:|---:|---:|---:| +| `computer_science` | 114.229 s | 11.906 | 58.719 ms | 0.6001 | -0.0009 | 0.7100 | +0.0000 | +| `energy` | 125.725 s | 17.697 | 65.943 ms | 0.5660 | -0.0120 | 0.5684 | -0.0156 | +| `finance_en` | 217.587 s | 13.521 | 61.112 ms | 0.4794 | -0.0186 | 0.5259 | -0.0221 | +| `finance_fr` | 185.335 s | 12.863 | 61.014 ms | 0.3082 | -0.0188 | 0.3313 | -0.0197 | +| `hr` | 68.946 s | 16.100 | 60.261 ms | 0.4480 | -0.0050 | 0.5234 | -0.0076 | +| `industrial` | 279.011 s | 18.795 | 65.456 ms | 0.3490 | +0.0010 | 0.3831 | +0.0011 | +| `pharmaceuticals` | 101.669 s | 22.750 | 69.020 ms | 0.5475 | -0.0015 | 0.6074 | +0.0004 | +| `physics` | 52.073 s | 32.147 | 65.281 ms | 0.3667 | -0.0033 | 0.4518 | -0.0002 | + +Total service ingest was 1,144.575 seconds (16.820 pages/s): **44.90% less +ingest time and 81.48% higher throughput than one-GPU local HF**, and 42.62% +less time and 74.27% higher throughput than the one-GPU NIM deployment. +Service query p50 stayed between 58.7 and 69.0 ms, versus 41.0-47.5 ms for +local HF and 198.4-201.5 ms for direct one-GPU NIM batch mode. + +The separately supplied current-main H100 nightly gives the upstream service +deployment reference that was missing from the first report. Its eight rounded +per-dataset ingest timers total approximately 1,284.02 seconds for 19,252 +pages, or 14.994 pages/s. The nightly topology provisions four distinct +one-GPU NIM pods on four H100s. + +| Service reference | Hardware | Ingest runtime | Throughput | +|---|---|---:|---:| +| Current-main nightly | 4 x H100 80GB | about 1,284.02 s | **about 14.994 pages/s** | +| PR compatibility run | 4 x H100 80GB | 1,144.575 s | **16.820 pages/s** | + +The nightly total is reconstructed from values rounded to two decimals, and +the two service runs were not a controlled patch A/B. Because this PR does not +change the remote OCR path, their difference is deployment/run context and is +not attributed to the patch. + +| Macro average | Service Recall@5 | Local HF | NIM | Nightly | Service nDCG@10 | Local HF | NIM | Nightly | +|---|---:|---:|---:|---:|---:|---:|---:|---:| +| English | 0.4795 | 0.4843 | 0.4751 | 0.485 | 0.5386 | 0.5445 | 0.5350 | 0.545 | +| All datasets | 0.4581 | 0.4645 | 0.4541 | 0.465 | 0.5127 | 0.5204 | 0.5095 | 0.521 | + +Service relevance was better than direct NIM batch by about 0.003-0.004 macro, +but remained 0.005-0.008 below local HF and nightly. Backend relevance parity +is therefore not established. + +The first sequential service attempt exposed an important isolation problem: +`overwrite=true` did not clear the persistent service collection between +runfiles. By `pharmaceuticals`, query hits included 58 PDF sources outside its +52-PDF corpus, and p50 had risen from 59 to 95 ms. A deterministic artifact +check failed that run. The valid suite above recreated only retriever and +vector-store volumes between datasets, retained all four NIM processes, and +asserted zero foreign sources after every run. All eight checks passed. The +accumulating-run relevance and latency values are excluded. + +The separately supplied latest `computer_science` service baseline was 10.11 +pages/s. The isolated run measured 11.906 pages/s with the same 0.6001/0.7100 +quality profile as the earlier service smoke. These deployment results validate +the current NIM/service stack; they do not attribute remote performance to this +local-only patch. + +#### BO767 four-GPU service result + +The same isolated service topology also completed BO767 with the four core +NIMs pinned one per GPU. Retriever and vector-store volumes were recreated +immediately before the run; all 767 files and 991 queries completed with zero +failed ingest jobs, container restarts, OCR errors, or OOMs. + +| Measurement | Result | Supplied baseline | Effect | +|---|---:|---:|---:| +| Ingest runtime | 623.545 s | 646.621 s implied | 3.57% lower | +| Throughput | **87.772 pages/s** | 84.64 pages/s | **3.70% higher** | +| Recall@5 | 0.856710 | - | - | +| Recall@10 | 0.907164 | - | - | +| nDCG@10 | 0.761228 | - | - | + +The client used an empty `CUDA_VISIBLE_DEVICES`, so all model inference went +through the NIM endpoints. Peak sampled memory on GPUs 0-3 was 1,929, 1,923, +3,865, and 9,698 MiB respectively; GPUs 4-7 remained unused. This is a service +deployment result and is not caused by the local-only batching change. + +Three current Compose portability/configuration issues required `/tmp`-only +workarounds; none is changed by this PR and all are tracked in +[#2424](https://github.com/NVIDIA/NeMo-Retriever/issues/2424): + +- this Docker daemon has no named `nvidia` runtime, so the override used + `runtime: runc` while retaining Compose GPU device reservations; +- the non-root NIM containers could not write the root-owned named model-store + volumes, so writable `/tmp` bind mounts were used; +- the generated service config includes `local_models.extract.use_graphic_elements`, + which the current `ServiceConfig` rejects, so only that invalid key was + removed from a temporary config before the service run. + +## Scope and validation + +This record supports the local Nemotron OCR v2 change only. The remote NIM and +service runs establish compatibility but do not change or attribute speedup to +the remote OCR path. JP20 was not run and byte-identical real OCR output is not +claimed. The one-GPU local-HF, one-GPU NIM, and four-GPU service runs sampled +their allocated devices continuously. The bounded local-HF batch of 8, four +colocated NIMs, and four-GPU service deployment were OOM-free on the tested +H100 80GB GPUs; smaller GPU classes require separate sizing evidence. + +The BO767 command was repeated against exported upstream and patch source +trees, changing only `PYTHONPATH` and the output/run identifiers: + +```bash +CUDA_VISIBLE_DEVICES=0 HF_HUB_OFFLINE=1 HF_DATASETS_OFFLINE=1 \ +VLLM_DEEP_GEMM_WARMUP=skip PYTHONPATH=/nemo_retriever/src \ +retriever harness run bo767_beir --mode batch --output-dir \ + --run-id \ + --set dataset.path=/localhome/local-jioffe/datasets/nv-ingest/bo767 \ + --set ingest.extract.extract_charts=true \ + --set ingest.extract.extract_tables=true \ + --set ingest.extract.batch.ocr_batch_size=8 --json +``` + +Validation on the rebased PR worktree: + +- expected upstream red: 1 failed in 0.08s; +- focused patch suite: 4 passed in 0.60s; +- related actor, graph, OCR, table, and video tests: 327 passed, 7 skipped; +- all pre-commit hooks and required PR validation checks passed. + +The expanded local ViDoRe run used: + +```bash +CUDA_VISIBLE_DEVICES=0 HF_HUB_OFFLINE=1 HF_DATASETS_OFFLINE=1 \ +VLLM_DEEP_GEMM_WARMUP=skip retriever harness run-files \ + nemo_retriever/harness/runfiles/vidore_v3_{computer_science,energy,finance_en,finance_fr,hr,industrial,pharmaceuticals,physics}_beir.json \ + --dataset-paths --mode batch \ + --output-dir --session-name issue-2323-vidore-local-batch --json +``` + +For the NIM checks, the same `computer_science` runfile was invoked first with +`--mode batch` and the four `localhost:8001` through `:8004` endpoint +overrides, then with `--mode service --service-endpoint http://localhost:7670`. +The full service image was built from the rebased PR commit. The one-GPU full +suite used the same eight runfiles shown above and these endpoint overrides: + +```bash +CUDA_VISIBLE_DEVICES= retriever harness run-files \ + --dataset-paths --mode batch \ + --output-dir --session-name issue-2323-one-gpu-nim-full-vidore \ + --set ingest.extract.page_elements_invoke_url=http://localhost:8001/v1/page-elements \ + --set ingest.extract.table_structure_invoke_url=http://localhost:8002/v1/table-structure \ + --set ingest.extract.ocr_invoke_url=http://localhost:8003/v1/ocr \ + --set ingest.embed.embed_invoke_url=http://localhost:8004/v1/embeddings \ + --set query.embed_invoke_url=http://localhost:8004/v1/embeddings --json +``` + +The valid service suite ran each runfile separately with the same four NIM +processes. Before each run, it recreated only the experiment's retriever and +vector-store volumes, waited for service health, and then invoked: + +```bash +CUDA_VISIBLE_DEVICES= retriever harness run-files \ + --dataset-paths --mode service \ + --service-endpoint http://localhost:7670 --output-dir \ + --session-name --json +``` + +Each result was followed by an assertion that every returned `source` belonged +to the current runfile's PDF corpus. + +BO767 service mode used the same healthy four-NIM stack and fresh service +storage: + +```bash +CUDA_VISIBLE_DEVICES= retriever harness run-files \ + nemo_retriever/harness/runfiles/bo767_beir.json \ + --dataset-paths --mode service \ + --service-endpoint http://localhost:7670 --output-dir \ + --session-name issue-2323-bo767-service-g4 --json +``` + +The fixed-query ViDoRe retrieval comparison used source commit +`611af594818342b655b5e9ae89c66aea2cbc3963`. +BO767 compared upstream `52886112cafab4c4bca1cda0d4f588785adfe4d3` +with patch `eaed9262780c45c1dce9e9a929357f2bcd886234`. Both used lock SHA-256 +`d9651104d0a10277642fa7e4794976948177f24c273da203e6bb694107d20bf6`. +Installed versions were `nemotron-ocr==2.0.1.dev20260720042916`, +`ray==2.55.1`, and `torch==2.11.0+cu130`. + +The expanded ViDoRe suite was measured at pre-rebase commit +`3c4ddef05ac8497855f346e68ef9c573e980fb0b`. Rebase added one upstream +service-only commit; the patched `shared.py` SHA-256 remained +`68cd70abbe1ef1beb1cac3fdd197053dfba946c63dd69f8c07623ff2b585ce72`. +The one-GPU ViDoRe upstream control used the patch series' exact parent, +`52886112cafab4c4bca1cda0d4f588785adfe4d3`, with the same lock SHA-256, +environment, H100, runfiles, and default worker policy as the patch run. +The NIM and service checks used rebased commit +`5a8ebd9e5468be4a72ae9888a7d0cba173e44e96`. The full one-GPU NIM suite used +commit `451ba127ea6fba72720c7f66753e2b73273eff6f`, whose merge base was current +`upstream/main` at `3d9e26f1a2d2fd73af499bb7a9ef7fe855739841`; the harness command took +5,456.73 seconds wall time including serial query evaluation. The isolated +full-service suite used PR commit +`3363bf663bbca03fd8300c07d7271a8000533694`; its runtime source was unchanged +from service-image commit `5a8ebd9e5468be4a72ae9888a7d0cba173e44e96`. +The eight isolated service invocations took 2,578.13 seconds wall time including +storage resets and query evaluation. Peak GPU memory was 1,927, 1,923, 3,865, +and 11,364 MiB on GPUs 0-3 respectively; GPUs 4-7 remained at zero. + +BO767 service used the same patch worktree and the unchanged service runtime +from image commit +`5a8ebd9e5468be4a72ae9888a7d0cba173e44e96`. All used lock SHA-256 +`d9651104d0a10277642fa7e4794976948177f24c273da203e6bb694107d20bf6` +and `nemotron-ocr==2.0.1.dev20260720042916`. Both Compose projects were stopped +after measurement; no NIM container was left running. diff --git a/nemo_retriever/developer_docs/ocr_cross_page_batching/proof-summary.svg b/nemo_retriever/developer_docs/ocr_cross_page_batching/proof-summary.svg new file mode 100644 index 0000000000..e17a324471 --- /dev/null +++ b/nemo_retriever/developer_docs/ocr_cross_page_batching/proof-summary.svg @@ -0,0 +1,109 @@ + + Issue 2323 local OCR batching evidence on NVIDIA H100 GPUs + Cross-page batching improves local OCR model and actor throughput. BO767 local Hugging Face whole-ingest throughput improves by 8.41 percent, while full ViDoRe v3 is flat because embedding remains the longer stage. Deployment context includes ViDoRe local, NIM, and service measurements plus the BO767 four-GPU service result. + + + + Issue #2323 · bounded cross-page local OCR batching + All GPU measurements use NVIDIA H100 80GB HBM3 · ingest timers exclude query evaluation + + + Local batching proof · 1× H100 + 128 fixed real crops · persistent Nemotron OCR v2 wrapper · crop lists capped at 8 + + + Model invocations + 128 → 16 + scalar calls → list-of-8 calls + + + OCR actor throughput + 34.4 → 55.8 + crops/s · 1.62× + + + Local model throughput + 40.7 → 73.4 + crops/s · 1.81× + + + Controlled whole-ingest A/B · local Hugging Face · 1× H100 + Same lock, profile, hardware, and default worker policy · extraction through LanceDB indexing included + + + BO767 · 54,730 pages + OCR/embedding frontier moved ≈128 s earlier + Upstream + + 34.04 + Patch + + 36.90 + +8.41% pages/s + ingest time −7.75% + Counterbalanced A/B · 1,608.019 s → 1,483.343 s + + + Full ViDoRe v3 · 19,252 pages + Embedding continues 60–381 s after OCR + Upstream + + 9.280 + Patch + + 9.268 + Flat · −0.12% + 8/8 datasets passed + Retrieval quality preserved · page-image embedding remains the longer stage + + + Deployment context · unchanged remote and service paths + These validate current deployment modes; they are not attributed to the local-only batching patch. + + + Full ViDoRe v3 throughput + pages/s · GPU count shown per deployment + Local HF patch · 1× + + 9.268 + Self-hosted NIM · 1× + + 9.652 + Service nightly · 4× + + ≈14.994 + Service PR run · 4× + + 16.820 + + + BO767 service · 4× H100 + pages/s · supplied service baseline + Baseline + + 84.64 + Measured + + 87.772 + +3.70% pages/s + 623.545 s · 767/767 files · context, not patch effect + + Local-HF bars are the controlled patch A/B · NIM/service measurements establish deployment compatibility + diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index 2530f2eb0c..2f21575bbb 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -157,8 +157,16 @@ It must always be greater than or equal to `--top-k`. Page deduplication and content-type filtering are applied after vector retrieval, preserving retriever ranking order and truncating the final output to -`--top-k`. When querying a local table ingested with an explicit embedding -model, pass the same `--embed-model-name` to `retriever query`. +`--top-k`. Local and batch ingest record the canonical embedding model on the +LanceDB table, and non-service query uses that model automatically. Use +`--embed-model-name` only as an explicit override or when querying a legacy or +third-party table without model metadata. Endpoint URLs and provider prefixes +remain runtime configuration, so continue to pass `--embed-invoke-url` and +`--embed-model-provider-prefix` when the selected model must be routed remotely. +For example, a table can store the canonical model +`nvidia/llama-nemotron-embed-vl-1b-v2` while a LiteLLM-routed request uses +`nvidia/nvidia/llama-nemotron-embed-vl-1b-v2`. The endpoint and routing prefix +are intentionally not persisted on the table. `--content-types` accepts comma-separated content types such as `text`, `table`, `chart`, `image`, and `infographic`. `images` is accepted as an alias for @@ -199,7 +207,8 @@ Unlike the dense path (which returns text-enriched hits), agentic mode returns the agent's ranked document IDs as JSON, each annotated with the source that produced it (`final_results`, `rrf`, or `selection_agent`). It reuses the same `--top-k`, `--lancedb-uri`, `--table-name`, `--embed-invoke-url`, and -`--embed-model-name` options as standard retrieval. +`--embed-model-name` options as standard retrieval. Agentic retrieval uses the +selected table's model automatically when `--embed-model-name` is omitted. **How it works.** Each agentic query runs `Query -> ReActAgentOperator -> (RRF fusion) -> SelectionAgentOperator -> ranked results`: @@ -218,17 +227,21 @@ Agentic-only knobs (apply only with `--agentic`): provided (`nemotron-8b` by default; `super-49b` also supported), or the remote model ID when `--agentic-invoke-url` is provided. - `--agentic-invoke-url` — OpenAI-compatible chat-completions endpoint for the - agent LLM. Providing it routes agent LLM calls to that remote endpoint. + agent LLM. Providing it routes agent LLM calls to that remote endpoint; omit it + to run the in-process local model. +- `--agentic-llm-client` (optional) — LLM client that builds the agent LLM. + Defaults to `callable`. It drives the in-process + adapter when `--agentic-invoke-url` is omitted, and the shared chat-completions + HTTP client when it is set. - `--agentic-reasoning-effort` (default `high`) — `reasoning_effort` forwarded on OpenAI-compatible agentic LLM calls; ignored by the local adapter. -- `--agentic-temperature` (default `0.0`) — sampling temperature for agent LLM - calls. Local and non-NVIDIA OpenAI-compatible endpoints allow up to `2.0`; - NVIDIA-hosted endpoints allow up to `1.0`. -- `--agentic-backend-top-k` (default `20`) — candidates pulled from the vector DB - per retrieval call. - `--agentic-react-max-steps` (default `50`) — maximum ReAct loop iterations. - `--agentic-text-truncation` (default `0`) — max characters of each candidate shown to the agent; `0` disables truncation. +- `--agentic-temperature` (default: unset) — sampling temperature for agent LLM + calls; omit to use the endpoint/model default (`0.0` = greedy). Local and + non-NVIDIA OpenAI-compatible endpoints allow up to `2.0`; NVIDIA-hosted + endpoints allow up to `1.0`. @@ -306,6 +319,65 @@ retriever ingest ./data/pdf_corpus \ --embed-model-name nvidia/llama-nemotron-embed-1b-v2 ``` +### Dense Nemotron embedding checkpoints + +`--embed-model-name` accepts a Hugging Face repository ID or an on-disk +checkpoint compatible with a supported dense Nemotron text or vision-language +embedding profile: + +```bash +retriever ingest ./data/pdf_corpus \ + --embed-model-name acme/my-finetuned-nemotron-embed +``` + +Tested official checkpoints: + +- `nvidia/llama-3.2-nv-embedqa-1b-v2` +- `nvidia/llama-nemotron-embed-1b-v2` +- `nvidia/llama-nemotron-embed-vl-1b-v2` +- `nvidia/llama-nemotron-embed-vl-1b-v2-fp8` +- `nvidia/llama-nv-embed-reasoning-3b` +- `nvidia/llama-embed-nemotron-8b` + +Equivalent local checkpoints and weight-only fine-tunes are supported. A +compatible checkpoint must be complete and loadable, use +`LlamaBidirectionalModel` or `LlamaNemotronVLModel`, and declare average +pooling with a positive output width. LanceDB infers the schema from the +produced vectors; the tested official checkpoints use 2048, 3072, and 4096 +dimensions. Query and document prompts are read from +`config_sentence_transformers.json` when the checkpoint supplies it. +Fine-tunes that require prefixes other than `query: ` and `passage: ` must +retain that prompt metadata. + +This does not add support for every model in the Nemotron RAG collection, +including rerankers, ColEmbed late-interaction models, Omni Embed, OCR, or +parsing models. Nemotron 3 Embed is also excluded because its Ministral3 +architecture requires a newer Transformers stack than this project currently +supports. Those models require different dependencies, outputs, modalities, +or operator contracts. + +Unregistered Hub repositories are resolved to an immutable commit and loaded +with `trust_remote_code=True`; only use repositories you trust. The resolved +model name and revision are recorded on the LanceDB table and reused by local +query. + +For a compatible ModelOpt checkpoint, including FP8 or NVFP4 variants, select +vLLM for ingest. Local query detects the ModelOpt configuration and selects +vLLM automatically: + +```bash +retriever ingest ./data/pdf_corpus \ + --embed-model-name /models/my-finetuned-nemotron-embed-fp8 \ + --local-ingest-embed-backend vllm + +retriever query "What is in this corpus?" \ + --table-name nemo-retriever +``` + +Hugging Face remains the local query backend for non-ModelOpt checkpoints. +Local directories must contain `config.json`, and their absolute path must be +available to every Ray worker or service replica that loads the model. + ### OCR language mode ```bash diff --git a/nemo_retriever/docs/cli/benchmarking.md b/nemo_retriever/docs/cli/benchmarking.md index bbf7c12fd0..078358d3ea 100644 --- a/nemo_retriever/docs/cli/benchmarking.md +++ b/nemo_retriever/docs/cli/benchmarking.md @@ -1,131 +1,42 @@ -# Benchmarking with the `retriever` CLI +# Benchmarking with the Retriever CLI -`retriever benchmark` and `retriever harness` are development and experimental subcommands -with no guarantees — refer to [Supported vs development / experimental subcommands](README.md#supported-vs-development--experimental-subcommands). +Retriever has two development benchmarking surfaces: -This page covers benchmark workflows for NeMo Retriever Library. Also refer to -[`nemo_retriever/harness/HANDOFF.md`](../../harness/HANDOFF.md) for operator-oriented -notes on `retriever harness`. +| Goal | Command | Documentation | +| --- | --- | --- | +| End-to-end ingest, query, and retrieval evaluation | `retriever harness` | [Retriever Harness](../../harness/README.md) | +| Throughput for one internal pipeline stage | `retriever benchmark` | [Stage micro-benchmarks](#stage-micro-benchmarks) | -Use `retriever harness` for benchmark orchestration and `retriever benchmark` for -per-stage micro-benchmarks. `retriever benchmark` remains callable but is hidden -from root help. +For product workflows on your own inputs, use `retriever ingest` and +`retriever query` instead. -## Harness (development / experimental) +## End-to-End Benchmarks -Run from the repository root or any directory. The harness uses code-owned -benchmark names from `nemo_retriever.harness.benchmark_registry`; use -`retriever harness list` to discover the available benchmarks and runsets. +The harness owns registered datasets, repeatable runfiles, metric gates, and +stable artifacts. Start with: ```bash -# List benchmark registry entries, optionally including runsets -retriever harness list retriever harness list --runsets - -# Inspect one concrete benchmark spec retriever harness show jp20_beir - -# Run one benchmark and write stable artifacts -retriever harness run jp20_beir - -# Run one benchmark in batch mode -retriever harness run bo767_beir --mode batch - -# Override a resolved config key for this run -retriever harness run bo767_beir --set query.top_k=5 - -# Expand and run a code-owned benchmark runset -retriever harness run-set jp20_core -``` - -Related commands: - -```bash -retriever harness --help -retriever harness list --help -retriever harness show --help -retriever harness run --help -retriever harness run-set --help -retriever harness diff --help ``` -### Agentic BEIR evaluation +Then choose the execution guide: -Harness runs use the standard dense retrieval path unless agentic retrieval is -enabled in the resolved benchmark query config. Set `query.agentic: true` in a -code-owned benchmark or runfile, or use repeatable `--set` overrides on the CLI. -The agentic harness path runs the same ReAct retrieval graph used by root query, -but only after ingest and only for BEIR evaluation (`evaluation.mode: beir`). +- [Local and batch library runs](../../harness/docs/library.md) +- [Existing-service and managed-Helm runs](../../harness/docs/service.md) +- [Recurring workstation nightly](../../../ops/retriever-nightly/README.md) -By default, agentic harness evaluation uses the in-process local vLLM backend -with `nemotron-8b`. Custom LLMs are not supported in process yet; run them behind -an OpenAI-compatible chat-completions endpoint and set `query.agentic_invoke_url`. +Agentic BEIR evaluation is documented under +[library harness runs](../../harness/docs/library.md#evaluate-agentic-retrieval). -Minimal BEIR override example: +## Stage Micro-Benchmarks -```bash -retriever harness run jp20_beir \ - --set query.agentic=true -``` - -Larger supported local profile: +`retriever benchmark` measures individual actors rather than an end-to-end +Retriever result. It remains callable for development compatibility but is +hidden from root help. ```bash -retriever harness run jp20_beir \ - --set query.agentic=true \ - --set query.agentic_llm_model=super-49b \ - --set query.agentic_local_tensor_parallel_size=2 -``` - -Custom/self-hosted OpenAI-compatible endpoint: - -```bash -retriever harness run jp20_beir \ - --set query.agentic=true \ - --set query.agentic_llm_model=custom-remote-model \ - --set query.agentic_invoke_url=http://localhost:9000/v1/chat/completions -``` - -Useful agentic query overrides: - -- `query.agentic_llm_model` — local profile alias/model ID when no invoke URL is - provided (`nemotron-8b` by default; `super-49b` also supported), or the remote - model ID when `query.agentic_invoke_url` is provided. -- `query.agentic_invoke_url` — OpenAI-compatible chat-completions endpoint. - Providing it routes agent LLM calls to that remote endpoint. -- `query.agentic_local_gpu_memory_utilization`, - `query.agentic_local_tensor_parallel_size`, `query.agentic_local_max_model_len`, - and `query.agentic_local_max_num_seqs` — harness-only local vLLM resource and - scheduling controls for benchmark runs. Use environment variables such as - `CUDA_VISIBLE_DEVICES` and the standard Hugging Face cache environment for - placement and model cache control. -- `query.agentic_backend_top_k` — backend candidate pool per ReAct retrieval - call. Must be at least the final requested metric depth (`max(evaluation.ks)`). -- `query.agentic_react_max_steps` — maximum ReAct loop iterations per query - (defaults to `50`). -- `query.agentic_text_truncation` — max characters of each candidate shown to - the agent; `0` disables truncation. -- `query.agentic_num_concurrent` — number of queries the agent batch runs - concurrently (defaults to `1`). -- `query.agentic_temperature` — defaults to `0.0`; local and non-NVIDIA - OpenAI-compatible endpoints allow `0.0..2.0`, while hosted/default NVIDIA - endpoints are validated as `0.0..1.0`. -- `query.agentic_reasoning_effort` — optional provider-specific field forwarded - only when configured. - -### Image storage - -For normal ingest, configure image persistence on `retriever ingest` with -`--store-images-uri ` (local path or fsspec URI). Stored assets follow -`--embed-granularity` (page vs element images). - -## Per-stage micro-benchmarks - -Stage throughput benchmarks remain callable for compatibility even though they -are hidden from root help: - -```bash -retriever benchmark --help # split, extract, audio-extract, page-elements, ocr, all +retriever benchmark --help retriever benchmark split --help retriever benchmark extract --help retriever benchmark audio-extract --help @@ -134,7 +45,7 @@ retriever benchmark ocr --help retriever benchmark all --help ``` -Example — PDF extraction actor: +Example: ```bash retriever benchmark extract ./data/pdf_corpus \ @@ -142,15 +53,6 @@ retriever benchmark extract ./data/pdf_corpus \ --pdf-extract-actors 4 ``` -Each benchmark reports rows/sec (or chunk rows/sec for audio) for its actor. - -## Notes - -- **Configuration:** `retriever harness` uses code-owned benchmarks/runsets from - `nemo_retriever.harness.benchmark_registry`; use `--set KEY=VALUE` for small - per-run config overrides. -- **Launcher:** for internal benchmarking, `retriever harness run BENCHMARK` and - `retriever harness run-set RUNSET` are the benchmark orchestration entry points - (development / experimental; no guarantees). -- **Stage benchmarks:** `retriever benchmark …` is specific to the retriever CLI and - covers per-stage throughput rather than full harness orchestration. +Stage commands report rows per second, or chunk rows per second for audio. They +do not produce the harness artifact contract and should not be used as retrieval +quality evidence. diff --git a/nemo_retriever/harness/EXPECTED_RESULTS.md b/nemo_retriever/harness/EXPECTED_RESULTS.md deleted file mode 100644 index ee9cb2349e..0000000000 --- a/nemo_retriever/harness/EXPECTED_RESULTS.md +++ /dev/null @@ -1,261 +0,0 @@ - - - -# Harness Expected Results - -Known dataset facts, canonical benchmark result ranges, and suggested integrity -gates for `retriever harness`. - -This file is documentation, not executable policy. Use dataset facts for -explicit `--require` gates. Use quality and performance ranges to judge whether -a result is in the expected ballpark. Update the references when datasets, -benchmark definitions, hardware, or retrieval behavior intentionally change. - -Only canonical benchmark expectations belong here. Exploratory runs, fast-text -fallbacks, chunking experiments, and failed attempts should stay in run -artifacts or handoff notes until the team chooses them as canonical benchmark -definitions. - -File counts, page counts, and query counts are portable dataset-integrity gates. -Recall and nDCG are reference ranges that help developers and agents identify -results outside the expected ballpark. Do not treat the quality ranges as -universal pass/fail policy across different hardware and runtime profiles. - -Performance observations such as ingest seconds, pages/sec, and query latency -are reference points for a specific environment. Treat them as hardware- and -configuration-sensitive unless the GPU SKU/count, CUDA driver, model backend, -vLLM/kernel settings, Ray worker layout, storage path, and dataset mount are -controlled. - -The dataset paths below are registry reference paths, not a portable filesystem -contract. Use `harness/dataset_paths.example.yaml` and `run-files ---dataset-paths` to point checked-in runfiles at the current machine. - -## JP20 - -Dataset: - -- Corpus path: `/datasets/nv-ingest/jp20` -- Query/qrels file: `data/jp20_query_gt.csv` -- Files: `20` -- Pages: `1940` - -Benchmarks: - -| Benchmark | Purpose | Ingest Profile | Queries | Expected Quality | -|-----------|---------|----------------|---------|------------------| -| `jp20_beir` | End-to-end retrieval quality | `auto` | `115` | `recall_5 >= 0.85`, `ndcg_10 >= 0.75` | - -Suggested full BEIR command: - -```bash -retriever harness run-files \ - --session-name jp20_beir \ - --output-dir /local/path/to/retriever-artifacts/jp20-beir \ - --dataset-paths /local/path/to/dataset_paths.yaml \ - --require 'files==20' \ - --require 'pages==1940' \ - --require 'query_count==115' \ - nemo_retriever/harness/runfiles/jp20_beir.json -``` - -Recent observed `jp20_beir` metrics on local hardware: - -- `rows_processed`: `3154` -- `ingest_secs`: about `215` to `223` -- `query_latency_p50_ms`: about `909` to `915` -- `query_latency_p95_ms`: about `953` to `1003` -- `recall_5`: about `0.878` to `0.887` -- `recall_10`: about `0.930` to `0.948` -- `ndcg_10`: about `0.793` to `0.802` - -Avoid hard-gating on latency or throughput unless the run environment is -controlled and recorded in the artifacts. - -## BO20 - -Dataset: - -- Corpus path: `/datasets/nv-ingest/bo20` -- Files: `20` -- BEIR qrels: not expected - -## BO767 - -Dataset: - -- Corpus path: `/datasets/nv-ingest/bo767` -- Query/qrels file: `data/bo767_query_gt.csv` -- Files: `767` -- Pages: `54730` - -Benchmark: - -| Benchmark | Purpose | Ingest Profile | Queries | Expected Quality | -|-----------|---------|----------------|---------|------------------| -| `bo767_beir` | End-to-end retrieval quality | `auto` | `991` | `recall_5 >= 0.84`, `ndcg_10 >= 0.74` | - -Suggested full BEIR command: - -```bash -retriever harness run-files \ - --session-name bo767_beir \ - --output-dir /local/path/to/retriever-artifacts/bo767-beir \ - --dataset-paths /local/path/to/dataset_paths.yaml \ - --require 'files==767' \ - --require 'pages==54730' \ - --require 'query_count==991' \ - nemo_retriever/harness/runfiles/bo767_beir.json -``` - -Observed `bo767_beir` metrics on an eight-H100 80GB HBM3 host: - -| Configuration | Workload GPUs | Rows | Ingest seconds | Pages/s | Recall@5 | Recall@10 | nDCG@10 | -|---------------|---------------|------|----------------|---------|----------|-----------|---------| -| RC26.05 Perflab | Not recorded | 79221 | 4036.847 | 13.56 | Not recorded | Not recorded | Not recorded | -| Automatic batch | 1 | 79229 | 1594.339 | 34.328 | 0.848638 | 0.897074 | 0.750110 | -| Automatic batch | 8 | 79230 | 764.946 | 71.548 | 0.849647 | 0.895055 | 0.748583 | -| Legacy worker-capped batch | 8 visible, effectively 1 used | 79230 | about 2265 | about 24.16 | 0.850656 | 0.896065 | 0.751507 | - -The RC26.05 Perflab artifact recorded eight physical GPUs but did not -distinguish the workload-visible GPU count. The checked-in BO767 runfile leaves -worker counts and batch sizes automatic so the batch planner can scale to the -GPUs available to the workload. Do not use -the legacy worker-capped result as evidence of eight-GPU scaling. The automatic -one- and eight-GPU runs differed by one output row; keep quality and row counts -visible when comparing performance results. - -## FinanceBench - -Dataset: - -- Corpus path: `/datasets/nv-ingest/foundation_rag/financebench` -- Query/qrels file: `data/financebench_train.json` -- Files: `369` -- Pages: `54057` - -Benchmark: - -| Benchmark | Purpose | Ingest Profile | Queries | Expected Quality | -|-----------|---------|----------------|---------|------------------| -| `financebench_beir` | End-to-end retrieval quality | `auto` | `150` | TBD after canonical run | - -## BO10K - -Dataset: - -- Corpus path: `/datasets/nv-ingest/bo10k` -- Query/qrels file: `data/digital_corpora_10k_annotations.csv` -- Files: `10000` - -Benchmark: - -| Benchmark | Purpose | Ingest Profile | Queries | Expected Quality | -|-----------|---------|----------------|---------|------------------| -| TBD | Canonical end-to-end retrieval quality | `auto` | TBD | TBD after canonical benchmark is defined | - -## Earnings Consulting - -Dataset: - -- Corpus path: `/datasets/nv-ingest/earnings_consulting_flattened` -- Query/qrels file: `data/earnings_consulting_multimodal.csv` -- Files: `514` -- Pages: `12988` - -Benchmark: - -| Benchmark | Purpose | Ingest Profile | Queries | Expected Quality | -|-----------|---------|----------------|---------|------------------| -| `earnings_beir` | End-to-end retrieval quality | `auto` | `628` | TBD after canonical run | - -## ViDoRe V3 - -The eight public ViDoRe v3 benchmarks use original PDFs for ingest and load -queries and qrels from the corresponding `vidore/` Hugging Face -dataset. The canonical configuration uses: - -- `nvidia/llama-nemotron-embed-vl-1b-v2` -- `text_image` embedding at page granularity -- page-image and infographic extraction -- page-level BEIR document IDs - -Dataset integrity gates: - -| Dataset | Files | Pages | Queries | -|---------|------:|------:|--------:| -| `vidore_v3_computer_science` | 2 | 1360 | 1290 | -| `vidore_v3_energy` | 41 | 2225 | 1848 | -| `vidore_v3_finance_en` | 6 | 2942 | 1854 | -| `vidore_v3_finance_fr` | 5 | 2384 | 1920 | -| `vidore_v3_hr` | 14 | 1110 | 1908 | -| `vidore_v3_industrial` | 27 | 5244 | 1698 | -| `vidore_v3_pharmaceuticals` | 52 | 2313 | 2184 | -| `vidore_v3_physics` | 42 | 1674 | 1812 | - -Canonical benchmarks: - -| Benchmark | Purpose | Ingest Profile | Expected Quality | -|-----------|---------|----------------|------------------| -| `vidore_v3_computer_science_beir` | Computer science page retrieval | `auto` | Observed `ndcg_10` about `0.708` to `0.709` | -| `vidore_v3_energy_beir` | Energy page retrieval | `auto` | Observed `ndcg_10` about `0.581` | -| `vidore_v3_finance_en_beir` | English finance page retrieval | `auto` | Observed `ndcg_10` about `0.547` | -| `vidore_v3_finance_fr_beir` | French finance page retrieval | `auto` | Observed `ndcg_10` about `0.345`; see coverage warning below | -| `vidore_v3_hr_beir` | Human-resources page retrieval | `auto` | Observed `ndcg_10` about `0.530` | -| `vidore_v3_industrial_beir` | Industrial page retrieval | `auto` | Observed `ndcg_10` about `0.381` | -| `vidore_v3_pharmaceuticals_beir` | Pharmaceuticals page retrieval | `auto` | Observed `ndcg_10` about `0.607` | -| `vidore_v3_physics_beir` | Physics page retrieval | `auto` | Observed `ndcg_10` about `0.451` | - -Use the checked-in runfiles for executable file, page, and query-count gates. -Add quality and performance observations here only after a complete run on the -canonical default configuration. - -Observed metrics from complete batch runs on an eight-H100 DGX with the default -benchmark configuration: - -| Dataset | Rows Processed | Indexed Rows | Ingest Seconds | Pages/s | Query p50 ms | Query p95 ms | Recall@5 | Recall@10 | nDCG@10 | -|---------|---------------:|-------------:|---------------:|--------:|-------------:|-------------:|---------:|----------:|--------:| -| Computer Science | 1360 | 1358 | 100.2-123.5 | 11.0-13.6 | 40.3-40.7 | 66.7-66.9 | 0.599-0.600 | 0.729-0.730 | 0.708-0.709 | -| Energy | 2225 | 2211 | 116.7 | 19.1 | 43.3 | 51.4 | 0.575 | 0.674 | 0.581 | -| Finance EN | 2942 | 2927 | 149.4 | 19.7 | 45.1 | 53.7 | 0.496 | 0.609 | 0.547 | -| Finance FR | 2384 | 2149 | 106.4 | 22.4 | 44.8 | 52.0 | 0.324 | 0.426 | 0.345 | -| HR | 1110 | 1091 | 82.6 | 13.4 | 39.7 | 46.3 | 0.452 | 0.574 | 0.530 | -| Industrial | 5244 | 5039 | 137.5 | 38.1 | 51.0 | 61.8 | 0.348 | 0.426 | 0.381 | -| Pharmaceuticals | 2313 | 2290 | 93.7 | 24.7 | 46.0 | 54.9 | 0.547 | 0.647 | 0.607 | -| Physics | 1674 | 1674 | 89.2 | 18.8 | 45.0 | 52.5 | 0.369 | 0.485 | 0.451 | - -Computer Science was run twice; the other domains have one complete observation -each. Computer Science quality was stable to within `0.0011` nDCG@10. The -eight-domain macro-average nDCG@10 was about `0.519`, using the mean of the two -Computer Science runs. - -The NeMo Retriever 26.05 image-plus-text release baseline reports average -Recall@5 of `0.490` over the English datasets and `0.465` over all datasets. -The default harness run reproduces the domain-level release results to within -`0.0032` absolute Recall@5: - -| Dataset | 26.05 Release Recall@5 | Observed Recall@5 | Delta | -|---------|------------------------:|------------------:|------:| -| Finance EN | 0.499 | 0.496 | -0.003 | -| Industrial | 0.348 | 0.348 | +0.000 | -| Computer Science | 0.600 | 0.599 | -0.001 | -| Pharmaceuticals | 0.549 | 0.547 | -0.002 | -| HR | 0.453 | 0.452 | -0.001 | -| Energy | 0.577 | 0.575 | -0.002 | -| Physics | 0.367 | 0.369 | +0.002 | -| Finance FR | 0.324 | 0.324 | +0.000 | - -The observed simple macro-average Recall@5 was `0.464` over all eight domains. - -The indexed-row audit found that every omitted page had empty corpus text. No -judged pages were omitted for seven of the eight datasets. Finance FR omitted -`235` empty-text pages, including `69` judged image-only pages, because the -current dense LanceDB write path drops records without text even when an image -embedding exists. Its Recall@5 still matches the 26.05 release baseline, which -suggests the release exercised the same behavior, but retaining those judged -image-only pages remains a correctness prerequisite before nightly scheduling. - -Do not apply default hard quality gates across machine profiles or GPU SKUs. -Keep the checked-in file, page, and query-count requirements as hard integrity -gates, record quality metrics on every run, and establish comparison ranges -from each nightly environment's own history. diff --git a/nemo_retriever/harness/HANDOFF.md b/nemo_retriever/harness/HANDOFF.md deleted file mode 100644 index 3efb6649c3..0000000000 --- a/nemo_retriever/harness/HANDOFF.md +++ /dev/null @@ -1,106 +0,0 @@ - - - -# Retriever Harness Maintainer Notes - -The user and agent contract is documented in [`README.md`](README.md). Keep this -file limited to implementation boundaries that maintainers need when changing -the harness. - -## Product Boundary - -- `retriever ingest` and `retriever query` are the direct product workflows. -- `retriever harness` is a developer benchmark/evaluation runner built on the - same planning and workflow APIs. -- `run` executes one registered benchmark using registry paths or explicit - overrides. -- `run-set` executes a code-owned benchmark group using registry paths. -- `run-files` executes one or more checked-in runfiles and can apply a - machine-local dataset path map. It is the portable session engine and owns - dry-run behavior for the complete session. -- `run-helm` is the supported optional provisioning wrapper around one - `run-files` session. It does not own benchmark execution semantics. -- `post-slack` only reads completed artifacts. It does not execute benchmarks - or mutate their results. -- Scheduling, retries, locking, and secret distribution are outside this - harness surface. Deployment is limited to the optional `run-helm` wrapper. -- `service` is a system-under-test mode that uses an endpoint supplied by the - caller; Helm is only an optional outer provisioning mechanism. - -The harness calls the shared ingest and query workflow modules directly. - -## Implementation Map - -- `runfile.py` parses one portable run request; `dataset_paths.py` resolves the - machine-local dataset map. -- `execution.py` preflights and executes one benchmark. -- `runsets.py` converts runsets or runfiles into prepared runs, then executes - both through one session loop. -- `json_io.py` atomically publishes JSON artifacts; `artifact_writer.py` owns - per-run status, events, logs, and artifact cleanup. -- `slack.py` reads completed artifacts and renders or posts a report. It does - not participate in benchmark execution. -- `helm_runner.py` and `HelmServiceManager` implement `run-helm` by provisioning - an immutable service, invoking the shared `run-files` CLI, collecting failure - logs, and tearing down. They do not implement benchmark sessions or reporting. - -## Configuration Ownership - -The Python registry owns benchmark and dataset semantics. Checked-in runfiles -own concrete modes, metric gates, and narrow overrides. A local -`dataset_paths.yaml` owns machine-specific document and query locations and -must remain outside source control. - -Resolution precedence is: - -1. Registry defaults. -2. Runfile overrides. -3. Machine-local dataset paths. -4. CLI `--set` overrides. - -Large checked-in benchmark runfiles use batch ingest. Do not silently change -their mode or hardware-sensitive worker tuning without fresh validation. - -## Artifact Contract - -Poll `status.json` while a run is active. Read `results.json` after one run is -terminal and `session_summary.json` after a multi-run session is terminal. -Those files contain summary metrics and relative pointers to detailed evidence. - -Detailed run evidence can include: - -- `events.jsonl` -- `runfile.json` -- `resolved_benchmark.json` -- `ingest_plan.json` -- `query_plan.json` -- `environment.json` -- `run.log` -- `beir_metrics.json` -- `beir_run.trec` -- `query_results.jsonl` -- `lancedb/` -- `service_logs/` - -New runs do not write `summary_metrics.json`. Compatibility readers may still -accept that file in older artifact directories. Failure summaries stay concise; -full tracebacks belong in `run.log`. - -## Validation - -At minimum, changes should cover: - -- CLI help and dry-run behavior for `run`, `run-set`, and `run-files`, plus - exit-code propagation through `run-helm`. -- One-run and multi-run terminal artifact shapes. -- Missing inputs, invalid overrides, and metric-gate failures. -- Dataset path precedence and secret redaction. -- Slack preview without a webhook, current-release reference rendering, and - transport errors without secret leakage. -- A real benchmark smoke run when execution behavior changes. - -Use `EXPECTED_RESULTS.md` for dataset facts and observed metric ranges. Do not -turn hardware-sensitive reference numbers into implicit global pass/fail policy. -A configured release snapshot is presentation input only: show it beside the -current nightly with hardware context, but do not maintain history or assign a -verdict in the harness. diff --git a/nemo_retriever/harness/README.md b/nemo_retriever/harness/README.md index 75600ee137..6ebaa8d25a 100644 --- a/nemo_retriever/harness/README.md +++ b/nemo_retriever/harness/README.md @@ -3,224 +3,128 @@ # Retriever Harness -Developer benchmark harness for repeatable Retriever ingest/query evaluation. +The Retriever Harness runs registered ingest and retrieval benchmarks with +repeatable configuration and machine-readable results. Use `retriever ingest` +and `retriever query` for your own data; use `retriever harness` for benchmark +and evaluation work. -Use `retriever ingest` and `retriever query` when you want to operate Retriever -directly on your own inputs. Use `retriever harness` when you want to run a -registered benchmark with reproducible settings, metric gates, and a stable -artifact contract. The harness does not install or operate a scheduler. +One harness supports two execution paths: -The harness is artifact-first. Poll `status.json` while one run is active, read -`results.json` when that run is terminal, and read `session_summary.json` for a -multi-run session. Those terminal files point to detailed evidence; agents do -not need to scan every file in the artifact directory. +- **Library execution** runs Retriever directly in `local` or Ray-backed + `batch` mode. +- **Service execution** tests an existing Retriever endpoint or a temporary + Helm deployment. -## Quick Start +Both paths use the same runfiles, metric gates, and artifact contract. -Run commands from the repository root through the `nemo_retriever` project: +## Start Here -```bash -uv run --project nemo_retriever retriever harness list --runsets -uv run --project nemo_retriever retriever harness show jp20_beir --json -``` - -The registry contains canonical benchmark definitions, but dataset mounts vary -by machine. Before executing a checked-in runfile, copy -[`dataset_paths.example.yaml`](dataset_paths.example.yaml) outside the repository -and replace the example document and query paths with paths available locally. - -Dry-run one checked-in dataset with that path map: - -```bash -uv run --project nemo_retriever retriever harness run-files \ - --session-name jp20_check \ - --output-dir /tmp/retriever-harness-jp20-check \ - --dataset-paths /local/path/to/dataset_paths.yaml \ - --dry-run \ - --json \ - nemo_retriever/harness/runfiles/jp20_beir.json -``` +| Goal | Guide | +| --- | --- | +| Run one small benchmark | [Library execution](docs/library.md#run-one-benchmark) | +| Run a larger or multi-GPU benchmark | [Library execution](docs/library.md#run-in-batch) | +| Test an existing Retriever service | [Service execution](docs/service.md#test-an-existing-service) | +| Provision and test a service with Helm | [Service execution](docs/service.md#provision-a-service-with-helm) | +| Run the workstation suite every day | [Nightly launcher](../../ops/retriever-nightly/README.md) | +| Review known dataset facts and observed results | [Expected results](docs/expected-results.md) | -After inspecting `session_summary.json` and the child run's resolved plans, -remove `--dry-run` to execute it. The same `run-files` command accepts multiple -runfiles for a collection. - -If the registry's default paths already exist on the machine, `run` is the -short form for one registered benchmark: +Run harness commands from the repository root: ```bash -uv run --project nemo_retriever retriever harness run jp20_beir \ - --output-dir /tmp/retriever-harness-jp20-beir \ - --require 'files==20' \ - --require 'pages==1940' \ - --require 'query_count==115' \ - --json +uv run --project nemo_retriever retriever harness list --runsets +uv run --project nemo_retriever retriever harness show jp20_beir --json ``` -Large checked-in BEIR runfiles such as BO767, FinanceBench, Earnings, and -ViDoRe use `mode: batch`. Keep JP20 local for quick smoke validation, and use -batch mode for larger canonical quality runs so Ray-backed ingest owns worker -parallelism and memory pressure. +Use `retriever harness --help` for the complete option list. -Use `mode: service` when the system under test is an already-running Retriever -service. Supply its machine-local URL with `--service-endpoint`; `run-files` -applies that URL only to service-mode children in a mixed session. Service mode -uses the product service APIs for ingest and query while preserving the same -`status.json`, `results.json`, metric gates, and session summary contract. +Service-mode benchmarks wait for remote document completion without downloading +retained result payloads. Their `rows_processed` value is the sum of +`result_rows` from successful document-completion events; the independent +VectorDB coverage checks remain the authoritative end-to-end validation. As a +result, service `ingest_secs` measures remote ingestion completion rather than +ingestion plus client-side result materialization. Service timing baselines +recorded before this behavior changed are not directly comparable. ## Commands -- `list`: list code-owned benchmarks and optional runsets. -- `show`: inspect one benchmark definition. -- `run`: run one registered benchmark using registry paths or explicit `--set` - overrides. -- `run-set`: expand a code-owned benchmark group using registry paths. -- `run-files`: execute one or more runfiles with an optional machine-local - dataset path map. Real children run sequentially in fresh processes. This is - the portable session engine for the checked-in suite and does not provision - infrastructure. -- `run-helm`: optionally provision a Helm service around one portable - `run-files` session, then collect failure logs and tear the service down. -- `check-vidore-access`: validate authenticated access to the queries, qrels, - and corpus objects for all eight ViDoRe v3 datasets without downloading them. -- `post-slack`: preview or post existing artifacts; it never executes a run. -- `diff`: compare two run artifact directories by `results.json` summary metrics. - -For the opinionated twelve-benchmark workstation workflow, including Git -selection, dataset defaults, Slack, and daily recurrence, use the -[Retriever nightly launcher](../../ops/retriever-nightly/README.md). - -Legacy sweep, recurring-job, runner, reporting-UI, and portal -commands are not part of this CLI surface. Scheduling and deployment belong to -separate infrastructure, not the benchmark harness. - -## Runfiles - -Runfiles are a small reproducibility helper for agents, handoffs, and -orchestrators. They describe one concrete run request: - -- registered `benchmark` -- optional `name`, `mode`, `run_id`, and `output_dir` -- optional `set` overrides -- optional `require` metric gates - -Runfiles cannot define new datasets or benchmarks. Add recurring benchmark -definitions to the Python registry instead. - -The harness accepts JSON, YAML, or YML runfiles. Runfiles use -`schema_version: 1`; unknown top-level runfile keys fail during resolution with -exit code `2`. The checked-in JP20 example is -[`runfiles/jp20_beir.json`](runfiles/jp20_beir.json). - -### Configure Machine-Local Dataset Paths - -Dataset locations vary between developer systems. Keep benchmark definitions -and checked-in runfiles independent of one machine's mount layout. Copy -[`dataset_paths.example.yaml`](dataset_paths.example.yaml) to an untracked -location, then set the document and query paths available on the machine that -runs the harness. - -The harness does not distribute or download private benchmark corpora or qrels. -The operator must have access to the datasets referenced by the selected -runfiles. - -Pass the local file with `--dataset-paths`. Relative paths in the file resolve -relative to the file itself. The harness writes the resolved absolute paths to -`expanded_runs.json` and each run's `resolved_benchmark.json`. - -Settings resolve in this order, from lowest to highest precedence: +| Command | Purpose | +| --- | --- | +| `list` | List registered benchmarks and optional runsets. | +| `show` | Show one resolved benchmark definition. | +| `run` | Run one registered benchmark. | +| `run-set` | Run a registered benchmark group using registry paths. | +| `run-files` | Run one or more portable runfiles as a session. | +| `run-helm` | Provision a service, run one session, collect failure logs, and tear it down. | +| `check-vidore-access` | Check remote ViDoRe queries, qrels, and corpora without downloading them. | +| `post-slack` | Preview or post completed artifacts without rerunning a benchmark. | +| `diff` | Compare summary metrics from two completed runs. | -1. Benchmark registry defaults. -2. Checked-in runfile overrides. -3. Machine-local dataset paths. -4. Command-line `--set` overrides. +`run-files` is the normal portable entrypoint. It accepts one runfile for one +benchmark or several runfiles for a suite. Real children run sequentially in +fresh processes so each run releases its model and Ray resources. -### Run One Or More Runfiles As A Session +## Configuration -For a single dataset, pass one runfile: +The Python registry owns benchmark definitions. A runfile selects one registered +benchmark and may set its mode, narrow overrides, and explicit metric gates: -```bash -uv run --project nemo_retriever retriever harness run-files \ - --session-name jp20_beir \ - --output-dir /local/path/to/retriever-artifacts/jp20-beir \ - --dataset-paths /local/path/to/dataset_paths.yaml \ - nemo_retriever/harness/runfiles/jp20_beir.json +```json +{ + "schema_version": 1, + "benchmark": "jp20_beir", + "mode": "local", + "require": ["files==20", "pages==1940"] +} ``` -For the four-dataset library suite, pass all four runfiles. This is still an -ordinary user-invoked harness session; the repository does not schedule it: +Runfiles may be JSON or YAML. They cannot define new benchmarks. For a +`run-files` session, set `--output-dir` and `--dry-run` on the command rather +than in individual runfiles. + +Dataset locations are machine-specific. Copy +[`dataset_paths.example.yaml`](dataset_paths.example.yaml) outside the +repository, edit it, and pass the file with `--dataset-paths`: ```bash -export RETRIEVER_SESSION_DIR=/local/path/to/retriever-artifacts/library-beir-$(date -u +%Y%m%d_%H%M%S_UTC) - -uv run --project nemo_retriever retriever harness run-files \ - --session-name library_beir \ - --output-dir "$RETRIEVER_SESSION_DIR" \ - --dataset-paths /local/path/to/dataset_paths.yaml \ - --json \ - nemo_retriever/harness/runfiles/jp20_beir.json \ - nemo_retriever/harness/runfiles/bo767_beir.json \ - nemo_retriever/harness/runfiles/earnings_beir.json \ - nemo_retriever/harness/runfiles/financebench_beir.json +cp nemo_retriever/harness/dataset_paths.example.yaml \ + /local/path/to/dataset_paths.yaml ``` -The ViDoRe v3 library follows the same runfile-first contract. Before GPU work, -export a Hugging Face read token and validate all remote evaluation partitions: +Relative paths in that file resolve from the file's directory. Settings resolve +from lowest to highest precedence: -```bash -export HF_TOKEN=... -uv run --project nemo_retriever retriever harness check-vidore-access -``` +1. Benchmark registry defaults. +2. Runfile values. +3. Machine-local dataset paths. +4. CLI `--set` values. -The check streams one byte from each queries, qrels, and corpus partition; it -does not download the complete parquet objects. Each benchmark uses the VL -embed model with `text_image` page embeddings and page-level BEIR scoring. Run -one domain while validating a machine or configuration: +Keep credentials, webhooks, and machine-local paths out of runfiles and source +control. -```bash -uv run --project nemo_retriever retriever harness run-files \ - --session-name vidore_v3_computer_science \ - --output-dir /local/path/to/retriever-artifacts/vidore-v3-computer-science \ - --dataset-paths /local/path/to/dataset_paths.yaml \ - --json \ - nemo_retriever/harness/runfiles/vidore_v3_computer_science_beir.json -``` +## Results and Artifacts -After single-domain validation, run all eight public ViDoRe v3 domains as one -portable session: +Use the process exit code and terminal JSON files—not console output—to decide +whether a run succeeded: -```bash -uv run --project nemo_retriever retriever harness run-files \ - --session-name vidore_v3_all \ - --output-dir /local/path/to/retriever-artifacts/vidore-v3-all \ - --dataset-paths /local/path/to/dataset_paths.yaml \ - --json \ - nemo_retriever/harness/runfiles/vidore_v3_computer_science_beir.json \ - nemo_retriever/harness/runfiles/vidore_v3_energy_beir.json \ - nemo_retriever/harness/runfiles/vidore_v3_finance_en_beir.json \ - nemo_retriever/harness/runfiles/vidore_v3_finance_fr_beir.json \ - nemo_retriever/harness/runfiles/vidore_v3_hr_beir.json \ - nemo_retriever/harness/runfiles/vidore_v3_industrial_beir.json \ - nemo_retriever/harness/runfiles/vidore_v3_pharmaceuticals_beir.json \ - nemo_retriever/harness/runfiles/vidore_v3_physics_beir.json -``` +| File | Read it when | +| --- | --- | +| `status.json` | A run is active and you need its current phase. | +| `results.json` | One run is terminal. | +| `session_summary.json` | A `run-files` or `run-set` session is terminal. | -The code-owned `vidore_v3_all` runset is also available when the registry's -default dataset paths are mounted. Prefer the checked-in runfiles for nightly -or other orchestrated sessions because they carry per-dataset integrity gates -and accept a machine-local path map. +Terminal files contain `success`, `exit_code`, summary metrics, and relative +paths to detailed evidence. Follow those paths only when needed: -Real `run-files` sessions execute children sequentially, each in a fresh spawned -process. The process boundary releases Ray and materialized dataframe memory -before the next benchmark while preserving one parent-owned -`session_summary.json`. A code-owned six-hour child deadline prevents one hung -benchmark from blocking the session forever. Dry-runs stay in the parent -process because they do not materialize batch data. +- `run.log` for full errors and lower-level output +- `resolved_benchmark.json`, `ingest_plan.json`, and `query_plan.json` for the + effective configuration +- `environment.json` for the source revision, GPU inventory, workload-visible + GPU count, and runtime context +- `beir_metrics.json`, `beir_run.trec`, and `query_results.jsonl` for evaluation + details +- `service_logs/` for a failed managed Helm deployment -`run-files` owns the session layout and execution mode. Runfiles passed to this -command cannot set their own `output_dir`, `run_id`, or `dry_run`; use the -session-level `--dry-run` flag instead. The session uses the following paths and -identifiers: +A multi-run session has this stable layout: ```text / @@ -228,71 +132,54 @@ identifiers: session_summary.json 001_/ 002_/ - -run ID: __ ``` -Session names and runfile names can contain letters, numbers, periods, -underscores, and hyphens. Other characters fail validation before execution. - -## Provision A Service With Helm +Dry-runs resolve configuration and write planning artifacts, but they are not +execution evidence. -Helm is one way to provision the service under test; it is not a benchmark -execution mode or part of the runfile schema. `run-helm` loads the non-secret -deployment settings in -[`examples/managed-helm-main.yaml`](examples/managed-helm-main.yaml), deploys an -explicit immutable image, waits for readiness and establishes a port-forward, -invokes the portable `run-files` engine with an existing runfile, collects -`service_logs/` on failure, and always tears the release down. +## Gates and Exit Codes -Set `HARNESS_HELM_SERVICE_IMAGE_REPOSITORY` and -`HARNESS_HELM_SERVICE_IMAGE_TAG` to an immutable image built from the checkout. -The external scheduler owns recurrence and the output directory. For JP20, run: +Metric gates compare values in `results.json.summary_metrics`: ```bash -export RETRIEVER_SESSION_DIR=/local/path/to/retriever-artifacts/helm-jp20-$(date -u +%Y%m%d_%H%M%S_UTC) - -uv run --project nemo_retriever retriever-harness run-helm \ - --config nemo_retriever/harness/examples/managed-helm-main.yaml \ - --output-dir "$RETRIEVER_SESSION_DIR" \ - --session-name helm_jp20 \ - --dataset-paths /local/path/to/dataset_paths.yaml \ - nemo_retriever/harness/runfiles/jp20_beir.json +--require 'files==20' +--require 'recall_5>=0.85' ``` -`run-helm` overrides the runfile mode to `service`; the shared JP20 runfile -continues to own its dataset-integrity gates. Recall and nDCG are recorded in -the standard artifacts without adding Helm-specific quality gates. The runner -never reads a Slack webhook. After the terminal session exists, read each -child's `results.json` for metrics and optionally invoke `post-slack --preview` -or `post-slack` as a separate operation. +During a dry-run, static facts can be evaluated; execution metrics are recorded +as skipped. The harness has no implicit quality or performance threshold. +[`expected-results.md`](docs/expected-results.md) contains human-readable observations +that can inform explicit gates. -The legacy module invocation remains supported for compatibility: +| Code | Meaning | +| ---: | --- | +| `0` | Success | +| `2` | Invalid benchmark, configuration, override, or gate | +| `3` | Missing dataset or input | +| `4` | Managed Helm lifecycle failure | +| `10` | Ingest failure | +| `11` | Query failure | +| `12` | Evaluation failure | +| `20` | Metric gate failure | +| `30` | Artifact write failure | +| `70` | Unexpected internal error | -```bash -uv run --project nemo_retriever \ - python -m nemo_retriever.harness.helm_runner --help -``` - -## Post Results to Slack +## Report Completed Results -Harness execution and Slack reporting are separate operations. `run-files` -writes local artifacts and never contacts Slack. `post-slack` reads an existing -session or run artifact, builds a summary, and sends that summary without -rerunning ingestion or queries. +Execution and reporting are separate. Preview a Slack report without a webhook: -Each new run records `gpu_sku` and `gpu_count` from the physical inventory -reported by `nvidia-smi`. It separately records `workload_gpu_count`, the number -of GPUs available to that process after an explicit `CUDA_VISIBLE_DEVICES` -constraint. Slack labels these as physical inventory versus GPUs available to -the workload, so a partition of a larger host is not mistaken for the whole -machine. +```bash +uv run --project nemo_retriever retriever harness post-slack \ + --preview \ + --title "Retriever benchmark results" \ + /path/to/session +``` -This separation lets you inspect a completed session before reporting it and -reuse the same artifacts when report formatting changes. +To post, export `SLACK_WEBHOOK_URL` and remove `--preview`. Each invocation +creates a new message and never changes the completed run. -To show the current release beside each nightly result, provide one external -release-reference snapshot: +A nightly can display the current release beside matching observations. Keep +one release snapshot outside the repository: ```json { @@ -307,290 +194,30 @@ release-reference snapshot: } ``` -Pass it with `post-slack --reference-file PATH`, or set -`RETRIEVER_HARNESS_REFERENCE_FILE` for recurring nightlies. Slack shows the -observed nightly and release values side by side with their GPU context. It -does not assign a verdict or use the reference for pass/fail. Keep the snapshot -outside the repository; advancing to the next RC only requires replacing its -release label and result values. The harness never appends to this file or -maintains reference history. - -### Prerequisites - -Before you post a report, verify the following: - -- The run completed far enough to write `session_summary.json` or - `results.json`. -- The environment includes the `requests` package. -- `SLACK_WEBHOOK_URL` contains an incoming webhook for the destination channel. - -Set the webhook in the process environment: - -```bash -export SLACK_WEBHOOK_URL="https://hooks.slack.com/services/..." -``` - -Do not put the webhook URL in a runfile, dataset paths file, shell argument, or -artifact. Load it from the process environment or a permissions-restricted -secret file outside the repository. - -### Post a Completed Session - -Pass the session directory to `post-slack`: - -```bash -uv run --project nemo_retriever retriever harness post-slack \ - --title "nemo-retriever library benchmarks" \ - "$RETRIEVER_SESSION_DIR" -``` - -You can also pass one or more run artifact directories or `results.json` files. -Each invocation sends a new Slack message; it does not modify the completed -harness artifacts. - -By default, the report includes file and page counts, ingest time, ingest -pages/sec, query count, recall, nDCG, and environment details when those values -are available. Use repeated `--metric-key` options to select a different metric -set. Use `--artifact-paths` when recipients can access the runner's local paths. - -ViDoRe v3 results always use a compact suite layout: total ingest time, -aggregate pages/sec, and a separate Recall@5/nDCG@10 table. A complete -eight-domain suite also reports macro averages across the seven English -datasets and across all datasets. Per-domain timing and other metadata remain -available in the session artifacts. This format also applies when previewing or -reposting completed artifacts; reporting never reruns ingestion or queries. +Pass it with `--reference-file` or set +`RETRIEVER_HARNESS_REFERENCE_FILE`. The report shows the two observations with +their GPU context; it does not assign a verdict, update the file, or maintain +history. -### Preview Report Formatting - -Use `--preview` to render the exact Slack payload without reading -`SLACK_WEBHOOK_URL` or making an HTTP request: +Compare two local runs without Slack: ```bash -uv run --project nemo_retriever retriever harness post-slack \ - --preview \ - --title "nemo-retriever library benchmarks" \ - "$RETRIEVER_SESSION_DIR" +uv run --project nemo_retriever retriever harness diff \ + /path/to/left/results.json \ + /path/to/right/results.json ``` -Preview the same completed session as often as needed while adjusting the -title, metric selection, or artifact-path setting. When the payload is ready, -run the command again without `--preview` to post it. Preview and posting use -the same artifact loader and payload formatter. - -`post-slack` has its own exit status and never changes the completed session's -status or artifacts. Any policy that combines run status, report status, -retries, locking, or recurrence belongs to the caller. - -## Controls And Overrides - -Benchmarks are code-owned defaults. Use `--set KEY=VALUE` for one-off -ablations, or put the same keys under `set` in a runfile for reproducible -agent/orchestrator runs. - -Examples: - -```bash -retriever harness run jp20_beir \ - --set query.top_k=20 \ - --set query.rerank=true \ - --set ingest.extract.batch.page_elements_workers=1 -``` - -Runfile equivalent: - -```json -{ - "schema_version": 1, - "benchmark": "bo767_beir", - "mode": "batch", - "set": { - "query.top_k": 10, - "ingest.extract.batch.pdf_extract_workers": 8, - "ingest.embed.batch.embed_batch_size": 64 - } -} -``` - -Supported override namespaces: - -- `dataset.*`: dataset path, query/qrels file, input type, BEIR loader, and - BEIR doc ID settings. -- `ingest.*`: profile, input type, Ray mode/address, extraction/media/caption, - dedup, chunk, embedding, image-store, storage, and batch worker settings. -- `query.*`: top-k, candidate-k, page dedup, content types, retrieval mode, - embedding endpoint/model, reranking, LanceDB URI, and table name. -- `evaluation.*`: evaluation mode, BEIR loader/dataset/split/language/doc ID - field, and metric cutoffs. - -Unknown override keys fail during resolution with exit code `2`. Values are -parsed as YAML scalars/lists/maps, so booleans, numbers, nulls, and lists can be -passed naturally. - -Use `retriever harness show --json` and `retriever harness run - --dry-run --json` to inspect the exact resolved benchmark and -plans before launching an expensive run. - -## Implementation Boundary - -The harness does not shell out to `retriever ingest` or `retriever query`. It -calls the same Python workflow/planning modules used by the CLI: +## Automation Contract -- ingest: `resolve_ingest_plan(...)` and `run_ingest_workflow(...)` -- query: `resolve_query_plan(...)` and shared query workflow objects -- BEIR: harness-owned query iteration over the resolved query plan +For agents and other callers: -For `mode: service`, the corresponding service ingest and query request APIs -replace the in-process plans. Helm deployment remains in `helm_runner.py`, -outside this benchmark contract. - -The harness controller calls those APIs in its Python process; this does not -force the ingest workload into local/in-process mode. A runfile with `mode: -batch` still resolves to Ray-backed batch ingest, while `mode: local` resolves -to local in-process ingest. Stdout remains diagnostic only; artifacts and exit -codes are the contract. - -## Artifacts - -Use one entrypoint for each lifecycle level instead of scanning the directory: - -- `status.json`: current phase and concise failure state while a run is active. -- `results.json`: authoritative terminal result, summary metrics, and relative - pointers to detailed run evidence. -- `session_summary.json`: authoritative terminal result for `run-set` and - `run-files` sessions, with relative pointers to each child run. - -The terminal files are deliberately compact. A successful run has this shape: - -```json -{ - "benchmark": "jp20_beir", - "dataset": "jp20", - "success": true, - "exit_code": 0, - "summary_metrics": {"files": 20, "pages": 1940, "recall_5": 0.887}, - "failure": null, - "artifacts": {"log": "run.log", "lancedb": "lancedb"} -} -``` - -A multi-run session summarizes its children without embedding their detailed -results: - -```json -{ - "session_type": "runfiles", - "session_name": "library_beir", - "run_commit": "0123456789abcdef0123456789abcdef01234567", - "working_tree_dirty": false, - "success": true, - "exit_code": 0, - "dry_run": false, - "isolate_runs": true, - "runs": [ - { - "benchmark": "jp20_beir", - "success": true, - "results_path": "001_jp20_beir/results.json" - } - ] -} -``` - -Follow the pointers in `results.json` only when deeper evidence is needed: - -- `events.jsonl`: phase transitions and harness events. -- `runfile.json`: original runfile payload, when a runfile was used. -- `resolved_benchmark.json`: exact effective benchmark spec. -- `ingest_plan.json`: redacted executable ingest plan. -- `query_plan.json`: executable query plan. -- `environment.json`: commit and runtime context. -- `run.log`: captured lower-level stdout/stderr and full exception tracebacks. -- `beir_metrics.json`: full BEIR metric family when evaluation executes. -- `beir_run.trec`: standard TREC runfile when evaluation executes. -- `query_results.jsonl`: per-query latency and ranked hits. -- `lancedb/`: the ingested table and index used by the run. - -Artifact manifest paths are relative to the run directory so a copied session -remains readable. Failure messages in `status.json` and `results.json` are kept -concise; use the listed debug artifacts, normally `run.log`, for full traces. -When an output directory is reused, the harness removes only its known generated -artifacts before starting; unrelated files in that directory are preserved. - -`environment.json` records an allowlisted set of reproducibility diagnostics, -including source revision, Python and package versions, accelerator information, -and selected Hugging Face, Ray, CUDA, and vLLM settings. Credentials and webhook -values are not recorded. - -New runs keep summary metrics inside `results.json`; they do not emit a separate -`summary_metrics.json`. `diff` and `post-slack` retain read compatibility with -older harness artifacts. - -Dry-runs write the terminal status/result manifests and planning artifacts. They -do not create empty `run.log`, `beir_metrics.json`, `beir_run.trec`, -`query_results.jsonl`, or `lancedb/`. -Treat `--dry-run` as configuration preflight, not execution evidence. When model -startup or runtime behavior changes, follow it with a real run of the smallest -appropriate registered benchmark. - -## Gates - -Use explicit `--require` gates. Gate expressions compare keys from -`results.json.summary_metrics`: - -```bash ---require 'files==20' ---require 'recall_5>=0.85' ---require 'query_latency_p95_ms<=1200' -``` +1. Discover benchmarks with `list --runsets --json`. +2. Inspect the target with `show --json`. +3. Supply machine-local dataset paths outside the repository. +4. Dry-run the exact command before expensive GPU work. +5. Use the exit code and terminal JSON file as the result. +6. Read detailed artifacts only when the terminal result points to them. -Gate failures exit with code `20` and still write artifacts. - -During `--dry-run`, gates for unavailable execution metrics are skipped and -listed in `results.json` as `skipped_metric_gates`. Static gates such as -`files==20` and `pages==1940` are still evaluated. - -Known dataset facts, observed result ranges, and suggested gates live in -[`EXPECTED_RESULTS.md`](EXPECTED_RESULTS.md). Keep threshold knowledge there, -not in benchmark Python code. - -## Agent Instructions - -For automated harness work: - -1. Use this harness only for registered benchmark/evaluation work. Use - `retriever ingest` and `retriever query` for direct product workflows. -2. Start with `retriever harness list --runsets --json`, then inspect the target - with `retriever harness show --json`. -3. Copy `dataset_paths.example.yaml` outside the repository and set the paths - available on the current machine. -4. Use `run-files --dataset-paths ...` with one runfile for one dataset or - multiple runfiles for a suite. -5. Always set `--output-dir`. Use `--dry-run` to preflight paths, overrides, and - gates before expensive execution, then use a small real benchmark when - runtime behavior needs validation. -6. Use explicit `--require` gates from `EXPECTED_RESULTS.md`. -7. Decide success from the process exit code and `results.json` for one run or - `session_summary.json` for a session. -8. Read `summary_metrics` from the applicable terminal JSON file. Follow its - pointers to `run.log` or other detailed evidence only when needed. -9. Do not parse progress bars, human CLI formatting, or raw stdout as the API. -10. Treat `post-slack` as optional post-processing. Previewing or posting never - executes a benchmark. - -## Exit Codes - -- `0`: success -- `2`: invalid benchmark/config/override/gate syntax -- `3`: dataset or input missing -- `10`: ingest failure -- `11`: query failure -- `12`: evaluation failure -- `20`: metric gate failure -- `30`: artifact write failure -- `70`: unexpected internal error - -## More Detail - -- [`EXPECTED_RESULTS.md`](EXPECTED_RESULTS.md): dataset facts, observed metrics, - and suggested explicit gates. -- [`HANDOFF.md`](HANDOFF.md): concise maintainer-oriented implementation notes. +Do not parse progress bars, human CLI formatting, or raw stdout. Scheduling, +retry policy, secret distribution, and result history belong to the caller. The +[nightly launcher](../../ops/retriever-nightly/README.md) is one such caller. diff --git a/nemo_retriever/harness/dataset_paths.example.yaml b/nemo_retriever/harness/dataset_paths.example.yaml index 81e22363b1..f4fe103964 100644 --- a/nemo_retriever/harness/dataset_paths.example.yaml +++ b/nemo_retriever/harness/dataset_paths.example.yaml @@ -1,10 +1,10 @@ schema_version: 1 -# Copy this file outside the checkout, replace the paths for the host, and pass -# the YAML file itself with --dataset-paths. Do not pass a dataset directory. +# Copy this file outside the checkout, replace the corpus paths for the host, +# and pass the YAML file itself with --dataset-paths. Keep query_file only for +# annotations that are not checked into data/. Do not pass a dataset directory. datasets: jp20: path: /path/to/datasets/nv-ingest/jp20 - query_file: /path/to/datasets/nv-ingest/jp20_query_gt.csv bo767: path: /path/to/datasets/nv-ingest/bo767 query_file: /path/to/NeMo-Retriever/data/bo767_query_gt.csv diff --git a/nemo_retriever/harness/docs/expected-results.md b/nemo_retriever/harness/docs/expected-results.md new file mode 100644 index 0000000000..f110447bf1 --- /dev/null +++ b/nemo_retriever/harness/docs/expected-results.md @@ -0,0 +1,132 @@ + + + +# Harness Expected Results + +This page records dataset facts and observed benchmark results. It is not an +executable baseline or universal pass/fail policy. + +Checked-in runfiles enforce only portable integrity facts such as file, page, +and query counts. Quality and performance vary with the Retriever revision, +hardware, runtime, and model configuration. Use those observations for review, +then add an explicit `--require` gate only when the environment and intended +policy justify it. + +## Dataset Facts + +Registry paths describe the standard internal mount. Use +[`dataset_paths.example.yaml`](../dataset_paths.example.yaml) to map these datasets +on another host. + +| Dataset | Benchmark | Standard path | Files | Pages | Queries | +| --- | --- | --- | ---: | ---: | ---: | +| JP20 | `jp20_beir` | `/datasets/nv-ingest/jp20` | 20 | 1,940 | 115 | +| BO20 | — | `/datasets/nv-ingest/bo20` | 20 | — | — | +| BO767 | `bo767_beir` | `/datasets/nv-ingest/bo767` | 767 | 54,730 | 991 | +| FinanceBench | `financebench_beir` | `/datasets/nv-ingest/foundation_rag/financebench` | 369 | 54,057 | 150 | +| BO10K | Not yet defined | `/datasets/nv-ingest/bo10k` | 10,000 | — | — | +| Earnings Consulting | `earnings_beir` | `/datasets/nv-ingest/earnings_consulting_flattened` | 514 | 12,988 | 628 | + +The query files are: + +| Dataset | Query or qrels file | +| --- | --- | +| JP20 | `data/jp20_query_gt.csv` | +| BO767 | `data/bo767_query_gt.csv` | +| FinanceBench | `data/financebench_train.json` | +| BO10K | `data/digital_corpora_10k_annotations.csv` | +| Earnings Consulting | `data/earnings_consulting_multimodal.csv` | + +## JP20 Observations + +Recent `jp20_beir` runs on local hardware: + +| Metric | Observed | +| --- | ---: | +| Rows processed | 3,154 | +| Ingest seconds | 215–223 | +| Query p50 | 909–915 ms | +| Query p95 | 953–1,003 ms | +| Recall@5 | 0.878–0.887 | +| Recall@10 | 0.930–0.948 | +| nDCG@10 | 0.793–0.802 | + +The original observations did not record the GPU SKU or workload-visible GPU +count. Treat their latency and throughput as context only. + +## BO767 Observations + +Runs on an eight-H100 80GB HBM3 host: + +| Configuration | Workload GPUs | Rows | Ingest seconds | Pages/s | Recall@5 | Recall@10 | nDCG@10 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| RC26.05 Perflab | Not recorded | 79,221 | 4,036.847 | 13.56 | — | — | — | +| Automatic batch | 1 | 79,229 | 1,594.339 | 34.328 | 0.848638 | 0.897074 | 0.750110 | +| Automatic batch | 8 | 79,230 | 764.946 | 71.548 | 0.849647 | 0.895055 | 0.748583 | +| Legacy worker-capped batch | 8 visible, effectively 1 used | 79,230 | about 2,265 | about 24.16 | 0.850656 | 0.896065 | 0.751507 | + +The RC26.05 artifact recorded eight physical GPUs but not the number visible to +the workload. Current runfiles leave worker counts and batch sizes automatic so +the planner can scale to available GPUs. The one- and eight-GPU automatic runs +differed by one output row; keep row counts and quality visible when comparing +throughput. + +## ViDoRe V3 + +The eight public ViDoRe v3 benchmarks use original PDFs with: + +- `nvidia/llama-nemotron-embed-vl-1b-v2` +- `text_image` page embeddings +- page-image and infographic extraction +- page-level BEIR document IDs + +Integrity facts: + +| Dataset | Files | Pages | Queries | +| --- | ---: | ---: | ---: | +| Computer Science | 2 | 1,360 | 1,290 | +| Energy | 41 | 2,225 | 1,848 | +| Finance EN | 6 | 2,942 | 1,854 | +| Finance FR | 5 | 2,384 | 1,920 | +| HR | 14 | 1,110 | 1,908 | +| Industrial | 27 | 5,244 | 1,698 | +| Pharmaceuticals | 52 | 2,313 | 2,184 | +| Physics | 42 | 1,674 | 1,812 | + +Observed complete batch runs on an eight-H100 DGX: + +| Dataset | Indexed rows | Ingest seconds | Pages/s | Recall@5 | Recall@10 | nDCG@10 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| Computer Science | 1,358 | 100.2–123.5 | 11.0–13.6 | 0.599–0.600 | 0.729–0.730 | 0.708–0.709 | +| Energy | 2,211 | 116.7 | 19.1 | 0.575 | 0.674 | 0.581 | +| Finance EN | 2,927 | 149.4 | 19.7 | 0.496 | 0.609 | 0.547 | +| Finance FR | 2,149 | 106.4 | 22.4 | 0.324 | 0.426 | 0.345 | +| HR | 1,091 | 82.6 | 13.4 | 0.452 | 0.574 | 0.530 | +| Industrial | 5,039 | 137.5 | 38.1 | 0.348 | 0.426 | 0.381 | +| Pharmaceuticals | 2,290 | 93.7 | 24.7 | 0.547 | 0.647 | 0.607 | +| Physics | 1,674 | 89.2 | 18.8 | 0.369 | 0.485 | 0.451 | + +Computer Science was run twice; the other domains have one complete observation. +The observed all-domain macro-average nDCG@10 was about `0.519`. The observations +predate separate physical and workload-visible GPU counts, so do not infer +scaling behavior from their throughput. + +### RC26.05 Recall Comparison + +| Dataset | RC26.05 Recall@5 | Observed Recall@5 | +| --- | ---: | ---: | +| Finance EN | 0.499 | 0.496 | +| Industrial | 0.348 | 0.348 | +| Computer Science | 0.600 | 0.599 | +| Pharmaceuticals | 0.549 | 0.547 | +| HR | 0.453 | 0.452 | +| Energy | 0.577 | 0.575 | +| Physics | 0.367 | 0.369 | +| Finance FR | 0.324 | 0.324 | + +The observed macro-average Recall@5 was `0.464` across all eight domains. + +Finance FR omitted 235 empty-text pages, including 69 judged image-only pages, +because the current dense LanceDB path drops records without text even when an +image embedding exists. Its Recall@5 still matched RC26.05, but preserving those +judged image-only pages remains correctness work. diff --git a/nemo_retriever/harness/docs/library.md b/nemo_retriever/harness/docs/library.md new file mode 100644 index 0000000000..4d204e39d2 --- /dev/null +++ b/nemo_retriever/harness/docs/library.md @@ -0,0 +1,159 @@ + + + +# Library Harness Runs + +Use library execution to benchmark the current NeMo Retriever checkout directly. +Choose the mode by workload size: + +| Mode | Use it for | Execution | +| --- | --- | --- | +| `local` | Small smoke tests such as JP20 | In-process ingest | +| `batch` | Larger corpora and multi-GPU runs | Ray-backed batch ingest | + +Both modes write the [same artifacts](../README.md#results-and-artifacts). + +## Prepare the Host + +You need the repository dependencies, NVIDIA drivers, the selected datasets, +and enough storage for artifacts. Run commands from the repository root through +the `nemo_retriever` project. + +Dataset mounts vary by machine: + +```bash +cp nemo_retriever/harness/dataset_paths.example.yaml \ + /local/path/to/dataset_paths.yaml +${EDITOR:-vi} /local/path/to/dataset_paths.yaml +``` + +Keep this file outside the repository. It maps registered dataset names to local +corpus paths and any annotations that are not checked in. JP20 uses the +repository's `data/jp20_query_gt.csv` unless `query_file` is explicitly +overridden. + +## Run One Benchmark + +Start with the checked-in JP20 runfile: + +```bash +uv run --project nemo_retriever retriever harness run-files \ + --session-name jp20_check \ + --output-dir /tmp/retriever-harness-jp20-check \ + --dataset-paths /local/path/to/dataset_paths.yaml \ + --dry-run \ + nemo_retriever/harness/runfiles/jp20_beir.json +``` + +Confirm that `session_summary.json` succeeds and inspect the child +`resolved_benchmark.json`. Then run the same command without `--dry-run`. + +If the registry paths already exist on the host, `run` is a shorter single-run +form: + +```bash +uv run --project nemo_retriever retriever harness run jp20_beir \ + --output-dir /tmp/retriever-harness-jp20 \ + --require 'files==20' \ + --require 'pages==1940' +``` + +## Run in Batch + +Checked-in runfiles for BO767, FinanceBench, Earnings, and ViDoRe already select +`batch`. Run one of them exactly as you ran JP20: + +```bash +uv run --project nemo_retriever retriever harness run-files \ + --session-name bo767_beir \ + --output-dir /local/path/to/retriever-artifacts/bo767-beir \ + --dataset-paths /local/path/to/dataset_paths.yaml \ + nemo_retriever/harness/runfiles/bo767_beir.json +``` + +Pass several runfiles to create one sequential session: + +```bash +uv run --project nemo_retriever retriever harness run-files \ + --session-name library_beir \ + --output-dir /local/path/to/retriever-artifacts/library-beir \ + --dataset-paths /local/path/to/dataset_paths.yaml \ + nemo_retriever/harness/runfiles/jp20_beir.json \ + nemo_retriever/harness/runfiles/bo767_beir.json \ + nemo_retriever/harness/runfiles/earnings_beir.json \ + nemo_retriever/harness/runfiles/financebench_beir.json +``` + +Children run sequentially in fresh processes. Ray and model resources are +released between datasets, while the parent writes one +`session_summary.json`. + +## Run ViDoRe + +ViDoRe queries, qrels, and corpus metadata require Hugging Face access. Check it +before GPU work: + +```bash +export HF_TOKEN=... +uv run --project nemo_retriever retriever harness check-vidore-access +``` + +Then run one domain with its checked-in runfile: + +```bash +uv run --project nemo_retriever retriever harness run-files \ + --session-name vidore_v3_computer_science \ + --output-dir /local/path/to/retriever-artifacts/vidore-v3-computer-science \ + --dataset-paths /local/path/to/dataset_paths.yaml \ + nemo_retriever/harness/runfiles/vidore_v3_computer_science_beir.json +``` + +The [nightly launcher](../../../ops/retriever-nightly/README.md) runs all eight +ViDoRe domains plus the four library benchmarks. + +## Apply a One-Off Override + +Use repeated `--set KEY=VALUE` options for an experiment: + +```bash +uv run --project nemo_retriever retriever harness run jp20_beir \ + --set query.top_k=20 \ + --set query.rerank=true +``` + +Supported namespaces are `dataset.*`, `ingest.*`, `query.*`, and +`evaluation.*`. Unknown keys fail before execution. Use `show --json` and +`run --dry-run --json` to inspect the resolved configuration. + +For reproducible recurring changes, put the values in a reviewed runfile or +benchmark definition rather than a shell command. + +## Evaluate Agentic Retrieval + +Agentic BEIR evaluation uses the same retrieval graph as `retriever query +--agentic` after ingest: + +```bash +uv run --project nemo_retriever retriever harness run jp20_beir \ + --set query.agentic=true +``` + +The default agent LLM runs locally. To use an OpenAI-compatible endpoint: + +```bash +uv run --project nemo_retriever retriever harness run jp20_beir \ + --set query.agentic=true \ + --set query.agentic_llm_model=custom-remote-model \ + --set query.agentic_invoke_url=http://localhost:9000/v1/chat/completions +``` + +Use `query.agentic_llm_client`, `query.agentic_react_max_steps`, +`query.agentic_num_concurrent`, and the local vLLM resource overrides only when +the experiment requires them. The resolved benchmark records every value. + +## Run on a Schedule + +The harness itself is one-shot. Use the +[nightly launcher](../../../ops/retriever-nightly/README.md) for the checked-in +twelve-benchmark suite, Git selection, Slack reporting, and a transparent daily +`tmux` loop. diff --git a/nemo_retriever/harness/docs/service.md b/nemo_retriever/harness/docs/service.md new file mode 100644 index 0000000000..ea191c6e41 --- /dev/null +++ b/nemo_retriever/harness/docs/service.md @@ -0,0 +1,108 @@ + + + +# Service Harness Runs + +Service execution runs the same registered benchmarks through Retriever service +ingest and query APIs. It writes the same artifacts and evaluates the same +metric gates as library execution. + +Choose how the service is provided: + +| Target | Command | +| --- | --- | +| An already-running service | `retriever harness run-files --mode service` | +| A temporary service provisioned by the harness | `retriever harness run-helm` | + +Helm is a provisioning wrapper, not a fourth benchmark mode. + +## Test an Existing Service + +The service must be reachable from the harness host. Export its bearer token +only when authentication is enabled: + +```bash +export HARNESS_SERVICE_API_TOKEN=... +``` + +Run a checked-in benchmark against the endpoint: + +```bash +uv run --project nemo_retriever retriever harness run-files \ + --mode service \ + --service-endpoint http://localhost:7670 \ + --session-name service_jp20 \ + --output-dir /local/path/to/retriever-artifacts/service-jp20 \ + --dataset-paths /local/path/to/dataset_paths.yaml \ + nemo_retriever/harness/runfiles/jp20_beir.json +``` + +`--mode service` overrides the runfile mode for the session. +`--service-endpoint` applies only to service-mode children, so mixed sessions +remain possible. If omitted, the endpoint defaults to +`http://localhost:7670`. + +The dataset paths identify documents and queries available to the harness +client. Service storage and infrastructure remain owned by the service. + +## Provision a Service with Helm + +`run-helm` deploys one immutable service image, waits for readiness, establishes +a local port-forward, runs the shared `run-files` session in service mode, +collects logs on failure, and tears the release down. + +Prerequisites: + +- working `helm` and `kubectl` commands, or equivalents selected in the config +- a Kubernetes cluster with the required image-pull and NGC secrets +- an immutable service image built from the checkout +- a non-secret Helm deployment file + +Start from +[`examples/managed-helm-main.yaml`](../examples/managed-helm-main.yaml). Select a +checked-in profile such as [`helm-profiles/core.yaml`](../helm-profiles/core.yaml) +with `helm_values_file`. + +Set the immutable image: + +```bash +export HARNESS_HELM_SERVICE_IMAGE_REPOSITORY=nvcr.io/example/nrl-service +export HARNESS_HELM_SERVICE_IMAGE_TAG= +``` + +Then run: + +```bash +uv run --project nemo_retriever retriever harness run-helm \ + --config nemo_retriever/harness/examples/managed-helm-main.yaml \ + --output-dir /local/path/to/retriever-artifacts/helm-jp20 \ + --session-name helm_jp20 \ + --dataset-paths /local/path/to/dataset_paths.yaml \ + nemo_retriever/harness/runfiles/jp20_beir.json +``` + +The deployment file may configure the chart, release, namespace, values file, +`helm_set` overrides, timeouts, local service port, command paths, and whether +Helm or kubectl requires `sudo`. Keep credentials out of it. + +Image tags named `latest`, `main`, or `nightly` are rejected because they are +not reproducible. `run-helm` returns the benchmark session status unless +deployment or teardown fails, in which case it returns `4`. + +## Inspect and Report Results + +Read `session_summary.json` first. Failed managed sessions may also contain +`service_logs/`. The artifact tree records the benchmark configuration and +runtime environment. Keep the immutable image reference and non-secret Helm +configuration with the surrounding job record. + +Reporting is a separate step: + +```bash +uv run --project nemo_retriever retriever harness post-slack \ + --preview \ + /local/path/to/retriever-artifacts/helm-jp20 +``` + +See the [shared artifact and reporting contract](../README.md#results-and-artifacts) +for exit codes, Slack posting, and release comparisons. diff --git a/nemo_retriever/harness/helm-profiles/no-nims-external.yaml b/nemo_retriever/harness/helm-profiles/no-nims-external.yaml index d59d4a04fa..1616b6ef30 100644 --- a/nemo_retriever/harness/helm-profiles/no-nims-external.yaml +++ b/nemo_retriever/harness/helm-profiles/no-nims-external.yaml @@ -6,9 +6,9 @@ topology: serviceConfig: nimEndpoints: - pageElementsInvokeUrl: "http://page-elements-nim.example.invalid:8000/v1/infer" - tableStructureInvokeUrl: "http://table-structure-nim.example.invalid:8000/v1/infer" - ocrInvokeUrl: "http://ocr-nim.example.invalid:8000/v1/infer" + pageElementsInvokeUrl: "http://page-elements-nim.example.invalid:8000/v1/page-elements" + tableStructureInvokeUrl: "http://table-structure-nim.example.invalid:8000/v1/table-structure" + ocrInvokeUrl: "http://ocr-nim.example.invalid:8000/v1/ocr" embedInvokeUrl: "http://embed-nim.example.invalid:8000/v1/embeddings" audioGrpcEndpoint: "audio-nim.example.invalid:50051" vectordb: diff --git a/nemo_retriever/helm/Chart.yaml b/nemo_retriever/helm/Chart.yaml index d711d0134a..8f377aac81 100644 --- a/nemo_retriever/helm/Chart.yaml +++ b/nemo_retriever/helm/Chart.yaml @@ -6,7 +6,7 @@ description: | This chart deploys the "service" run mode of nemo-retriever as a single scalable web endpoint that orchestrates document ingestion against - remote NIM endpoints (page-elements, table-structure, OCR, VLM embed). + remote NIM endpoints (object detection, OCR, VLM embed). In-cluster NIMs are provisioned via NVIDIA NIM Operator custom resources (NIMCache + NIMService, apps.nvidia.com/v1alpha1). The diff --git a/nemo_retriever/helm/README.md b/nemo_retriever/helm/README.md index 6e48784a19..0cbf88db33 100644 --- a/nemo_retriever/helm/README.md +++ b/nemo_retriever/helm/README.md @@ -3,7 +3,7 @@ A Kubernetes Helm chart for running the **service** mode of [`nemo-retriever`](../README.md): a FastAPI document ingestion server that streams uploads through a set of NVIDIA NIM microservices -(page-elements, table-structure, OCR, VLM embed by default) and exposes +(object detection, OCR, VLM embed by default) and exposes result + status APIs over HTTP / SSE. Use **Helm** (this chart and/or the **additional Library charts** documented in the @@ -67,9 +67,9 @@ nemo_retriever/helm/ ├── pvc.yaml # general persistence PVC ├── secrets.yaml # ngc-secret + ngc-api └── nims/ - ├── nemotron-page-elements-v3.yaml # NIMCache + NIMService + ├── nemotron-page-elements-v3.yaml # NIMCache + NIMService ├── nemotron-table-structure-v1.yaml # NIMCache + NIMService - ├── nemotron-ocr-v2.yaml # NIMCache + NIMService (OCR) + ├── nemotron-ocr-v2.yaml # NIMCache + NIMService ├── llama-nemotron-embed-vl-1b-v2.yaml # NIMCache + NIMService (VLM embed) ├── llama-nemotron-rerank-vl-1b-v2.yaml # NIMCache + NIMService (optional; not auto-wired) ├── nemotron-parse.yaml # NIMCache + NIMService (optional; not auto-wired) @@ -83,7 +83,7 @@ nemo_retriever/helm/ ### 1. Service image { #1-service-image } -The chart defaults to the GA image published to NGC: +The chart defaults to the image published to NGC: ``` nvcr.io/nvidia/nemo-microservices/nrl-service:26.5.0 @@ -100,6 +100,7 @@ then override `service.image.repository` / `service.image.tag`: # from the repo root: docker build \ --target service \ + --build-arg DOWNLOAD_DEFAULT_TOKENIZER=True \ -t /nemo-retriever-service: . docker push /nemo-retriever-service: ``` @@ -170,12 +171,12 @@ helm install retriever ./nemo_retriever/helm \ --set ngcImagePullSecret.password=$NGC_API_KEY \ --set ngcApiSecret.create=true \ --set ngcApiSecret.password=$NGC_API_KEY \ - --set serviceConfig.nimEndpoints.pageElementsInvokeUrl=http://page-elements.svc:8000/v1/infer \ - --set serviceConfig.nimEndpoints.tableStructureInvokeUrl=http://table-structure.svc:8000/v1/infer \ - --set serviceConfig.nimEndpoints.ocrInvokeUrl=http://ocr.svc:8000/v1/infer \ + --set serviceConfig.nimEndpoints.pageElementsInvokeUrl=http://page-elements.svc:8000/v1/page-elements \ + --set serviceConfig.nimEndpoints.tableStructureInvokeUrl=http://table-structure.svc:8000/v1/table-structure \ + --set serviceConfig.nimEndpoints.ocrInvokeUrl=http://ocr.svc:8000/v1/ocr \ --set serviceConfig.nimEndpoints.embedInvokeUrl=http://embed.svc:8000/v1/embeddings -``` +``` `ngcApiSecret` materialises an `ngc-api` Secret containing both `NGC_API_KEY` and `NGC_CLI_API_KEY` keys; the service container reads it via `optional: true` `secretKeyRef`, so the install still succeeds when @@ -191,7 +192,7 @@ NIM (the VL reranker `rerankqa`, Nemotron Parse, Omni 30B, and the Parakeet `audio` ASR NIM) is **disabled by default** to honor the "optional and disabled by default" contract in [deployment-options.md](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/deployment-options.md); -refer to [Recommended minimal install](#recommended-minimal-install-2605) +refer to [Recommended minimal install](#recommended-minimal-install-2608) for the opt-in `--set` flags that turn any of them on. ```bash @@ -202,7 +203,7 @@ helm install retriever ./nemo_retriever/helm \ --set ngcApiSecret.password=$NGC_API_KEY ``` -### Recommended minimal install (26.05) { #recommended-minimal-install-2605 } +### Recommended minimal install (26.08) { #recommended-minimal-install-2608 } Deploy only the four core NIMs that the retriever service auto-wires (`page_elements`, `table_structure`, `ocr`, `vlm_embed`): @@ -211,7 +212,8 @@ helm install retriever ./nemo_retriever/helm \ --set ngcImagePullSecret.create=true \ --set ngcImagePullSecret.password=$NGC_API_KEY \ --set ngcApiSecret.create=true \ - --set ngcApiSecret.password=$NGC_API_KEY + --set ngcApiSecret.password=$NGC_API_KEY \ + --set service.image.tag=26.8.0 ``` > The VL reranker (`rerankqa`), Nemotron Parse, the Nemotron 3 Nano Omni 30B caption NIM, the generic answer-generation LLM (`answer_llm`, Super-49B defaults), and the Parakeet `audio` ASR NIM are **all off by default** — they only reconcile when you explicitly opt in. Opt-in flags: @@ -224,21 +226,21 @@ helm install retriever ./nemo_retriever/helm \ > > This matches the "optional and disabled by default" contract in [deployment-options.md](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/deployment-options.md) and avoids silently pulling ≈ 62 GiB of Omni weights, loading a large two-GPU LLM, or claiming extra dedicated GPUs on a "default" install. Refer to the [model hardware requirements](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/prerequisites-support-matrix.md#model-hardware-requirements) table for per-NIM GPU and disk costs. -The chart auto-wires the operator-managed in-cluster URLs of the four +The chart auto-wires the operator-managed in-cluster URLs of the three "core" NIMs into the service's `nim_endpoints` block: | key | operator-managed Service | invoke path | | --- | ------------------------ | ----------- | -| `nimOperator.page_elements` | `nemotron-page-elements-v3` | `/v1/infer` | -| `nimOperator.table_structure` | `nemotron-table-structure-v1` | `/v1/infer` | -| `nimOperator.ocr` | `nemotron-ocr-v2` | `/v1/infer` | +| `nimOperator.page_elements` | `nemotron-page-elements-v3` | `/v1/page-elements` | +| `nimOperator.table_structure` | `nemotron-table-structure-v1` | `/v1/table-structure` | +| `nimOperator.ocr` | `nemotron-ocr-v2` | `/v1/ocr` | | `nimOperator.vlm_embed` | `llama-nemotron-embed-vl-1b-v2` | `/v1/embeddings` | Track operator reconciliation with: ```bash kubectl get nimcache,nimservice -n -kubectl describe nimservice nemotron-page-elements-v3 -n +kubectl describe nimservice nemotron-object-detection -n ``` First-time NIMCache reconciliation downloads model weights to a PVC. By @@ -290,7 +292,7 @@ short list of knobs you'll touch first. | Path | Default | Notes | |-------------------------------|------------------------------------|-------| -| `service.image.repository` | `nvcr.io/nvidia/nemo-microservices/nrl-service` | GA NGC image; override to pin a different build or use a local registry. | +| `service.image.repository` | `nvcr.io/nvidia/nemo-microservices/nrl-service` | NGC image; override to pin a different build or use a local registry. | | `service.image.tag` | `26.5.0` | | | `service.replicas` | `1` | Keep at 1 because standalone job and scheduler state are process-local. | | `service.installFfmpeg` | `false` | Install `ffmpeg`/`ffprobe` at container startup by setting `INSTALL_FFMPEG=true`. Requires network egress, writable root filesystem, and sudo/setuid allowed. Not for air-gapped clusters — use a custom image instead. | @@ -307,7 +309,7 @@ For air-gapped clusters, refer to [Deployment options — Air-gapped and disconn To run self-hosted Parakeet for [audio and video extraction](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/audio-video.md): -1. Set `nimOperator.audio.enabled=true` (it is on by default; disable other optional NIMs you do not need per [Recommended minimal install](#recommended-minimal-install-2605)). +1. Set `nimOperator.audio.enabled=true` (it is on by default; disable other optional NIMs you do not need per [Recommended minimal install](#recommended-minimal-install-2608)). 2. Pin the ASR `NIMService` to a **dedicated GPU** with `nimOperator.audio.resources`, `nodeSelector`, or `tolerations` (refer to [NIM Operator](https://docs.nvidia.com/nim-operator/latest/index.html)). 3. Confirm the GPU SKU in [Model hardware requirements](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/prerequisites-support-matrix.md#model-hardware-requirements) (footnote ⁴ lists Blackwell limitations). 4. Set `service.installFfmpeg=true` when the retriever service will process audio or video on clusters that allow runtime package install (refer to `service.installFfmpeg` above). On **OpenShift restricted-v2**, use a [prebuilt service image](./openshift.md#audio-and-video-ffmpeg-on-restricted-openshift) instead. @@ -331,6 +333,10 @@ The retriever service picks up the in-cluster ASR endpoint when `nimOperator.aud | `serviceConfig.llm.model` | `""` | Optional explicit LiteLLM model id. Leave empty to inherit `nimOperator.answer_llm.model` when using the operator-managed answer LLM; set it for external endpoints. | | `serviceConfig.llm.ragSystemPromptPrefix` | `""` | Optional explicit RAG prompt prefix. Leave empty unless an endpoint needs model-specific prompt directives. | | `serviceConfig.llm.reasoningEnabled` | `true` | Request-level reasoning toggle for `/v1/answer`. Defaults to true for external OpenAI-compatible providers; set false for Nemotron endpoints that should receive portable no-reasoning controls. | +| `serviceConfig.agentic.enabled` | `false` | Enables `POST /v1/query` with `agentic=true` and the additive `agentic_query` MCP tool. | +| `serviceConfig.agentic.llmModel` | `""` | Chat model used by the inner agentic retrieval loop. Required when `invokeUrl` is set. | +| `serviceConfig.agentic.invokeUrl` | `""` | OpenAI-compatible chat completions endpoint used by agentic retrieval. | +| `serviceConfig.agentic.requestTimeoutS` | `1800` | Gateway and MCP timeout for the multi-step agentic retrieval call. | | `serviceConfig.vectordb.enabled` | `true` | Deploy the LanceDB vectordb Pod. When `true` the chart **requires** a resolvable embed endpoint (refer to [VectorDB and the embed endpoint](#vectordb-and-the-embed-endpoint)); `helm install` / `helm upgrade` fails fast otherwise. | | `serviceConfig.vectordb.lancedbUri` | `/data/vectordb` | LanceDB on the vectordb Pod's PVC. | | `serviceConfig.vectordb.embedModel` | `nvidia/llama-nemotron-embed-vl-1b-v2` | Passed to vectordb + worker `embed_model_name`. | @@ -456,8 +462,8 @@ environment variable instead of writing the key into the ConfigMap. ### NIM Operator sub-stack -Each NIM block under `nimOperator.` renders a `NIMCache` + `NIMService` -pair gated on three conditions ALL holding: +Each enabled NIM block under `nimOperator.` renders operator resources +gated on three conditions ALL holding: 1. The `apps.nvidia.com/v1alpha1` CRDs are installed in the cluster. 2. The master switch `nims.enabled` is `true`. @@ -466,15 +472,17 @@ pair gated on three conditions ALL holding: | Path | Default | Notes | |----------------------------------------|---------|-------| | `nims.enabled` | `true` | Master switch. Set false to render no NIM resources. | -| `nimOperator.page_elements.enabled` | `true` | Page-elements detector NIM. | -| `nimOperator.table_structure.enabled` | `true` | Table-structure detector NIM. | +| `nimOperator.page_elements.enabled` | `true` | Page Elements 2.0 service; auto-wired to `/v1/page-elements`. | +| `nimOperator.table_structure.enabled` | `true` | Table Structure 2.0 service; auto-wired to `/v1/table-structure`. | +| `nimOperator..image` | `nvcr.io/nim/nvidia/nemotron-object-detection:2.0.1` | Both services use the combined image but select distinct models. | | `nimOperator.ocr.enabled` | `true` | OCR NIM. | -| `nimOperator.ocr.image` | `nvcr.io/nim/nvidia/nemotron-ocr-v2:1.4.0` | Default OCR NIM image. | +| `nimOperator.ocr.image` | `nvcr.io/nim/nvidia/nemotron-ocr-v2:2.0.1` | Default OCR NIM image. | | `nimOperator.vlm_embed.enabled` | `true` | Multimodal embedding NIM (also used by the vectordb Pod). | | `nimOperator.vlm_embed.nimServiceName` | `llama-nemotron-embed-vl-1b-v2` | NIMService / in-cluster DNS name. | -| `nimOperator.vlm_embed.image` | `nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2:1.12.0` | Default VLM embed NIM image. | +| `nimOperator.vlm_embed.image` | `nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2:2.3.0` | Default VLM embed NIM image. | | `nimOperator.rerankqa.enabled` | `false` | VL reranker NIM (optional; not auto-wired). Set `true` to opt in. Default `false` so chart installs honor the "optional and disabled by default" contract in [deployment-options.md](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/deployment-options.md) and do not silently provision an extra ≈ 3.1 GiB GPU NIM. The image points at the **VL** SKU (`llama-nemotron-rerank-vl-1b-v2`) per [prerequisites-support-matrix.md](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/prerequisites-support-matrix.md#default-helm-nims) — the text-only `llama-nemotron-rerank-1b-v2` silently degrades multimodal reranking and is not the documented POR. | -| `nimOperator.nemotron_parse.enabled` | `false` | Structured-parse NIM (optional). Set `true` when using `extract_method="nemotron_parse"`. Default `false` so chart installs honor the "optional and disabled by default" contract in [deployment-options.md](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/deployment-options.md). Image tag follows the [image tag conventions](#image-tag-conventions). | +| `nimOperator.rerankqa.image` | `nvcr.io/nim/nvidia/llama-nemotron-rerank-vl-1b-v2:2.3.0` | Default optional VL reranker NIM image. | +| `nimOperator.nemotron_parse.enabled` | `false` | Structured-parse NIM (optional). Set `true` when using `method="nemotron_parse"`. Default `false` so chart installs honor the "optional and disabled by default" contract in [deployment-options.md](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/deployment-options.md). Image tag follows the [image tag conventions](#image-tag-conventions). | | `nimOperator.nemotron_3_nano_omni_30b_a3b_reasoning.enabled` | `false` | Omni 30B caption NIM (optional). Set `true` to enable image captioning — refer to [Image captioning (Omni 30B)](#image-captioning-omni-30b). Default `false` so chart installs do not silently pull ≈ 62 GiB of BF16 weights or claim a second dedicated GPU. Image tag follows the [image tag conventions](#image-tag-conventions). | | `nimOperator.answer_llm.enabled` | `false` | Generic answer-generation LLM NIM (optional; Super-49B defaults). Set `true` to enable `/v1/answer` — refer to [Answer generation (operator-managed LLM)](#answer-generation-llm). Default `false` so installs do not silently claim answer-generation GPUs. | | `nimOperator.answer_llm.model` | `openai/nvidia/llama-3.3-nemotron-super-49b-v1.5` | LiteLLM/OpenAI model id inherited by `serviceConfig.llm.model` when the operator-managed answer LLM is enabled and no explicit service model is set. | @@ -496,7 +504,7 @@ pair gated on three conditions ALL holding: > are auto-wired into the retriever-service config. Optional NIMs may reconcile > when `nimOperator..enabled` is `true` in `values.yaml`, but the > retriever-service won't call them unless you wire your pipeline to use them. -> For minimal installs, prefer the [minimal install](#recommended-minimal-install-2605) overrides. +> For minimal installs, prefer the [minimal install](#recommended-minimal-install-2608) overrides. #### Filtering cached GPU profiles { #filtering-cached-gpu-profiles } @@ -566,7 +574,7 @@ Every NIM in this chart pins an exact NGC image tag in `values.yaml` | Family | Example | Meaning | | ------ | ------- | ------- | -| Plain semver | `nemotron-page-elements-v3:1.8.0` | A standard NIM release, identical bytes on every pull. Used by the four core NIMs and the reranker / ASR NIMs. | +| Plain semver | `nemotron-object-detection:2.0.1` | A standard NIM release, identical bytes on every pull. Used by the four core NIMs and the reranker / ASR NIMs. | | `-variant` | `nemotron-parse-v1.2:1.7.0-variant`, `nemotron-3-nano-omni-30b-a3b-reasoning:1.7.0-variant` | The Nemotron Parse and Nemotron 3 Nano Omni 30B builds that ship per-GPU TensorRT engine variants the NIM Operator selects from at reconciliation time (refer to the Omni and Parse rows in the [model hardware requirements](https://github.com/NVIDIA/NeMo-Retriever/blob/main/docs/docs/extraction/prerequisites-support-matrix.md#model-hardware-requirements) table). The `-variant` suffix is the NGC tag that ships alongside this chart and matches footnote ³ of the support matrix. | For air-gapped mirror pipelines: mirror the *exact* tag — both the @@ -585,7 +593,7 @@ helm upgrade --install retriever ./nemo_retriever/helm \ and validate against the same release of the retriever service before production rollout. -**Charts and captioning.** Charts and infographics use **page_elements** +**Charts and captioning.** Charts and infographics use **Page Elements, Table Structure** and **ocr**. For image captioning, set `nimOperator.nemotron_3_nano_omni_30b_a3b_reasoning.enabled=true` — refer to [Image captioning (Omni 30B)](#image-captioning-omni-30b) for the @@ -642,7 +650,7 @@ different VLM SKU. The chart defaults to **`nimOperator.nimServiceGpuLimit: 1`**, which renders `spec.resources.limits.nvidia.com/gpu: 1` on every NIMService unless a per-NIM `resources` map overrides it. This is required on -NIM Operator **v3.1.1** (and other versions tested on A100/H100): when +NIM Operator **v3.1.2** (and other versions tested on A100/H100): when the chart omits the `resources` block entirely, the operator often **does not** populate GPU limits from the model profile, and NIM pods start without GPU access (`The NVIDIA Driver was not detected`). @@ -740,6 +748,11 @@ custom service configuration files. | `ngcApiSecret.name` | `ngc-api` | Name referenced by NIMCache/NIMService `authSecret`. | | `ngcApiSecret.password` | `""` | NGC API key (populates `NGC_API_KEY` + `NGC_CLI_API_KEY`). | | `imagePullSecrets` | `[]` | Extra pre-existing pull secrets appended to every Pod. | +| `serviceConfig.vectordb.internalAuth.enabled` | `false` | Enable dedicated Secret-backed Retriever-to-VectorDB authentication. | +| `serviceConfig.vectordb.internalAuth.existingSecret.name` | `""` | Existing Secret shared by Retriever and VectorDB pods. | +| `serviceConfig.auth.scopeTokenSecret.name` | `""` | Existing Secret containing the public scope-token JSON file. | +| `serviceConfig.auth.enabled` | `false` | Require bearer authentication for the public gateway. | +| `serviceConfig.auth.allowInsecureInlineApiToken` | `false` | Explicit development-only gate for ConfigMap-backed `apiToken`. | ### Optional features @@ -776,6 +789,33 @@ The chart will skip Secret creation. Make sure `my-org-ngc-pull` exists as `kubernetes.io/dockerconfigjson` and `my-org-ngc-api` as `Opaque` with an `NGC_API_KEY` key, in the release namespace. +Protect the public gateway and its dedicated VectorDB hop with two separate +pre-existing Secrets: + +```yaml +serviceConfig: + auth: + scopeTokenSecret: + name: nrl-public-auth + key: scope-tokens.json + enabled: true + vectordb: + internalAuth: + enabled: true + existingSecret: + name: nrl-internal-vdb-auth + key: token +``` + +`nrl-public-auth` must contain a JSON document such as +`{"tokens":[{"token":"","scopes":["workspace-123"]}]}` under the +configured key. `nrl-internal-vdb-auth` must contain a distinct, high-entropy +credential. Internal authentication is opt-in for local compatibility; enable +it for production deployments. When enabled, a missing Secret or key prevents +the pods from starting instead of falling back to unauthenticated VectorDB +access. Inline `serviceConfig.auth.apiToken` is rejected unless +`allowInsecureInlineApiToken=true`, and must never be used for production. + ### Disable one NIM and supply an external URL for it ```yaml @@ -1176,11 +1216,11 @@ your release tag). Defaults below match | Role | `nimOperator` key | Default image (`repository:tag`) | |------|-------------------|----------------------------------| | Retriever service | — | `service.image.repository`:`service.image.tag` (override for production) | -| Page elements | `page_elements` | `nvcr.io/nim/nvidia/nemotron-page-elements-v3:1.8.0` | -| Table structure | `table_structure` | `nvcr.io/nim/nvidia/nemotron-table-structure-v1:1.8.0` | -| OCR | `ocr` | `nvcr.io/nim/nvidia/nemotron-ocr-v2:1.4.0` | -| VL embed | `vlm_embed` | `nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2:1.12.0` | -| VL reranker (optional) | `rerankqa` | `nvcr.io/nim/nvidia/llama-nemotron-rerank-vl-1b-v2:1.10.0` | +| Page Elements | `page_elements` | `nvcr.io/nim/nvidia/nemotron-object-detection:2.0.1` | +| Table Structure | `table_structure` | `nvcr.io/nim/nvidia/nemotron-object-detection:2.0.1` | +| OCR | `ocr` | `nvcr.io/nim/nvidia/nemotron-ocr-v2:2.0.1` | +| VL embed | `vlm_embed` | `nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2:2.3.0` | +| VL reranker (optional) | `rerankqa` | `nvcr.io/nim/nvidia/llama-nemotron-rerank-vl-1b-v2:2.3.0` | | Nemotron Parse (optional) | `nemotron_parse` | `nvcr.io/nim/nvidia/nemotron-parse-v1.2:1.7.0-variant` | | Omni caption (optional) | `nemotron_3_nano_omni_30b_a3b_reasoning` | `nvcr.io/nim/nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:1.7.0-variant` | | Answer LLM (optional, Super-49B default) | `answer_llm` | `nvcr.io/nim/nvidia/llama-3.3-nemotron-super-49b-v1.5:2.0.5` | @@ -1219,8 +1259,8 @@ ngcImagePullSecret: nimOperator: page_elements: image: - repository: /nemotron-page-elements-v3 - tag: "1.8.0" + repository: /nemotron-object-detection + tag: "2.0.1" pullPolicy: IfNotPresent # Repeat for table_structure, ocr, vlm_embed, and any optional keys you enable. ``` @@ -1239,10 +1279,10 @@ nimOperator: ```bash docker login nvcr.io -u '$oauthtoken' -p "$NGC_API_KEY" -docker pull nvcr.io/nim/nvidia/nemotron-page-elements-v3:1.8.0 -docker tag nvcr.io/nim/nvidia/nemotron-page-elements-v3:1.8.0 \ - /nemotron-page-elements-v3:1.8.0 -docker push /nemotron-page-elements-v3:1.8.0 +docker pull nvcr.io/nim/nvidia/nemotron-object-detection:2.0.1 +docker tag nvcr.io/nim/nvidia/nemotron-object-detection:2.0.1 \ + /nemotron-object-detection:2.0.1 +docker push /nemotron-object-detection:2.0.1 ``` For bulk sync, prefer [skopeo](https://github.com/containers/skopeo) or diff --git a/nemo_retriever/helm/openshift.md b/nemo_retriever/helm/openshift.md index c15925077b..88977542b1 100644 --- a/nemo_retriever/helm/openshift.md +++ b/nemo_retriever/helm/openshift.md @@ -190,7 +190,7 @@ helm install retriever ./nemo_retriever/helm -n nemo-retriever \ --set persistence.enabled=false \ --set retrieverResults.enabled=false \ --set service.image.repository=nvcr.io/nvidia/nemo-microservices/nrl-service \ - --set service.image.tag=26.5.0 + --set service.image.tag=26.8.0 ``` Verify pods: @@ -212,7 +212,7 @@ helm install retriever ./nemo_retriever/helm -n nemo-retriever \ --set ngcImagePullSecret.create=false \ --set ngcApiSecret.create=false \ --set service.image.repository=nvcr.io/nvidia/nemo-microservices/nrl-service \ - --set service.image.tag=26.5.0 + --set service.image.tag=26.8.0 ``` After install, confirm workloads reach Ready before you run ingest: diff --git a/nemo_retriever/helm/templates/NOTES.txt b/nemo_retriever/helm/templates/NOTES.txt index 78c57a8fdd..114fcfb6f2 100644 --- a/nemo_retriever/helm/templates/NOTES.txt +++ b/nemo_retriever/helm/templates/NOTES.txt @@ -61,13 +61,15 @@ Services: the following NIMServices were submitted for reconciliation (each lands as a Deployment + Service owned by the operator): {{- if .Values.nimOperator.page_elements.enabled }} - - nemotron-page-elements-v3 → http://nemotron-page-elements-v3:{{ .Values.nimOperator.page_elements.expose.service.port }}/v1/infer + - nemotron-page-elements-v3 → http://nemotron-page-elements-v3:{{ .Values.nimOperator.page_elements.expose.service.port }}/v1/page-elements + (nemotron-object-detection image; auto-wired into nim_endpoints.page_elements_invoke_url) {{- end }} {{- if .Values.nimOperator.table_structure.enabled }} - - nemotron-table-structure-v1 → http://nemotron-table-structure-v1:{{ .Values.nimOperator.table_structure.expose.service.port }}/v1/infer + - nemotron-table-structure-v1 → http://nemotron-table-structure-v1:{{ .Values.nimOperator.table_structure.expose.service.port }}/v1/table-structure + (nemotron-object-detection image; auto-wired into nim_endpoints.table_structure_invoke_url) {{- end }} {{- if .Values.nimOperator.ocr.enabled }} - - {{ .Values.nimOperator.ocr.nimServiceName }} → http://{{ .Values.nimOperator.ocr.nimServiceName }}:{{ .Values.nimOperator.ocr.expose.service.port }}/v1/infer + - {{ .Values.nimOperator.ocr.nimServiceName }} → http://{{ .Values.nimOperator.ocr.nimServiceName }}:{{ .Values.nimOperator.ocr.expose.service.port }}/v1/ocr {{- end }} {{- if .Values.nimOperator.vlm_embed.enabled }} - {{ .Values.nimOperator.vlm_embed.nimServiceName }} → http://{{ .Values.nimOperator.vlm_embed.nimServiceName }}:{{ .Values.nimOperator.vlm_embed.expose.service.port }}/v1/embeddings @@ -96,8 +98,9 @@ Services: kubectl --namespace {{ .Release.Namespace }} get nimcache,nimservice kubectl --namespace {{ .Release.Namespace }} describe nimservice - First-time NIMCache reconciliation downloads model weights to a PVC; - allow several minutes before each NIMService reports ready. + Direct-PVC NIMServices download their selected model during service startup. + NIMCache-mode services wait for their cache job first. Allow several minutes + before each NIMService reports ready. {{- else }} diff --git a/nemo_retriever/helm/templates/_helpers.tpl b/nemo_retriever/helm/templates/_helpers.tpl index e92b7c4d65..a73158d72b 100644 --- a/nemo_retriever/helm/templates/_helpers.tpl +++ b/nemo_retriever/helm/templates/_helpers.tpl @@ -487,7 +487,7 @@ NIMService GPU resources By default the chart sets ``spec.resources.limits.nvidia.com/gpu`` on every NIMService (see ``nimOperator.nimServiceGpuLimit``) because the NIM Operator does **not** reliably populate that field from the model -profile on all tested versions (for example v3.1.1 on A100/H100), which +profile on all tested versions (for example v3.1.2 on A100/H100), which otherwise leaves NIM pods without GPU access. Helm and the operator may both server-side-apply the same field; a @@ -525,9 +525,9 @@ file name under templates/nims/.yaml) so the retriever-service config can address each NIM as `http://:`. Mapping (key -> Service name, default invokePath): - page_elements -> nemotron-page-elements-v3 /v1/infer - table_structure -> nemotron-table-structure-v1 /v1/infer - ocr -> nemotron-ocr-v2 /v1/infer + page_elements -> nemotron-page-elements-v3 /v1/page-elements + table_structure -> nemotron-table-structure-v1 /v1/table-structure + ocr -> nemotron-ocr-v2 /v1/ocr vlm_embed -> llama-nemotron-embed-vl-1b-v2 /v1/embeddings nemotron_3_nano_omni_30b_a3b_reasoning -> nemotron-3-nano-omni-30b-a3b-reasoning /v1/chat/completions answer_llm -> Values.nimOperator.answer_llm.nimServiceName /v1 @@ -627,7 +627,7 @@ nemo-retriever.nimOperator.url "context" $ "key" "page_elements" "serviceName" "nemotron-page-elements-v3" - "invokePath" "/v1/infer") }} + "invokePath" "/v1/page-elements") }} */}} {{- define "nemo-retriever.nimOperator.url" -}} {{- $ctx := .context -}} @@ -659,7 +659,7 @@ nemo-retriever.nim.endpointURL "key" "page_elements" "serviceName" "nemotron-page-elements-v3" "configKey" "pageElementsInvokeUrl" - "invokePath" "/v1/infer") }} + "invokePath" "/v1/page-elements") }} */}} {{- define "nemo-retriever.nim.endpointURL" -}} {{- $ctx := .context -}} diff --git a/nemo_retriever/helm/templates/configmap.yaml b/nemo_retriever/helm/templates/configmap.yaml index 3ed7a77ed6..4a170089b2 100644 --- a/nemo_retriever/helm/templates/configmap.yaml +++ b/nemo_retriever/helm/templates/configmap.yaml @@ -8,17 +8,44 @@ auto-wire the NIM Operator-managed in-cluster Service URL when the apps.nvidia.com/v1alpha1 CRDs are present and the corresponding `nimOperator..enabled` flag is true. Each operator-managed Service inherits the NIMService resource name, so the mapping is fixed: - page_elements -> nemotron-page-elements-v3 /v1/infer - table_structure -> nemotron-table-structure-v1 /v1/infer - ocr -> nemotron-ocr-v2 /v1/infer + page_elements -> nemotron-page-elements-v3 /v1/page-elements + table_structure -> nemotron-table-structure-v1 /v1/table-structure + ocr -> nemotron-ocr-v2 /v1/ocr + nemotron_parse -> nemotron-parse /v1/chat/completions vlm_embed -> llama-nemotron-embed-vl-1b-v2 /v1/embeddings nemotron_3_nano_omni_30b_a3b_reasoning -> nemotron-3-nano-omni-30b-a3b-reasoning /v1/chat/completions answer_llm -> Values.nimOperator.answer_llm.nimServiceName /v1 */}} {{- $ctx := . -}} -{{- $pageElementsURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "page_elements" "serviceName" "nemotron-page-elements-v3" "configKey" "pageElementsInvokeUrl" "invokePath" "/v1/infer") -}} -{{- $tableStructureURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "table_structure" "serviceName" "nemotron-table-structure-v1" "configKey" "tableStructureInvokeUrl" "invokePath" "/v1/infer") -}} -{{- $ocrURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "ocr" "serviceName" $ctx.Values.nimOperator.ocr.nimServiceName "configKey" "ocrInvokeUrl" "invokePath" "/v1/infer") -}} +{{- $auth := $ctx.Values.serviceConfig.auth -}} +{{- $scopeTokenSecret := $auth.scopeTokenSecret -}} +{{- $internalAuth := $ctx.Values.serviceConfig.vectordb.internalAuth -}} +{{- if and $auth.enabled (not $auth.apiToken) (not $scopeTokenSecret.name) -}} +{{- fail "serviceConfig.auth.enabled=true requires serviceConfig.auth.scopeTokenSecret.name or serviceConfig.auth.apiToken." -}} +{{- end -}} +{{- if and (not $auth.enabled) (or $auth.apiToken $scopeTokenSecret.name) -}} +{{- fail "serviceConfig.auth.apiToken and serviceConfig.auth.scopeTokenSecret.name require serviceConfig.auth.enabled=true." -}} +{{- end -}} +{{- if and $auth.apiToken (not $auth.allowInsecureInlineApiToken) -}} +{{- fail "serviceConfig.auth.apiToken writes a credential to a ConfigMap; set serviceConfig.auth.allowInsecureInlineApiToken=true only for explicit insecure development use, or configure serviceConfig.auth.scopeTokenSecret.name." -}} +{{- end -}} +{{- if and $auth.apiToken $scopeTokenSecret.name -}} +{{- fail "serviceConfig.auth.apiToken and serviceConfig.auth.scopeTokenSecret.name are mutually exclusive." -}} +{{- end -}} +{{- if and $scopeTokenSecret.name (not $scopeTokenSecret.key) -}} +{{- fail "serviceConfig.auth.scopeTokenSecret.name requires serviceConfig.auth.scopeTokenSecret.key." -}} +{{- end -}} +{{- if and $internalAuth.enabled (not $internalAuth.existingSecret.name) -}} +{{- fail "serviceConfig.vectordb.internalAuth.enabled=true requires serviceConfig.vectordb.internalAuth.existingSecret.name." -}} +{{- end -}} +{{- if and $internalAuth.enabled (not $internalAuth.existingSecret.key) -}} +{{- fail "serviceConfig.vectordb.internalAuth.enabled=true requires serviceConfig.vectordb.internalAuth.existingSecret.key." -}} +{{- end -}} +{{- $pageElementsURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "page_elements" "serviceName" "nemotron-page-elements-v3" "configKey" "pageElementsInvokeUrl" "invokePath" "/v1/page-elements") -}} +{{- $tableStructureURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "table_structure" "serviceName" "nemotron-table-structure-v1" "configKey" "tableStructureInvokeUrl" "invokePath" "/v1/table-structure") -}} +{{- $ocrURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "ocr" "serviceName" $ctx.Values.nimOperator.ocr.nimServiceName "configKey" "ocrInvokeUrl" "invokePath" "/v1/ocr") -}} +{{- $nemotronParseURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "nemotron_parse" "serviceName" "nemotron-parse" "configKey" "nemotronParseInvokeUrl" "invokePath" "/v1/chat/completions") -}} +{{- $nemotronParseModel := $ctx.Values.serviceConfig.nimEndpoints.nemotronParseModel | default "" -}} {{- $embedURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "vlm_embed" "serviceName" $ctx.Values.nimOperator.vlm_embed.nimServiceName "configKey" "embedInvokeUrl" "invokePath" "/v1/embeddings") -}} {{- $captionURL := include "nemo-retriever.nim.endpointURL" (dict "context" $ctx "key" "nemotron_3_nano_omni_30b_a3b_reasoning" "serviceName" "nemotron-3-nano-omni-30b-a3b-reasoning" "configKey" "captionInvokeUrl" "invokePath" "/v1/chat/completions") -}} {{- /* @@ -73,9 +100,15 @@ nim_endpoints: page_elements_invoke_url: {{ .pageElementsURL | quote }} table_structure_invoke_url: {{ .tableStructureURL | quote }} ocr_invoke_url: {{ .ocrURL | quote }} + nemotron_parse_invoke_url: {{ if .nemotronParseURL }}{{ .nemotronParseURL | quote }}{{ else }}null{{ end }} + nemotron_parse_model: {{ if .nemotronParseModel }}{{ .nemotronParseModel | quote }}{{ else }}null{{ end }} embed_invoke_url: {{ .embedURL | quote }} embed_model_name: {{ .Values.serviceConfig.vectordb.embedModel | quote }} - embed_model_provider_prefix: {{ if .Values.serviceConfig.vectordb.embedModelProviderPrefix }}{{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }}{{ else }}null{{ end }} +{{- if .Values.serviceConfig.vectordb.embedModelProviderPrefix }} + embed_model_provider_prefix: {{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }} +{{- else }} + embed_model_provider_prefix: null +{{- end }} caption_invoke_url: {{ if .captionURL }}{{ .captionURL | quote }}{{ else }}null{{ end }} caption_model_name: {{ if .captionModelName }}{{ .captionModelName | quote }}{{ else }}null{{ end }} audio_grpc_endpoint: {{ if .audioGrpcEndpoint }}{{ .audioGrpcEndpoint | quote }}{{ else }}null{{ end }} @@ -106,6 +139,7 @@ mcp: path: {{ .Values.serviceConfig.mcp.path | quote }} base_url: {{ if .Values.serviceConfig.mcp.baseUrl }}{{ .Values.serviceConfig.mcp.baseUrl | quote }}{{ else }}null{{ end }} enable_write_tools: {{ .Values.serviceConfig.mcp.enableWriteTools }} + query_methods: {{ .Values.serviceConfig.mcp.queryMethods | default "classic" | quote }} max_concurrency: {{ .Values.serviceConfig.mcp.maxConcurrency }} request_timeout_s: {{ .Values.serviceConfig.mcp.requestTimeoutS }} ingest_timeout_s: {{ .Values.serviceConfig.mcp.ingestTimeoutS }} @@ -126,6 +160,17 @@ llm: rag_system_prompt_prefix: {{ if .llmRagSystemPromptPrefix }}{{ .llmRagSystemPromptPrefix | quote }}{{ else }}null{{ end }} reasoning_enabled: {{ .Values.serviceConfig.llm.reasoningEnabled }} +agentic: + enabled: {{ .Values.serviceConfig.agentic.enabled }} + llm_model: {{ if .Values.serviceConfig.agentic.llmModel }}{{ .Values.serviceConfig.agentic.llmModel | quote }}{{ else }}null{{ end }} + invoke_url: {{ if .Values.serviceConfig.agentic.invokeUrl }}{{ .Values.serviceConfig.agentic.invokeUrl | quote }}{{ else }}null{{ end }} + reasoning_effort: {{ if .Values.serviceConfig.agentic.reasoningEffort }}{{ .Values.serviceConfig.agentic.reasoningEffort | quote }}{{ else }}null{{ end }} + backend_top_k: {{ .Values.serviceConfig.agentic.backendTopK }} + react_max_steps: {{ .Values.serviceConfig.agentic.reactMaxSteps }} + text_truncation: {{ .Values.serviceConfig.agentic.textTruncation }} + temperature: {{ .Values.serviceConfig.agentic.temperature }} + request_timeout_s: {{ .Values.serviceConfig.agentic.requestTimeoutS }} + pipeline: realtime_workers: {{ .Values.serviceConfig.pipeline.realtimeWorkers }} realtime_queue_size: {{ .Values.serviceConfig.pipeline.realtimeQueueSize }} @@ -150,7 +195,10 @@ resources: max_upload_bytes: {{ $maxUploadBytes }} auth: + enabled: {{ .Values.serviceConfig.auth.enabled }} api_token: {{ .Values.serviceConfig.auth.apiToken | default "null" }} + default_scope: {{ .Values.serviceConfig.auth.defaultScope | quote }} + allow_unscoped_dev: {{ .Values.serviceConfig.auth.allowUnscopedDev }} header_name: {{ .Values.serviceConfig.auth.headerName | quote }} bypass_paths: {{ toJson .Values.serviceConfig.auth.bypassPaths }} @@ -160,7 +208,11 @@ vectordb: lancedb_uri: {{ .Values.serviceConfig.vectordb.lancedbUri | quote }} table_name: {{ .Values.serviceConfig.vectordb.tableName | quote }} embed_model: {{ .Values.serviceConfig.vectordb.embedModel | quote }} - embed_model_provider_prefix: {{ if .Values.serviceConfig.vectordb.embedModelProviderPrefix }}{{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }}{{ else }}null{{ end }} +{{- if .Values.serviceConfig.vectordb.embedModelProviderPrefix }} + embed_model_provider_prefix: {{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }} +{{- else }} + embed_model_provider_prefix: null +{{- end }} vectordb_url: "http://{{ .vectordbSvc }}:{{ .vectordbPort }}" {{- else }} vectordb: @@ -184,7 +236,7 @@ metadata: data: retriever-service.yaml: | mode: standalone -{{ include "nemo-retriever.configBody" (dict "Values" .Values "gatewaySvc" $gatewaySvc "pageElementsURL" $pageElementsURL "tableStructureURL" $tableStructureURL "ocrURL" $ocrURL "embedURL" $embedURL "captionURL" $captionURL "captionModelName" $captionModelName "audioGrpcEndpoint" $audioGrpcEndpoint "llmEnabled" $llmEnabled "llmModel" $llmModel "llmAPIBase" $llmAPIBase "llmRagSystemPromptPrefix" $llmRagSystemPromptPrefix "vectordbSvc" $vectordbSvc "vectordbPort" $vectordbPort) | indent 4 }} +{{ include "nemo-retriever.configBody" (dict "Values" .Values "gatewaySvc" $gatewaySvc "pageElementsURL" $pageElementsURL "tableStructureURL" $tableStructureURL "ocrURL" $ocrURL "nemotronParseURL" $nemotronParseURL "nemotronParseModel" $nemotronParseModel "embedURL" $embedURL "captionURL" $captionURL "captionModelName" $captionModelName "audioGrpcEndpoint" $audioGrpcEndpoint "llmEnabled" $llmEnabled "llmModel" $llmModel "llmAPIBase" $llmAPIBase "llmRagSystemPromptPrefix" $llmRagSystemPromptPrefix "vectordbSvc" $vectordbSvc "vectordbPort" $vectordbPort) | indent 4 }} {{- else }} # ========================================================================= # Split mode — one ConfigMap per role with the appropriate mode + gateway @@ -212,6 +264,6 @@ data: timeout_s: 300.0 max_connections: 100 {{- end }} -{{ include "nemo-retriever.configBody" (dict "Values" $.Values "gatewaySvc" $gatewaySvc "pageElementsURL" $pageElementsURL "tableStructureURL" $tableStructureURL "ocrURL" $ocrURL "embedURL" $embedURL "captionURL" $captionURL "captionModelName" $captionModelName "audioGrpcEndpoint" $audioGrpcEndpoint "llmEnabled" $llmEnabled "llmModel" $llmModel "llmAPIBase" $llmAPIBase "llmRagSystemPromptPrefix" $llmRagSystemPromptPrefix "vectordbSvc" $vectordbSvc "vectordbPort" $vectordbPort) | indent 4 }} +{{ include "nemo-retriever.configBody" (dict "Values" $.Values "gatewaySvc" $gatewaySvc "pageElementsURL" $pageElementsURL "tableStructureURL" $tableStructureURL "ocrURL" $ocrURL "nemotronParseURL" $nemotronParseURL "nemotronParseModel" $nemotronParseModel "embedURL" $embedURL "captionURL" $captionURL "captionModelName" $captionModelName "audioGrpcEndpoint" $audioGrpcEndpoint "llmEnabled" $llmEnabled "llmModel" $llmModel "llmAPIBase" $llmAPIBase "llmRagSystemPromptPrefix" $llmRagSystemPromptPrefix "vectordbSvc" $vectordbSvc "vectordbPort" $vectordbPort) | indent 4 }} {{- end }} {{- end }} diff --git a/nemo_retriever/helm/templates/deployment-vectordb.yaml b/nemo_retriever/helm/templates/deployment-vectordb.yaml index 19c4efc009..1df0a86a3a 100644 --- a/nemo_retriever/helm/templates/deployment-vectordb.yaml +++ b/nemo_retriever/helm/templates/deployment-vectordb.yaml @@ -6,6 +6,8 @@ {{- $embedURL := include "nemo-retriever.nim.endpointURL" (dict "context" . "key" "vlm_embed" "serviceName" .Values.nimOperator.vlm_embed.nimServiceName "configKey" "embedInvokeUrl" "invokePath" "/v1/embeddings") -}} {{- $localEmbed := include "nemo-retriever.localEmbed.enabled" . | eq "true" -}} {{- $localModels := .Values.serviceConfig.localModels -}} +{{- $agentic := .Values.serviceConfig.agentic -}} +{{- $internalAuth := .Values.serviceConfig.vectordb.internalAuth -}} {{- /* Fail-fast guard: rendering a vectordb Deployment without any query-time embedding backend produces a "healthy" Pod whose first /v1/query request @@ -99,6 +101,31 @@ spec: - --embed-model-provider-prefix - {{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }} {{- end }} + {{- if $agentic.enabled }} + - --agentic + {{- if $agentic.llmModel }} + - --agentic-llm-model + - {{ $agentic.llmModel | quote }} + {{- end }} + {{- if $agentic.invokeUrl }} + - --agentic-invoke-url + - {{ $agentic.invokeUrl | quote }} + {{- end }} + {{- if $agentic.reasoningEffort }} + - --agentic-reasoning-effort + - {{ $agentic.reasoningEffort | quote }} + {{- end }} + - --agentic-backend-top-k + - {{ $agentic.backendTopK | quote }} + - --agentic-react-max-steps + - {{ $agentic.reactMaxSteps | quote }} + - --agentic-text-truncation + - {{ $agentic.textTruncation | quote }} + - --agentic-temperature + - {{ $agentic.temperature | quote }} + - --agentic-request-timeout + - {{ $agentic.requestTimeoutS | quote }} + {{- end }} - --port - {{ $vdb.port | quote }} {{- if $embedURL }} @@ -110,7 +137,15 @@ spec: containerPort: {{ $vdb.port }} protocol: TCP env: - {{- if $embedURL }} + {{- if $internalAuth.enabled }} + - name: NRL_INTERNAL_VDB_TOKEN + valueFrom: + secretKeyRef: + name: {{ $internalAuth.existingSecret.name | quote }} + key: {{ $internalAuth.existingSecret.key | quote }} + optional: false + {{- end }} + {{- if or $embedURL (and $agentic.enabled $agentic.invokeUrl) }} - name: NVIDIA_API_KEY valueFrom: secretKeyRef: diff --git a/nemo_retriever/helm/templates/deployment.yaml b/nemo_retriever/helm/templates/deployment.yaml index a695bdc898..891d861c91 100644 --- a/nemo_retriever/helm/templates/deployment.yaml +++ b/nemo_retriever/helm/templates/deployment.yaml @@ -18,6 +18,9 @@ {{- end -}} {{- end -}} {{- $svc := .Values.service -}} +{{- $internalAuth := .Values.serviceConfig.vectordb.internalAuth -}} +{{- $scopeTokenSecret := .Values.serviceConfig.auth.scopeTokenSecret -}} +{{- $scopeTokenMountPath := "/var/run/secrets/nemo-retriever/auth" -}} {{- if and (eq .Values.topology.mode "split") (ne (int .Values.topology.gateway.replicas) 1) -}} {{- fail "split mode requires topology.gateway.replicas=1 (process-local scheduler)" -}} {{- end -}} @@ -107,6 +110,18 @@ spec: env: - name: NEMO_RETRIEVER_SERVICE_CONFIG value: /etc/nemo-retriever/retriever-service.yaml + {{- if $internalAuth.enabled }} + - name: NRL_INTERNAL_VDB_TOKEN + valueFrom: + secretKeyRef: + name: {{ $internalAuth.existingSecret.name | quote }} + key: {{ $internalAuth.existingSecret.key | quote }} + optional: false + {{- end }} + {{- if $scopeTokenSecret.name }} + - name: NRL_SCOPE_TOKEN_FILE + value: {{ printf "%s/scope-tokens.json" $scopeTokenMountPath | quote }} + {{- end }} {{- if .Values.retrieverResults.enabled }} - name: NEMO_RETRIEVER_RESULTS_DIR value: {{ .Values.retrieverResults.mountPath | quote }} @@ -164,6 +179,11 @@ spec: {{- end }} - name: tmp mountPath: /tmp + {{- if $scopeTokenSecret.name }} + - name: scope-token + mountPath: {{ $scopeTokenMountPath | quote }} + readOnly: true + {{- end }} {{- with $svc.extraVolumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -185,6 +205,15 @@ spec: name: {{ include "nemo-retriever.configMapName" . }} - name: tmp emptyDir: {} + {{- if $scopeTokenSecret.name }} + - name: scope-token + secret: + secretName: {{ $scopeTokenSecret.name | quote }} + defaultMode: 288 + items: + - key: {{ $scopeTokenSecret.key | quote }} + path: scope-tokens.json + {{- end }} {{- if .Values.persistence.enabled }} - name: data persistentVolumeClaim: @@ -281,6 +310,18 @@ spec: env: - name: NEMO_RETRIEVER_SERVICE_CONFIG value: /etc/nemo-retriever/retriever-service.yaml + {{- if $internalAuth.enabled }} + - name: NRL_INTERNAL_VDB_TOKEN + valueFrom: + secretKeyRef: + name: {{ $internalAuth.existingSecret.name | quote }} + key: {{ $internalAuth.existingSecret.key | quote }} + optional: false + {{- end }} + {{- if and (eq $role "gateway") $scopeTokenSecret.name }} + - name: NRL_SCOPE_TOKEN_FILE + value: {{ printf "%s/scope-tokens.json" $scopeTokenMountPath | quote }} + {{- end }} {{- if $.Values.retrieverResults.enabled }} - name: NEMO_RETRIEVER_RESULTS_TTL_SECONDS value: {{ $.Values.retrieverResults.ttlSeconds | quote }} @@ -347,6 +388,11 @@ spec: {{- end }} - name: tmp mountPath: /tmp + {{- if and (eq $role "gateway") $scopeTokenSecret.name }} + - name: scope-token + mountPath: {{ $scopeTokenMountPath | quote }} + readOnly: true + {{- end }} {{- if $svc.startupProbe.enabled }} startupProbe: {{- omit $svc.startupProbe "enabled" | toYaml | nindent 12 }} @@ -368,6 +414,15 @@ spec: {{- if eq $role "gateway" }} sizeLimit: {{ int64 $.Values.serviceConfig.workQueue.spoolLimitBytes | quote }} {{- end }} + {{- if and (eq $role "gateway") $scopeTokenSecret.name }} + - name: scope-token + secret: + secretName: {{ $scopeTokenSecret.name | quote }} + defaultMode: 288 + items: + - key: {{ $scopeTokenSecret.key | quote }} + path: scope-tokens.json + {{- end }} {{- if and $.Values.persistence.enabled (eq $role "gateway") }} - name: data persistentVolumeClaim: diff --git a/nemo_retriever/helm/templates/nims/nemotron-page-elements-v3.yaml b/nemo_retriever/helm/templates/nims/nemotron-page-elements-v3.yaml index 6d3b936abc..daf73b8564 100644 --- a/nemo_retriever/helm/templates/nims/nemotron-page-elements-v3.yaml +++ b/nemo_retriever/helm/templates/nims/nemotron-page-elements-v3.yaml @@ -30,7 +30,7 @@ metadata: spec: image: repository: {{ .Values.nimOperator.page_elements.image.repository }} - tag: {{ .Values.nimOperator.page_elements.image.tag }} + tag: {{ .Values.nimOperator.page_elements.image.tag | toString | quote }} pullPolicy: {{ .Values.nimOperator.page_elements.image.pullPolicy }} pullSecrets: {{ toYaml .Values.nimOperator.page_elements.image.pullSecrets | indent 6 }} diff --git a/nemo_retriever/helm/templates/nims/nemotron-table-structure-v1.yaml b/nemo_retriever/helm/templates/nims/nemotron-table-structure-v1.yaml index fee6dff7c0..2a13ed8670 100644 --- a/nemo_retriever/helm/templates/nims/nemotron-table-structure-v1.yaml +++ b/nemo_retriever/helm/templates/nims/nemotron-table-structure-v1.yaml @@ -29,7 +29,7 @@ metadata: spec: image: repository: {{ .Values.nimOperator.table_structure.image.repository }} - tag: {{ .Values.nimOperator.table_structure.image.tag }} + tag: {{ .Values.nimOperator.table_structure.image.tag | toString | quote }} pullPolicy: {{ .Values.nimOperator.table_structure.image.pullPolicy }} pullSecrets: {{ toYaml .Values.nimOperator.table_structure.image.pullSecrets | indent 6 }} diff --git a/nemo_retriever/helm/values.yaml b/nemo_retriever/helm/values.yaml index 188ebd851e..043433a3d7 100644 --- a/nemo_retriever/helm/values.yaml +++ b/nemo_retriever/helm/values.yaml @@ -544,6 +544,13 @@ serviceConfig: pageElementsInvokeUrl: "" tableStructureInvokeUrl: "" ocrInvokeUrl: "" + # Optional Nemotron Parse chat-completions endpoint. When the + # operator-managed Parse NIM is enabled, its in-cluster URL is used + # unless this explicit endpoint is set. + nemotronParseInvokeUrl: "" + # Optional model override. Leave empty to infer the hosted + # nvidia/nemotron-parse or self-hosted v1.2 contract from the URL. + nemotronParseModel: "" embedInvokeUrl: "" # Optional remote VLM endpoint for image captioning (Nemotron 3 Nano # Omni). Auto-wired from the in-cluster Service when @@ -590,6 +597,8 @@ serviceConfig: path: "/mcp" baseUrl: "" enableWriteTools: true + # classic | agentic | all — which retrieval MCP tools to register. + queryMethods: classic maxConcurrency: 8 requestTimeoutS: 60.0 ingestTimeoutS: 1800.0 @@ -634,6 +643,20 @@ serviceConfig: # portable no-reasoning controls. reasoningEnabled: true + # Agentic (ReAct) retrieval exposed as POST /v1/query with agentic=true and the + # agentic_query MCP tool. The LLM runs inside the vectordb pod because that + # pod owns the LanceDB volume. + agentic: + enabled: false + llmModel: "" + invokeUrl: "" + reasoningEffort: high + backendTopK: 20 + reactMaxSteps: 50 + textTruncation: 0 + temperature: 0.0 + requestTimeoutS: 1800.0 + # Pipeline worker pools. Workers are abstract dispatchers — sizing # depends on whether they do local GPU work or fan out to remote NIMs. # For CPU-only NIM-forwarding nodes, higher worker counts are fine. @@ -670,11 +693,29 @@ serviceConfig: tableName: "nemo_retriever" embedModel: "nvidia/llama-nemotron-embed-vl-1b-v2" embedModelProviderPrefix: "" + # Optional dedicated gateway/worker-to-VectorDB authentication. When + # enabled, all Retriever and VectorDB pods read the same existing Secret + # key and VectorDB rejects missing or invalid internal credentials. + internalAuth: + enabled: false + existingSecret: + name: "" + key: token - # Optional bearer-token authentication. When apiToken is set, every - # request must carry "Authorization: Bearer ". + # Public bearer-token authentication. Production deployments should mount + # a Secret-backed scope-token JSON file. It is disabled by default so a + # standalone deployment accepts gateway requests without a bearer token. + # Inline apiToken is retained only as an explicitly enabled insecure + # development fallback. auth: + enabled: false apiToken: null + allowInsecureInlineApiToken: false + defaultScope: "default" + allowUnscopedDev: false + scopeTokenSecret: + name: "" + key: scope-tokens.json headerName: "Authorization" bypassPaths: - "/v1/health" @@ -876,7 +917,7 @@ nims: # # When auto-resolution is in effect, the chart populates # `serviceConfig.nimEndpoints.*` with the NIMService's in-cluster URL -# (e.g. http://nemotron-page-elements-v3:8000/v1/infer). An explicit value +# (e.g. http://nemotron-page-elements-v3:8000/v1/page-elements). An explicit value # in `serviceConfig.nimEndpoints.*` always wins. # # NIMCache resources may carry `helm.sh/resource-policy: keep` so model @@ -894,8 +935,11 @@ nimOperator: keepOnUninstall: true # NIMCache/NIMService names retired from the chart. deploy.sh deletes these # after each successful helm reconcile so they do not linger with keep-on-uninstall. + # nemotron-object-detection was a short-lived single-NIM chart name; page_elements + # and table_structure now share that container image under their legacy names. pruneRetiredNimResources: - nemotron-ocr-v1 + - nemotron-object-detection pvc: create: true # If set, applies to every per-NIM PVC that doesn't override @@ -998,13 +1042,14 @@ nimOperator: # ``1.7.0-variant`` Parse / Omni tags mean. # --------------------------------------------------------------------------- - # Page-elements detector (YOLOX-style). Used by the page-element extraction - # stage of the ingest pipeline. + # Page-element detector. Uses the combined nemotron-object-detection NIM image + # but keeps the legacy in-cluster Service name so the retriever pipeline's + # page_elements_invoke_url continues to resolve independently. page_elements: enabled: true image: - repository: nvcr.io/nim/nvidia/nemotron-page-elements-v3 - tag: "1.8.0" + repository: nvcr.io/nim/nvidia/nemotron-object-detection + tag: "2.0.1" pullPolicy: IfNotPresent pullSecrets: - ngc-secret @@ -1035,25 +1080,32 @@ nimOperator: endpoint: "" env: {} env: - - name: NIM_HTTP_API_PORT - value: "8000" - - name: NIM_TRITON_LOG_VERBOSE + - name: NIM_SERVER_BIND_ADDR + value: "0.0.0.0:8000" + - name: NIM_PERFORMANCE_MODE + value: "0" + - name: NIM_SERVER_MODE + value: latency + - name: NIM_SERVER_MAX_WAIT_MS + value: "0" + - name: NIM_ENGINE_COUNT value: "1" - - name: NIM_TRITON_MAX_BATCH_SIZE - value: "32" - - name: NIM_TRITON_CPU_THREADS_PRE_PROCESSOR - value: "2" - - name: NIM_TRITON_CPU_THREADS_POST_PROCESSOR + - name: NIM_PIPELINE_MAX_BATCH_SIZE value: "1" - - name: OMP_NUM_THREADS - value: "2" + - name: NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER + value: ngc + - name: NIM_ENGINE_MODEL_NAME + value: nvidia/nemotron-page-elements-v3 + - name: NIM_ENGINE_MODEL_PATH + value: /model-store/page-elements - # Table structure detector. Used by the table extraction stage. + # Table structure detector. Uses the same nemotron-object-detection NIM image + # as page_elements but keeps the legacy Service name for table_structure_invoke_url. table_structure: enabled: true image: - repository: nvcr.io/nim/nvidia/nemotron-table-structure-v1 - tag: "1.8.0" + repository: nvcr.io/nim/nvidia/nemotron-object-detection + tag: "2.0.1" pullPolicy: IfNotPresent pullSecrets: - ngc-secret @@ -1084,18 +1136,24 @@ nimOperator: endpoint: "" env: {} env: - - name: NIM_HTTP_API_PORT - value: "8000" - - name: NIM_TRITON_LOG_VERBOSE + - name: NIM_SERVER_BIND_ADDR + value: "0.0.0.0:8000" + - name: NIM_PERFORMANCE_MODE + value: "0" + - name: NIM_SERVER_MODE + value: latency + - name: NIM_SERVER_MAX_WAIT_MS + value: "0" + - name: NIM_ENGINE_COUNT value: "1" - - name: NIM_TRITON_RATE_LIMIT - value: "3" - - name: NIM_TRITON_MAX_BATCH_SIZE - value: "32" - - name: NIM_TRITON_CUDA_MEMORY_POOL_MB - value: "2048" - - name: OMP_NUM_THREADS + - name: NIM_PIPELINE_MAX_BATCH_SIZE value: "1" + - name: NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER + value: ngc + - name: NIM_ENGINE_MODEL_NAME + value: nvidia/nemotron-table-structure-v1 + - name: NIM_ENGINE_MODEL_PATH + value: /model-store/table-structure # Nemotron OCR v2. Used by the OCR stage of the pipeline. ocr: @@ -1104,7 +1162,7 @@ nimOperator: nimServiceName: nemotron-ocr-v2 image: repository: nvcr.io/nim/nvidia/nemotron-ocr-v2 - tag: "1.4.0" + tag: "2.0.1" pullPolicy: IfNotPresent pullSecrets: - ngc-secret @@ -1135,14 +1193,26 @@ nimOperator: endpoint: "" env: {} env: - - name: OMP_NUM_THREADS - value: "8" - - name: NIM_HTTP_API_PORT - value: "8000" - - name: NIM_TRITON_LOG_VERBOSE + - name: NIM_SERVER_BIND_ADDR + value: "0.0.0.0:8000" + - name: NIM_PERFORMANCE_MODE + value: "0" + - name: NIM_SERVER_MODE + value: latency + - name: NIM_SERVER_MAX_WAIT_MS + value: "0" + - name: NIM_ENGINE_COUNT + value: "1" + - name: NIM_PIPELINE_MAX_BATCH_SIZE value: "1" - - name: NIM_TRITON_MAX_BATCH_SIZE - value: "32" + - name: NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER + value: ngc + - name: NIM_ENGINE_MODEL_NAME + value: nvidia/nemotron-ocr-v2 + - name: NIM_ENGINE_MODEL_PATH + value: /model-store/ocr + - name: NIM_ENGINE_MODEL_VARIANT + value: multilingual # Llama Nemotron multimodal (VL) embed 1B v2. Used by the embedding stage, # the vectordb query path, and multimodal ingest. @@ -1152,7 +1222,7 @@ nimOperator: nimServiceName: llama-nemotron-embed-vl-1b-v2 image: repository: nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2 - tag: "1.12.0" + tag: "2.3.0" pullPolicy: IfNotPresent pullSecrets: - ngc-secret @@ -1288,7 +1358,7 @@ nimOperator: enabled: false image: repository: nvcr.io/nim/nvidia/llama-nemotron-rerank-vl-1b-v2 - tag: "1.11.0" + tag: "2.3.0" pullPolicy: IfNotPresent pullSecrets: - ngc-secret diff --git a/nemo_retriever/pyproject.toml b/nemo_retriever/pyproject.toml index 543f07f7ef..df143c15b7 100644 --- a/nemo_retriever/pyproject.toml +++ b/nemo_retriever/pyproject.toml @@ -16,6 +16,20 @@ required-environments = [ # fastparquet 2026.3.0 has no wheel for Linux x86_64 (manylinux); constrain to last version that does. override-dependencies = [ "fastparquet>=2024.11.0,<2026", + # Security minimums for transitive dependencies. + "langsmith>=0.8.18", + "msgpack>=1.2.1", + "cryptography>=48.0.1", + "aiohttp>=3.14.3", + "requests>=2.33.0", + "idna>=3.15", + "pydantic-settings>=2.14.2", + "vllm==0.25.1", + "transformers>=5.14.1", + "ray>=2.56.1", + "nltk>=3.10.1", + "litellm>=1.95.0rc3", + "pillow>=12.3.0", # Suppress opencv-python from transitive deps. nemotron wheels pull in both # opencv-python and opencv-python-headless, which stomp on the same cv2/ dir. # The 'never' marker makes uv treat this as unsatisfiable on all platforms. @@ -41,7 +55,7 @@ classifiers = [ dependencies = [ # Core orchestration and framework "backoff", - "ray[data,serve]>=2.49.0", + "ray[data,serve]>=2.56.1", "ffmpeg-python", "pandas>=2.0,<3", "sqlglot>=30.0.0", @@ -49,11 +63,13 @@ dependencies = [ "typer>=0.12.0", "click>=8.2.0", "pyyaml>=6.0", + # Agent prompt templates are rendered at runtime. + "jinja2>=3.1", # Service layer "fastapi>=0.114.0", "uvicorn[standard]>=0.30.0", "python-multipart>=0.0.9", - "prometheus-fastapi-instrumentator>=7.0,<8", + "prometheus-fastapi-instrumentator>=8.0,<9", "opentelemetry-api>=1.41.1", "opentelemetry-sdk>=1.41.1", "opentelemetry-exporter-otlp-proto-grpc>=1.41.1", @@ -61,7 +77,7 @@ dependencies = [ "fastmcp>=2.0.0", # HTTP clients "httpx>=0.27.0", - "requests>=2.32.5", + "requests>=2.33.0", "urllib3==2.7.0", # Utilities "pydantic>=2.8.0", @@ -74,8 +90,8 @@ dependencies = [ # Core ingest packages # Document parsing and NIM client libs "pypdfium2==4.30.0", - "pillow==12.2.0", - "nltk==3.9.4", + "pillow>=12.3.0", + "nltk>=3.10.1", "markitdown", "langchain-nvidia-ai-endpoints>=1.4.0", # Default VDB solution @@ -83,6 +99,9 @@ dependencies = [ # gRPC client for Parakeet/Riva ASR. Required for ASRCPUActor when it # targets the public NVCF Parakeet endpoint (the default) or any remote NIM. "nvidia-riva-client>=2.25.1", + # Exact tokenizer-only TXT/HTML chunking (no transformers, torch, or weights). + "tokenizers>=0.21.1", + "huggingface-hub>=0.34.0", ] [project.optional-dependencies] @@ -101,7 +120,7 @@ service = [ # HTTP retry decorator used by NIM model interfaces "backoff", # Remote answer generation client used by serviceConfig.llm and /v1/answer - "litellm>=1.86.0,<2", + "litellm>=1.95.0rc3,<2", # Lightweight utilities used by pipeline operators "easydict", "addict", @@ -117,8 +136,7 @@ service = [ # ── Local model inference (GPU assumed; torch resolves to CUDA on Linux) ───── # Stable Nemotron extraction package selection for published local installs. local = [ - "transformers>=4.57.6,<5", - "tokenizers>=0.21.1", + "transformers>=5.14.1,<6", "accelerate==1.12.0", "opencv-python-headless>=4.8.0", "torch==2.11.0; sys_platform == 'linux'", @@ -143,10 +161,10 @@ local = [ # vLLM compiles CUDA kernels at runtime via torch inductor (requires Python.h). # Use a uv-managed Python (`uv python install 3.12`) so headers are available; # system Python installs typically omit them and will fail with InductorError. - "vllm==0.20.0; sys_platform == 'linux'", - # flashinfer and cubin versions must match - "flashinfer-cubin==0.6.8.post1; sys_platform == 'linux'", - "flashinfer-python==0.6.8.post1; sys_platform == 'linux'", + "vllm==0.25.1; sys_platform == 'linux'", + # Keep the CUDA inference stack reproducible; vLLM 0.25.1 requires FlashInfer 0.6.13. + "flashinfer-cubin==0.6.13; sys_platform == 'linux'", + "flashinfer-python==0.6.13; sys_platform == 'linux'", ] # ── Multimedia — audio/ASR and SVG rendering ──────────────────────────────── @@ -192,7 +210,7 @@ benchmarks = [ # or construct an ``LLMJudge`` / ``LiteLLMClient`` directly. Powers both the # live-RAG SDK and the batch evaluation framework. llm = [ - "litellm>=1.86.0,<2", + "litellm>=1.95.0rc3,<2", ] dev = [ @@ -242,6 +260,7 @@ explicit = true where = ["src"] [tool.setuptools.package-data] +"nemo_retriever._agentic.nemo_agent.prompts" = ["templates/**/*.j2"] "nemo_retriever.harness.portal" = ["static/**/*"] "nemo_retriever.service" = ["retriever-service.yaml"] "nemo_retriever.tools.skill_eval" = ["prompts/*.j2", "configs/*.yaml"] diff --git a/nemo_retriever/src/nemo_retriever/__init__.py b/nemo_retriever/src/nemo_retriever/__init__.py index 5815a1bbcb..58cf6ececc 100644 --- a/nemo_retriever/src/nemo_retriever/__init__.py +++ b/nemo_retriever/src/nemo_retriever/__init__.py @@ -30,6 +30,18 @@ "ingestor", "retriever", "RetrieverServiceCompatibilityError", + "RetrieverServiceClient", + "RetrieverServiceError", + "RetrieverServiceNotFoundError", + "RetrieverServiceConflictError", + "RetrieverServiceValidationError", + "CollectionInfo", + "CollectionDeleteResult", + "CollectionPage", + "DocumentInfo", + "DocumentPage", + "DocumentDeleteResult", + "QueryHit", ] retriever = _retriever_cls() @@ -56,8 +68,48 @@ def __getattr__(name: str): from nemo_retriever.ingestor.graph_ingestor import GraphIngestionError return GraphIngestionError - if name == "RetrieverServiceCompatibilityError": - from nemo_retriever.service.client import RetrieverServiceCompatibilityError + if name in { + "RetrieverServiceClient", + "RetrieverServiceCompatibilityError", + }: + from nemo_retriever.service.client import RetrieverServiceClient, RetrieverServiceCompatibilityError - return RetrieverServiceCompatibilityError + return { + "RetrieverServiceClient": RetrieverServiceClient, + "RetrieverServiceCompatibilityError": RetrieverServiceCompatibilityError, + }[name] + if name in { + "RetrieverServiceError", + "RetrieverServiceNotFoundError", + "RetrieverServiceConflictError", + "RetrieverServiceValidationError", + }: + from nemo_retriever.service.errors import ( + RetrieverServiceConflictError, + RetrieverServiceError, + RetrieverServiceNotFoundError, + RetrieverServiceValidationError, + ) + + return locals()[name] + if name in { + "CollectionInfo", + "CollectionDeleteResult", + "CollectionPage", + "DocumentInfo", + "DocumentPage", + "DocumentDeleteResult", + "QueryHit", + }: + from nemo_retriever.common.schemas.collections import ( + CollectionDeleteResult, + CollectionInfo, + CollectionPage, + DocumentDeleteResult, + DocumentInfo, + DocumentPage, + QueryHit, + ) + + return locals()[name] raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/nemo_retriever/src/nemo_retriever/_agentic/__init__.py b/nemo_retriever/src/nemo_retriever/_agentic/__init__.py new file mode 100644 index 0000000000..9b6077faf1 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Private namespace for agentic-retrieval implementations. + +Internal API. Users should drive agentic retrieval through the public graph +operators (see :mod:`nemo_retriever.operators.graph_ops`) rather than importing +from this package directly. Names and layout here may change without notice. +""" diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/__init__.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/__init__.py new file mode 100644 index 0000000000..00bc4e2cb1 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/__init__.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Agentic retrieval core: a self-contained agent loop with pluggable LLM backends and tools. + +Internal API. This package is driven by the public graph operators (see +:mod:`nemo_retriever.operators.graph_ops`); prefer those over importing from here +directly. Typical integration:: + + from nemo_retriever._agentic.nemo_agent import Agent, AgentConfig, create_retrieve_tool + from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config + + llm = create_llm(create_llm_config("litellm", model=...)) + retrieve = create_retrieve_tool("default", my_retriever_fn) + agent = Agent(config=AgentConfig(mode="select"), llm=llm, retrieve_tool=retrieve) + result = agent.run_sync("what is ...", raw_log_dir=None) + result.final_doc_ids # [] + result.error when the run failed + +Standalone selection over candidate documents (an explicit ``scores`` ranking +arms the context-overflow shrink retry; without it there is one attempt):: + + sel = SelectionAgent(config=SelectionAgentConfig(target_top_k=10), llm=llm) + top = sel.select_sync("what is ...", documents, scores=rrf_scores) +""" + +from .agent import Agent, AgentConfig, ToolExecutionError +from .loop import KNOWN_LLM_ERRORS +from .results import ( + ERROR_BAD_FINISH_REASON, + ERROR_CONTENT_POLICY, + ERROR_CONTEXT_LIMIT, + ERROR_LLM_CALL_FAILED, + ERROR_MAX_STEPS, + ERROR_TOOL_FAILED, + ERROR_UNEXPECTED, + AgentError, + AgentRunResult, +) +from .selection_agent import SelectionAgent, SelectionAgentConfig +from .tools import ( + BaseEndTool, + BaseRetrieveTool, + BaseTool, + FinalResults, + LogSelectedDocs, + ReasoningAugmentedRetrieveTool, + RetrieveContext, + RetrieveTool, + SelectionThinkTool, + ThinkTool, + ToolContractError, + ToolError, + create_retrieve_tool, +) + +__all__ = [ + "ERROR_BAD_FINISH_REASON", + "ERROR_CONTENT_POLICY", + "ERROR_CONTEXT_LIMIT", + "ERROR_LLM_CALL_FAILED", + "ERROR_MAX_STEPS", + "ERROR_TOOL_FAILED", + "ERROR_UNEXPECTED", + "KNOWN_LLM_ERRORS", + "Agent", + "AgentConfig", + "AgentError", + "AgentRunResult", + "BaseEndTool", + "BaseRetrieveTool", + "BaseTool", + "FinalResults", + "LogSelectedDocs", + "ReasoningAugmentedRetrieveTool", + "RetrieveContext", + "RetrieveTool", + "SelectionAgent", + "SelectionAgentConfig", + "SelectionThinkTool", + "ThinkTool", + "ToolContractError", + "ToolError", + "ToolExecutionError", + "create_retrieve_tool", +] diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/agent.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/agent.py new file mode 100644 index 0000000000..5f977e596b --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/agent.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The agentic-retrieval loop: the Agent class for document retrieval. + +The generic loop mechanics (step loop, tool dispatch, error policy, raw-IO +logging, result building) live in ``loop.py`` and are shared with +``selection_agent.py``; this module adds everything retrieval-specific — +construction of the standard tool set, the user-message bootstrap, and the +over-fetch/dedup/exclusion bookkeeping around retrieve tools. +""" + +from __future__ import annotations + +import asyncio +import logging +from contextlib import nullcontext +from pathlib import Path +from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union + +from .cache_propagation import PropagationPacer +from .llm import BaseLLMBackend, bind_query_id +from .loop import ( + BaseAgentLoopConfig, + ToolExecutionError, + _awrite_json, + _BaseAgentLoop, + _RunState, + build_auto_continue_msg, +) +from .prompts import render_system_prompt +from .results import AgentRunResult +from .tools import ( + BaseEndTool, + BaseRetrieveTool, + BaseTool, + FinalResults, + RetrieveContext, + ThinkTool, +) +from .tools.base_tool import tool_error_text +from .tools.retrieve import retrieve_output_to_msg_content + +logger = logging.getLogger(__name__) + + +class AgentConfig(BaseAgentLoopConfig): + """Configuration for the agent. + + Pure data — LLM and tool *instances* are constructor arguments, not + config. ``system_prompt`` is a packaged prompt name or a filesystem path; + ``None`` selects the default. ``mode`` is the extension point for + additional agent behaviors; ``"select"`` (document selection) is the only + one implemented, and any other value is rejected. The loop-policy + fields (``max_steps``, ``on_error``, logging + and pacing knobs) are inherited from + :class:`~nemo_agent.loop.BaseAgentLoopConfig`. + """ + + mode: Literal["select"] = "select" + system_prompt: Optional[str] = None + enforce_top_k: bool = True + target_top_k: Optional[int] = 10 + extended_relevance: bool = True + enable_think: bool = False + end_tool_with_msg: bool = True + user_msg_type: Literal["simple", "with_results"] = "with_results" + ensure_new_docs: bool = True + + +class Agent(_BaseAgentLoop): + """LLM agent loop with retrieval tools. + + Construction wires everything the loop needs; per-query state lives in a + private run state, so one instance safely serves many (including + concurrent) runs. Standard tools are assembled internally: the think tool + (``config.enable_think``), the end tool (``final_results``; overridable + via the ``end_tool`` argument), and the required primary ``retrieve_tool`` + (which also powers the ``user_msg_type="with_results"`` bootstrap). + ``tool_overrides`` adds/replaces tools by name — extra + :class:`BaseRetrieveTool` instances get the same retrieval bookkeeping as + the primary one. + + Error behavior is driven by ``config.on_error``; see + :data:`~nemo_agent.loop.KNOWN_LLM_ERRORS` and + :class:`~nemo_agent.results.AgentError`. + """ + + def __init__( + self, + config: AgentConfig, + llm: BaseLLMBackend, + retrieve_tool: BaseRetrieveTool, + *, + end_tool: Optional[BaseEndTool] = None, + tool_overrides: Optional[Dict[str, BaseTool]] = None, + ) -> None: + if not isinstance(config, AgentConfig): + raise TypeError(f"config must be an AgentConfig, got {type(config).__name__}.") + super().__init__(config=config, llm=llm) + if not isinstance(retrieve_tool, BaseRetrieveTool): + raise TypeError( + "retrieve_tool must be a BaseRetrieveTool (build one with " + f"create_retrieve_tool or subclass it), got {type(retrieve_tool).__name__}." + ) + + if end_tool is not None and not isinstance(end_tool, BaseEndTool): + raise TypeError(f"end_tool must be a BaseEndTool, got {type(end_tool).__name__}.") + default_prompt = "06_select_lean_v1.j2" + end_payload_phrase = "with your selected doc_ids" + if end_tool is None: + top_k = int(config.target_top_k) if config.enforce_top_k and config.target_top_k else None + end_tool = FinalResults(top_k=top_k, include_msg=config.end_tool_with_msg) + self._retrieve_tool = retrieve_tool + self._end_tool = end_tool + + tools: Dict[str, BaseTool] = {} + # Primary retrieve tool first (it also powers the with_results + # bootstrap retrieve). + tools[retrieve_tool.name] = retrieve_tool + # Extra/override tools next (e.g. a keyword-search retrieve tool) so + # they sit between the primary retrieve tool and the end tool. + for key, tool in (tool_overrides or {}).items(): + if not isinstance(tool, BaseTool): + raise TypeError(f"tool_overrides[{key!r}] must be a BaseTool, got {type(tool).__name__}.") + if key != tool.name: + raise ValueError(f"tool_overrides key {key!r} does not match the tool's spec name {tool.name!r}.") + if key == retrieve_tool.name: + raise ValueError( + f"tool_overrides may not replace the primary retrieve tool {key!r}; " + "pass it as the retrieve_tool argument instead." + ) + if key == end_tool.name: + raise ValueError( + f"tool_overrides may not replace the end tool {key!r}; " "pass it as the end_tool argument instead." + ) + tools[key] = tool + # End tool after the retrieve tools. + if end_tool.name in tools: + raise ValueError(f"Duplicate name {end_tool.name!r} among the agent's tools.") + tools[end_tool.name] = end_tool + # Think tool LAST (optional scratchpad). An override that already + # supplied a same-named tool keeps its slot. + if config.enable_think: + think = ThinkTool(extended_relevance=config.extended_relevance) + if think.name in tools: + logger.info("tool_overrides replaces the think tool %r.", think.name) + else: + tools[think.name] = think + self._tool_map = tools + # Specs are static per instance. Order is deterministic (retrieve, + # then overrides in insertion order, then the end tool, then the think + # tool last) — it matters both for model behavior and because + # Anthropic-style prompt-cache markers land on the last tool spec. + self._tool_specs = [t.spec for t in tools.values()] + + prompt_name = config.system_prompt or default_prompt + system_prompt = render_system_prompt( + prompt_name, + with_init_docs=config.user_msg_type == "with_results", + enforce_top_k=config.enforce_top_k, + top_k=config.target_top_k, + extended_relevance=config.extended_relevance, + ) + self._system_msg = {"role": "system", "content": [{"type": "text", "text": system_prompt}]} + self._auto_user_msg = build_auto_continue_msg(end_tool.name, end_payload_phrase) + + # ------------------------------------------------------------------ + # Entry points. + # ------------------------------------------------------------------ + + async def run( + self, + query: str, + *, + query_id: Optional[str] = None, + task_instruction: Optional[str] = None, + task_info: Optional[Any] = None, + exclude_docids: Optional[Set[str]] = None, + raw_log_dir: Optional[Union[str, Path]] = None, + ) -> AgentRunResult: + """Run the agent for one query. + + Parameters + ---------- + query: + The user's question; also the retrieve tools' ``global_query``. + query_id: + Optional id for this run. When given, the agent binds it itself, + so token usage lands under ``llm.get_usage(query_id)`` and + progress/error logs are labeled — callers then don't need to + manage ``bind_query_id`` at all. When omitted, an ambient + ``bind_query_id(...)`` established by the caller still applies. + task_instruction: + Optional retrieval instruction, prefixed onto the user message + (``Instruct: ...``) unless it already starts with one. + task_info: + Arbitrary JSON-serializable info (e.g. the query id) written to + ``extra_info.json`` when ``raw_log_dir`` is set — it associates + the log directory with the query. + exclude_docids: + Document ids retrieve tools must never surface this run. + raw_log_dir: + Per-query directory for raw LLM IO artifacts, built by the + caller. ``None`` discards the artifacts. See + ``config.write_all_llm_io_logs`` for all-steps vs last-step; + per-step extras are always written at run end as + ``api_response_extras.json``. + """ + binding = bind_query_id(query_id) if query_id is not None else nullcontext() + with binding: + state = _RunState( + query=str(query), + raw_log_dir=Path(raw_log_dir) if raw_log_dir is not None else None, + exclude_docs=set(exclude_docids) if exclude_docids is not None else set(), + pacer=PropagationPacer(target_s=self.config.cache_propagation_target_s), + tool_map=self._tool_map, + tool_specs=self._tool_specs, + auto_user_msg=self._auto_user_msg, + stage="main_agent", + message_history=[self._system_msg], + ) + if state.raw_log_dir is not None and task_info is not None: + await _awrite_json(task_info, state.raw_log_dir, "extra_info.json") + return await self._run_state_to_result( + state, prologue=lambda: self._append_user_message(state, task_instruction) + ) + + def run_sync(self, query: str, **kwargs: Any) -> AgentRunResult: + """Synchronous facade over :meth:`run` (one ``asyncio.run`` per call). + + For thread-based callers; must not be invoked from inside a running + event loop — ``await run(...)`` there instead. + """ + return asyncio.run(self.run(query, **kwargs)) + + # ------------------------------------------------------------------ + # Retrieval-specific pieces on top of the shared loop. + # ------------------------------------------------------------------ + + async def _append_user_message(self, state: _RunState, task_instruction: Optional[str]) -> None: + instruction = (task_instruction or "").strip() + if instruction and not instruction.lower().startswith("instruct"): + instruction = f"Instruct: {instruction}" + if instruction: + instruction += "\n" + task_inst_query = f"{instruction}Query:\n{state.query}" + + if self.config.user_msg_type == "simple": + state.message_history.append({"role": "user", "content": [{"type": "text", "text": task_inst_query}]}) + return + # "with_results": bootstrap the conversation with an initial retrieve + # through the primary retrieve tool. + try: + content = await self._execute_retrieve( + state, tool=self._retrieve_tool, llm_kwargs={"query": state.query}, query_type="main" + ) + except Exception as e: + raise ToolExecutionError(self._retrieve_tool.name, e) from e + state.message_history.append( + { + "role": "user", + "content": [ + {"type": "text", "text": task_inst_query}, + {"type": "text", "text": "Retrieved Documents:"}, + ] + + content, + } + ) + + async def _dispatch_tool_call( + self, state: _RunState, fn_name: str, fn_kwargs: Dict[str, Any] + ) -> Tuple[List[Dict[str, Any]], bool]: + """Intercept retrieve tools for the retrieval bookkeeping; defer the rest.""" + tool = state.tool_map.get(fn_name) + if isinstance(tool, BaseRetrieveTool): + try: + content = await self._execute_retrieve(state, tool=tool, llm_kwargs=fn_kwargs, query_type="agent") + except Exception as e: + raise ToolExecutionError(fn_name, e) from e + return content, False + return await super()._dispatch_tool_call(state, fn_name, fn_kwargs) + + async def _execute_retrieve( + self, + state: _RunState, + *, + tool: BaseRetrieveTool, + llm_kwargs: Dict[str, Any], + query_type: str, + ) -> List[Dict[str, Any]]: + """Run one retrieval with over-fetch/dedup/exclusion bookkeeping. + + Returns tool-message content blocks; LLM-recoverable problems (bad + arguments, ToolError from the tool) come back as error-text content. + """ + + def _llm_error(exc: BaseException) -> List[Dict[str, Any]]: + return [{"type": "text", "text": tool_error_text(tool.name, exc)}] + + unexpected = set(llm_kwargs) - {"query", "top_k"} + if unexpected: + return _llm_error(TypeError(f"unexpected argument(s): {', '.join(sorted(map(str, unexpected)))}.")) + if "query" not in llm_kwargs: + return _llm_error(TypeError("the 'query' argument is required.")) + query = str(llm_kwargs["query"]) + top_k_raw = llm_kwargs.get("top_k", tool.default_top_k) + try: + top_k = int(top_k_raw) + except (TypeError, ValueError): + return _llm_error(TypeError(f"'top_k' must be an integer, got {top_k_raw!r}.")) + if top_k <= 0: + return _llm_error(TypeError(f"'top_k' must be a positive integer, got {top_k}.")) + + context = RetrieveContext(global_query=state.query, reasoning=state.last_reasoning) + # Over-fetch so that, after dropping excluded ids and counting only + # docs new to this run (when ensure_new_docs), top_k new docs remain. + seen = state.retrieved_docs if self.config.ensure_new_docs else set() + fetch_k = top_k + len(seen) + len(state.exclude_docs) + result = await tool.acall(query, top_k=fetch_k, context=context) + if isinstance(result, str): + return [{"type": "text", "text": result}] + + result = sorted(result, key=lambda d: d["score"], reverse=True) + output: List[Dict[str, Any]] = [] + num_new = 0 + call_seen: Set[str] = set() + for item in result: + if item["id"] in state.exclude_docs: + continue + if item["id"] in call_seen: + continue + call_seen.add(item["id"]) + rec = dict(item) + if rec["id"] not in seen: + num_new += 1 + output.append(rec) + if num_new >= top_k: + break + + # Repeats keep their slot but drop their content (the LLM saw it + # already); must run BEFORE the ids below join retrieved_docs. + for rec in output: + if rec["id"] in state.retrieved_docs: + rec.pop("image", None) + rec.pop("text", None) + rec["note"] = ( + "This document is retrieved before. See previous retrieval results " + f"for the content of this document (id: {rec['id']})." + ) + for rec in output: + state.retrieved_docs.add(rec["id"]) + + state.retrieval_log.append( + { + "input": {"query": query, "top_k": top_k}, + "tool_name": tool.name, + "query_type": query_type, + "output": output, + } + ) + return retrieve_output_to_msg_content(output) diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/cache_propagation.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/cache_propagation.py new file mode 100644 index 0000000000..98a4696402 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/cache_propagation.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Adaptive propagation pacing for back-to-back LLM calls. + +Anthropic prompt-cache writes via Bedrock-via-NIM propagate across cache +replicas with a small lag (empirically ~5-15s on this stack). When a +downstream LLM call needs to read a cache entry written by an immediately +preceding upstream call, the natural inter-call gap may be too short and +the read misses. + +:class:`PropagationPacer` lets a caller declare "ensure at least N seconds +between consecutive marked calls" without inserting a blind sleep on every +call: it sleeps only the *remaining* time relative to the most recent mark, +so calls that take longer than ``target_s`` on their own incur no extra +wait at all. + +Set ``target_s = 0`` (the default for non-Anthropic providers) to disable +pacing entirely; both ``await_propagation`` and ``mark`` become cheap +no-ops. +""" + +from __future__ import annotations + +import asyncio +import time +from typing import Awaitable, Callable, Optional + + +class PropagationPacer: + """Maintain a minimum-wall-time gap between consecutive LLM calls. + + Typical usage:: + + pacer = PropagationPacer(target_s=30.0) + for k in topk_list: + await pacer.await_propagation() + result = await llm_call(k) + pacer.mark() + """ + + def __init__( + self, + target_s: float, + *, + now: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + self.target_s = max(0.0, float(target_s)) + self._last_ts: Optional[float] = None + self._now = now + self._sleep = sleep + + @property + def primed(self) -> bool: + """True once :meth:`mark` has been called at least once.""" + return self._last_ts is not None + + def reset(self) -> None: + """Clear the most recent mark (next ``await_propagation`` is a no-op).""" + self._last_ts = None + + def mark(self) -> None: + """Record the current time as the latest call's finish timestamp. + + Always call this AFTER the LLM call returns; ``await_propagation`` + reads back from this timestamp on the next iteration. + """ + if self.target_s <= 0: + return + self._last_ts = self._now() + + async def await_propagation(self) -> float: + """Sleep just long enough to meet the target gap. Returns the seconds slept. + + No-op (returns 0.0) when: + - ``target_s`` is zero (pacing disabled), OR + - the pacer has not been marked yet (first call in the sequence), OR + - the natural gap since the last mark already meets or exceeds + ``target_s``. + """ + if self.target_s <= 0 or self._last_ts is None: + return 0.0 + elapsed = self._now() - self._last_ts + remaining = self.target_s - elapsed + if remaining <= 0: + return 0.0 + await self._sleep(remaining) + return remaining diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/__init__.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/__init__.py new file mode 100644 index 0000000000..d51a5c5358 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/__init__.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal LLM backend components for the private agent implementation. + +Internal integration surface: + +- :class:`BaseLLMBackend` — the ABC to implement for a custom backend + (one required method: ``_completion_impl``). +- :class:`CompletionResult` — the envelope every completion call returns. +- :class:`BaseLLMConfig` / :class:`CallableLLMConfig` / :class:`LiteLLMConfig` — + configuration models, each colocated with its backend class. +- :func:`create_llm` — config-driven backend factory. +- :func:`create_llm_config` — kwargs-filtering config factory: builds a + backend's config, dropping (with a warning) any field that config does not + support. +- :func:`get_available_backends` — the registered backend names; the source of + truth callers should read. +- ``errors`` — the :class:`LLMCallError` hierarchy callers branch on. +- ``error_classification`` — :func:`classify_call_exception`, which maps a + completion callable's failure onto that hierarchy without depending on any + HTTP library. +- ``usage`` binders — :func:`bind_query_id` / :func:`bind_stage` for + per-(query, stage) token-usage attribution. +- ``helpers`` — reusable pieces for backend authors + (:func:`extract_reasoning_from_message`, :func:`extract_text_content`, + :func:`normalize_messages_for_api`, :func:`strip_private_message_keys`, + :func:`redact_url`, :func:`excerpt`, :func:`resolve_api_key`). +""" + +import logging +from typing import Any + +from .base_backend import BaseLLMBackend, BaseLLMConfig +from .callable_backend import CallableLLMBackend, CallableLLMConfig +from .error_classification import classify_call_exception +from .errors import ContentPolicyError, ContextLimitError, LLMCallError, RateLimitError +from .helpers import ( + excerpt, + extract_reasoning_from_message, + extract_text_content, + normalize_messages_for_api, + redact_url, + resolve_api_key, + strip_private_message_keys, +) +from .litellm_backend import LiteLLMBackend, LiteLLMConfig +from .result import CompletionResult +from .usage import ( + UNSET_QUERY, + UNSET_STAGE, + bind_query_id, + bind_stage, + coerce_usage_to_dict, + deep_merge_usage, + deep_merge_usage_breakdown, + get_query_id, + get_stage, + sum_usage_breakdown, +) + +logger = logging.getLogger(__name__) + +_BACKEND_REGISTRY = { + "litellm": LiteLLMBackend, + "callable": CallableLLMBackend, +} + + +def get_available_backends() -> tuple[str, ...]: + """Return the names of all registered LLM backends, sorted.""" + return tuple(sorted(_BACKEND_REGISTRY)) + + +def create_llm(config: BaseLLMConfig, **kwargs: Any) -> BaseLLMBackend: + """Instantiate the LLM backend selected by ``config.backend``. + + Each backend validates that it received its own config subclass (e.g. + ``backend="litellm"`` requires a :class:`LiteLLMConfig`). + + Extra keyword arguments are forwarded verbatim to the backend constructor — + this is how injection-only backends receive their runtime dependency (e.g. + ``create_llm(config, completion_fn=fn)`` for :class:`CallableLLMBackend`). The + factory does not inspect or filter them: passing a kwarg a backend does not + accept is a caller error and surfaces as a ``TypeError`` from the constructor. + """ + try: + backend_cls = _BACKEND_REGISTRY[config.backend] + except KeyError: + raise ValueError( + f"Unrecognized LLM backend {config.backend!r}. " f"Available backends: {sorted(_BACKEND_REGISTRY)}." + ) from None + return backend_cls(config, **kwargs) + + +def create_llm_config(backend: str, **kwargs: Any) -> BaseLLMConfig: + """Build the configuration for ``backend``, dropping unsupported fields. + + ``backend`` selects the config subclass (via the same registry as + :func:`create_llm`). Any keyword not declared on that config subclass is + dropped with a ``WARNING`` naming the dropped keys (values are never logged); + required-field and per-field type validation still apply to the kept fields. + + ``backend`` is the selector and is never forwarded into the config — the + chosen subclass's ``Literal`` default sets ``config.backend``. Building a + config never instantiates the backend. + """ + try: + backend_cls = _BACKEND_REGISTRY[backend] + except KeyError: + raise ValueError( + f"Unrecognized LLM backend {backend!r}. " f"Available backends: {sorted(_BACKEND_REGISTRY)}." + ) from None + config_cls = backend_cls.config_cls + supported = set(config_cls.model_fields) + kept = {k: v for k, v in kwargs.items() if k in supported} + dropped = sorted(k for k in kwargs if k not in supported) + if dropped: + logger.warning( + "create_llm_config: backend %r (%s) does not support config field(s) " "%s; ignoring them.", + backend, + config_cls.__name__, + dropped, + ) + return config_cls(**kept) + + +__all__ = [ + "BaseLLMBackend", + "BaseLLMConfig", + "CallableLLMBackend", + "CallableLLMConfig", + "CompletionResult", + "ContentPolicyError", + "ContextLimitError", + "LLMCallError", + "LiteLLMBackend", + "LiteLLMConfig", + "RateLimitError", + "UNSET_QUERY", + "UNSET_STAGE", + "bind_query_id", + "bind_stage", + "classify_call_exception", + "coerce_usage_to_dict", + "create_llm", + "create_llm_config", + "deep_merge_usage", + "deep_merge_usage_breakdown", + "excerpt", + "extract_reasoning_from_message", + "extract_text_content", + "get_available_backends", + "get_query_id", + "get_stage", + "normalize_messages_for_api", + "redact_url", + "resolve_api_key", + "strip_private_message_keys", + "sum_usage_breakdown", +] diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/base_backend.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/base_backend.py new file mode 100644 index 0000000000..29a9caa714 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/base_backend.py @@ -0,0 +1,321 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Abstract base class every LLM backend implements, and its base configuration.""" + +from __future__ import annotations + +import asyncio +import logging +import threading +import time +from abc import ABC, abstractmethod +from copy import deepcopy +from typing import Any, ClassVar, Dict, List, Optional + +from pydantic import BaseModel, ConfigDict + +from .errors import RateLimitError +from .result import CompletionResult +from .usage import ( + UNSET_QUERY, + UNSET_STAGE, + deep_merge_usage, + deep_merge_usage_breakdown, + get_query_id, + get_stage, + sum_usage_breakdown, +) + +logger = logging.getLogger(__name__) + +# Ceiling for provider-communicated `RateLimitError.retry_after` values, so a +# bogus Retry-After header can never stall the retry loop for e.g. an hour. +_RETRY_AFTER_CAP_S = 300.0 + + +class BaseLLMConfig(BaseModel): + """Backend-agnostic LLM configuration. + + Each backend subclasses this with its own fields, colocated with the + backend class (see ``LiteLLMConfig`` in ``litellm_backend.py``). Note what + is deliberately absent: raw-IO log paths and error-handling policy belong + to the caller (agent/pipeline), not the LLM — do not add them back here. + + Attributes + ---------- + backend: + Discriminator used by :func:`nemo_agent.llm.create_llm` to pick the + backend class (kept as a string so yaml/hydra-driven config works). + Defaults to ``"callable"``, the library's default backend. Note each + backend subclass pins this by ``Literal``, so the default here only + applies to a bare ``BaseLLMConfig`` — which cannot build a backend, since + every backend requires its own config subclass. + model: + Model identifier, in whatever form the backend expects. + api_key: + API key, or the indirection ``"os.environ/VAR_NAME"`` to resolve the + key from the environment at backend construction. + base_url: + Endpoint base URL, if the backend needs one. + tool_choice: + OpenAI-style tool-choice policy; sent only when tools are provided. + max_completion_tokens: + Per-request completion budget, if set. Overridable per call. + reasoning_effort: + Provider reasoning-effort knob (e.g. ``"low"`` / ``"medium"`` / + ``"high"``, provider-dependent). Backend-agnostic *declaration*; each + backend applies or ignores it. + temperature: + Sampling temperature. Backend-agnostic *declaration*; each backend + forwards it when set and leaves it unset (provider default) when ``None``. + parallel_tool_calls: + Whether the model may emit multiple tool calls in one turn. Backend-agnostic + *declaration*; forwarded when set, left unset (provider default) when ``None``. + capture_raw_io: + When true, backends populate ``CompletionResult.raw_request`` / + ``raw_response`` (JSON-serializable, credential-redacted). Off by + default — the dumps cost memory and nobody should pay for artifacts + they don't persist. + rate_limit_max_retries / rate_limit_retry_sleep_s: + Policy for the base class's retry-on-:class:`RateLimitError` loop + wrapped around every call. An exception-provided ``retry_after`` wins + over the configured sleep but never exceeds a hardcoded cap. + """ + + model_config = ConfigDict(extra="forbid") + + backend: str = "callable" + model: str + api_key: Optional[str] = None + base_url: Optional[str] = None + tool_choice: str = "auto" + max_completion_tokens: Optional[int] = None + reasoning_effort: Optional[str] = None + temperature: Optional[float] = None + parallel_tool_calls: Optional[bool] = None + capture_raw_io: bool = False + rate_limit_max_retries: int = 3 + rate_limit_retry_sleep_s: float = 60.0 + + +class BaseLLMBackend(ABC): + """Chat-completion client with built-in usage tracking and rate-limit retry. + + Subclassing + ----------- + Implement :meth:`_completion_impl` (sync). Backends with a native async + client should also override :meth:`_acompletion_impl`; otherwise the + default bridges via ``asyncio.to_thread`` (contextvars propagate, so usage + attribution survives the bridge). NEVER override :meth:`completion` / + :meth:`acompletion` — they are the templates that guarantee usage + recording and retry policy run on every call. + + Contract for implementations + ---------------------------- + - Return a fully-assembled :class:`CompletionResult` (see its docstring for + per-field contracts, including ``capture_raw_io`` gating and + ``raw_request`` credential redaction). + - Raise only :class:`~nemo_agent.llm.errors.LLMCallError` subclasses for + call failures (``raise ... from e``, wrapping only the client call); + anything else escaping is treated as a bug. See ``errors`` module. + - Never mutate ``messages`` / ``tools`` in place — callers pass live + history. + - ``**overrides`` are per-call, backend-interpreted request overrides + (e.g. ``max_completion_tokens=64, num_retries=0`` for a preflight + probe); apply what you understand, tolerate what you don't. + + Usage tracking + -------------- + After every successful call the template deep-merges ``result.usage`` into + ``usage[query_id][stage]``, keyed by the ambient + :func:`~nemo_agent.llm.usage.bind_query_id` / + :func:`~nemo_agent.llm.usage.bind_stage` bindings (falling back to + ``UNSET_QUERY`` / ``UNSET_STAGE``). One backend instance is typically + shared across all queries and stages; accumulation is guarded by a lock + and is safe under both threads and asyncio. A failure to record usage is + logged loudly but never fails a completed call. + """ + + #: Each concrete backend MUST set this to its paired config subclass. Read by + #: ``create_llm_config`` WITHOUT instantiating the backend, and used by + #: ``__init__`` below for the generic config-type check. + config_cls: ClassVar[type[BaseLLMConfig]] + + def __init__(self, config: BaseLLMConfig) -> None: + expected = type(self).config_cls + if not isinstance(config, expected): + raise TypeError(f"{type(self).__name__} requires a {expected.__name__}, " f"got {type(config).__name__}") + self.config = config + # Name-mangled on purpose: subclasses own the plain `_` namespace, and + # accidentally clobbering the accumulator must be impossible. Access + # goes through the public usage accessors below. + self.__usage: Dict[str, Dict[str, Dict[str, Any]]] = {} + self.__usage_lock = threading.Lock() + + # ------------------------------------------------------------------ + # Public templates — do not override. + # ------------------------------------------------------------------ + + def completion( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + **overrides: Any, + ) -> CompletionResult: + """Make a chat-completion call (sync). See class docstring for the contract.""" + attempt = 0 + while True: + try: + result = self._completion_impl(messages, tools=tools, **overrides) + break + except RateLimitError as e: + delay = self._rate_limit_retry_delay(e, attempt) + if delay is None: + raise + time.sleep(delay) + attempt += 1 + result = self._require_completion_result(result, "_completion_impl") + self._record_usage(result.usage) + return result + + async def acompletion( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + **overrides: Any, + ) -> CompletionResult: + """Make a chat-completion call (async). See class docstring for the contract.""" + attempt = 0 + while True: + try: + result = await self._acompletion_impl(messages, tools=tools, **overrides) + break + except RateLimitError as e: + delay = self._rate_limit_retry_delay(e, attempt) + if delay is None: + raise + await asyncio.sleep(delay) + attempt += 1 + result = self._require_completion_result(result, "_acompletion_impl") + self._record_usage(result.usage) + return result + + # ------------------------------------------------------------------ + # Implementation hooks. + # ------------------------------------------------------------------ + + @abstractmethod + def _completion_impl( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + **overrides: Any, + ) -> CompletionResult: + """Backend-specific synchronous completion. See class docstring.""" + + async def _acompletion_impl( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + **overrides: Any, + ) -> CompletionResult: + """Backend-specific async completion. + + Default bridges to :meth:`_completion_impl` in a worker thread. + Override with a native async client when available. + """ + return await asyncio.to_thread(self._completion_impl, messages, tools=tools, **overrides) + + def _require_completion_result(self, result: Any, impl_name: str) -> CompletionResult: + """Make a contract violation self-diagnosing instead of a downstream AttributeError.""" + if not isinstance(result, CompletionResult): + raise TypeError( + f"{type(self).__name__}.{impl_name} must return a CompletionResult, " f"got {type(result).__name__}" + ) + return result + + # ------------------------------------------------------------------ + # Rate-limit retry policy (shared by both templates so they can't drift). + # ------------------------------------------------------------------ + + def _rate_limit_retry_delay(self, e: RateLimitError, attempt: int) -> Optional[float]: + """Seconds to sleep before retrying ``attempt`` (0-based), or None to re-raise. + + A provider-communicated ``retry_after`` wins over the configured sleep + but is capped by ``_RETRY_AFTER_CAP_S``; garbage values (negative, NaN, + non-numeric) fall back to the configured sleep. + """ + if attempt >= self.config.rate_limit_max_retries: + return None + retry_after: Optional[float] = None + if e.retry_after is not None: + try: + retry_after = float(e.retry_after) + except (TypeError, ValueError): + retry_after = None + if retry_after is not None and 0 <= retry_after < float("inf"): + delay = min(retry_after, _RETRY_AFTER_CAP_S) + else: + delay = self.config.rate_limit_retry_sleep_s + logger.info( + "Rate limited; sleeping %.1fs before retry %d/%d.", + delay, + attempt + 1, + self.config.rate_limit_max_retries, + ) + return float(delay) + + # ------------------------------------------------------------------ + # Usage tracking. + # ------------------------------------------------------------------ + + def _record_usage(self, usage: Optional[Dict[str, Any]]) -> None: + if not usage: + return + query_id = get_query_id() or UNSET_QUERY + stage = get_stage() or UNSET_STAGE + try: + with self.__usage_lock: + bucket = self.__usage.setdefault(query_id, {}).setdefault(stage, {}) + deep_merge_usage(bucket, usage) + except Exception: + logger.warning( + "Failed to record LLM token usage for query_id=%r stage=%r; " "usage totals will undercount.", + query_id, + stage, + exc_info=True, + ) + + def get_usage(self, query_id: Optional[str] = None) -> Dict[str, Any]: + """Deep copy of accumulated usage. + + With ``query_id``: that query's ``{stage: usage}`` breakdown (``{}`` if + unknown). Without: the full ``{query_id: {stage: usage}}`` mapping. + """ + with self.__usage_lock: + if query_id is None: + return deepcopy(self.__usage) + return deepcopy(self.__usage.get(query_id, {})) + + def pop_query_usage(self, query_id: str) -> Dict[str, Any]: + """Remove and return one query's ``{stage: usage}`` breakdown (``{}`` if unknown).""" + with self.__usage_lock: + return self.__usage.pop(query_id, {}) + + def get_usage_by_stage(self) -> Dict[str, Any]: + """Derived run totals: ``{stage: usage}`` merged across all queries.""" + with self.__usage_lock: + totals: Dict[str, Any] = {} + for stage_breakdown in self.__usage.values(): + deep_merge_usage_breakdown(totals, stage_breakdown) + return totals + + def get_total_usage(self) -> Dict[str, Any]: + """Derived grand total across all queries and stages.""" + return sum_usage_breakdown(self.get_usage_by_stage()) + + def reset_usage(self) -> None: + """Drop all accumulated usage.""" + with self.__usage_lock: + self.__usage = {} diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/callable_backend.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/callable_backend.py new file mode 100644 index 0000000000..cb697b7b23 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/callable_backend.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LLM backend that adapts an external chat-completion *callable* to the backend contract. + +This is the library's default backend. It carries no HTTP client or LLM serving logic of its own: +everything related to LLM serving or HTTP client logic lives inside the injected callable. + +The callable contract +--------------------- +The callable is keyword-only and returns an OpenAI-compatible ``chat.completion`` +dict. It must accept the full keyword set this backend sends:: + + invoke_url, messages, model, api_key, tools, tool_choice, timeout_s, + temperature, max_tokens, extra_body, max_retries, max_429_retries + +Every keyword is sent on every call. + +``temperature`` may be ``None``, meaning *unset*. Callables interpret that +themselves — a hosted endpoint should omit the field so the provider default +applies, while an in-process engine that has no provider to defer to should pick +a concrete value. That asymmetry is intentional; see the callables' docstrings. + +Unlike the registry-buildable backends, :class:`CallableLLMBackend` needs a live +``completion_fn`` and therefore cannot be constructed from config alone. Callers +inject it: ``create_llm(config, completion_fn=fn)``. +""" + +from __future__ import annotations + +import logging +from copy import deepcopy +from typing import Any, Callable, Dict, List, Literal, Optional + +from pydantic import Field + +from .base_backend import BaseLLMBackend, BaseLLMConfig +from .error_classification import classify_call_exception +from .errors import LLMCallError +from .helpers import ( + extract_reasoning_from_message, + extract_text_content, + normalize_messages_for_api, + redact_url, + resolve_api_key, + strip_private_message_keys, +) +from .result import CompletionResult +from .usage import coerce_usage_to_dict + +logger = logging.getLogger(__name__) + +#: The completion callable contract. See the module docstring for the required +#: keyword set — ``Callable[..., Dict]`` cannot express it. +CompletionFn = Callable[..., Dict[str, Any]] + +_REDACTED = "***REDACTED***" +#: Request keys whose values are credentials and must be scrubbed from a captured +#: ``raw_request``. Exact-match (not substring) so e.g. ``max_tokens`` is never +#: caught by a naive "token" check. +_SENSITIVE_REQUEST_KEYS = frozenset({"api_key"}) +#: Request keys holding a URL, which may carry credentials in userinfo or query. +_URL_REQUEST_KEYS = frozenset({"invoke_url"}) + + +def _redacted_request(call_kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Deep-copied, credential-redacted snapshot of the request.""" + out: Dict[str, Any] = {} + for key, value in call_kwargs.items(): + if key in _SENSITIVE_REQUEST_KEYS: + out[key] = _REDACTED + elif key in _URL_REQUEST_KEYS and value is not None: + out[key] = redact_url(str(value)) + else: + out[key] = deepcopy(value) + return out + + +class CallableLLMConfig(BaseLLMConfig): + """Configuration for :class:`CallableLLMBackend`. + + ``backend`` is pinned by type, so a ``CallableLLMConfig`` is only ever routed + to :class:`CallableLLMBackend`. Request knobs are inherited from + :class:`BaseLLMConfig` (``model``, ``temperature``, ``tool_choice``, + ``max_completion_tokens``, ...). + + ``base_url`` (inherited, optional) is the callable's ``invoke_url``, forwarded + **verbatim**. It stays optional and unvalidated on purpose: an in-process + callable has no endpoint, and neither this config nor the backend can tell + a remote callable from an in-process one. A remote callable handed ``None`` + fails on its first call with its own error. + + Attributes + ---------- + timeout_s: + Per-request wire timeout handed to the callable. + max_retries / max_429_retries: + Retry budget handed to the callable, which owns retrying. Both are + forwarded rather than acted on here. + rate_limit_max_retries: + Redeclared as ``0``. The callable owns retries, so the base template must + make exactly one attempt; retrying on top would multiply the two budgets. + """ + + backend: Literal["callable"] = "callable" + + timeout_s: float = Field(default=120.0, gt=0) + #: A TOTAL ATTEMPT count in the callable's own spelling, not a retry count: + #: 3 means at most 3 requests. ``ge=1`` is load-bearing — a callable that + #: loops ``while attempt < max_retries`` issues ZERO requests at 0 and then + #: reports retries-exhausted, which reads as an endpoint failure. + max_retries: int = Field(default=3, ge=1) + #: Responses tolerated before giving up on a rate-limited endpoint. Sized so a + #: sustained 429 fails in roughly two minutes instead of stalling the agent. + max_429_retries: int = Field(default=6, ge=1) + #: The callable retries internally, so the base template must not retry too. + #: Currently reached only when the callable surfaces a rate limit as a typed + #: ``RateLimitError``; kept explicit so a future callable that reports rate + #: limits WITHOUT retrying can opt back in by overriding it. + rate_limit_max_retries: int = 0 + + +class CallableLLMBackend(BaseLLMBackend): + """Adapter from an OpenAI-compatible completion callable to :class:`BaseLLMBackend`. + + The base class owns usage recording and the rate-limit retry template; this + subclass implements :meth:`_completion_impl`, shapes the request, parses the + response, and translates failures. + + Error translation + ----------------- + An exception raised by the callable is classified by + :func:`~nemo_agent.llm.error_classification.classify_call_exception` and + re-raised as the matching :class:`~nemo_agent.llm.errors.LLMCallError` + subclass, with the original chained on ``__cause__``. Classification is + best-effort and degrades to a plain ``LLMCallError``: a failure carrying a + response object is classified from its status and body, anything else from + its message text. A malformed response dict is surfaced as an explicit + ``LLMCallError`` rather than a raw ``KeyError``. + + Not populated + ------------- + :attr:`CompletionResult.extra_response_info` is always ``{}``. The callable + returns a decoded response *body*, so response headers, HTTP status, and + attempt counts do not survive the call boundary and cannot be recovered here. + """ + + config_cls = CallableLLMConfig + + def __init__(self, config: CallableLLMConfig, completion_fn: Optional[CompletionFn] = None) -> None: + super().__init__(config) + self.config: CallableLLMConfig + if completion_fn is None: + raise ValueError( + "CallableLLMBackend requires a completion_fn; it cannot be built from config " + "alone. Pass it via create_llm(config, completion_fn=...)." + ) + self._completion_fn = completion_fn + self._api_key = resolve_api_key(config.api_key) or None + + def _completion_impl( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + **overrides: Any, + ) -> CompletionResult: + # Order is load-bearing: private-key stripping runs BEFORE normalization, + # which collapses text-only content-block lists and discards block metadata. + prepared = normalize_messages_for_api(strip_private_message_keys(messages)) + call_kwargs = self._build_call_kwargs(prepared, tools, overrides) + + # Client call: the ONLY statement wrapped for error translation. We catch + # Exception (never BaseException), so cancellation / KeyboardInterrupt + # still propagate. + try: + response = self._completion_fn(**call_kwargs) + except LLMCallError: + # Already one of ours — a backend-shaped callable, or a classifier that + # ran closer to the wire. Re-wrapping would erase the subclass the + # agent branches on. + raise + except Exception as e: + raise classify_call_exception(e) from e + + # Response parsing stays OUTSIDE the try (a bug here is ours, not the + # wire's); a malformed response is surfaced as an explicit LLMCallError. + return self._build_result(response, call_kwargs) + + def _build_call_kwargs( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + overrides: Dict[str, Any], + ) -> Dict[str, Any]: + """Assemble the keyword arguments for the completion callable.""" + overrides = dict(overrides) + + extra_body: Dict[str, Any] = {} + parallel_tool_calls = overrides.pop("parallel_tool_calls", self.config.parallel_tool_calls) + if parallel_tool_calls is not None: + extra_body["parallel_tool_calls"] = parallel_tool_calls + reasoning_effort = overrides.pop("reasoning_effort", self.config.reasoning_effort) + if reasoning_effort: + extra_body["reasoning_effort"] = reasoning_effort + override_extra_body = overrides.pop("extra_body", None) + if isinstance(override_extra_body, dict): + extra_body.update(override_extra_body) + + # The FULL contract, always, in the callable's own spelling. Nothing is + # conditionally omitted (see the module docstring). + call_kwargs: Dict[str, Any] = { + "invoke_url": self.config.base_url, + "messages": messages, + "model": self.config.model, + "api_key": self._api_key, + "tools": tools, + "tool_choice": "none" if not tools else self.config.tool_choice, + "timeout_s": self.config.timeout_s, + # None means "unset" and is forwarded as such; the callable decides + # whether that means "omit the field" or a concrete default. + "temperature": overrides.pop("temperature", self.config.temperature), + "max_tokens": overrides.pop("max_completion_tokens", self.config.max_completion_tokens), + "extra_body": extra_body, + "max_retries": self.config.max_retries, + "max_429_retries": self.config.max_429_retries, + } + # Remaining overrides are passed through as keyword arguments so nothing is + # silently ignored. + call_kwargs.update(overrides) + return call_kwargs + + def _build_result(self, response: Any, call_kwargs: Dict[str, Any]) -> CompletionResult: + if not isinstance(response, dict): + raise LLMCallError(f"Callable returned {type(response).__name__}, expected an OpenAI chat.completion dict.") + choices = response.get("choices") + if not isinstance(choices, list) or not choices: + raise LLMCallError("Callable response is missing a non-empty 'choices' list.") + choice = choices[0] + if not isinstance(choice, dict): + raise LLMCallError(f"Callable response choice must be a dict, got {type(choice).__name__}.") + raw_message = choice.get("message") + if not isinstance(raw_message, dict): + raise LLMCallError("Callable response choice is missing a 'message' object.") + + # Spec says `content` is a string, but some OpenAI-compatible endpoints + # return a block list; coerce so one exotic response cannot fail a run. + message: Dict[str, Any] = {"role": "assistant", "content": extract_text_content(raw_message.get("content"))} + tool_calls = raw_message.get("tool_calls") + if tool_calls: + # Already OpenAI-shaped with ``arguments`` as a JSON string — pass through verbatim. + message["tool_calls"] = tool_calls + + raw_request: Optional[Dict[str, Any]] = None + raw_response: Optional[Dict[str, Any]] = None + if self.config.capture_raw_io: + # Best-effort: capturing artifacts must never fail an otherwise-good call. + try: + raw_request = _redacted_request(call_kwargs) + raw_response = deepcopy(response) + except Exception: + logger.warning("Failed to capture raw LLM IO; continuing without it.", exc_info=True) + + # `.strip()` is not cosmetic: the agent loop treats any finish reason + # outside ("stop", "tool_calls") as terminal, so a padded " tool_calls " + # would end an otherwise healthy run. + raw_finish = choice.get("finish_reason") + finish_reason = raw_finish.strip() if isinstance(raw_finish, str) and raw_finish.strip() else "stop" + + return CompletionResult( + message=message, + finish_reason=finish_reason, + reasoning=extract_reasoning_from_message(raw_message), + usage=coerce_usage_to_dict(response.get("usage")), + raw_request=raw_request, + raw_response=raw_response, + ) diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/error_classification.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/error_classification.py new file mode 100644 index 0000000000..644a0d8bba --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/error_classification.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Map completion-callable failures onto the error hierarchy the agent branches on. + +Why this module exists +---------------------- +:class:`~nemo_agent.llm.callable_backend.CallableLLMBackend` receives an injected +completion callable and therefore cannot know which client library raised. The +agent, however, *does* branch on the distinction: ``loop.py`` records +:class:`ContextLimitError` / :class:`ContentPolicyError` as expected outcomes, +and ``selection_agent.py`` retries a shrunken candidate list only when the run +failed with a context limit. Flattening every failure to a bare +:class:`LLMCallError` silently disables both. + +Deliberately HTTP-library-free +------------------------------ +This package must import neither ``requests`` nor ``httpx``. An exception that +carries a response is therefore **duck-typed** on ``.status_code`` / ``.text`` / +``.headers`` / ``.url`` — a surface both ``requests.Response`` and +``httpx.Response`` satisfy. + +Sole entry point is :func:`classify_call_exception`, which is **total**: it +always returns an :class:`LLMCallError` subclass and never raises. + +Caveat on message content +------------------------- +Some clients embed the raw endpoint URL and the full response body in the +exception message. :func:`~nemo_agent.llm.helpers.excerpt` and +:func:`~nemo_agent.llm.helpers.redact_url` bound and scrub the message this +module *builds*, but the original text survives on ``__cause__``, which +``loop.py`` logs with ``exc_info=True``. That is tidiness, not containment. +""" + +from __future__ import annotations + +import logging +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime +from typing import Any, Mapping, Optional, Tuple + +from .errors import ContentPolicyError, ContextLimitError, LLMCallError, RateLimitError +from .helpers import excerpt, redact_url, redact_urls_in_text + +logger = logging.getLogger(__name__) + +#: ``Retry-After`` values above this are treated as unusable; the base class's +#: configured sleep is a better answer than a 15-minute stall. +_RETRY_AFTER_MAX_S = 300.0 + +_CONTEXT_LIMIT_CODES = frozenset({"context_length_exceeded", "string_above_max_length"}) +_CONTENT_POLICY_CODES = frozenset({"content_filter", "content_policy_violation"}) + +_CONTENT_POLICY_MARKERS = ( + "content filter", + "content_filter", + "content policy", + "content_policy", + "content management policy", + "guardrail", + "responsible ai", +) + + +# ---------------------------------------------------------------------- +# Pure helpers. None of these log — their callers do. +# ---------------------------------------------------------------------- + + +def _header(headers: Optional[Mapping[str, Any]], name: str) -> Optional[Any]: + """Case-insensitive header lookup that works on any mapping.""" + if not headers: + return None + for key, value in headers.items(): + if str(key).lower() == name: + return value + return None + + +def _parse_retry_after(headers: Optional[Mapping[str, Any]]) -> Optional[float]: + """Seconds from a ``Retry-After`` header, or None when unusable. + + Accepts delta-seconds and RFC-9110 HTTP-date. Returns None for ``<= 0`` and + for implausibly large values: the base class accepts any finite ``>= 0`` and + caps at its own ceiling, so ``Retry-After: 0`` would burn the whole + rate-limit budget instantly and a skewed date would stall every retry. + """ + raw = _header(headers, "retry-after") + if raw is None: + return None + text = str(raw).strip() + if not text: + return None + + try: + seconds = float(text) + except ValueError: + try: + when = parsedate_to_datetime(text) + except (TypeError, ValueError): + return None + if when is None: + return None + if when.tzinfo is None: + when = when.replace(tzinfo=timezone.utc) + seconds = (when - datetime.now(timezone.utc)).total_seconds() + + if seconds <= 0 or seconds > _RETRY_AFTER_MAX_S: + return None + return seconds + + +def _error_fields(body_json: Any) -> Tuple[str, str]: + """``(code, type)`` from an OpenAI-style error envelope, lowercased; ``("", "")`` if absent.""" + if not isinstance(body_json, dict): + return "", "" + error = body_json.get("error") + if not isinstance(error, dict): + return "", "" + code = str(error.get("code") or "").strip().lower() + type_ = str(error.get("type") or "").strip().lower() + return code, type_ + + +def _looks_like_context_limit(lowered: str) -> bool: + """Prose fallback shared by the response and no-response paths. + + Provider wording drifts, so these markers are a heuristic layered *under* the + structured ``error.code`` check — never a replacement for it. + """ + return ( + "contextwindowexceedederror" in lowered + or ("context" in lowered and "window" in lowered) + or ("context" in lowered and "reduce" in lowered) + or "maximum context length" in lowered + or "is longer than the maximum model length" in lowered + or "please reduce the length" in lowered + # Prompt so long the completion budget went negative — a context-overflow + # symptom rather than a literal context-window message. + or ("max_tokens must be at least 1" in lowered and "got -" in lowered) + ) + + +def _looks_like_content_policy(lowered: str) -> bool: + return any(marker in lowered for marker in _CONTENT_POLICY_MARKERS) + + +# ---------------------------------------------------------------------- +# Response duck-typing. Every accessor is total. +# ---------------------------------------------------------------------- + + +def _status_code(response: Any) -> Optional[int]: + """Status code from a duck-typed response, or None when there isn't one. + + Compares ``is None`` and NEVER truthiness: ``requests.Response.__bool__`` + returns ``status_code < 400``, so every 4xx/5xx response object — exactly the + ones this module exists to classify — is falsy. + """ + if response is None: + return None + try: + return int(getattr(response, "status_code", None)) + except (TypeError, ValueError): + return None + + +def _response_text(response: Any) -> str: + # getattr INSIDE the try: `.text` is a property that decodes `.content` and can + # itself raise (chunked-encoding errors, "content already consumed", ...). + try: + value = getattr(response, "text", None) + return "" if value is None else str(value) + except Exception: + return "" + + +def _response_json(response: Any) -> Any: + """Best-effort decode; None when the body is not JSON. + + Broad on purpose: ``.json()`` raises a ``ValueError`` subclass, but the + ``.content`` read underneath it can raise non-``ValueError``s. + """ + getter = getattr(response, "json", None) + if not callable(getter): + return None + try: + return getter() + except Exception: + return None + + +def _response_headers(response: Any) -> Mapping[str, Any]: + try: + headers = getattr(response, "headers", None) + return headers if isinstance(headers, Mapping) else {} + except Exception: + return {} + + +def _response_url(response: Any) -> str: + try: + return str(getattr(response, "url", "") or "") + except Exception: + return "" + + +def _safe_str(exc: BaseException) -> str: + """``str(exc)`` that cannot itself raise. + + Not paranoia: a custom exception with a broken ``__str__`` would otherwise + defeat this module's totality guarantee at the one moment it matters most — + while something is already going wrong. + """ + try: + return str(exc) + except Exception: + return f"" + + +# ---------------------------------------------------------------------- +# Classification. +# ---------------------------------------------------------------------- + + +def classify_http_error( + status_code: int, + body_text: str, + body_json: Any, + headers: Optional[Mapping[str, Any]], + url: str, +) -> LLMCallError: + """Map an HTTP error response to the exception class the agent branches on. + + Structured ``error.code`` / ``error.type`` are checked first because they are + stable; prose markers are a fallback because provider wording drifts. + ``body_json`` may be None (an error body is not always JSON) — every prose + check runs against ``body_text``, which is a superset of it. + + Pure: an unclassified error is returned as a plain ``LLMCallError`` and the + caller decides whether to log it. + """ + # ``body_text`` is scrubbed for the message only; the marker matching below + # runs against the raw text so redaction can never cost a classification. + message = f"HTTP {status_code} from {redact_url(url)}: {excerpt(redact_urls_in_text(body_text))}" + + if status_code == 429: + return RateLimitError(message, retry_after=_parse_retry_after(headers)) + + code, type_ = _error_fields(body_json) + if code in _CONTEXT_LIMIT_CODES or type_ in _CONTEXT_LIMIT_CODES: + return ContextLimitError(message) + if code in _CONTENT_POLICY_CODES or type_ in _CONTENT_POLICY_CODES: + return ContentPolicyError(message) + + lowered = str(body_text or "").lower() + if _looks_like_context_limit(lowered): + return ContextLimitError(message) + if _looks_like_content_policy(lowered): + return ContentPolicyError(message) + + return LLMCallError(message) + + +def classify_prose_error(message: str, detail: str) -> LLMCallError: + """Classify from an exception's text alone, for failures carrying no response. + + ``message`` is what the returned exception says; ``detail`` is the raw text + the markers are matched against. They are separate arguments so a caller can + scrub/bound the former without narrowing the latter. + """ + lowered = str(detail or "").lower() + if _looks_like_context_limit(lowered): + return ContextLimitError(message) + if _looks_like_content_policy(lowered): + return ContentPolicyError(message) + return LLMCallError(message) + + +def classify_call_exception(exc: Exception, *, context: str = "completion callable failed") -> LLMCallError: + """Total exception -> :class:`LLMCallError` subclass. NEVER raises. + + Classification quality degrades gracefully by design: a failure carrying a + response object gets full HTTP classification; anything else falls back to + prose matching and, failing that, to a plain ``LLMCallError``. In-process + backends land in the latter buckets, which is acceptable — the agent treats + an unclassified failure as a terminal call error, which it is. + """ + try: + return _classify(exc, context) + except Exception: + logger.warning("LLM error classification failed; degrading to LLMCallError.", exc_info=True) + # ``_safe_str`` again, not ``str(exc)``: a raising ``__str__`` is one of the + # few things that can land us here, so re-reading it would raise from the + # handler itself. Scrubbed like every other path — a degraded message is + # still a message we surface. + return LLMCallError(f"{context}: {redact_urls_in_text(_safe_str(exc))}") + + +def _classify(exc: Exception, context: str) -> LLMCallError: + if isinstance(exc, LLMCallError): + # Already one of ours (a backend-shaped callable). Re-wrapping would erase + # the subclass the agent branches on. + return exc + + detail = _safe_str(exc) + response = getattr(exc, "response", None) + status = _status_code(response) + if status is None: + # No response to inspect: timeouts, transport errors, retries-exhausted + # wrappers, and every in-process backend failure. Prose is all we have. + # + # The message is scrubbed but ``detail`` is passed through raw: the + # message is what we surface, while ``detail`` is only ever matched + # against markers, and narrowing it could lose a classification. + return classify_prose_error(f"{context}: {excerpt(redact_urls_in_text(detail))}", detail) + + return classify_http_error( + status, + _response_text(response), + _response_json(response), + _response_headers(response), + _response_url(response), + ) diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/errors.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/errors.py new file mode 100644 index 0000000000..fb54ed0fec --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/errors.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exception hierarchy for LLM backends. + +Contract +-------- +``BaseLLMBackend.completion`` / ``acompletion`` raise only :class:`LLMCallError` +subclasses **for call failures**. Any other exception escaping them is, by +definition, a bug (in the backend implementation or in the caller's inputs) and +must propagate untranslated so it surfaces as one. + +Rules for backend implementations: + +- Translate provider exceptions at the raise site with ``raise (...) + from e`` so the original exception stays chained on ``__cause__``. +- Wrap ONLY the client call in the translating ``try/except`` — never your own + request-prep or response-parsing code. Errors from the wire get translated; + errors from your code propagate raw. +- Never catch or wrap ``BaseException`` (``asyncio.CancelledError``, + ``KeyboardInterrupt``); wrapping ``CancelledError`` breaks task cancellation. +- Add a new subclass here only when a caller *branches* on it. Nicer error + messages alone do not justify a class — format the chained ``__cause__``. +""" + + +class LLMCallError(Exception): + """An LLM API call failed. + + Instantiable catch-all for call failures that don't fit a more specific + subclass (auth errors, timeouts, 5xx, unrecognized bad requests, ...). + The provider exception, when there is one, is chained on ``__cause__``. + """ + + +class ContextLimitError(LLMCallError): + """The request exceeded the model's context window. + + Also raised for indirect symptoms of an oversized prompt, e.g. a provider + rejecting the request because the remaining completion budget is negative + ("max_tokens must be at least 1, got -N"). + """ + + +class ContentPolicyError(LLMCallError): + """The provider refused the request on content-policy grounds.""" + + +class RateLimitError(LLMCallError): + """The provider rate-limited the request. + + ``BaseLLMBackend``'s public templates catch this and retry with a pause + (see ``rate_limit_max_retries`` / ``rate_limit_retry_sleep_s`` on the + config). Backends may set ``retry_after`` (seconds) when the provider + communicates one; the retry loop honors it over the configured sleep, + capped by a hardcoded ceiling (``base_backend._RETRY_AFTER_CAP_S``). + """ + + def __init__(self, message: str = "", retry_after: "float | None" = None) -> None: + super().__init__(message) + self.retry_after = retry_after diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/helpers.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/helpers.py new file mode 100644 index 0000000000..faff4ded57 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/helpers.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable, backend-agnostic helpers for LLM backend implementations. + +Nothing here is required by the ``BaseLLMBackend`` contract; these exist so +custom backends don't re-invent the fiddly parts (most OpenAI-compatible +backends can use them verbatim). +""" + +from __future__ import annotations + +import os +import re +from typing import Any, Dict, List, Mapping, Optional +from urllib.parse import urlsplit, urlunsplit + +from .errors import LLMCallError + +_THINK_BLOCK_RE = re.compile(r"(.*?)", flags=re.DOTALL) + +#: Cap on provider text echoed into exception messages and logs, so a multi-megabyte +#: error body never lands in an agent trajectory verbatim. +BODY_EXCERPT_CHARS = 2000 + +#: Indirection prefix for reading an API key from the environment at construction. +_ENV_PREFIX = "os.environ/" + + +def excerpt(text: Any, *, limit: int = BODY_EXCERPT_CHARS) -> str: + """Bound provider text before it reaches an exception message or a log line.""" + body = str(text or "") + return body if len(body) <= limit else body[:limit] + "..." + + +def _redact_one_url(url: str) -> str: + try: + parts = urlsplit(str(url)) + netloc = parts.hostname or "" + if parts.port: + netloc = f"{netloc}:{parts.port}" + return urlunsplit((parts.scheme, netloc, parts.path, "", "")) + except Exception: + return "" + + +def redact_url(url: str) -> str: + """Drop userinfo, query, and fragment so a URL is safe to log or persist. + + Handles the comma-separated multi-endpoint form some callers accept: a bare + ``urlsplit`` would leave every URL after the first sitting in ``.path`` with + its userinfo intact, so each segment is redacted independently. + """ + return ",".join(_redact_one_url(part.strip()) for part in str(url or "").split(",")) + + +#: A full URL embedded in free text. Excludes whitespace and the delimiters that +#: commonly close a URL in prose, so a trailing quote or bracket is not swallowed. +_URL_IN_TEXT_RE = re.compile(r"https?://[^\s\"'<>)\]]+", re.IGNORECASE) +#: A bare ``?key=value`` query fragment, for messages that carry only a path. +#: Requires an ``=`` so ordinary prose ending in a question mark is left alone. +_QUERY_IN_TEXT_RE = re.compile(r"\?[^\s\"'<>)\]]*=[^\s\"'<>)\]]*") +#: Punctuation that ends a sentence rather than the URL it follows. +_TRAILING_PUNCTUATION = ".,;:!" + + +def _redact_url_in_text(match: "re.Match[str]") -> str: + raw = match.group(0) + trailing = "" + while raw and raw[-1] in _TRAILING_PUNCTUATION: + raw, trailing = raw[:-1], raw[-1] + trailing + return redact_url(raw) + trailing + + +def redact_urls_in_text(text: str) -> str: + """Scrub URLs and query strings out of an arbitrary message.""" + scrubbed = _URL_IN_TEXT_RE.sub(_redact_url_in_text, str(text or "")) + return _QUERY_IN_TEXT_RE.sub("?", scrubbed) + + +def resolve_api_key(api_key: Optional[str]) -> str: + """Resolve the configured key, following an ``os.environ/VAR`` indirection. + + Backends call this once at construction so a typo'd variable name fails at + build time rather than on the first request. + """ + raw = (api_key or "").strip() + if not raw.startswith(_ENV_PREFIX): + return raw + var = raw[len(_ENV_PREFIX) :].strip() + try: + return os.environ[var].strip() + except KeyError: + raise ValueError(f"Environment variable '{var}' is not set. Set it with: export {var}=") from None + + +def strip_private_message_keys(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Return copies of ``messages`` without ``__``-prefixed top-level keys. + + Agent code stashes non-API metadata on history messages under dunder-style + keys (e.g. ``"__reasoning__"``). Backends should strip these before the + wire rather than rely on provider tolerance. Never mutates the input. + """ + out: List[Dict[str, Any]] = [] + for msg in messages: + if isinstance(msg, dict): + out.append({k: v for k, v in msg.items() if not (isinstance(k, str) and k.startswith("__"))}) + else: + out.append(msg) + return out + + +def normalize_messages_for_api(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Normalize message content from list-of-content-blocks to plain strings. + + Some OpenAI-compatible endpoints only accept string content for certain + roles. This converts text-only ``content`` lists (e.g. + ``[{"type": "text", "text": "..."}]``) into a plain string. Messages with + non-text blocks (e.g. ``image_url``) are left as-is. Block-level metadata + (e.g. ``cache_control``) is discarded when a text-only list is collapsed, + so run this BEFORE adding block-level markers. Never mutates the input. + """ + normalized: List[Dict[str, Any]] = [] + for msg in messages: + msg = dict(msg) + content = msg.get("content") + if isinstance(content, list): + text_parts: List[str] = [] + all_text = True + for item in content: + if isinstance(item, dict) and item.get("type") == "text": + text_parts.append(str(item.get("text", ""))) + else: + all_text = False + break + if all_text: + if len(text_parts) == 0: + msg["content"] = None + elif len(text_parts) == 1: + msg["content"] = text_parts[0] + else: + msg["content"] = "\n".join(text_parts) + normalized.append(msg) + return normalized + + +def extract_text_content(content: Any) -> Optional[str]: + """Coerce assistant message content to ``str | None`` per the envelope contract. + + Providers occasionally return content as a list of blocks; extract and join + the text blocks, skipping non-text blocks (thinking/tool blocks — reasoning + is surfaced separately on ``CompletionResult.reasoning``). Any other shape + is a malformed provider response and raises :class:`LLMCallError`. + """ + if content is None or isinstance(content, str): + return content + if isinstance(content, list): + parts: List[str] = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + if block.get("type", "text") == "text" and "text" in block: + parts.append(str(block.get("text", ""))) + else: + text = getattr(block, "text", None) + if text: + parts.append(str(text)) + return "\n".join(parts) if parts else None + raise LLMCallError(f"Unexpected assistant message content type from provider: {type(content).__name__}") + + +def extract_reasoning_from_message(message: object) -> Optional[str]: + """Best-effort extraction of the per-turn reasoning trace. + + Accepts either a provider message *object* (attribute access) or a plain + message *dict* (key access). Handles three exposure shapes observed across + providers: + + 1. ``reasoning_content`` (gpt-oss via NIM, GLM, DeepSeek-R1) + 2. ``thinking_blocks`` (Anthropic extended thinking) + 3. Inline ``...`` blocks in ``content`` + (Tongyi-DeepResearch and similar ReAct-style models) + + Returns ``None`` when no reasoning channel is populated. + """ + if message is None: + return None + + reasoning_content = _coerce_str(_get_field(message, "reasoning_content")) + if reasoning_content: + return reasoning_content + + thinking_blocks = _get_field(message, "thinking_blocks") + if thinking_blocks: + parts: List[str] = [] + for block in thinking_blocks: + if isinstance(block, dict): + text = block.get("thinking") or block.get("text") or "" + else: + text = getattr(block, "thinking", None) or getattr(block, "text", None) or "" + if text: + parts.append(str(text).strip()) + if parts: + return "\n".join(parts) + + content = _get_field(message, "content") + if isinstance(content, str) and "" in content and "" in content: + # Capture the last ... block; if multiple are present + # the final one reflects the model's most recent reasoning state. + matches = _THINK_BLOCK_RE.findall(content) + if matches: + tail = matches[-1].strip() + if tail: + return tail + + return None + + +def _get_field(message: object, name: str) -> Any: + if isinstance(message, Mapping): + return message.get(name) + return getattr(message, name, None) + + +def _coerce_str(value: object) -> Optional[str]: + if value is None: + return None + if isinstance(value, str): + v = value.strip() + return v if v else None + try: + v = str(value).strip() + return v if v else None + except Exception: + return None diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/litellm_backend.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/litellm_backend.py new file mode 100644 index 0000000000..24f8a296c0 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/litellm_backend.py @@ -0,0 +1,687 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LiteLLM-backed implementation of :class:`BaseLLMBackend`. + +litellm is an OPTIONAL dependency: this module must never import it at top +level. The import happens in ``LiteLLMBackend.__init__`` so missing litellm +fails fast at construction with a clear error, while ``import nemo_agent.llm`` +stays safe for users of other backends. +""" + +from __future__ import annotations + +import logging +import os +from copy import deepcopy +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional + +from pydantic import field_validator + +from .base_backend import BaseLLMBackend, BaseLLMConfig +from .errors import ContentPolicyError, ContextLimitError, LLMCallError, RateLimitError +from .helpers import ( + extract_reasoning_from_message, + extract_text_content, + normalize_messages_for_api, + strip_private_message_keys, +) +from .result import CompletionResult +from .usage import coerce_usage_to_dict + +if TYPE_CHECKING: # pragma: no cover - typing only, litellm may be absent + from litellm.types.utils import ModelResponse + +logger = logging.getLogger(__name__) + +_REDACTED = "***REDACTED***" +_SENSITIVE_REQUEST_KEYS = frozenset({"api_key"}) +_HEADER_CONTAINER_KEYS = frozenset({"headers", "extra_headers", "default_headers"}) +_SENSITIVE_HEADER_MARKERS = ("authorization", "api-key", "api_key", "token", "secret") + +_TAGGABLE_BLOCK_TYPES = frozenset({"text", "tool_result", "tool_use", "image", "document"}) + + +def _patch_litellm_nvidia_nim_cache_passthrough() -> bool: + """Stop LiteLLM's ``nvidia_nim`` provider from stripping ``cache_control``. + + LiteLLM's ``OpenAIGPTConfig.(async_)transform_request`` calls + ``remove_cache_control_flag_from_messages_and_tools`` on every outgoing + request, which deletes any ``cache_control`` field from ``messages`` and + ``tools``. Because ``NvidiaNimConfig`` extends ``OpenAIGPTConfig`` without + overriding that method, every Anthropic-style cache-control marker is + silently scrubbed before the request leaves the process, even though NIM + itself relays the markers to Bedrock/Anthropic correctly. The Databricks + provider has the same problem and ships an explicit no-op override; we + mirror that fix by replacing the method with a no-op on the NIM config + class. + + Applied at backend construction (idempotent). Disable by exporting + ``RB_PATCH_LITELLM_CACHE_CONTROL=0`` (e.g. to A/B against an upstream + LiteLLM release that ships this fix natively). + + Returns ``True`` if the patch was applied this call, ``False`` otherwise + (disabled via env, already applied, or the import failed on an + incompatible LiteLLM version). + """ + if os.environ.get("RB_PATCH_LITELLM_CACHE_CONTROL", "1").lower() in ("0", "false", "no", "off"): + logger.info("RB_PATCH_LITELLM_CACHE_CONTROL=0 — skipping cache_control passthrough patch.") + return False + + try: + from litellm.llms.nvidia_nim.chat.transformation import NvidiaNimConfig + except Exception as e: # pragma: no cover - depends on installed litellm + logger.warning( + "Could not import NvidiaNimConfig to patch cache_control passthrough: %s. " + "Anthropic prompt caching may be silently disabled for nvidia_nim/* models.", + e, + ) + return False + + if getattr(NvidiaNimConfig, "_agent_core_cache_control_patched", False): + return False + + def _noop_remove_cache_control_flag(self, model, messages, tools=None): + # Same signature as OpenAIGPTConfig.remove_cache_control_flag_from_messages_and_tools + # but a no-op so cache_control survives all the way to NIM. + return messages, tools + + NvidiaNimConfig.remove_cache_control_flag_from_messages_and_tools = ( # type: ignore[assignment] + _noop_remove_cache_control_flag + ) + NvidiaNimConfig._agent_core_cache_control_patched = True # type: ignore[attr-defined] + logger.info( + "Patched litellm NvidiaNimConfig.remove_cache_control_flag_from_messages_and_tools " + "to no-op so Anthropic cache_control markers reach NIM (set " + "RB_PATCH_LITELLM_CACHE_CONTROL=0 to disable)." + ) + return True + + +def _block_has_cache_control(block: Any) -> bool: + """True if ``block`` is a dict carrying an Anthropic ``cache_control`` marker.""" + return isinstance(block, dict) and isinstance(block.get("cache_control"), dict) + + +def _is_taggable_block(block: Any) -> bool: + """True if ``block`` is a content block that Anthropic accepts a marker on.""" + return isinstance(block, dict) and block.get("type") in _TAGGABLE_BLOCK_TYPES + + +def _find_last_markable_msg_idx(messages: List[Dict[str, Any]]) -> Optional[int]: + """Return the index of the most-recent message we can safely add a marker to. + + Returns ``None`` when either no message is markable OR the most-recent + markable message already carries a deliberate ``cache_control`` marker + (see :meth:`LiteLLMBackend._with_last_message_cache_control_marker` for the + dedup rationale). + """ + for i in range(len(messages) - 1, -1, -1): + msg = messages[i] + if not isinstance(msg, dict): + continue + content = msg.get("content") + if isinstance(content, str): + if content.strip() == "": + # Empty assistant tool-call turn: skip, look further back. + continue + return i # plain-string content cannot carry a pre-existing marker + if isinstance(content, list) and any(_is_taggable_block(p) for p in content): + if any(_block_has_cache_control(p) for p in content): + return None # caller already marked this message; preserve their intent + return i + return None + + +def _apply_cache_control_to_last_block(message: Dict[str, Any], cache_control: Dict[str, Any]) -> Dict[str, Any]: + """Return a copy of ``message`` with ``cache_control`` on its last taggable block. + + Handles both content shapes: + - ``content: str``: wrap into a single-block text list carrying the marker. + - ``content: [blocks]``: copy the list, attach the marker to the last + taggable block (text / tool_result / tool_use / image / document); leave + the rest untouched. If no taggable block exists the message is returned + verbatim. + """ + out = dict(message) + content = out.get("content") + cc = dict(cache_control) + if isinstance(content, str): + out["content"] = [{"type": "text", "text": content, "cache_control": cc}] + return out + if not isinstance(content, list): + return out + last_taggable_idx = max((j for j, p in enumerate(content) if _is_taggable_block(p)), default=-1) + if last_taggable_idx < 0: + return out + new_content: List[Any] = [] + for j, part in enumerate(content): + if j == last_taggable_idx and isinstance(part, dict): + marked = dict(part) + marked["cache_control"] = cc + new_content.append(marked) + else: + new_content.append(part) + out["content"] = new_content + return out + + +def _redacted_request(request_kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Deep-copied, credential-redacted snapshot of the request for ``raw_request``.""" + out: Dict[str, Any] = {} + for k, v in request_kwargs.items(): + if k in _SENSITIVE_REQUEST_KEYS: + out[k] = _REDACTED + continue + if k in _HEADER_CONTAINER_KEYS and isinstance(v, dict): + out[k] = { + hk: (_REDACTED if any(m in str(hk).lower() for m in _SENSITIVE_HEADER_MARKERS) else deepcopy(hv)) + for hk, hv in v.items() + } + continue + out[k] = deepcopy(v) + return out + + +class LiteLLMConfig(BaseLLMConfig): + """Configuration for :class:`LiteLLMBackend`. + + Attributes + ---------- + thinking: + Provider extended-thinking knob (e.g. Anthropic-style + ``{"type": "enabled", "budget_tokens": ...}``); routed through litellm's + supported-params check and forwarded via ``extra_body`` when the provider + adapter does not recognize it (so ``drop_params=True`` cannot silently + strip it). ``reasoning_effort`` (inherited from ``BaseLLMConfig``) is + routed the same way. + api_version: + Provider API version (e.g. Azure-style endpoints). + num_retries: + litellm-internal retry count for transient errors. Distinct from the + base class's rate-limit pause loop. Overridable per call + (e.g. ``num_retries=0`` for a fail-fast preflight probe). + drop_params / allowed_openai_params: + litellm parameter-filtering controls, passed through as-is. + ``drop_params`` defaults to ``True`` so an OpenAI-compatible endpoint that + rejects a param litellm's adapter emits does not 400 the whole run; set it + ``False`` to surface those rejections instead. Note this is exactly why + provider-specific params (``thinking`` / ``reasoning_effort`` / + ``cache_control``) are routed via ``extra_body`` below — top-level params + litellm does not recognize are silently stripped when this is on. + cache_control: + Prompt-cache marker (e.g. ``{"type": "ephemeral"}``). For Anthropic + models this becomes per-message/tool ``cache_control`` block markers; + for other providers it is sent as a request param. + prompt_cache_key / prompt_cache_retention: + OpenAI-style prompt-cache controls; ignored (with a warning) for + Anthropic models, which use ``cache_control`` markers instead. + + ``backend`` is pinned by type: a ``LiteLLMConfig`` cannot be constructed + with (and therefore never routed to) a different backend. + + ``base_url`` (inherited) is normalized on construction: a trailing + ``/chat/completions`` is stripped. + """ + + backend: Literal["litellm"] = "litellm" + thinking: Optional[Dict[str, Any]] = None + api_version: Optional[str] = None + num_retries: Optional[int] = 4 + drop_params: bool = True + allowed_openai_params: Optional[List[str]] = None + cache_control: Optional[Dict[str, Any]] = None + prompt_cache_key: Optional[str] = None + prompt_cache_retention: Optional[str] = None + + @field_validator("base_url", mode="after") + @classmethod + def _strip_chat_completions_suffix(cls, base_url: Optional[str]) -> Optional[str]: + """Normalize ``base_url`` to the endpoint base litellm expects. + + litellm appends ``/chat/completions`` itself, so a caller-supplied + endpoint that already ends in ``/chat/completions`` must + have that suffix removed. + """ + if base_url is None: + return base_url + cleaned = base_url.rstrip("/").removesuffix("/chat/completions") + return cleaned or base_url + + +class LiteLLMBackend(BaseLLMBackend): + """LLM backend that makes chat-completion calls through litellm. + + Handles the provider-specific details: param routing via ``extra_body`` + for params the provider adapter doesn't recognize, Anthropic prompt-cache + block markers (system / last message / last tool), the NIM cache-control + passthrough patch, and ``os.environ/`` api-key resolution. + """ + + config_cls = LiteLLMConfig + + def __init__(self, config: LiteLLMConfig) -> None: + super().__init__(config) + self.config: LiteLLMConfig + + try: + import litellm + except ImportError as e: + raise RuntimeError( + "litellm is not installed but LiteLLMBackend was requested. " + "Install it (`pip install litellm`) or use a different backend." + ) from e + self._litellm = litellm + + # Reduce noisy provider/help banners on handled API errors. + if hasattr(litellm, "suppress_debug_info"): + litellm.suppress_debug_info = True + _patch_litellm_nvidia_nim_cache_passthrough() + + self._cache_control_message_marker: Optional[Dict[str, Any]] = None + + self.completion_kwargs: Dict[str, Any] = dict( + model=config.model, + tool_choice=config.tool_choice, + base_url=config.base_url, + api_version=config.api_version, + num_retries=config.num_retries, + max_completion_tokens=config.max_completion_tokens, + ) + if config.drop_params: + self.completion_kwargs["drop_params"] = config.drop_params + if config.allowed_openai_params: + self.completion_kwargs["allowed_openai_params"] = config.allowed_openai_params + if config.temperature is not None: + self.completion_kwargs["temperature"] = config.temperature + if config.parallel_tool_calls is not None: + self.completion_kwargs["parallel_tool_calls"] = config.parallel_tool_calls + + supported_params = None + get_supported_openai_params = getattr(litellm, "get_supported_openai_params", None) + if callable(get_supported_openai_params): + try: + supported_params = get_supported_openai_params(model=config.model) + except Exception: + supported_params = None + + def _param_supported(name: str) -> bool: + return isinstance(supported_params, list) and name in supported_params + + def _set_param_or_extra_body(name: str, value: Any) -> None: + if value is None: + return + if _param_supported(name): + self.completion_kwargs[name] = value + return + # Some OpenAI-compatible endpoints only accept provider-specific + # parameters via `extra_body`; keep top-level kwargs clean. + extra_body = self.completion_kwargs.get("extra_body") + if not isinstance(extra_body, dict): + extra_body = {} + extra_body[name] = value + self.completion_kwargs["extra_body"] = extra_body + logger.info("Forwarding `%s` via extra_body for model=%s.", name, config.model) + + model_lower = str(config.model).lower() + use_anthropic_message_cache_markers = "anthropic" in model_lower + + # `reasoning_effort` must go through the same supported/extra_body + # router as `thinking` / `cache_control`: several provider adapters + # (e.g. nvidia_nim) ship static supported-params lists that predate + # reasoning_effort, and with drop_params=True litellm silently strips + # unsupported top-level params before the wire. extra_body keys pass + # through opaquely. + if isinstance(config.reasoning_effort, str) and config.reasoning_effort.strip(): + _set_param_or_extra_body("reasoning_effort", config.reasoning_effort.strip()) + + if isinstance(config.thinking, dict) and len(config.thinking) != 0: + _set_param_or_extra_body("thinking", dict(config.thinking)) + if isinstance(config.cache_control, dict) and len(config.cache_control) != 0: + effective_cache_control = dict(config.cache_control) + if use_anthropic_message_cache_markers: + # Anthropic prompt caching is exclusively driven by `cache_control` + # markers on content blocks (system/messages/tools). There is no + # top-level `cache_control` API field; sending one via extra_body + # makes Bedrock-proxied routes reject the request. + self._cache_control_message_marker = dict(effective_cache_control) + logger.info( + "Applying `cache_control` as message-block marker for model=%s " + "(no request-level cache_control).", + config.model, + ) + else: + _set_param_or_extra_body("cache_control", effective_cache_control) + if isinstance(config.prompt_cache_key, str) and config.prompt_cache_key.strip(): + if use_anthropic_message_cache_markers: + logger.warning( + "Ignoring `prompt_cache_key` for anthropic model=%s; using cache_control markers instead.", + config.model, + ) + else: + _set_param_or_extra_body("prompt_cache_key", config.prompt_cache_key.strip()) + if isinstance(config.prompt_cache_retention, str) and config.prompt_cache_retention.strip(): + if use_anthropic_message_cache_markers: + logger.warning( + "Ignoring `prompt_cache_retention` for anthropic model=%s; use cache_control ttl instead.", + config.model, + ) + else: + _set_param_or_extra_body("prompt_cache_retention", config.prompt_cache_retention.strip()) + + if config.api_key is not None: + self.completion_kwargs["api_key"] = config.api_key + self._resolved_api_key: Optional[str] = None + if config.api_key is not None and config.api_key.strip().startswith("os.environ/"): + self._resolved_api_key = os.environ[config.api_key.strip().removeprefix("os.environ/")] + + # ------------------------------------------------------------------ + # Impls. + # ------------------------------------------------------------------ + + def _completion_impl( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + **overrides: Any, + ) -> CompletionResult: + request_kwargs = self._prepare_request(messages, tools=tools, **overrides) + try: + response = self._litellm.completion(**request_kwargs) + except Exception as e: + translated = self._translate_provider_error(e) + if translated is None: + raise + raise translated from e + return self._build_result(response, request_kwargs) + + async def _acompletion_impl( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + **overrides: Any, + ) -> CompletionResult: + request_kwargs = self._prepare_request(messages, tools=tools, **overrides) + try: + response = await self._litellm.acompletion(**request_kwargs) + except Exception as e: + translated = self._translate_provider_error(e) + if translated is None: + raise + raise translated from e + return self._build_result(response, request_kwargs) + + # ------------------------------------------------------------------ + # Request preparation (shared by both impls). + # ------------------------------------------------------------------ + + def _prepare_request( + self, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + **overrides: Any, + ) -> Dict[str, Any]: + """Build litellm request kwargs. Never mutates ``messages`` / ``tools``. + + Order matters: normalization collapses text-only content-block lists to + strings (discarding block-level metadata), so it must run BEFORE cache + markers are attached. + """ + request_messages = strip_private_message_keys(messages) + request_messages = normalize_messages_for_api(request_messages) + request_tools = tools + if isinstance(self._cache_control_message_marker, dict): + # Multi-turn caching strategy (Anthropic allows up to 4 breakpoints + # per request; longest prefix match wins): + # 1. Mark the SYSTEM message so caching works from the first turn + # even if a provider does not forward per-tool `cache_control`. + # 2. Mark the LAST message so the entire prior history is in the + # cache prefix; next turn the breakpoint moves forward and + # Anthropic's cascading window picks up the longest matching + # prefix from the prior breakpoint. + # 3. Mark the LAST tools entry so the stable `tools + system` + # prefix is always cached. + request_messages = self._with_cache_control_marker( + messages=request_messages, + cache_control=self._cache_control_message_marker, + ) + request_messages = self._with_last_message_cache_control_marker( + messages=request_messages, + cache_control=self._cache_control_message_marker, + ) + if isinstance(request_tools, list) and len(request_tools) > 0: + request_tools = self._with_tool_cache_control_marker( + tools=request_tools, + cache_control=self._cache_control_message_marker, + ) + + # Merge in steps so per-call overrides can intentionally win over + # config defaults without raising `TypeError` on duplicate keys. + # completion_kwargs holds shared mutables (extra_body, allowed lists): + # deep-copy per request so downstream mutation of one request can never + # corrupt the template for subsequent calls. + request_kwargs: Dict[str, Any] = {"messages": request_messages} + request_kwargs.update(deepcopy(self.completion_kwargs)) + request_kwargs.update(overrides) + if request_tools is not None: + request_kwargs["tools"] = request_tools + else: + # No tools this call: drop tool-only params so providers that reject + # them without a `tools` array don't 400. + request_kwargs.pop("tool_choice", None) + request_kwargs.pop("parallel_tool_calls", None) + if self._resolved_api_key is not None: + request_kwargs["api_key"] = self._resolved_api_key + return request_kwargs + + # ------------------------------------------------------------------ + # Error translation (client-call failures only; see errors module). + # ------------------------------------------------------------------ + + def _translate_provider_error(self, e: Exception) -> Optional[LLMCallError]: + """Map a client-call exception to a library error, or None to re-raise raw.""" + exceptions = self._litellm.exceptions + if isinstance(e, exceptions.RateLimitError): + return RateLimitError(str(e)) + # NOTE: ContextWindowExceededError and ContentPolicyViolationError are + # BadRequestError subclasses in litellm — check them first. + if isinstance(e, exceptions.ContextWindowExceededError): + return ContextLimitError(str(e)) + if isinstance(e, exceptions.ContentPolicyViolationError): + return ContentPolicyError(str(e)) + if isinstance(e, exceptions.BadRequestError): + err_str = str(e).lower() + if "contentpolicyviolationerror" in err_str: + return ContentPolicyError(str(e)) + if ( + "contextwindowexceedederror" in err_str + or ("context" in err_str and "window" in err_str) + or ("context" in err_str and "reduce" in err_str) + # Prompt so long the completion budget went negative — a + # context-overflow symptom. + or ("max_tokens must be at least 1" in err_str and "got -" in err_str) + ): + return ContextLimitError(str(e)) + return LLMCallError(str(e)) + # Any other litellm/openai exception type is still an API-call failure + # (auth, timeout, connection, 5xx, ...). Exceptions from other modules + # (KeyError, TypeError, ...) are bugs and propagate raw. + module = (type(e).__module__ or "").split(".")[0] + if module in ("litellm", "openai"): + return LLMCallError(str(e)) + return None + + # ------------------------------------------------------------------ + # Result assembly. + # ------------------------------------------------------------------ + + def _build_result(self, response: "ModelResponse", request_kwargs: Dict[str, Any]) -> CompletionResult: + choices = getattr(response, "choices", None) or [] + if len(choices) != 1: + raise LLMCallError(f"Expected exactly 1 choice in the API response, got {len(choices)}.") + choice = choices[0] + message_obj = choice.message + + message: Dict[str, Any] = {"role": "assistant", "content": extract_text_content(message_obj.content)} + tool_calls = getattr(message_obj, "tool_calls", None) + if tool_calls: + message["tool_calls"] = [tc.model_dump() if hasattr(tc, "model_dump") else dict(tc) for tc in tool_calls] + + raw_request: Optional[Dict[str, Any]] = None + raw_response: Optional[Dict[str, Any]] = None + if self.config.capture_raw_io: + try: + raw_request = _redacted_request(request_kwargs) + raw_response = response.model_dump() + except Exception: + logger.warning("Failed to capture raw LLM IO; continuing without it.", exc_info=True) + + return CompletionResult( + message=message, + finish_reason=choice.finish_reason or "unknown", + reasoning=extract_reasoning_from_message(message_obj), + usage=coerce_usage_to_dict(getattr(response, "usage", None)), + extra_response_info=self._build_extra_response_info(response), + raw_request=raw_request, + raw_response=raw_response, + ) + + def _build_extra_response_info(self, response: "ModelResponse") -> Dict[str, Any]: + """Everything the backend knows about the response except the message itself. + + Schema-unstable, log-only (see :class:`CompletionResult`). Best-effort: + degrades to ``{}`` rather than failing a successful call. + """ + info: Dict[str, Any] = {} + try: + info = response.model_dump(exclude={"choices"}) + except Exception: + info = {} + try: + additional_headers = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {} + info["ratelimit"] = { + "TPM": additional_headers.get("llm_provider-x-ratelimit-remaining-tokens"), + "RQ": additional_headers.get("llm_provider-x-ratelimit-remaining-requests"), + } + except Exception: + pass + return info + + # ------------------------------------------------------------------ + # Anthropic cache-control markers. + # ------------------------------------------------------------------ + + @staticmethod + def _with_cache_control_marker( + messages: List[Dict[str, Any]], cache_control: Dict[str, Any] + ) -> List[Dict[str, Any]]: + """Return a copy of messages with one explicit cache marker block.""" + if not isinstance(messages, list) or len(messages) == 0: + return messages + + out_messages: List[Dict[str, Any]] = [] + for msg in messages: + out_messages.append(dict(msg)) + + if isinstance(out_messages[0].get("role"), str) and out_messages[0].get("role") == "system": + target_idx = 0 + else: + target_idx = len(out_messages) - 1 + + target = dict(out_messages[target_idx]) + content = target.get("content") + + cc = dict(cache_control) + if isinstance(content, str): + target["content"] = [{"type": "text", "text": content, "cache_control": cc}] + elif isinstance(content, list): + replaced = False + new_parts = [] + for part in content: + if isinstance(part, dict): + p = dict(part) + if not replaced and p.get("type") == "text": + p["cache_control"] = cc + replaced = True + new_parts.append(p) + else: + new_parts.append(part) + if not replaced: + new_parts.append({"type": "text", "text": "", "cache_control": cc}) + target["content"] = new_parts + else: + return out_messages + + out_messages[target_idx] = target + return out_messages + + @staticmethod + def _with_last_message_cache_control_marker( + messages: List[Dict[str, Any]], cache_control: Dict[str, Any] + ) -> List[Dict[str, Any]]: + """Attach a ``cache_control`` marker on the LAST eligible message. + + This caches the entire conversation history up to and including the + most recent assistant/tool turn. Combined with the static system + + tools breakpoint, multi-turn agents reuse the maximum possible prefix + on each call. + + Content-shape behaviour + ----------------------- + - ``content: str`` → wrap into ``[{type: "text", text, cache_control}]``. + - ``content: [blocks]`` → set ``cache_control`` on the LAST taggable + block (text / tool_result / tool_use / image / document). + - ``content: None`` or empty string (typical of assistant tool-call + turns) → skip and fall back to the previous non-empty message. + + Pre-existing breakpoint dedup + ----------------------------- + If the first eligible message (scanning from the tail) already carries + a ``cache_control`` marker on at least one block, this function returns + the input unchanged: the caller has placed an explicit breakpoint at a + specific prefix boundary we must not shadow. Two breakpoints in one + message empirically break cache_read on Bedrock-via-NIM — only the + byte-identical explicit one ever matches on subsequent requests. + """ + if not isinstance(messages, list) or len(messages) == 0: + return messages + target_idx = _find_last_markable_msg_idx(messages) + if target_idx is None: + return messages + out_messages: List[Dict[str, Any]] = [dict(m) if isinstance(m, dict) else m for m in messages] + out_messages[target_idx] = _apply_cache_control_to_last_block( + message=out_messages[target_idx], cache_control=cache_control + ) + return out_messages + + @staticmethod + def _with_tool_cache_control_marker( + tools: List[Dict[str, Any]], cache_control: Dict[str, Any] + ) -> List[Dict[str, Any]]: + """Return a copy of ``tools`` with a cache marker on the last tool definition. + + For Anthropic models, ``cache_control`` may be attached either at the + top of a tool entry or on its nested ``function`` object; we attach it + to the nested ``function`` since LiteLLM's anthropic adapter also reads + ``tool["function"]["cache_control"]`` and Bedrock's tool-cache markers + are passed through similarly. + """ + if not isinstance(tools, list) or len(tools) == 0: + return tools + out_tools: List[Dict[str, Any]] = [] + for t in tools: + out_tools.append(dict(t) if isinstance(t, dict) else t) + last_idx = len(out_tools) - 1 + last = out_tools[last_idx] + if not isinstance(last, dict): + return out_tools + last = dict(last) + func = last.get("function") + if isinstance(func, dict): + func = dict(func) + func["cache_control"] = dict(cache_control) + last["function"] = func + else: + last["cache_control"] = dict(cache_control) + out_tools[last_idx] = last + return out_tools diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/result.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/result.py new file mode 100644 index 0000000000..c64dc2b473 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/result.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The result envelope every LLM backend returns from a completion call.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + + +@dataclass +class CompletionResult: + """Normalized result of one chat-completion call. + + The stable/unstable boundary: anything a consumer needs to *branch* on is a + typed field with a fixed meaning; :attr:`extra_response_info` is explicitly + schema-unstable and may only be logged. + + Attributes + ---------- + message: + OpenAI-chat-format assistant message dict, directly appendable to the + conversation history and valid to send back on the next call. + ``content`` is ``str | None``. ``tool_calls`` (when present) follow the + OpenAI shape ``{"id", "type", "function": {"name", "arguments"}}`` with + ``arguments`` kept as the raw JSON *string* — parsing (and the error + message when parsing fails) belongs to the caller. Exactly one + assistant message per call. + finish_reason: + Normalized finish reason; ``"stop"`` and ``"tool_calls"`` are the two + values agent loops branch on. Backends map anything unknown to a + non-empty descriptive string. + reasoning: + The model's reasoning trace for this turn, if the backend can extract + one (see ``helpers.extract_reasoning_from_message``), else ``None``. + usage: + Token usage as a JSON-ish nested dict with int leaves, or ``None`` when + the backend/server does not report usage. No fixed schema is required: + a minimal backend may return just ``{"prompt_tokens", "completion_tokens", + "total_tokens"}``; richer backends may nest arbitrarily (cache tiers, + reasoning tokens, ...). Consumers must tolerate missing keys. The base + class also merges this into its per-(query, stage) usage tracker. + extra_response_info: + Always populated (possibly ``{}``), built best-effort — an exotic + response must degrade this to ``{}``, never fail a successful call. + **Backend-defined and schema-unstable: consumers may log it, never + branch on it.** Anything decision-worthy must be promoted to a typed + field. (The litellm backend puts the full response dump minus + ``choices`` here, plus rate-limit headers under ``"ratelimit"``.) + raw_request: + The request as sent, as a plain JSON-serializable dict — populated only + when ``config.capture_raw_io`` is true, else ``None``. MUST be + credential-redacted by the backend (no ``api_key`` / authorization + material). + raw_response: + Full JSON-serializable dump of the provider response — populated only + when ``config.capture_raw_io`` is true, else ``None``. + """ + + message: Dict[str, Any] + finish_reason: str + reasoning: Optional[str] = None + usage: Optional[Dict[str, Any]] = None + extra_response_info: Dict[str, Any] = field(default_factory=dict) + raw_request: Optional[Dict[str, Any]] = None + raw_response: Optional[Dict[str, Any]] = None + + def __post_init__(self) -> None: + if not isinstance(self.message, dict): + raise TypeError( + f"CompletionResult.message must be an OpenAI-format assistant message dict, " + f"got {type(self.message).__name__}" + ) + if self.message.get("role") != "assistant": + raise ValueError(f"CompletionResult.message must have role='assistant', got {self.message.get('role')!r}") + if not isinstance(self.finish_reason, str) or not self.finish_reason: + raise ValueError(f"CompletionResult.finish_reason must be a non-empty string, got {self.finish_reason!r}") + if not isinstance(self.extra_response_info, dict): + raise TypeError( + f"CompletionResult.extra_response_info must be a dict, " + f"got {type(self.extra_response_info).__name__}" + ) diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py new file mode 100644 index 0000000000..da3a287bc6 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Token-usage attribution and aggregation for LLM backends. + +Attribution +----------- +Two independent ContextVars label every LLM call: + +- **query id** — bound once per query by the pipeline/driver + (:func:`bind_query_id`). +- **stage** — bound per call site by agent code (:func:`bind_stage`), e.g. + ``"main_agent"`` or ``"top5_agent"``. + +``BaseLLMBackend`` reads both after every successful call and deep-merges +``CompletionResult.usage`` into ``usage[query_id][stage]``. Unbound values fall +back to the public sentinels :data:`UNSET_QUERY` / :data:`UNSET_STAGE` (these +names appear in persisted traces — treat them as schema). + +Always prefer the context-manager binders over raw ``ContextVar.set``: they +reset to the *previous* value on exit, so nested scopes compose and nothing +leaks across queries within a task. + +Threading caveat +---------------- +``contextvars`` propagate into ``asyncio`` tasks and ``asyncio.to_thread`` +automatically, but **worker threads do NOT inherit the spawning thread's +context** (e.g. ``ThreadPoolExecutor``). Bind inside the thread/task that makes +the LLM call, or wrap submissions with ``contextvars.copy_context().run``. +Otherwise usage silently lands in the ``UNSET_*`` buckets. + +Aggregation rules (contract of :func:`deep_merge_usage`) +--------------------------------------------------------- +Usage dicts across providers are heterogeneous (nested ``*_tokens_details``, +cache tiers that appear mid-run, ``None`` placeholders). For each key ``k`` +across the running aggregate and an incoming sample: + +1. ``sample[k]`` is ``None`` and ``acc[k]`` missing or ``None`` -> store/keep ``None``. +2. ``sample[k]`` is ``None`` and ``acc[k]`` already typed (int or dict) -> no-op. +3. ``sample[k]`` non-``None`` and ``acc[k]`` missing or ``None`` -> deep-copy + ``sample[k]`` into ``acc[k]`` (the type-lock event). +4. dict + dict -> recurse. +5. int + int -> sum. +6. Type mismatch (locked-int meets dict or vice versa) -> :class:`TypeError` + carrying the dotted key path. A real schema regression should be loud. + +Booleans are int subclasses in Python; they are explicitly excluded from the +"int sum" branch so a ``True`` never becomes ``1`` of someone's +``prompt_tokens``. Other scalar types are out of contract; the helper keeps the +most recent value rather than crashing on a provider extension. +""" + +from __future__ import annotations + +import contextvars +from contextlib import contextmanager +from copy import deepcopy +from typing import Any, Dict, Iterator, Mapping, Optional + +UNSET_QUERY = "" +"""Usage bucket for calls made with no bound query id (e.g. preflight probes). + +Angle-bracketed so it cannot collide with a real query id; reserved — do not +bind it as an actual query id. +""" + +UNSET_STAGE = "" +"""Usage bucket for calls made with no bound stage. Reserved, like UNSET_QUERY.""" + +_MISSING = object() +"""Sentinel distinguishing 'key not yet seen' from 'key seen as None'.""" + +_QUERY_ID: "contextvars.ContextVar[Optional[str]]" = contextvars.ContextVar( + "agent_core_llm_query_id", + default=None, +) +_STAGE: "contextvars.ContextVar[Optional[str]]" = contextvars.ContextVar( + "agent_core_llm_stage", + default=None, +) + + +def get_query_id() -> Optional[str]: + """Return the ambient usage query id, or ``None`` when unbound.""" + return _QUERY_ID.get() + + +def get_stage() -> Optional[str]: + """Return the ambient usage stage, or ``None`` when unbound.""" + return _STAGE.get() + + +@contextmanager +def bind_query_id(query_id: Optional[str]) -> Iterator[None]: + """Scoped binding of the usage query id (typically once per query).""" + token = _QUERY_ID.set(query_id) + try: + yield + finally: + _QUERY_ID.reset(token) + + +@contextmanager +def bind_stage(stage: Optional[str]) -> Iterator[None]: + """Scoped binding of the usage stage (typically around each LLM call site).""" + token = _STAGE.set(stage) + try: + yield + finally: + _STAGE.reset(token) + + +def deep_merge_usage( + acc: Optional[Dict[str, Any]], + sample: Optional[Mapping[str, Any]], + *, + _path: str = "", +) -> Dict[str, Any]: + """Aggregate ``sample`` into ``acc`` in place; return ``acc``. + + See the module docstring for the full rule contract. + + Parameters + ---------- + acc: + Running aggregate. ``None`` is treated as an empty dict (a fresh one is + created and returned). + sample: + New usage snapshot. ``None`` and empty mappings are no-ops. + + Raises + ------ + TypeError + On a type-lock conflict; the message includes the dotted key path so + schema regressions are easy to localize. + """ + if acc is None: + acc = {} + if sample is None: + return acc + for k, v in sample.items(): + path = f"{_path}.{k}" if _path else k + cur = acc.get(k, _MISSING) + + if v is None: + # Rule 1: nothing seen yet -> remember the slot exists, value None. + # Rule 2: already typed -> no-op. + if cur is _MISSING: + acc[k] = None + continue + + if isinstance(v, dict): + if cur is _MISSING or cur is None: + # Rule 3: type-lock event for a dict slot. + acc[k] = deepcopy(v) + elif isinstance(cur, dict): + # Rule 4: dict + dict -> recurse. + deep_merge_usage(cur, v, _path=path) + else: + # Rule 6: int slot, sample is dict. + raise TypeError( + f"deep_merge_usage: type lock conflict at {path!r}: " f"acc is {type(cur).__name__}, sample is dict" + ) + continue + + # int (excluding bool, which is an int subclass). + if isinstance(v, int) and not isinstance(v, bool): + if cur is _MISSING or cur is None: + # Rule 3: type-lock event for an int slot. + acc[k] = int(v) + elif isinstance(cur, int) and not isinstance(cur, bool): + # Rule 5: int + int -> sum. + acc[k] = cur + int(v) + else: + # Rule 6: dict (or other) slot, sample is int. + raise TypeError( + f"deep_merge_usage: type lock conflict at {path!r}: " f"acc is {type(cur).__name__}, sample is int" + ) + continue + + # Out-of-contract scalar (float, str, bool, list, ...): keep the latest + # value as a safety net rather than crashing on a provider extension. + acc[k] = v + + return acc + + +def deep_merge_usage_breakdown( + acc: Optional[Dict[str, Any]], + sample: Optional[Mapping[str, Any]], +) -> Dict[str, Any]: + """Deep-merge a stage-keyed usage breakdown (``{stage: usage}``) into ``acc``.""" + if acc is None: + acc = {} + if sample is None: + return acc + for stage, usage in sample.items(): + if not isinstance(stage, str) or not stage: + continue + if not isinstance(usage, Mapping) or not usage: + continue + stage_acc = acc.get(stage) + if not isinstance(stage_acc, dict): + stage_acc = {} + acc[stage] = stage_acc + deep_merge_usage(stage_acc, usage) + return acc + + +def sum_usage_breakdown(usage_by_stage: Optional[Mapping[str, Any]]) -> Dict[str, Any]: + """Collapse a stage-keyed usage breakdown into one total usage dict.""" + total: Dict[str, Any] = {} + if usage_by_stage is None: + return total + for usage in usage_by_stage.values(): + if isinstance(usage, Mapping) and usage: + deep_merge_usage(total, usage) + return total + + +def coerce_usage_to_dict(usage: Any) -> Optional[Dict[str, Any]]: + """Best-effort conversion of a provider usage object to a plain dict. + + Handles pydantic v2 models (``model_dump``), pydantic v1 / namedtuple-ish + objects (``dict``), and ``SimpleNamespace`` / plain objects (``__dict__``), + in that order. ``None`` and unsupported shapes return ``None`` so callers + can short-circuit. + """ + if usage is None: + return None + if isinstance(usage, dict): + return dict(usage) + # pydantic v2 + md = getattr(usage, "model_dump", None) + if callable(md): + try: + out = md() + if isinstance(out, dict): + return out + except Exception: + pass + # pydantic v1 / namedtuple-ish + d = getattr(usage, "dict", None) + if callable(d): + try: + out = d() + if isinstance(out, dict): + return out + except Exception: + pass + # SimpleNamespace / plain object + raw = getattr(usage, "__dict__", None) + if isinstance(raw, dict): + return dict(raw) + return None diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/loop.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/loop.py new file mode 100644 index 0000000000..583efedfa7 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/loop.py @@ -0,0 +1,485 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal agent-loop engine shared by :class:`~nemo_agent.agent.Agent` and +:class:`~nemo_agent.selection_agent.SelectionAgent`. + +This module owns everything both agents have in common: the per-run state, the +LLM step loop (finish-reason branching, auto-continue, reasoning annotation), +tool-call dispatch (unknown-tool and malformed-arguments recovery, structural +:class:`~nemo_agent.tools.BaseEndTool` termination), the ``on_error`` policy, +raw-IO capture/flush, progress logging, and result building. + +Subclasses provide assembly, not loop mechanics: they seed a :class:`_RunState` +(system/user messages, per-run tool map and specs, auto-continue text, usage +stage label) and call :meth:`_BaseAgentLoop._run_state_to_result`. The one +dispatch hook is :meth:`_BaseAgentLoop._dispatch_tool_call`, which +:class:`~nemo_agent.agent.Agent` overrides to intercept retrieve tools. + +Everything here is private implementation detail within +:mod:`nemo_retriever._agentic.nemo_agent`; names used by sibling modules are not +part of the supported external interface. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Set, Tuple, Union + +from pydantic import BaseModel, ConfigDict + +from .cache_propagation import PropagationPacer +from .llm import ( + BaseLLMBackend, + CompletionResult, + ContentPolicyError, + ContextLimitError, + LLMCallError, + bind_stage, + get_query_id, +) +from .results import ( + ERROR_BAD_FINISH_REASON, + ERROR_CONTENT_POLICY, + ERROR_CONTEXT_LIMIT, + ERROR_LLM_CALL_FAILED, + ERROR_MAX_STEPS, + ERROR_TOOL_FAILED, + ERROR_UNEXPECTED, + AgentError, + AgentRunResult, +) +from .tools import BaseEndTool, BaseTool + +logger = logging.getLogger(__name__) + +# LLM call failures that are expected, model-input-dependent, and +# non-actionable. Under ``on_error="raise_unknown"`` ONLY these end the run +# with an error record; everything else (bare LLMCallError — auth/timeout/5xx, +# a RateLimitError that exhausted the backend's retries, tool crashes, bugs) +# raises so it gets seen and fixed. +KNOWN_LLM_ERRORS: Tuple[type, ...] = (ContextLimitError, ContentPolicyError) + + +class ToolExecutionError(Exception): + """A tool raised an unexpected exception while the agent executed it. + + The original exception is chained on ``__cause__``. Never sent to the LLM; + the agent's error policy maps it to the ``tool_failed`` category. + """ + + def __init__(self, tool_name: str, original: BaseException) -> None: + super().__init__(f"Tool '{tool_name}' failed. {type(original).__name__}: {original}") + self.tool_name = tool_name + + +def build_auto_continue_msg(end_tool_name: str, end_payload_phrase: str) -> str: + """The user message appended when the model stops without calling a tool. + + One template for every agent, interpolating the actual end tool's name so + the instruction can never point at a tool that doesn't exist in the run. + """ + return ( + "Please continue on whatever approach you think is suitable.\n" + f"If you think you have solved the task, you MUST call the {end_tool_name} tool " + f"{end_payload_phrase}. Saying the task is complete in text only does NOT " + f"end the interaction; you must call {end_tool_name}.\n" + "IMPORTANT: YOU SHOULD NEVER ASK FOR HUMAN RESPONSE.\n" + ) + + +class BaseAgentLoopConfig(BaseModel): + """Config fields the shared loop engine reads; both agent configs inherit it. + + Pure data — see the subclasses (``AgentConfig``, ``SelectionAgentConfig``) + for the agent-specific fields and the field semantics they add. + + Attributes + ---------- + max_steps: + Maximum number of LLM calls per run; ``None`` = unlimited (callers set + their own limit). An end-tool call made on the last allowed step still + counts as success (the limit is checked before each LLM call). + on_error: + Error policy. ``raise_unknown`` (default): expected + model-input-dependent failures (:data:`KNOWN_LLM_ERRORS`) become a + terminal error record; everything else raises. ``never_raise`` + (production): every ``Exception`` becomes an error record (full + traceback logged). ``raise_all`` (debugging): everything raises except + the normal terminations (max-steps, bad finish_reason). + write_all_llm_io_logs: + When a run gets a ``raw_log_dir``: ``True`` writes every step's + prompt/response pair immediately; ``False`` (default) writes only the + most recent pair, once, at run end. + verbose_progress_logs: + Per-step INFO progress line (step counter + rate-limit info), labeled + with the bound query id. + cache_propagation_target_s: + Minimum wall-time gap between consecutive LLM calls in a run, for + Anthropic prompt-cache propagation. ``0.0`` (default) disables pacing. + """ + + model_config = ConfigDict(extra="forbid") + + max_steps: Optional[int] = None + on_error: Literal["never_raise", "raise_unknown", "raise_all"] = "raise_unknown" + write_all_llm_io_logs: bool = False + verbose_progress_logs: bool = True + cache_propagation_target_s: float = 0.0 + + +@dataclass +class _RunState: + """All per-run mutable state plus the per-run inputs the loop consumes. + + One instance per run (for the selection agent: per shrink attempt). Agent + instances themselves are immutable after construction, so a single agent + safely serves many (including concurrent) runs. ``tool_map`` / + ``tool_specs`` / ``auto_user_msg`` / ``stage`` are per-run *inputs* seeded + by the agent: static copies for ``Agent``, built per call for + ``SelectionAgent`` (whose end tool and prompt depend on the candidates). + """ + + query: str + raw_log_dir: Optional[Path] + exclude_docs: Set[str] + pacer: PropagationPacer + tool_map: Dict[str, BaseTool] + tool_specs: List[Dict[str, Any]] + auto_user_msg: str + stage: str + # When True, a ContextLimitError is recorded as an error result even under + # on_error="raise_all" — set only by SelectionAgent's non-final shrink + # attempts, whose retry loop must inspect the typed result. + record_context_limit: bool = False + message_history: List[Dict[str, Any]] = field(default_factory=list) + steps: int = 0 + retrieved_docs: Set[str] = field(default_factory=set) + retrieval_log: List[Dict[str, Any]] = field(default_factory=list) + extra_data: Dict[str, Any] = field(default_factory=dict) + extra_response_infos: List[Dict[str, Any]] = field(default_factory=list) + last_reasoning: Optional[str] = None + last_raw_io: Optional[Tuple[int, Optional[Dict[str, Any]], Optional[Dict[str, Any]]]] = None + end_payload: Optional[Dict[str, Any]] = None + # Best-effort payload from the most recent INVALID end-tool call, kept so a + # run that fails without ever ending validly can still surface the agent's + # last attempt. The last attempt wins. + last_end_attempt: Optional[Dict[str, Any]] = None + error: Optional[AgentError] = None + warned_missing_raw_io: bool = False + + +class _BaseAgentLoop: + """The loop engine. Subclasses assemble runs; this class executes them. + + Mirrors the ``BaseLLMBackend`` template pattern: the run/step/dispatch + methods here are the templates — subclasses must not override them, with + the single documented exception of :meth:`_dispatch_tool_call` (extend by + intercepting a tool kind, then defer to ``super()``). + """ + + def __init__(self, config: BaseAgentLoopConfig, llm: BaseLLMBackend) -> None: + if not isinstance(config, BaseAgentLoopConfig): + raise TypeError(f"config must be a BaseAgentLoopConfig, got {type(config).__name__}.") + if not isinstance(llm, BaseLLMBackend): + raise TypeError(f"llm must be a BaseLLMBackend, got {type(llm).__name__}.") + self.config = config + self.llm = llm + + # ------------------------------------------------------------------ + # Run skeleton. + # ------------------------------------------------------------------ + + async def _run_state_to_result( + self, + state: _RunState, + *, + prologue: Optional[Callable[[], Awaitable[None]]] = None, + ) -> AgentRunResult: + """Drive a seeded run state through the loop under the error policy. + + ``prologue`` (when given) runs inside the try block, so its failures — + e.g. the main agent's bootstrap retrieve — hit the same error policy + as loop failures. Deferred log artifacts flush in ``finally`` so they + land on error and raise paths too. + """ + try: + if prologue is not None: + await prologue() + await self._loop(state) + except Exception as e: + if not self._handle_run_exception(state, e): + raise + finally: + await self._flush_deferred_logs(state) + return self._build_result(state) + + # ------------------------------------------------------------------ + # Loop. + # ------------------------------------------------------------------ + + async def _loop(self, state: _RunState) -> None: + while True: + if self.config.max_steps is not None and state.steps >= self.config.max_steps: + self._record_error( + state, + AgentError(category=ERROR_MAX_STEPS, message="Agent reached maximum allowed iterations"), + ) + return + await self._step(state) + if state.error is not None: + return + tool_calls = state.message_history[-1].get("tool_calls") or [] + if len(tool_calls) == 0: + state.message_history.append( + {"role": "user", "content": [{"type": "text", "text": state.auto_user_msg}]} + ) + continue + ended = await self._process_tool_calls(state) + if ended: + return + + async def _step(self, state: _RunState) -> None: + """One LLM call: append the assistant message or record a terminal error.""" + await state.pacer.await_propagation() + with bind_stage(state.stage): + result = await self.llm.acompletion(messages=state.message_history, tools=state.tool_specs) + state.pacer.mark() + step_idx = state.steps + state.steps += 1 + state.extra_response_infos.append(result.extra_response_info) + await self._capture_raw_io(state, step_idx, result) + self._log_progress(state, result) + + # Overwritten every step — None when this turn exposed no reasoning, + # so stale reasoning never leaks into the next retrieve. + state.last_reasoning = result.reasoning + + if result.finish_reason not in ("stop", "tool_calls"): + self._record_error( + state, + AgentError( + category=ERROR_BAD_FINISH_REASON, + message=f"LLM failed with finish_reason '{result.finish_reason}'", + ), + ) + return + message = dict(result.message) + if result.reasoning is not None: + message["__reasoning__"] = result.reasoning # backends strip __-prefixed keys + state.message_history.append(message) + + async def _process_tool_calls(self, state: _RunState) -> bool: + """Execute the last assistant message's tool calls; True when the run ended.""" + ended = False + tool_messages: List[Dict[str, Any]] = [] + for call_info in state.message_history[-1]["tool_calls"]: + fn_name = str((call_info.get("function") or {}).get("name")) + content: Optional[List[Dict[str, Any]]] = None + fn_kwargs: Optional[Dict[str, Any]] = None + try: + fn_kwargs = json.loads(call_info["function"]["arguments"]) + if not isinstance(fn_kwargs, dict): + raise TypeError("tool arguments must decode to an object") + except Exception: + content = [ + {"type": "text", "text": "Error parsing tool arguments. Tool arguments not correctly formatted."} + ] + if content is None: + assert fn_kwargs is not None + content, call_ended = await self._dispatch_tool_call(state, fn_name, fn_kwargs) + ended = ended or call_ended + tool_messages.append( + {"content": content, "role": "tool", "tool_call_id": call_info.get("id"), "name": fn_name} + ) + state.message_history.extend(tool_messages) + return ended + + async def _dispatch_tool_call( + self, state: _RunState, fn_name: str, fn_kwargs: Dict[str, Any] + ) -> Tuple[List[Dict[str, Any]], bool]: + """Return ``(tool_message_content, run_ended)`` for one tool call. + + Handles unknown tools, end tools, and generic tools. ``Agent`` + overrides this to intercept retrieve tools first, then defers here. + """ + tool = state.tool_map.get(fn_name) + if tool is None: + # LLMs occasionally hallucinate tool names; an error result lets + # the model self-correct instead of aborting the whole run. + available = sorted(state.tool_map) + state.extra_data.setdefault("unknown_tool_calls", []).append( + {"requested": str(fn_name), "available": available} + ) + text = ( + f"Error: tool '{fn_name}' is not available. " + f"Available tools: {', '.join(available) if available else '(none)'}. " + "Please retry using one of the available tool names exactly as listed." + ) + return [{"type": "text", "text": text}], False + try: + if isinstance(tool, BaseEndTool): + payload, text = tool.try_end(**fn_kwargs) + if payload is not None: + state.end_payload = payload + return [{"type": "text", "text": text}], True + # Invalid end call: the agent retries (error text below), but keep + # its best-effort attempt so a run that later fails (e.g. max + # steps) can still surface the model's last answer/doc_ids + # instead of nothing. Overwrite so the LAST attempt wins. + salvaged = tool.salvage_payload(fn_kwargs) + if salvaged is not None: + state.last_end_attempt = salvaged + return [{"type": "text", "text": text}], False + output = await tool.acall(**fn_kwargs) + if not isinstance(output, str): + output = json.dumps(output) + return [{"type": "text", "text": output}], False + except Exception as e: + raise ToolExecutionError(fn_name, e) from e + + # ------------------------------------------------------------------ + # Errors. + # ------------------------------------------------------------------ + + def _handle_run_exception(self, state: _RunState, exc: Exception) -> bool: + """Apply ``config.on_error``; True when recorded (caller must not re-raise).""" + category, exc_class = _classify_exception(exc) + if self.config.on_error == "never_raise": + record = True + elif self.config.on_error == "raise_unknown": + record = isinstance(exc, KNOWN_LLM_ERRORS) + else: # raise_all + record = False + if state.record_context_limit and isinstance(exc, ContextLimitError): + record = True + if not record: + return False + if isinstance(exc, KNOWN_LLM_ERRORS): + logger.warning("Agent run (query_id=%r) ended with expected error (%s): %s", get_query_id(), category, exc) + else: + logger.error( + "Agent run (query_id=%r) ended with error (%s): %s", + get_query_id(), + category, + exc, + exc_info=True, + ) + self._record_error(state, AgentError(category=category, message=str(exc), exception_class=exc_class)) + return True + + @staticmethod + def _record_error(state: _RunState, error: AgentError) -> None: + state.error = error + state.message_history.append({"role": "agent_error", "content": f"[{error.category}] {error.message}"}) + + # ------------------------------------------------------------------ + # Raw-IO logging (paths come from the caller; the LLM only captures). + # ------------------------------------------------------------------ + + async def _capture_raw_io(self, state: _RunState, step_idx: int, result: CompletionResult) -> None: + if state.raw_log_dir is None: + return + if result.raw_request is None and result.raw_response is None: + if not state.warned_missing_raw_io: + logger.warning( + "raw_log_dir=%s was provided but the LLM backend returned no raw IO; construct " + "the backend with capture_raw_io=True to get per-step prompt/response logs.", + state.raw_log_dir, + ) + state.warned_missing_raw_io = True + return + if self.config.write_all_llm_io_logs: + await _write_raw_pair(state.raw_log_dir, step_idx, result.raw_request, result.raw_response) + else: + state.last_raw_io = (step_idx, result.raw_request, result.raw_response) + + async def _flush_deferred_logs(self, state: _RunState) -> None: + """Write deferred artifacts at run end (also on error/raise paths). + + Best-effort: a failing log write must never mask the run's outcome. + """ + if state.raw_log_dir is None: + return + try: + if not self.config.write_all_llm_io_logs and state.last_raw_io is not None: + step_idx, raw_request, raw_response = state.last_raw_io + await _write_raw_pair(state.raw_log_dir, step_idx, raw_request, raw_response) + if state.extra_response_infos: + await _awrite_json(state.extra_response_infos, state.raw_log_dir, "api_response_extras.json") + except Exception: + logger.exception("Failed to write LLM IO logs under %s.", state.raw_log_dir) + + # ------------------------------------------------------------------ + # Progress + result. + # ------------------------------------------------------------------ + + def _log_progress(self, state: _RunState, result: CompletionResult) -> None: + if not self.config.verbose_progress_logs: + return + parts = [f"S: {state.steps}"] + ratelimit = result.extra_response_info.get("ratelimit") + if isinstance(ratelimit, dict): + parts.extend(f"{k}: {v}" for k, v in ratelimit.items()) + qid = get_query_id() + prefix = f"[{qid}] " if qid else "" + logger.info("%s%s", prefix, " ".join(parts)) + + def _build_result(self, state: _RunState) -> AgentRunResult: + payload = state.end_payload + if payload is None and state.last_end_attempt is not None: + # The run ended in error without a valid end call. Fall back to the + # agent's last (invalid) end-tool attempt so callers still get its + # best-effort doc_ids. The run still counts as failed + # (``error`` set, ``succeeded`` False). + payload = state.last_end_attempt + final_doc_ids: List[str] = [] + if payload is not None: + doc_ids = payload.get("doc_ids") + if isinstance(doc_ids, list): + final_doc_ids = [str(d) for d in doc_ids] + return AgentRunResult( + final_doc_ids=final_doc_ids, + end_payload=payload, + error=state.error, + trajectory=state.message_history, + retrieval_log=state.retrieval_log, + extra_data=state.extra_data, + ) + + +def _classify_exception(exc: Exception) -> Tuple[str, str]: + """Map an exception to an ``AgentError`` category + exception class name.""" + if isinstance(exc, ContextLimitError): + return ERROR_CONTEXT_LIMIT, type(exc).__name__ + if isinstance(exc, ContentPolicyError): + return ERROR_CONTENT_POLICY, type(exc).__name__ + if isinstance(exc, LLMCallError): + return ERROR_LLM_CALL_FAILED, type(exc).__name__ + if isinstance(exc, ToolExecutionError): + cause = exc.__cause__ + return ERROR_TOOL_FAILED, type(cause).__name__ if cause is not None else type(exc).__name__ + return ERROR_UNEXPECTED, type(exc).__name__ + + +async def _write_raw_pair( + log_dir: Path, step_idx: int, raw_request: Optional[Dict[str, Any]], raw_response: Optional[Dict[str, Any]] +) -> None: + if raw_request is not None: + await _awrite_json(raw_request, log_dir, f"{step_idx}_prompt.json") + if raw_response is not None: + await _awrite_json(raw_response, log_dir, f"{step_idx}_response.json") + + +async def _awrite_json(obj: Any, log_dir: Union[str, Path], filename: str) -> None: + def _write() -> None: + path = Path(log_dir, filename) + path.parent.mkdir(exist_ok=True, parents=True) + with open(path, "w") as f: + json.dump(obj, f, indent=2) + + await asyncio.to_thread(_write) diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/__init__.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/__init__.py new file mode 100644 index 0000000000..6715c39b8f --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/__init__.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prompt rendering for the agent's system prompts. + +The agent renders its system prompt through :func:`render_system_prompt` and +never touches a template engine. Resolution order: an existing filesystem path +wins (so a caller can point at any file); otherwise the name must be a +packaged template under ``templates/``. + +jinja2 is imported lazily, only when a prompt is actually rendered. The render +*variables* and the templates' conditional logic are the contract — if jinja2 +ever becomes unacceptable as a dependency, this module's internals swap +template files for plain Python functions and nothing outside it changes. + +Render variables used by the packaged system prompts: ``with_init_docs``, +``enforce_top_k``, ``top_k``, ``extended_relevance``. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, List + +_TEMPLATES_DIR = Path(__file__).parent / "templates" + + +def available_prompts() -> List[str]: + """Sorted relative names of the packaged prompt templates. + + Templates in subdirectories keep their subpath, e.g. the selection-agent + set is listed (and rendered) as ``"selection/01_v0.j2"``. + """ + if not _TEMPLATES_DIR.is_dir(): + return [] + return sorted(p.relative_to(_TEMPLATES_DIR).as_posix() for p in _TEMPLATES_DIR.rglob("*.j2")) + + +def render_system_prompt(name_or_path: str, **variables: Any) -> str: + """Render a system prompt by packaged template name or filesystem path. + + An existing file path is rendered directly; otherwise ``name_or_path`` + must be a packaged template filename. The rendered prompt is stripped of + surrounding whitespace. + """ + path = Path(name_or_path) + if not path.is_file(): + path = _TEMPLATES_DIR / name_or_path + if not path.is_file(): + raise ValueError( + f"Unknown prompt {name_or_path!r}: not an existing file path and not one of " + f"the packaged templates: {available_prompts()}." + ) + return _render_template_text(path.read_text(), variables) + + +def _render_template_text(text: str, variables: Dict[str, Any]) -> str: + try: + import jinja2 + except ImportError as e: # pragma: no cover - depends on the environment + raise RuntimeError("Rendering prompts requires the 'jinja2' package.") from e + return jinja2.Template(text.strip()).render(**variables).strip() diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/00_default.j2 b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/00_default.j2 new file mode 100644 index 0000000000..90ce03438f --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/00_default.j2 @@ -0,0 +1 @@ +You are a helpful assistant. Use the "retrieve" tool to find all the documents related to the given query. diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/01_v0.j2 b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/01_v0.j2 new file mode 100644 index 0000000000..f054bb1638 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/01_v0.j2 @@ -0,0 +1,19 @@ +You are a retrieval agent, which uses retrieval to help users find the documents they need. + + +Your goal is to help users find documents they related to their search queries. You should be thorough and find all documents relevant to the user's query. If the user's query is a question, you should not answer the question yourself. Instead, you should find the related documents for the given query. + + + +* You are given a retrieval tool, powered by a dense embedding model, that takes a text query and returns the most similar documents. +* You can call the search tool multiple times. +* Search for related documents to the user's query from different angles. +* If needed, revise your search queries based on the documents you find in previous steps. +* Once you find the relevant documents, report the ID of the relevant documents by calling the "final_results" tool, which also ends the interaction. + + + +* You should be thorough and find all related documents. +* While you can use the retrieval tools as many times as you want, it is an expensive tool. So, try to be efficient and find the documents in as few searches as possible. +* The goal is to increase the **Recall** of your search attempt. So, if multiple documents are relevant to the given query, you should find and report all of them even if only a subset of them is enough for answering the query. + diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/02_v1.j2 b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/02_v1.j2 new file mode 100644 index 0000000000..d4608693ff --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/02_v1.j2 @@ -0,0 +1,47 @@ +You are a retrieval agent that finds all documents related to a given query. + + +You are given a search query and a list of documents retrieved for that query. Your task is to write new queries and use the given search tool to find *ALL* the related and somewhat related documents to the given query (i.e., maximize recall). +If the user's query is a question, you should not answer the question yourself. Instead, you should find the related documents for the given query. + + +{% if extended_relevance %} + +- You should be careful, in the context of this task, what it means to be a "query", "document", and "relevant" can sometimes be very complex and might not follow the traditional definition of these terms in standard information retrieval. +- In standard retrieval, a query is usually a user question (like a web search query), the document is some sort of content that provides information (e.g., a web page), and these two are considered relevant if the document provides information that answers the user's query. +- However, in our setting, this could be different. Here are some examples: + * the query is a programming problem and documents are programming language syntax references. A document is relevant if it contains the reference for the programming syntax used for solving the problem. + * both query and documents are descriptions programming problems and a query and document are relevant if the same approach is used to solve them. + * the query is a math problem and documents are theorems. Relevant documents (theorems) are the ones that are useful for solving the math problem. + * the query and document are both math problems. A query and a document are relevant if the same theorem is used for solving them. + * the query is a task description (e.g., for an API programmer) and documents are descriptions of available APIs. Relevant documents (e.g., APIs) are the ones needed for completing the task. +- This is not an exhaustive list. These are just some examples to show you the complexity of queries, documents, and the concept of relevance in this task. +- Note that even here, the relevant documents are still the ones that are useful for a user who is searching for the given query. But the relation is more nuanced. +- You should analyze the query and some of the available documents. And then reason about what could be a meaningful definition of relevance in this case, and what the user could be looking for. +- Moreover, sometimes, the query could be even a prompt that is given to a Large Language Model (LLM) and the user wants to find the useful documents for the LLM that help answering/solving this prompt. + + +{% endif %} + +- You are given a retrieval tool, powered by a dense embedding model, that takes a text query and returns the most similar documents. +{%- if extended_relevance %} +- As explained above, reason and figure out what the meaning of relevance is in this case, and what could be relevant and useful information for the given query. +{%- endif %} +- You can call the search tool multiple times. +- Search for related documents to the user's query from different angles. +- If needed, revise your search queries based on the documents you find in previous steps. +- Once you are confident that you have found all the related and somewhat related documents and there are no more related documents in the corpus, call the "final_results" tool to finish the task. The interaction only ends when you call this tool; saying you are done in text does not end it. +{%- if enforce_top_k %} +- When calling "final_results", you must select exactly the {{ top_k }} most relevant documents among all documents you have retrieved. +{%- endif %} +- When calling the "final_results" tool, the list of documents must be sorted in the decreasing level of relevance to the query. I.e., the first document is the most relevant to the query, the second document is the second most relevant to the query, and so on. + + + + +- You should be thorough and find all related and somewhat related documents. +- The goal is to increase the **Recall** of your search attempt. So, if multiple documents are relevant to the given query, you should find and report all of them even if only a subset of them is enough for answering the query. +{%- if with_init_docs %} +- **TIP**: you can look at the list of documents retrieved using the original query and think what other queries you can use to find the potentially related documents that are missing in these results. +{%- endif %} + diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/06_select_lean_v1.j2 b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/06_select_lean_v1.j2 new file mode 100644 index 0000000000..c52ad33122 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/06_select_lean_v1.j2 @@ -0,0 +1,47 @@ +You are a retrieval agent that finds all documents related to a given query. + + +You are given a search query and a list of documents retrieved for that query. Your task is to write new queries and use the given search tool to find *ALL* the related and somewhat related documents to the given query (i.e., maximize recall). +If the user's query is a question, you should not answer the question yourself. Instead, you should find the related documents for the given query. + + +{% if extended_relevance %} + +- You should be careful, in the context of this task, what it means to be a "query", "document", and "relevant" can sometimes be very complex and might not follow the traditional definition of these terms in standard information retrieval. +- In standard retrieval, a query is usually a user question (like a web search query), the document is some sort of content that provides information (e.g., a web page), and these two are considered relevant if the document provides information that answers the user's query. +- However, in our setting, this could be different. Here are some examples: + * the query is a programming problem and documents are programming language syntax references. A document is relevant if it contains the reference for the programming syntax used for solving the problem. + * both query and documents are descriptions programming problems and a query and document are relevant if the same approach is used to solve them. + * the query is a math problem and documents are theorems. Relevant documents (theorems) are the ones that are useful for solving the math problem. + * the query and document are both math problems. A query and a document are relevant if the same theorem is used for solving them. + * the query is a task description (e.g., for an API programmer) and documents are descriptions of available APIs. Relevant documents (e.g., APIs) are the ones needed for completing the task. +- This is not an exhaustive list. These are just some examples to show you the complexity of queries, documents, and the concept of relevance in this task. +- Note that even here, the relevant documents are still the ones that are useful for a user who is searching for the given query. But the relation is more nuanced. +- You should analyze the query and some of the available documents. And then reason about what could be a meaningful definition of relevance in this case, and what the user could be looking for. +- Moreover, sometimes, the query could be even a prompt that is given to a Large Language Model (LLM) and the user wants to find the useful documents for the LLM that help answering/solving this prompt. + + +{% endif %} + +- You are given a retrieval tool, powered by a dense embedding model, that takes a text query and returns the most similar documents. +{%- if extended_relevance %} +- As explained above, reason and figure out what the meaning of relevance is in this case, and what could be relevant and useful information for the given query. +{%- endif %} +- You can call the search tool multiple times. +- Search for related documents to the user's query from different angles. +- If needed, revise your search queries based on the documents you find in previous steps. +- Once you are confident that you have found all the related and somewhat related documents and there are no more related documents in the corpus, call the "final_results" tool to finish the task. +{%- if enforce_top_k %} +- When calling "final_results", you must select exactly the {{ top_k }} most relevant documents among all documents you have retrieved. +{%- endif %} +- When calling the "final_results" tool, the list of documents must be sorted in the decreasing level of relevance to the query. I.e., the first document is the most relevant to the query, the second document is the second most relevant to the query, and so on. + + + + +- You should be thorough and find all related and somewhat related documents. +- The goal is to increase the **Recall** of your search attempt. So, if multiple documents are relevant to the given query, you should find and report all of them even if only a subset of them is enough for answering the query. +{%- if with_init_docs %} +- **TIP**: you can look at the list of documents retrieved using the original query and consider what other queries you can use to find the potentially related documents that are missing in these results. +{%- endif %} + diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/selection/00_demo.j2 b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/selection/00_demo.j2 new file mode 100644 index 0000000000..a2233c05bb --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/selection/00_demo.j2 @@ -0,0 +1 @@ +You are given a query and a list of documents. Select the {{ top_k }} most relevant documents for the given query. diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/selection/01_v0.j2 b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/selection/01_v0.j2 new file mode 100644 index 0000000000..6438ededb0 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/prompts/templates/selection/01_v0.j2 @@ -0,0 +1,43 @@ +You are a document re-ranker agent, which is the final stage in an information retrieval pipeline. + + +You are given a search query and a list of retrieved candidate documents that are potentially relevant to the given query. Your goal is to help the users identify the most relevant documents to the given query from the list of candidate documents. + + +{% if extended_relevance %} + +- You should be careful, in the context of this task, what it means to be a "query", "document", and "relevant" can sometimes be very complex and might not follow the traditional definition of these terms in standard re-ranking and retrieval. +- In standard re-ranking/retrieval, a query is usually a user question (like a web search query), the document is some sort of content that provides information (e.g., a web page), and these two are considered relevant if the document provides information that answers the user's query. +- However, in our setting, this could be different. Here are some examples: + * the query is a programming problem and documents are programming language syntax references. A document is relevant if it contains the reference for the programming syntax used for solving the problem. + * both query and documents are descriptions programming problems and a query and document are relevant if the same approach is used to solve them. + * the query is a math problem and documents are theorems. Relevant documents (theorems) are the ones that are useful for solving the math problem. + * the query and document are both math problems. A query and a document are relevant if the same theorem is used for solving them. + * the query is a task description (e.g., for an API programmer) and documents are descriptions of available APIs. Relevant documents (e.g., APIs) are the ones needed for completing the task. +- This is not an exhaustive list. These are just some examples to show you the complexity of queries, documents, and the concept of relevance in this task. +- Note that even here, the relevant documents are still the ones that are useful for a user who is searching for the given query. But the relation is more nuanced. +- You should analyze the query and the available documents. And then reason about what could be a meaningful definition of relevance in this case, and what the user could be looking for. +- Moreover, sometimes, the query could be even a prompt that is given to a Large Language Model (LLM) and the user wants to find the useful documents for the LLM that help answering/solving this prompt. + + +{% endif %} + +* You are given a search query and a list of candidate documents. You have access to the ID and content of each candidate document. +* You should read the query carefully and understand it. +{%- if extended_relevance %} +* As explained above, reason and figure out what the meaning of relevance is in this case, and what could be relevant and useful information for the given query. +{%- endif %} +* Then you should compare the query with each one of the candidate documents. In this comparison, you want to identify if the document is relevant/useful for the given query and to what extent. +* Select the {{ top_k }} most relevant candidate documents for the given query. +* Note that just selecting the most relevant documents is not enough. You should identify the relative level of relevance between the query and selected documents. This helps you sort the selected documents later based on how relevant they are to the query. +* Once you have this information, you should call the "log_selected_documents" function to report the final results and signal the completion of the task. +* Note that the selected document IDs must be reported in the decreasing level of relevance. I.e., The first document in the list is the most relevant, the second is the second most relevant, and so on. This is similar to what a search engine (e.g., Google Search) does (it shows you the relevant results in a sorted order, where the most relevant results appear on top of the list). + + + + +* you have access to a "think" tool that you can use for complex thinking and analysis. Here are examples of cases where the think tool might be useful: + - complex analysis and thinking to understand the meaning and intent of the query. E.g., what is the user trying to find with this query? what kind of information is helpful for the user? + - extended thinking to analyze how each candidate document could or could not be relevant to the given query. + - reasoning to identify the relative level of relevance between the query and selected documents. It helps you sort the documents correctly when reporting the final answer. + diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/results.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/results.py new file mode 100644 index 0000000000..e0a93e2d72 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/results.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed result and error records returned by the agent.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +# Error categories. Downstream code branches on these constants (via +# ``AgentRunResult.error.category``), never on error-message text. +ERROR_MAX_STEPS = "max_steps" +ERROR_BAD_FINISH_REASON = "bad_finish_reason" +ERROR_CONTEXT_LIMIT = "context_limit" +ERROR_CONTENT_POLICY = "content_policy" +ERROR_LLM_CALL_FAILED = "llm_call_failed" +ERROR_TOOL_FAILED = "tool_failed" +ERROR_UNEXPECTED = "unexpected" + + +@dataclass +class AgentError: + """Why an agent run ended without a successful end-tool call. + + Attributes + ---------- + category: + One of the ``ERROR_*`` constants in this module. + message: + Human-readable description (also appended to the trajectory as a + ``role="agent_error"`` message). Do not branch on it. + exception_class: + Class name of the underlying exception, or ``None`` for the two + normal terminations (``max_steps``, ``bad_finish_reason``). + """ + + category: str + message: str + exception_class: Optional[str] = None + + +@dataclass +class AgentRunResult: + """Everything one agent run produces. + + ``final_doc_ids`` is a convenience extracted from ``end_payload`` (its + ``doc_ids`` key). On a successful run ``end_payload`` is the full + *validated* end-tool arguments (e.g. ``message``, ``search_successful``). + When a run FAILS without ever making a valid end call, ``end_payload`` + falls back to the agent's last invalid end-tool attempt (a lenient + best-effort subset of what the model supplied) so callers still get the + model's final intent; ``error`` still says why the run failed and + ``succeeded`` stays ``False``. ``end_payload`` is ``None`` only when the + run failed and no usable end attempt was ever made. + + The verbose fields (``trajectory``, ``retrieval_log``, ``extra_data``) + are always populated; callers that don't need them simply don't persist + them. + """ + + final_doc_ids: List[str] = field(default_factory=list) + end_payload: Optional[Dict[str, Any]] = None + error: Optional[AgentError] = None + trajectory: List[Dict[str, Any]] = field(default_factory=list) + retrieval_log: List[Dict[str, Any]] = field(default_factory=list) + extra_data: Dict[str, Any] = field(default_factory=dict) + + @property + def succeeded(self) -> bool: + """True iff the run ended via a successful end-tool call.""" + return self.error is None diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/selection_agent.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/selection_agent.py new file mode 100644 index 0000000000..485c8f0ee0 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/selection_agent.py @@ -0,0 +1,368 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Standalone selection agent: pick the top-k most relevant candidate documents. + +The agent takes a query plus a list of candidate documents and runs an LLM +loop (on the shared engine in ``loop.py``) that ends with a +``log_selected_documents`` call naming the ``target_top_k`` best candidates. +It knows nothing about retrieval, RRF, or the main agent. + +Two modes, selected explicitly by the ``scores`` argument of :meth:`SelectionAgent.select`: + +- ``scores=None`` — single attempt; a context-window overflow surfaces per + ``config.on_error`` like any other known LLM error. +- ``scores={doc_id: priority}`` — passing a priority ranking is the permission + to shrink: on a context-window overflow the agent drops the lowest-priority + quarter of the candidates and retries (up to ``config.shrink_attempts`` + total attempts). + +The mode is never inferred from the documents themselves: candidate docs that +came out of retrieval carry per-subquery ``score`` keys which are NOT +comparable across subqueries — truncating by them would be wrong, so they are +ignored (and never shown to the LLM). The ``scores`` side table is a globally +comparable ranking the caller owns (e.g. RRF fusion scores). +""" + +from __future__ import annotations + +import asyncio +import math +from contextlib import nullcontext +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +from pydantic import Field + +from .cache_propagation import PropagationPacer +from .llm import BaseLLMBackend, bind_query_id +from .loop import ( + BaseAgentLoopConfig, + _awrite_json, + _BaseAgentLoop, + _RunState, + build_auto_continue_msg, +) +from .prompts import render_system_prompt +from .results import ERROR_CONTEXT_LIMIT, AgentRunResult +from .tools import BaseTool, LogSelectedDocs, SelectionThinkTool + +# Fraction of candidates dropped (lowest priority first) before a retry after +# a context-window overflow. Deliberately a constant, not a config knob: +# 0.25 drops a quarter of the candidates (``floor(n * 0.25)``) per shrink step. +_SHRINK_DROP_FRACTION = 0.25 + +_DEFAULT_SELECTION_PROMPT = "selection/01_v0.j2" +_END_TOOL_NAME = "log_selected_documents" +_END_PAYLOAD_PHRASE = "with your selected doc_ids" + + +class SelectionAgentConfig(BaseAgentLoopConfig): + """Configuration for the selection agent. + + Pure data, like ``AgentConfig``; the LLM instance is a constructor + argument. The loop-policy fields (``max_steps``, ``on_error``, logging and + pacing knobs) are inherited from + :class:`~nemo_agent.loop.BaseAgentLoopConfig`. + + Attributes + ---------- + system_prompt: + Packaged prompt name or filesystem path; ``None`` selects the default + selection prompt. Render variables: ``extended_relevance`` and the + per-run ``top_k``. + target_top_k: + How many documents a run selects — one k per agent instance. Runs with + fewer candidates than ``target_top_k`` self-adjust to the candidate + count. + extended_relevance: + Render the extended relevance-definition guidance into the system + prompt (and the think tool's description). + enable_think: + Register the selection think tool. Off by default, consistent with + ``AgentConfig``. + end_tool_with_msg: + Require (and pass through) the end tool's diagnostic ``message`` + argument. + shrink_attempts: + Total attempts (including the first) when ``scores`` are provided and + a context-window overflow occurs. Defaults to 2: one shrink retry + after the initial attempt. + """ + + system_prompt: Optional[str] = None + target_top_k: int = Field(default=10, ge=1) + extended_relevance: bool = True + enable_think: bool = False + end_tool_with_msg: bool = True + shrink_attempts: int = Field(default=2, ge=1) + + +class SelectionAgent(_BaseAgentLoop): + """LLM agent that selects the top-k candidates for a query from a given list. + + Construction wires the static pieces (stage label, auto-continue message, + optional think tool); everything that depends on a run's candidate set — + the system prompt (``top_k`` baked in), the + :class:`~nemo_agent.tools.LogSelectedDocs` end tool (count + allowed ids), + and the documents user message — is assembled per attempt inside + :meth:`select`. The instance is immutable after construction and safely + serves many (including concurrent) runs. + + Results reuse :class:`~nemo_agent.results.AgentRunResult`: the selected + ids land in ``final_doc_ids`` (with the end tool's ``message`` in + ``end_payload``), ``retrieval_log`` stays empty, and when the shrink retry + fired, ``extra_data["context_shrink"]`` records what was dropped. + + Token usage is attributed to the stage ``f"top{target_top_k}_agent"`` — + the *configured* k, even when a run's feasible k is smaller. + """ + + def __init__(self, config: SelectionAgentConfig, llm: BaseLLMBackend) -> None: + if not isinstance(config, SelectionAgentConfig): + raise TypeError(f"config must be a SelectionAgentConfig, got {type(config).__name__}.") + super().__init__(config=config, llm=llm) + self._stage = f"top{config.target_top_k}_agent" + self._auto_user_msg = build_auto_continue_msg(_END_TOOL_NAME, _END_PAYLOAD_PHRASE) + self._prompt_name = config.system_prompt or _DEFAULT_SELECTION_PROMPT + self._think_tool: Optional[SelectionThinkTool] = ( + SelectionThinkTool(extended_relevance=config.extended_relevance) if config.enable_think else None + ) + + # ------------------------------------------------------------------ + # Entry points. + # ------------------------------------------------------------------ + + async def select( + self, + query: str, + documents: List[Dict[str, Any]], + *, + scores: Optional[Dict[str, float]] = None, + query_id: Optional[str] = None, + task_info: Optional[Any] = None, + raw_log_dir: Optional[Union[str, Path]] = None, + ) -> AgentRunResult: + """Select the ``config.target_top_k`` most relevant documents for ``query``. + + Parameters + ---------- + query: + The user's question. + documents: + Candidate documents: dicts with a required ``id`` (coerced to + ``str``) and optional ``text`` / ``image`` strings. Duplicate ids + are deduplicated (first occurrence wins); any other keys — + including retrieval ``score`` keys — are ignored and never shown + to the LLM. + scores: + Optional ``{doc_id: priority}`` ranking covering every candidate + id (a superset is fine). Passing it arms the context-overflow + shrink retry: intermediate context-limit failures are consumed by + the retry loop regardless of ``config.on_error`` — dropping the + lowest-priority quarter each time — and only the final attempt's + outcome is subject to the configured policy. Every failure that is + not a context-window overflow follows the policy immediately, on + any attempt. Without ``scores`` there is exactly one attempt. + query_id: + Optional id bound for the run, exactly like ``Agent.run``: token + usage lands under ``llm.get_usage(query_id)`` and logs are + labeled. An ambient caller-side ``bind_query_id`` also applies. + task_info: + Arbitrary JSON-serializable info written once to + ``extra_info.json`` when ``raw_log_dir`` is set. + raw_log_dir: + Per-call directory for raw LLM IO artifacts (caller-built; + ``None`` discards them). The first attempt writes at the root — + identical layout to an ``Agent`` run — and shrink retries write + under ``attempt_{i}/`` subdirectories (``i`` >= 2), so nothing + overwrites. + """ + candidates = _validated_unique_documents(documents) + if scores is not None: + _validate_scores_cover(scores, candidates) + + binding = bind_query_id(query_id) if query_id is not None else nullcontext() + with binding: + base_log_dir = Path(raw_log_dir) if raw_log_dir is not None else None + if base_log_dir is not None and task_info is not None: + await _awrite_json(task_info, base_log_dir, "extra_info.json") + + if scores is None: + return await self._run_attempt( + query=str(query), + candidates=candidates, + raw_log_dir=base_log_dir, + record_context_limit=False, + ) + + # Scored mode: retry on context overflow, shrinking the candidates. + initial_count = len(candidates) + dropped_ids: List[str] = [] + attempt = 1 + while True: + final = attempt >= self.config.shrink_attempts or len(candidates) == 1 + attempt_dir = base_log_dir + if base_log_dir is not None and attempt > 1: + attempt_dir = base_log_dir / f"attempt_{attempt}" + result = await self._run_attempt( + query=str(query), + candidates=candidates, + raw_log_dir=attempt_dir, + record_context_limit=not final, + ) + if final or result.error is None or result.error.category != ERROR_CONTEXT_LIMIT: + break + candidates, newly_dropped = _drop_lowest_priority(candidates, scores) + dropped_ids.extend(newly_dropped) + attempt += 1 + if dropped_ids: + result.extra_data["context_shrink"] = { + "attempts_run": attempt, + "dropped_doc_ids": dropped_ids, + "initial_candidates": initial_count, + "final_candidates": len(candidates), + } + return result + + def select_sync(self, query: str, documents: List[Dict[str, Any]], **kwargs: Any) -> AgentRunResult: + """Synchronous facade over :meth:`select` (one ``asyncio.run`` per call). + + For thread-based callers; must not be invoked from inside a running + event loop — ``await select(...)`` there instead. + """ + return asyncio.run(self.select(query, documents, **kwargs)) + + # ------------------------------------------------------------------ + # Per-attempt assembly. + # ------------------------------------------------------------------ + + async def _run_attempt( + self, + *, + query: str, + candidates: List[Dict[str, Any]], + raw_log_dir: Optional[Path], + record_context_limit: bool, + ) -> AgentRunResult: + feasible_topk = min(self.config.target_top_k, len(candidates)) + system_prompt = render_system_prompt( + self._prompt_name, + extended_relevance=self.config.extended_relevance, + top_k=feasible_topk, + ) + end_tool = LogSelectedDocs( + top_k=feasible_topk, + candidate_docids=[d["id"] for d in candidates], + include_msg=self.config.end_tool_with_msg, + ) + tool_map: Dict[str, BaseTool] = {} + if self._think_tool is not None: + tool_map[self._think_tool.name] = self._think_tool + tool_map[end_tool.name] = end_tool + + state = _RunState( + query=query, + raw_log_dir=raw_log_dir, + exclude_docs=set(), + pacer=PropagationPacer(target_s=self.config.cache_propagation_target_s), + tool_map=tool_map, + tool_specs=[t.spec for t in tool_map.values()], + auto_user_msg=self._auto_user_msg, + stage=self._stage, + record_context_limit=record_context_limit, + message_history=[ + {"role": "system", "content": [{"type": "text", "text": system_prompt}]}, + _documents_user_message(query, candidates), + ], + ) + return await self._run_state_to_result(state) + + +# --------------------------------------------------------------------------- +# Input validation + per-run message building. +# --------------------------------------------------------------------------- + + +def _validated_unique_documents(documents: Any) -> List[Dict[str, Any]]: + """Validate the candidate documents and dedup by id (first occurrence wins). + + Returns copies with ``id`` coerced to ``str``; never mutates the caller's + dicts. Violations raise ``ValueError`` naming the offending item — these + are caller bugs, not run outcomes. + """ + if not isinstance(documents, list) or len(documents) == 0: + raise ValueError("documents must be a non-empty list of document dicts.") + out: List[Dict[str, Any]] = [] + seen: set = set() + for idx, doc in enumerate(documents): + if not isinstance(doc, dict): + raise ValueError(f"documents[{idx}] must be a dict, got {type(doc).__name__}.") + if doc.get("id") is None: + raise ValueError(f"documents[{idx}] is missing the required 'id' key.") + doc_id = str(doc["id"]) + for key in ("text", "image"): + val = doc.get(key) + if val is not None and not isinstance(val, str): + raise ValueError(f"documents[{idx}] (id {doc_id!r}) has a non-string {key!r}: {type(val).__name__}.") + if doc_id in seen: + continue + seen.add(doc_id) + out.append({**doc, "id": doc_id}) + return out + + +def _validate_scores_cover(scores: Any, candidates: List[Dict[str, Any]]) -> None: + """Require a numeric score for every candidate id (supersets tolerated).""" + if not isinstance(scores, dict): + raise ValueError(f"scores must be a dict of doc_id -> priority, got {type(scores).__name__}.") + missing = [d["id"] for d in candidates if d["id"] not in scores] + if missing: + raise ValueError( + "scores must cover every candidate document id; missing: " + ", ".join(repr(i) for i in missing) + "." + ) + bad = [d["id"] for d in candidates if not isinstance(scores[d["id"]], (int, float))] + if bad: + raise ValueError( + "scores values must be numbers; non-numeric score for id(s): " + ", ".join(repr(i) for i in bad) + "." + ) + + +def _drop_lowest_priority( + candidates: List[Dict[str, Any]], scores: Dict[str, float] +) -> Tuple[List[Dict[str, Any]], List[str]]: + """Drop the lowest-priority quarter (at least 1 doc, never all of them). + + Deterministic: ascending score, original position as the tiebreak (ties + drop the earlier doc first, via a stable sort over first-seen order). The + kept documents preserve their original relative order, so the retry's user + message renders them in the same sequence minus the dropped ones. + """ + n = len(candidates) + drop_n = min(max(1, math.floor(n * _SHRINK_DROP_FRACTION)), n - 1) + by_priority = sorted(range(n), key=lambda i: (scores[candidates[i]["id"]], i)) + drop_idx = set(by_priority[:drop_n]) + dropped_ids = [candidates[i]["id"] for i in by_priority[:drop_n]] + kept = [candidates[i] for i in range(n) if i not in drop_idx] + return kept, dropped_ids + + +def _documents_user_message(query: str, candidates: List[Dict[str, Any]]) -> Dict[str, Any]: + """Build the candidates user message. + + Only ``id`` / ``text`` / ``image`` are rendered; other keys (including + retrieval ``score``s) are never shown to the LLM. + """ + content: List[Dict[str, Any]] = [ + {"type": "text", "text": f"Query:\n{query}"}, + {"type": "text", "text": "Candidate Documents:"}, + ] + for doc in candidates: + content.append({"type": "text", "text": f"Doc ID: {doc['id']}"}) + text = doc.get("text") + if isinstance(text, str) and text.strip() != "": + content.append({"type": "text", "text": f"Doc Text: {text}"}) + image = doc.get("image") + if image is not None and str(image).strip() != "": + content.append({"type": "image_url", "image_url": {"url": image}}) + return {"role": "user", "content": content} diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/__init__.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/__init__.py new file mode 100644 index 0000000000..6e8cc9d456 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/__init__.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tools the agent can call, and the bases for plugging in your own. + +One module per tool family: + +- ``base_tool.py`` — :class:`BaseTool`; :class:`ToolError` (LLM-recoverable) + and :class:`ToolContractError` (integration bug). +- ``retrieve.py`` — :class:`BaseRetrieveTool` / :class:`RetrieveContext`, the + concrete :class:`RetrieveTool` and :class:`ReasoningAugmentedRetrieveTool` + (wrap a plain ``retriever_fn(query, top_k)``), and + :func:`create_retrieve_tool` picking between them by name. +- ``end_tools.py`` — :class:`BaseEndTool` plus the standard end tools: + :class:`FinalResults` (the agent's end tool) and :class:`LogSelectedDocs` + (the selection agent's per-run end tool). +- ``think_tool.py`` — :class:`ThinkTool`, the optional scratchpad, and + :class:`SelectionThinkTool`, its selection-flavored variant. +""" + +from .base_tool import BaseTool, ToolContractError, ToolError +from .end_tools import BaseEndTool, FinalResults, LogSelectedDocs +from .retrieve import ( + BaseRetrieveTool, + ReasoningAugmentedRetrieveTool, + RetrieveContext, + RetrieveTool, + create_retrieve_tool, +) +from .think_tool import SelectionThinkTool, ThinkTool + +__all__ = [ + "BaseEndTool", + "BaseRetrieveTool", + "BaseTool", + "FinalResults", + "LogSelectedDocs", + "ReasoningAugmentedRetrieveTool", + "RetrieveContext", + "RetrieveTool", + "SelectionThinkTool", + "ThinkTool", + "ToolContractError", + "ToolError", + "create_retrieve_tool", +] diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/base_tool.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/base_tool.py new file mode 100644 index 0000000000..54a48e8353 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/base_tool.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Base tool abstraction shared by every tool the agent can call.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class ToolError(Exception): + """Deliberate, LLM-recoverable tool failure. + + Raise from a tool implementation when the *call* was invalid (bad argument + values, unusable input). The public wrappers turn it into an + ``"Error calling ..."`` string returned as the tool result, so the LLM can + see what it did wrong and retry on the next turn. + """ + + +class ToolContractError(Exception): + """A tool implementation violated a base-class contract. + + Deliberately NOT caught by the ``call``/``acall`` wrappers: a contract + violation is an integration bug, not something the LLM can correct. It + propagates to the agent's error policy so it surfaces as "your tool + returned X" instead of a ``KeyError`` deep inside the agent loop. + """ + + +def tool_error_text(tool_name: str, exc: BaseException) -> str: + """The LLM-visible error string for a recoverable tool failure.""" + return f"Error calling '{tool_name}' tool. {type(exc).__name__}: {str(exc)}" + + +class BaseTool(ABC): + """Define a tool to be passed to the LLM. + + Subclasses implement :meth:`_spec` (OpenAI function-tool spec dict) plus + ``_call`` and/or ``_acall``. The public ``call``/``acall`` wrappers convert + ``TypeError`` (signature mismatch from LLM-supplied kwargs) and + :class:`ToolError` into LLM-visible error strings; every other exception + propagates untranslated. + """ + + @abstractmethod + def _spec(self) -> dict: + raise NotImplementedError + + def _call(self, **kwargs: Any) -> Any: + raise NotImplementedError + + async def _acall(self, **kwargs: Any) -> Any: + raise NotImplementedError + + def call(self, **kwargs: Any) -> Any: + try: + output = self._call(**kwargs) + except (TypeError, ToolError) as e: + output = tool_error_text(self.name, e) + return output + + async def acall(self, **kwargs: Any) -> Any: + try: + output = await self._acall(**kwargs) + except (TypeError, ToolError) as e: + output = tool_error_text(self.name, e) + return output + + def __call__(self, **kwargs: Any) -> Any: + return self.call(**kwargs) + + @property + def spec(self) -> dict: + return self._spec() + + @property + def name(self) -> str: + return str(self.spec["function"]["name"]) diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/end_tools.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/end_tools.py new file mode 100644 index 0000000000..0f9044b3c7 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/end_tools.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tools that terminate the agent loop: the base plus the standard end tools. + +:class:`FinalResults` ends an agent run with the chosen document IDs; +:class:`LogSelectedDocs` ends a selection-agent run with the top-k subset of +the run's candidate documents. +""" + +# flake8: noqa: E501 +# The prompt text below exceeds the line limit. Two of the long lines live inside a +# triple-quoted prompt string, where a trailing ``# noqa`` would be interpolated into +# the prompt the model sees, so the exemption is applied per file rather than per line. + +from __future__ import annotations + +from abc import abstractmethod +from typing import Any, Dict, List, Optional, Tuple + +from .base_tool import BaseTool, ToolError, tool_error_text + + +class BaseEndTool(BaseTool): + """A tool whose successful call ends the agent run. + + The agent identifies end tools structurally (``isinstance(tool, + BaseEndTool)``) and calls :meth:`try_end` instead of ``acall``: a valid + call hands the agent the validated payload at call time; an invalid call + produces the usual LLM-visible error text so the model retries. + + Subclasses implement :meth:`_validate_payload` with an explicit signature + matching their spec (so unexpected LLM-supplied kwargs raise ``TypeError`` + naturally). Raise ``TypeError`` / :class:`ToolError` for invalid calls; + return the normalized payload dict on success. The agent's result exposes + the payload's ``doc_ids`` key as a typed convenience, with the full + payload available verbatim. + """ + + success_message: str = "The results have been successfully logged and the interaction ended." + + @abstractmethod + def _validate_payload(self, **kwargs: Any) -> Dict[str, Any]: + """Validate the end-call arguments; return the normalized payload.""" + raise NotImplementedError + + def try_end(self, **kwargs: Any) -> Tuple[Optional[Dict[str, Any]], str]: + """Return ``(payload, tool_message_text)``; payload is ``None`` for invalid calls.""" + try: + payload = self._validate_payload(**kwargs) + except (TypeError, ToolError) as e: + return None, tool_error_text(self.name, e) + return payload, self.success_message + + def salvage_payload(self, kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Best-effort payload from an INVALID (non-terminating) end call. + + A fallback only: when a run ends in error without ever making a valid + end call, the agent's last invalid attempt still carries the model's + intended output. Returns a lenient subset — a non-empty ``doc_ids`` + list of strings and/or a non-empty ``answer`` string, plus a non-blank + ``message`` — or ``None`` when nothing usable was supplied. The strict + contract lives in :meth:`_validate_payload`; this never raises and is + deliberately generic (each end tool only ever supplies its own keys). + """ + out: Dict[str, Any] = {} + doc_ids = kwargs.get("doc_ids") + if isinstance(doc_ids, list) and doc_ids and all(isinstance(i, str) for i in doc_ids): + out["doc_ids"] = list(doc_ids) + answer = kwargs.get("answer") + if isinstance(answer, str) and answer.strip(): + out["answer"] = answer + # Only worth salvaging when there is a primary result; a bare + # ``message`` with no doc_ids/answer carries nothing usable. + if not out: + return None + message = kwargs.get("message") + if isinstance(message, str) and message.strip(): + out["message"] = message + return out + + def _call(self, **kwargs: Any) -> str: + self._validate_payload(**kwargs) + return self.success_message + + async def _acall(self, **kwargs: Any) -> str: + return self._call(**kwargs) + + +class FinalResults(BaseEndTool): + """Tool for logging selected document IDs and signaling the end of the interaction.""" + + def __init__(self, top_k: Optional[int] = None, include_msg: bool = True): + self.top_k = top_k + self.include_msg = include_msg + + tk_ins = "" + if top_k is not None: + tk_ins = f"- You must choose exactly {top_k} document IDs when calling this function." + + required: List[str] = [] + properties: Dict[str, Any] = {} + + desc = """Signals the completion of the search process for the current query. + +Use this tool when: +- You have found all the relevant documents to the query. +- Despite several attempts, you cannot find good documents for the given query. + +""" + if self.include_msg: + desc += """The message should include: +- A brief summary of your exploration and the results +- Explanation if the search was unsuccessful + +""" + desc += f"""When reporting the selected document IDs, make sure: +- the list of document IDs is sorted in the decreasing level of relevance to the query. I.e., the first document in the list is the most relevant to the query, the second is the second most relevant to the query, and so on. +{tk_ins} + +The successful_search field should be set to true if you believed you have found the most relevant documents to the user's query, and false otherwise. And partial if it is in between.""" + + if self.include_msg: + required.append("message") + properties["message"] = { + "type": "string", + "description": "A message for the user to explain why you think you found all the related documents and there is no related document is missing. Also, include a short description of your exploration process. If your attempts to find related documents were unsuccessful, explain why.", + } + required.append("doc_ids") + required.append("search_successful") + properties["doc_ids"] = { + "type": "array", + "items": {"type": "string"}, + "description": "List of document IDs that are relevant to the user's query sorted descending by their level of relevance to the user's query. I.e., the first document is the most relevant to the query, the second is the second most relevant to the query, and so on.", + } + properties["search_successful"] = { + "type": "string", + "enum": ["true", "false", "partial"], + "description": "Whether you managed to find all the related documents to the query.", + } + + self.spec_dict = { + "type": "function", + "function": { + "name": "final_results", + "description": desc, + "parameters": { + "type": "object", + "required": required, + "properties": properties, + }, + }, + } + + def _spec(self) -> dict: + return self.spec_dict + + def _validate_payload( + self, doc_ids: List[str], search_successful: str, message: Optional[str] = None + ) -> Dict[str, Any]: + if self.include_msg: + if message is None: + raise TypeError("The `message` argument is required.") + if not isinstance(message, str): + raise TypeError(f"The `message` argument must be a string. Got `{type(message)}` type.") + if not isinstance(doc_ids, list): + raise TypeError(f"The `doc_ids` argument must be a list. Got `{type(doc_ids)}` type.") + if len(doc_ids) == 0: + raise ToolError("`doc_ids` cannot be empty. You must choose at least one relevant document.") + if not all(isinstance(i, str) for i in doc_ids): + raise TypeError("Items in `doc_ids` must be of type string (i.e., python's `str` type).") + if not isinstance(search_successful, str): + raise TypeError(f"The `search_successful` argument must be a string. Got `{type(search_successful)}` type.") + if search_successful not in ["true", "false", "partial"]: + raise ToolError( + f"`search_successful` must be one of `true`, `false`, or `partial`. Got `{search_successful}` instead." + ) + if self.top_k is not None and len(doc_ids) != self.top_k: + raise ToolError( + f"`doc_ids` must contain exactly {self.top_k} documents. But got {len(doc_ids)} document IDs instead." + ) + payload: Dict[str, Any] = {"doc_ids": list(doc_ids), "search_successful": search_successful} + if message is not None: + payload["message"] = message + return payload + + +class LogSelectedDocs(BaseEndTool): + """Tool for reporting the selected document IDs and ending a selection run. + + This is the end tool of the :class:`~nemo_agent.selection_agent.SelectionAgent`. + It is constructed per run, not per agent: the required count (``top_k``) + and the allowed ids (``candidate_docids``) depend on the run's candidate + documents, and per-run construction keeps concurrent runs on one agent + instance isolated. + """ + + def __init__(self, top_k: int, candidate_docids: List[str], include_msg: bool = True): + self.top_k = int(top_k) + self.allowed_doc_ids = {str(i) for i in candidate_docids} + self.include_msg = include_msg + + desc = ( + "Records the selected documents and signals the end of the task.\n\n" + "Use this tool when you have carefully considered the candidate " + f"documents and have selected exactly {self.top_k} of the most " + "relevant documents to the query.\n\n" + ) + if self.include_msg: + desc += ( + "The message argument should explain your reasoning and " + "justification for selecting this specific set of documents as " + "the most relevant to the query.\n\n" + ) + desc += ( + "**Note**: the list of document IDs passed as the `doc_ids` " + "argument must be sorted in the decreasing level of relevance. " + "In other words, the first document in `doc_ids` list is the " + "most relevant to the query, the second document is the second " + "most relevant document, and so on." + ) + + required: List[str] = [] + properties: Dict[str, Any] = {} + if self.include_msg: + required.append("message") + properties["message"] = { + "type": "string", + "description": "A message for the user to explain why you think the selected are the most relevant to the query. Also, explain why this specific order of document IDs satisfies the most to least relevant ordering criteria.", + } + required.append("doc_ids") + properties["doc_ids"] = { + "type": "array", + "items": {"type": "string"}, + "description": ( + f"The IDs of the {self.top_k} most relevant documents to the given query. " + "The IDs must be sorted in the decreasing " + "level of relevance. I.e., the first document must be the " + "most relevant to the query." + ), + } + + self.spec_dict = { + "type": "function", + "function": { + "name": "log_selected_documents", + "description": desc, + "parameters": { + "type": "object", + "required": required, + "properties": properties, + }, + }, + } + + def _spec(self) -> dict: + return self.spec_dict + + def _validate_payload(self, doc_ids: List[str], message: Optional[str] = None) -> Dict[str, Any]: + if self.include_msg: + if message is None: + raise TypeError("The `message` argument is required.") + if not isinstance(message, str): + raise TypeError(f"The `message` argument must be a string. Got `{type(message)}` type.") + if not isinstance(doc_ids, list): + raise TypeError(f"The `doc_ids` argument must be a list. Got `{type(doc_ids)}` type.") + if len(doc_ids) != self.top_k: + raise ToolError(f"You must select exactly {self.top_k} documents. Got {len(doc_ids)} documents.") + if not all(isinstance(i, str) for i in doc_ids): + raise TypeError("Items in `doc_ids` must be of type string (i.e., python's `str` type).") + for i in doc_ids: + if i not in self.allowed_doc_ids: + raise ToolError(f"Document with ID `{i}` is not among the candidate documents.") + payload: Dict[str, Any] = {"doc_ids": list(doc_ids)} + if message is not None: + payload["message"] = message + return payload diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/retrieve.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/retrieve.py new file mode 100644 index 0000000000..fbde635a0d --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/retrieve.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Retrieve tools: the base contract plus ready-made function-backed tools. + +- :class:`BaseRetrieveTool` — subclass for full control; the agent detects + retrieve tools by isinstance and applies its over-fetch/dedup/exclusion/ + logging machinery around :meth:`~BaseRetrieveTool.acall`. +- :class:`RetrieveTool` — wraps a plain ``retriever_fn(query, top_k)``. +- :class:`ReasoningAugmentedRetrieveTool` — same, but prepends the run's + original question and/or the LLM's latest reasoning trace to the query. +- :func:`create_retrieve_tool` — picks one of the two classes by name. +""" + +from __future__ import annotations + +import asyncio +import inspect +import json +from abc import abstractmethod +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Union + +from .base_tool import BaseTool, ToolContractError, ToolError, tool_error_text + + +@dataclass(frozen=True) +class RetrieveContext: + """Per-call context the agent hands to every retrieve tool. + + ``global_query`` is the run's original question. ``reasoning`` is the most + recent turn's reasoning trace (``None`` when the turn exposed none; the + agent overwrites it every step, so it is never stale). + """ + + global_query: str + reasoning: Optional[str] = None + + +class BaseRetrieveTool(BaseTool): + """A tool that retrieves documents. + + Implement :meth:`_spec` and :meth:`_acall` (async-only; the standard tool + hook, here with the fixed signature ``(query, top_k, context)``). ``top_k`` + is the number of documents the caller needs back (the agent's value + already accounts for over-fetching); ``context`` carries the run's global + query and latest reasoning — ignore it if unused. + + Result contract (validated here; violations raise + :class:`ToolContractError` naming your tool): a list of dicts, each with + ``id`` (coerced to str), ``score`` (coerced to float), and ``text`` + (coerced to str; may be empty); ``image`` optional; extra keys are shown + to the LLM verbatim. + Raise :class:`ToolError` for LLM-recoverable failures; let genuine + failures propagate. Tools may be called from concurrent runs — keep them + stateless. + """ + + default_top_k: int = 20 + + @abstractmethod + async def _acall(self, query: str, top_k: int, context: RetrieveContext) -> List[Dict[str, Any]]: + raise NotImplementedError + + async def acall( + self, query: str, top_k: Optional[int] = None, context: Optional[RetrieveContext] = None + ) -> Union[List[Dict[str, Any]], str]: + """Retrieve and validate — do not override; implement :meth:`_acall`. + + Narrows :meth:`BaseTool.acall` to the retrieve signature. ``top_k`` + defaults to ``default_top_k``; ``context`` defaults to an empty + context built from the query. Unlike the generic wrapper, only + :class:`ToolError` becomes LLM-visible error text — the agent + validates LLM-supplied arguments before calling, so a ``TypeError`` + here is a tool bug and propagates. + """ + if top_k is None: + top_k = self.default_top_k + if context is None: + context = RetrieveContext(global_query=str(query)) + try: + result = await self._acall(query, top_k, context) + except ToolError as e: + return tool_error_text(self.name, e) + return self._validate_result(result) + + def _validate_result(self, result: Any) -> List[Dict[str, Any]]: + cls_name = type(self).__name__ + if not isinstance(result, list): + raise ToolContractError( + f"{cls_name}: retrieve result must be a list of document dicts, " f"got {type(result).__name__}." + ) + out: List[Dict[str, Any]] = [] + for i, item in enumerate(result): + if not isinstance(item, dict): + raise ToolContractError(f"{cls_name}: result item {i} must be a dict, got {type(item).__name__}.") + missing = [k for k in ("id", "score", "text") if k not in item] + if missing: + raise ToolContractError( + f"{cls_name}: result item {i} is missing required key(s) {missing}. " + "Each document must carry 'id', 'score', and 'text'." + ) + doc = dict(item) + doc["id"] = str(doc["id"]) + try: + doc["score"] = float(doc["score"]) + except (TypeError, ValueError): + raise ToolContractError( + f"{cls_name}: result item {i} has a non-numeric 'score': {item['score']!r}." + ) from None + if not isinstance(doc["text"], str): + doc["text"] = "" if doc["text"] is None else str(doc["text"]) + out.append(doc) + return out + + +class RetrieveTool(BaseRetrieveTool): + """Retrieve tool backed by a plain ``retriever_fn(query, top_k)``. + + ``retriever_fn`` may be sync (runs via ``asyncio.to_thread`` — it may + block on network/disk/inference) or async, and is called positionally. + """ + + def __init__( + self, + retriever_fn: Callable[..., Any], + *, + name: str = "retrieve", + description: str = "Search for documents related to a query using dense retrieval.", + default_top_k: int = 20, + ) -> None: + if not callable(retriever_fn): + raise TypeError(f"retriever_fn must be callable, got {type(retriever_fn).__name__}.") + self._fn = retriever_fn + self._fn_is_async = inspect.iscoroutinefunction(retriever_fn) + self._name = str(name) + self._description = str(description) + self.default_top_k = int(default_top_k) + + def _spec(self) -> Dict[str, Any]: + return { + "type": "function", + "function": { + "name": self._name, + "description": self._description, + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query.", + }, + "top_k": { + "type": "integer", + "description": "Number of documents to retrieve.", + "default": self.default_top_k, + }, + }, + "required": ["query"], + }, + }, + } + + async def _acall(self, query: str, top_k: int, context: RetrieveContext) -> List[Dict[str, Any]]: + return await self._call_retriever(query, top_k) + + async def _call_retriever(self, query: str, top_k: int) -> Any: + if self._fn_is_async: + return await self._fn(query, top_k) + return await asyncio.to_thread(self._fn, query, top_k) + + +class ReasoningAugmentedRetrieveTool(RetrieveTool): + """RetrieveTool that adds run context to the query before retrieving. + + Composes the query body from ``body_template`` (or the ``context_mode`` + default) with the run's original question (``{global_query}``) and the + LLM's latest reasoning trace (``{reasoning}``); the retriever function + stays a plain ``(query, top_k)``. When a turn exposed no reasoning, + ``empty_reasoning_policy`` either substitutes a placeholder + (``"substitute"``, default) or sends the raw query unchanged (``"skip"``). + """ + + BODY_TEMPLATES: Dict[str, str] = { + "reasoning": "Reasoning: {reasoning}\n\nQuery: {query}", + "global_query": "Original question: {global_query}\n\nQuery: {query}", + "both": "Original question: {global_query}\nReasoning: {reasoning}\n\nQuery: {query}", + } + EMPTY_REASONING_PLACEHOLDER = "Empty" + + def __init__( + self, + retriever_fn: Callable[..., Any], + *, + context_mode: str = "both", + body_template: Optional[str] = None, + empty_reasoning_policy: str = "substitute", + **kwargs: Any, + ) -> None: + super().__init__(retriever_fn, **kwargs) + if body_template is None: + if context_mode not in self.BODY_TEMPLATES: + raise ValueError(f"context_mode must be one of {sorted(self.BODY_TEMPLATES)}, got {context_mode!r}.") + body_template = self.BODY_TEMPLATES[context_mode] + try: + body_template.format(query="", reasoning="", global_query="") + except (KeyError, IndexError) as e: + raise ValueError( + f"body_template has an unsupported placeholder ({e!s}); " + "supported placeholders are {query}, {reasoning}, {global_query}." + ) from e + if empty_reasoning_policy not in ("substitute", "skip"): + raise ValueError(f"empty_reasoning_policy must be 'substitute' or 'skip', got {empty_reasoning_policy!r}.") + self._body_template = body_template + self._empty_reasoning_policy = empty_reasoning_policy + + async def _acall(self, query: str, top_k: int, context: RetrieveContext) -> List[Dict[str, Any]]: + return await self._call_retriever(self._build_query(query, context), top_k) + + def _build_query(self, query: str, context: RetrieveContext) -> str: + reasoning = (context.reasoning or "").strip() + if "{reasoning}" in self._body_template and not reasoning: + if self._empty_reasoning_policy == "skip": + return str(query) + reasoning = self.EMPTY_REASONING_PLACEHOLDER + return self._body_template.format( + query=str(query), + reasoning=reasoning, + global_query=(context.global_query or "").strip(), + ) + + +def create_retrieve_tool(kind: str, retriever_fn: Callable[..., Any], **kwargs: Any) -> BaseRetrieveTool: + """Create a retrieve tool from a plain retriever function. + + ``kind="default"`` builds a :class:`RetrieveTool`; + ``kind="reasoning_augmented"`` builds a + :class:`ReasoningAugmentedRetrieveTool`. ``kwargs`` go to the class + constructor. For anything fancier (e.g. reading ``context`` raw), subclass + :class:`RetrieveTool` (or :class:`BaseRetrieveTool`) and override ``_acall``. + """ + if kind == "default": + return RetrieveTool(retriever_fn, **kwargs) + if kind == "reasoning_augmented": + return ReasoningAugmentedRetrieveTool(retriever_fn, **kwargs) + raise ValueError(f"kind must be 'default' or 'reasoning_augmented', got {kind!r}.") + + +def retrieve_output_to_msg_content(output: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Convert a validated retrieve output list into LLM message content blocks.""" + content_list: List[Dict[str, Any]] = [] + for doc_in in output: + doc = {**doc_in} + if str(doc.get("text") or "").strip() == "": + doc.pop("text", None) + img = doc.pop("image", None) + content_list.append({"type": "text", "text": json.dumps(doc)}) + if img is not None: + content_list.append({"type": "image_url", "image_url": {"url": img}}) + return content_list diff --git a/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/think_tool.py b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/think_tool.py new file mode 100644 index 0000000000..16c156f369 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/tools/think_tool.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Scratchpad tool that lets the LLM think with output tokens. + +:class:`ThinkTool` carries retrieval-flavored use-case bullets (the main +agent's version); :class:`SelectionThinkTool` swaps in selection-flavored +bullets. Both descriptions are LLM-visible; only the bullet list differs. +""" + +from __future__ import annotations + +from typing import List + +from .base_tool import BaseTool + + +class ThinkTool(BaseTool): + """Tool that allows the LLM to think with output tokens.""" + + _use_case_lines: List[str] = [ + "- When processing a complex query, use this tool to organize your thoughts and think about the sub queries that you need to search for to find the relevant information", # noqa: E501 + "- If a query is vague is very difficult to find information for it, you can use this tool to think about clues in the query that you can use to narrow down the search and spot relevant pieces of information.", # noqa: E501 + "- When finding related documents that help you create better search queries in the next step, use this tool to think about what pieces of information from these documents are helpful to search for.", # noqa: E501 + "- When you fail to find any related information to the query, use this tool to think about other search strategies that you can take to retrieve the related documents", # noqa: E501 + ] + + def __init__(self, extended_relevance: bool = False): + if extended_relevance: + ext_lines = [ + "- When it is difficult to understand what is the intent of the user and what they are trying to find with this query, use this tool to think about potential definitions of relevance that could be meaningful/useful to the user for this task.", # noqa: E501 + "- If the intention of the user is vague especially given the available documents, use this tool to think how you should decide what documents are relevant and what the metric of relevance is.", # noqa: E501 + ] + ext = "\n".join(ext_lines) + "\n" + else: + ext = "" + description = ( + "Use the tool to think about something. It will not obtain new information or make any changes, " + "but just log the thought. Use it when complex reasoning or brainstorming is needed.\n" + "\n" + "Common use cases:\n" + f"{ext}" + "\n".join(self._use_case_lines) + "\n" + "\n" + "The tool simply logs your thought process for better transparency and does not make any changes." + ) + self.spec_dict = { + "type": "function", + "function": { + "name": "think", + "description": description, + "parameters": { + "type": "object", + "properties": { + "thought": { + "type": "string", + "description": "The thought to log.", + } + }, + "required": ["thought"], + }, + }, + } + + def _spec(self) -> dict: + return self.spec_dict + + def _call(self, thought: str) -> str: + return "Your thought has been logged." + + async def _acall(self, thought: str) -> str: + return "Your thought has been logged." + + +class SelectionThinkTool(ThinkTool): + """The selection agent's think tool: same contract, selection-flavored bullets.""" + + _use_case_lines: List[str] = [ + "- When processing a complex query, use this tool to organize your thoughts and think about how each document might be related to the given query.", # noqa: E501 + "- If a query is vague or hard to understand, you can use this tool to think about clues in the query that help you identify the connections between a document and the query.", # noqa: E501 + "- You can use this tool to think what pieces of information in each document are the most important or relevant for the given query.", # noqa: E501 + ] diff --git a/nemo_retriever/src/nemo_retriever/cli/ingest/options.py b/nemo_retriever/src/nemo_retriever/cli/ingest/options.py index da6d8dc0bc..8aae3b5ec8 100644 --- a/nemo_retriever/src/nemo_retriever/cli/ingest/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/ingest/options.py @@ -241,6 +241,7 @@ str | None, typer.Option( "--embed-invoke-url", + envvar="EMBED_INVOKE_URL", help=( "Embedding endpoint override. On CPU-only hosts, ingest automatically uses NVIDIA's hosted " "embedding endpoint with NVIDIA_API_KEY or NGC_API_KEY; pass this only for another endpoint." @@ -251,6 +252,7 @@ str | None, typer.Option( "--embed-model-name", + envvar="EMBED_MODEL_NAME", help=f"Optional embedding model name override. Defaults to {DEFAULT_EMBED_MODEL} when omitted.", ), ] @@ -258,6 +260,7 @@ str | None, typer.Option( "--embed-model-provider-prefix", + envvar="EMBED_MODEL_PROVIDER_PREFIX", help="Optional LiteLLM provider prefix prepended to the remote embedding model name.", ), ] diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 23af798552..b41c4759ef 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -16,7 +16,10 @@ from nemo_retriever.cli.query import options as opts from nemo_retriever.cli.query_workflow import agentic_query_documents as query_agentic_documents from nemo_retriever.cli.query_workflow import query_documents_with_metadata as query_local_documents_with_metadata -from nemo_retriever.query.agentic_options import agentic_backend_top_k_error, agentic_temperature_error +from nemo_retriever.query.agentic_options import ( + agentic_llm_client_error, + agentic_temperature_error, +) from nemo_retriever.cli.shared import ( ROOT_CLI_ERRORS, quiet_capture, @@ -158,7 +161,8 @@ def _retrieval_options( hidden=True, help=( "Query a LanceDB index produced by local or batch ingest; retrieval mode auto-detects the index.\n\n" - f"Default embedding model: {opts.DEFAULT_EMBED_MODEL}.\n\n" + "Embedding model: read from the selected table when available; " + f"legacy tables fall back to {opts.DEFAULT_EMBED_MODEL}.\n\n" f"Default local reranker model when reranking: {opts.DEFAULT_RERANK_MODEL}.\n\n" "For a service deployment, use retriever query service --help." ), @@ -186,19 +190,19 @@ def _local_command( agentic_llm_model: opts.AgenticLlmModelOption = None, agentic_invoke_url: opts.AgenticInvokeUrlOption = None, agentic_reasoning_effort: opts.AgenticReasoningEffortOption = "high", - agentic_backend_top_k: opts.AgenticBackendTopKOption = 20, agentic_react_max_steps: opts.AgenticReactMaxStepsOption = 50, agentic_text_truncation: opts.AgenticTextTruncationOption = 0, - agentic_temperature: opts.AgenticTemperatureOption = 0.0, + agentic_temperature: opts.AgenticTemperatureOption = None, + agentic_llm_client: opts.AgenticLlmClientOption = None, ) -> None: _validate_output_options(output_format, max_text_chars) if reranker_invoke_url is None: reranker_invoke_url = os.environ.get("RERANKER_INVOKE_URL") or None - if embed_invoke_url is None: - embed_invoke_url = os.environ.get("EMBED_INVOKE_URL") or None rerank = rerank or bool(reranker_invoke_url) or bool(reranker_model_name) or bool(reranker_backend) silence_noisy_libraries() if agentic: + # Relaxed model gating: an explicit model is required only for the remote + # (invoke_url) path; in-process runs default to the local model. if agentic_invoke_url and not agentic_llm_model: typer.echo( "Error: --agentic-invoke-url requires --agentic-llm-model.", @@ -208,15 +212,31 @@ def _local_command( if not agentic_invoke_url and not agentic_llm_model: agentic_llm_model = "nemotron-8b" - backend_error = agentic_backend_top_k_error(agentic_backend_top_k, target_top_k=top_k) - if backend_error: - typer.echo(f"Error: {backend_error}", err=True) - raise typer.Exit(1) - temperature_invoke_url = agentic_invoke_url or "local://in-process" - temperature_error = agentic_temperature_error(agentic_temperature, invoke_url=temperature_invoke_url) - if temperature_error: - typer.echo(f"Error: {temperature_error}", err=True) - raise typer.Exit(1) + if agentic_temperature is not None: + # Use a local sentinel so in-process runs get the in-process bound. + temperature_invoke_url = agentic_invoke_url or "local://in-process" + temperature_error = agentic_temperature_error(agentic_temperature, invoke_url=temperature_invoke_url) + if temperature_error: + typer.echo(f"Error: {temperature_error}", err=True) + raise typer.Exit(1) + + # Fast-fail with a flag-named message; AgenticRetrievalConfig.__post_init__ + # is the authoritative resolver. `callable` drives either an in-process + # engine or the shared HTTP client, so it is valid with and without an + # invoke_url; every other client is remote-only. + if agentic_llm_client is not None: + client_error = agentic_llm_client_error(agentic_llm_client, field_name="agentic_llm_client") + if client_error: + typer.echo(f"Error: {client_error}", err=True) + raise typer.Exit(1) + client = agentic_llm_client.strip().lower() + if not agentic_invoke_url and client != "callable": + typer.echo( + "Error: a remote LLM client requires --agentic-invoke-url; " + "omit --agentic-llm-client to run the local in-process model.", + err=True, + ) + raise typer.Exit(1) try: reranker_api_key = _api_key_from_env_option(reranker_api_key_env) if reranker_invoke_url else None @@ -253,10 +273,10 @@ def _local_command( llm_model=agentic_llm_model, invoke_url=agentic_invoke_url, reasoning_effort=agentic_reasoning_effort, - backend_top_k=agentic_backend_top_k, react_max_steps=agentic_react_max_steps, text_truncation=agentic_text_truncation, temperature=agentic_temperature, + llm_client=agentic_llm_client, ), ) with quiet_capture(): diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index b47d1f3646..3fabce18d3 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -8,11 +8,16 @@ import typer +from nemo_retriever._agentic.nemo_agent.llm import get_available_backends from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL DEFAULT_EMBED_MODEL = VL_EMBED_MODEL DEFAULT_RERANK_MODEL = VL_RERANK_MODEL +# Advertised in --agentic-llm-client help; sourced from the registry so a newly +# registered client shows up without editing this string. +_AGENTIC_LLM_CLIENT_CHOICES = ", ".join(get_available_backends()) + QueryArgument = Annotated[str, typer.Argument(..., help="Query text.")] TopKOption = Annotated[ @@ -63,19 +68,24 @@ ] EmbedInvokeUrlOption = Annotated[ str | None, - typer.Option("--embed-invoke-url", help="Embedding NIM endpoint URL."), + typer.Option("--embed-invoke-url", envvar="EMBED_INVOKE_URL", help="Embedding NIM endpoint URL."), ] EmbedModelNameOption = Annotated[ str | None, typer.Option( "--embed-model-name", - help=f"Optional embedding model name override. Defaults to {DEFAULT_EMBED_MODEL} when omitted.", + envvar="EMBED_MODEL_NAME", + help=( + "Embedding model override. When omitted, use the model recorded on the selected table, " + f"then fall back to {DEFAULT_EMBED_MODEL} for a legacy table without metadata." + ), ), ] EmbedModelProviderPrefixOption = Annotated[ str | None, typer.Option( "--embed-model-provider-prefix", + envvar="EMBED_MODEL_PROVIDER_PREFIX", help="Optional LiteLLM provider prefix prepended to the remote embedding model name.", ), ] @@ -179,14 +189,6 @@ help="reasoning_effort forwarded on agentic LLM calls.", ), ] -AgenticBackendTopKOption = Annotated[ - int, - typer.Option( - "--agentic-backend-top-k", - min=1, - help="Backend retrieve-pool depth per agentic retrieval call.", - ), -] AgenticReactMaxStepsOption = Annotated[ int, typer.Option( @@ -204,11 +206,26 @@ ), ] AgenticTemperatureOption = Annotated[ - float, + float | None, typer.Option( "--agentic-temperature", min=0.0, - help="Sampling temperature for agentic LLM calls (0.0 = greedy).", + help=( + "Sampling temperature for agentic LLM calls. " + "Omit to leave it unset (endpoint/model default; 0.0 = greedy)." + ), + ), +] +AgenticLlmClientOption = Annotated[ + str | None, + typer.Option( + "--agentic-llm-client", + help=( + "LLM client that builds the agent LLM in agentic mode. Optional: defaults to " + "'callable' for both in-process local runs and remote (--agentic-invoke-url) runs. " + f"Registered clients: {_AGENTIC_LLM_CLIENT_CHOICES}. Any client other than 'callable' " + "is remote-only and requires --agentic-invoke-url." + ), ), ] ServiceUrlOption = Annotated[ diff --git a/nemo_retriever/src/nemo_retriever/common/api/util/converters/datetools.py b/nemo_retriever/src/nemo_retriever/common/api/util/converters/datetools.py index bff3215af2..7e515e9ff1 100644 --- a/nemo_retriever/src/nemo_retriever/common/api/util/converters/datetools.py +++ b/nemo_retriever/src/nemo_retriever/common/api/util/converters/datetools.py @@ -11,6 +11,43 @@ from nemo_retriever.common.api.util.exception_handlers.converters import datetools_exception_handler +def normalize_timezone_aware_iso8601_to_utc(date_string: str) -> str: + """Normalize a timezone-aware ISO-8601 timestamp to UTC. + + Parameters + ---------- + date_string : str + An ISO-8601 timestamp containing ``Z`` or an explicit UTC offset. + + Returns + ------- + str + The same instant represented as an ISO-8601 timestamp in UTC. + + Raises + ------ + ValueError + If ``date_string`` is invalid or does not include timezone information. + + Notes + ----- + This helper deliberately rejects naive timestamps instead of assuming the + host timezone. That keeps persisted lifecycle and scheduling timestamps + consistent across services running in different regions. + """ + + if not isinstance(date_string, str): + raise ValueError("timestamp must be ISO-8601") + + try: + parsed = datetime.fromisoformat(date_string.replace("Z", "+00:00")) + except ValueError as exc: + raise ValueError("timestamp must be ISO-8601") from exc + if parsed.tzinfo is None or parsed.utcoffset() is None: + raise ValueError("timestamp must include a timezone offset") + return parsed.astimezone(timezone.utc).isoformat() + + @datetools_exception_handler def datetimefrompdfmeta(pdf_formated_date: str, keep_tz: bool = False) -> str: """ diff --git a/nemo_retriever/src/nemo_retriever/common/inline_text.py b/nemo_retriever/src/nemo_retriever/common/inline_text.py new file mode 100644 index 0000000000..ab957f8437 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/common/inline_text.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared validation and logical identities for inline text sources.""" + +from __future__ import annotations + +from typing import Sequence + +INLINE_TEXT_SOURCE_PREFIX = "inline://" + + +def normalize_inline_texts(texts: str | Sequence[str]) -> list[str]: + """Copy and validate raw inline documents.""" + if isinstance(texts, str): + values = [texts] + elif isinstance(texts, Sequence): + values = list(texts) + else: + raise TypeError(f"texts must be a string or sequence of strings, got {type(texts).__name__}") + + for index, value in enumerate(values): + if not isinstance(value, str): + raise TypeError(f"texts[{index}] must be a string, got {type(value).__name__}") + return values + + +def inline_text_source_id(index: int) -> str: + """Return the deterministic ordinal identity for an inline document.""" + return f"{INLINE_TEXT_SOURCE_PREFIX}{index:08d}" + + +def is_inline_text_source(source_id: str) -> bool: + """Return whether a source identifier denotes inline text.""" + return source_id.startswith(INLINE_TEXT_SOURCE_PREFIX) + + +def is_blank_inline_corpus(texts: Sequence[str] | None) -> bool: + """Return whether a configured inline corpus contains only blank text.""" + return texts is not None and not any(text.strip() for text in texts) diff --git a/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py b/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py index 9295e122f2..81f2e0951b 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py @@ -12,6 +12,7 @@ by PDFium in the PDF extraction stage. """ +from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Sequence, Tuple import base64 @@ -637,6 +638,369 @@ def _find_ts_detections_for_bbox( return None +@dataclass +class _OCRRowResult: + """Mutable OCR output accumulated for one source page row.""" + + table_items: List[Dict[str, Any]] = field(default_factory=list) + chart_items: List[Dict[str, Any]] = field(default_factory=list) + infographic_items: List[Dict[str, Any]] = field(default_factory=list) + text_blocks: List[Dict[str, Any]] = field(default_factory=list) + error: Any = None + + +@dataclass(frozen=True) +class _PreparedOCRRow: + """Validated source row and the labels that should be cropped from it.""" + + row_index: int + row: Any + page_image_b64: str + detections: List[Dict[str, Any]] + wanted_labels: set[str] + + +@dataclass(frozen=True) +class _OCRCropJob: + """One local crop plus the address needed to stitch its result.""" + + row_index: int + row: Any + label_name: str + bbox: List[float] + crop_array: np.ndarray + + +def _record_ocr_error(row_result: _OCRRowResult, exc: BaseException) -> None: + print(f"Warning: OCR failed: {type(exc).__name__}: {exc}") + row_result.error = { + "stage": "ocr_page_elements", + "type": exc.__class__.__name__, + "message": str(exc), + "traceback": "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)), + } + + +def _prepare_ocr_rows( + batch_df: pd.DataFrame, + *, + wanted_labels: set[str], + extract_text: bool, +) -> Tuple[List[_PreparedOCRRow], List[_OCRRowResult]]: + """Validate page inputs and retain one result slot per source row.""" + + prepared_rows: List[_PreparedOCRRow] = [] + row_results = [_OCRRowResult() for _ in range(len(batch_df.index))] + + for row_index, row in enumerate(batch_df.itertuples(index=False)): + row_result = row_results[row_index] + try: + page_elements = getattr(row, "page_elements_v3", None) + detections: List[Dict[str, Any]] = [] + if isinstance(page_elements, dict): + detections = page_elements.get("detections") or [] + if not isinstance(detections, list): + detections = [] + + page_image = getattr(row, "page_image", None) or {} + page_image_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None + if not isinstance(page_image_b64, str) or not page_image_b64: + metadata = getattr(row, "metadata", None) or {} + upstream_error = metadata.get("error") if isinstance(metadata, dict) else None + page_num = getattr(row, "page_number", "?") + path = getattr(row, "path", "?") + if upstream_error: + _logger.warning( + "OCR skipping page %s of %s — no page image (upstream error: %s)", + page_num, + path, + upstream_error, + ) + else: + _logger.debug( + "OCR skipping page %s of %s — no page image (text-only or raster not requested)", + page_num, + path, + ) + row_result.error = upstream_error + continue + + row_wanted_labels = wanted_labels + if extract_text: + metadata = getattr(row, "metadata", None) or {} + needs_ocr = metadata.get("needs_ocr_for_text", False) if isinstance(metadata, dict) else False + if needs_ocr: + row_wanted_labels = wanted_labels | _TEXT_LABELS + + prepared_rows.append( + _PreparedOCRRow( + row_index=row_index, + row=row, + page_image_b64=page_image_b64, + detections=detections, + wanted_labels=row_wanted_labels, + ) + ) + except BaseException as exc: + _record_ocr_error(row_result, exc) + + return prepared_rows, row_results + + +def _append_ocr_prediction( + row_result: _OCRRowResult, + *, + row: Any, + label_name: str, + bbox: List[float], + preds: Any, + crop_hw: Tuple[int, int], + use_table_structure: bool, +) -> None: + """Parse one prediction and stitch it into its source-row result.""" + + blocks = _parse_ocr_result(preds) + if label_name == "table": + text = "" + if use_table_structure: + ts_match = _find_ts_detections_for_bbox(row, bbox) + if ts_match is not None: + ts_dets, ts_hw = ts_match + text = join_table_structure_and_ocr_output(ts_dets, preds, ts_hw or crop_hw) + if not text: + text = _blocks_to_pseudo_markdown(blocks, crop_hw=crop_hw) + if not text: + text = _blocks_to_text(blocks) + else: + text = _blocks_to_text(blocks) + + entry = {"bbox_xyxy_norm": bbox, "text": text} + if label_name == "table": + row_result.table_items.append(entry) + elif label_name == "chart": + row_result.chart_items.append(entry) + elif label_name == "infographic": + row_result.infographic_items.append(entry) + elif label_name in _TEXT_LABELS: + row_result.text_blocks.extend(blocks) + + +def _remote_crop_shape(crop_b64: str) -> Tuple[int, int]: + """Return ``(height, width)`` for a remote crop when it can be decoded.""" + + try: + raw = base64.b64decode(crop_b64) + with Image.open(io.BytesIO(raw)) as crop_image: + width, height = crop_image.size + return (height, width) + except Exception: + return (0, 0) + + +def _run_remote_ocr( + prepared_rows: List[_PreparedOCRRow], + row_results: List[_OCRRowResult], + *, + invoke_url: str, + api_key: Optional[str], + request_timeout_s: float, + max_batch_size: int, + retry: RemoteRetryParams, + nim_client: NIMClient | None, + use_table_structure: bool, +) -> None: + """Invoke the existing per-page remote OCR path and stitch its results.""" + + for prepared in prepared_rows: + row_result = row_results[prepared.row_index] + try: + crops = _crop_all_from_page( + prepared.page_image_b64, + prepared.detections, + prepared.wanted_labels, + as_b64=True, + ) + crop_b64s: List[str] = [crop_b64 for _label, _bbox, crop_b64 in crops] + crop_metadata: List[Tuple[str, List[float]]] = [(label_name, bbox) for label_name, bbox, _crop_b64 in crops] + if not crop_b64s: + continue + + invoke_kwargs = dict( + invoke_url=invoke_url, + image_b64_list=crop_b64s, + api_key=api_key, + timeout_s=float(request_timeout_s), + max_batch_size=max_batch_size, + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + if nim_client is not None: + response_items = nim_client.invoke_image_inference_batches(**invoke_kwargs) + else: + response_items = invoke_image_inference_batches( + **invoke_kwargs, + max_pool_workers=int(retry.remote_max_pool_workers), + ) + if len(response_items) != len(crop_metadata): + raise RuntimeError(f"Expected {len(crop_metadata)} OCR responses, got {len(response_items)}") + + for index, (label_name, bbox) in enumerate(crop_metadata): + preds = _extract_remote_ocr_item(response_items[index]) + crop_hw = _remote_crop_shape(crop_b64s[index]) if label_name == "table" else (0, 0) + _append_ocr_prediction( + row_result, + row=prepared.row, + label_name=label_name, + bbox=bbox, + preds=preds, + crop_hw=crop_hw, + use_table_structure=use_table_structure, + ) + except BaseException as exc: + _record_ocr_error(row_result, exc) + + +def _collect_local_crop_jobs( + prepared_rows: List[_PreparedOCRRow], + row_results: List[_OCRRowResult], +) -> Dict[str, List[_OCRCropJob]]: + """Collect compatible local crops across all prepared page rows.""" + + jobs_by_merge_level: Dict[str, List[_OCRCropJob]] = {"word": [], "paragraph": []} + for prepared in prepared_rows: + try: + crops = _crop_all_from_page( + prepared.page_image_b64, + prepared.detections, + prepared.wanted_labels, + ) + for label_name, bbox, crop_array in crops: + merge_level = "word" if label_name == "table" else "paragraph" + jobs_by_merge_level[merge_level].append( + _OCRCropJob( + row_index=prepared.row_index, + row=prepared.row, + label_name=label_name, + bbox=bbox, + crop_array=crop_array, + ) + ) + except BaseException as exc: + _record_ocr_error(row_results[prepared.row_index], exc) + return jobs_by_merge_level + + +def _run_local_ocr_batches( + model: Any, + jobs_by_merge_level: Dict[str, List[_OCRCropJob]], + row_results: List[_OCRRowResult], + *, + batch_size: int, + use_table_structure: bool, +) -> None: + """Run bounded local crop lists, falling back per crop when required.""" + + for merge_level, jobs in jobs_by_merge_level.items(): + for start in range(0, len(jobs), batch_size): + batch_jobs = jobs[start : start + batch_size] + batch_crops = [job.crop_array for job in batch_jobs] + + try: + batch_preds = model.invoke(batch_crops, merge_level=merge_level) + except Exception: + batch_preds = None + + if not isinstance(batch_preds, list) or len(batch_preds) != len(batch_jobs): + for job in batch_jobs: + try: + preds = model.invoke(job.crop_array, merge_level=merge_level) + _append_ocr_prediction( + row_results[job.row_index], + row=job.row, + label_name=job.label_name, + bbox=job.bbox, + preds=preds, + crop_hw=(job.crop_array.shape[0], job.crop_array.shape[1]), + use_table_structure=use_table_structure, + ) + except BaseException as exc: + _record_ocr_error(row_results[job.row_index], exc) + continue + + for job, preds in zip(batch_jobs, batch_preds): + try: + _append_ocr_prediction( + row_results[job.row_index], + row=job.row, + label_name=job.label_name, + bbox=job.bbox, + preds=preds, + crop_hw=(job.crop_array.shape[0], job.crop_array.shape[1]), + use_table_structure=use_table_structure, + ) + except BaseException as exc: + _record_ocr_error(row_results[job.row_index], exc) + + +def _build_ocr_page_elements_output( + batch_df: pd.DataFrame, + row_results: List[_OCRRowResult], + *, + extract_text: bool, + extract_tables: bool, + extract_charts: bool, + extract_infographics: bool, + elapsed_s: float, +) -> pd.DataFrame: + """Finalize per-row metadata and add OCR output columns to the batch.""" + + all_table = [result.table_items for result in row_results] + all_chart = [result.chart_items for result in row_results] + all_infographic = [result.infographic_items for result in row_results] + all_text: List[Optional[str]] = [] + all_ocr_meta: List[Dict[str, Any]] = [] + + for result in row_results: + all_text.append(_blocks_to_text(result.text_blocks) if extract_text and result.text_blocks else None) + + counts_by_label: Dict[str, int] = {} + if result.table_items: + counts_by_label["table"] = len(result.table_items) + if result.chart_items: + counts_by_label["chart"] = len(result.chart_items) + if result.infographic_items: + counts_by_label["infographic"] = len(result.infographic_items) + if result.text_blocks: + counts_by_label["text"] = len(result.text_blocks) + + all_ocr_meta.append( + { + "timing": {"seconds": float(elapsed_s)}, + "error": result.error, + "num_detections": sum(counts_by_label.values()), + "counts_by_label": counts_by_label, + } + ) + + out = batch_df.copy() + if extract_tables or "table" not in out.columns: + out["table"] = all_table + if extract_charts or "chart" not in out.columns: + out["chart"] = all_chart + if extract_infographics or "infographic" not in out.columns: + out["infographic"] = all_infographic + if extract_text and "text" in out.columns: + for index, ocr_text in enumerate(all_text): + if ocr_text is not None: + out.iat[index, out.columns.get_loc("text")] = ocr_text + elif extract_text: + out["text"] = [text if text is not None else "" for text in all_text] + out["ocr"] = all_ocr_meta + out["ocr_v1_num_detections"] = [metadata["num_detections"] for metadata in all_ocr_meta] + out["ocr_v1_counts_by_label"] = [metadata["counts_by_label"] for metadata in all_ocr_meta] + return out + + # --------------------------------------------------------------------------- # Core function # --------------------------------------------------------------------------- @@ -659,11 +1023,6 @@ def ocr_page_elements( nim_client: NIMClient | None = None, **kwargs: Any, ) -> Any: - retry = remote_retry or RemoteRetryParams( - remote_max_pool_workers=int(kwargs.get("remote_max_pool_workers", 16)), - remote_max_retries=int(kwargs.get("remote_max_retries", 10)), - remote_max_429_retries=int(kwargs.get("remote_max_429_retries", 5)), - ) """ Run Nemotron OCR on cropped regions detected by PageElements v3. @@ -688,6 +1047,11 @@ def ocr_page_elements( Original columns plus ``table``, ``chart``, ``infographic``, and ``ocr``. """ + retry = remote_retry or RemoteRetryParams( + remote_max_pool_workers=int(kwargs.get("remote_max_pool_workers", 16)), + remote_max_retries=int(kwargs.get("remote_max_retries", 10)), + remote_max_429_retries=int(kwargs.get("remote_max_429_retries", 5)), + ) if not isinstance(batch_df, pd.DataFrame): raise NotImplementedError("ocr_page_elements currently only supports pandas.DataFrame input.") @@ -706,254 +1070,54 @@ def ocr_page_elements( if extract_infographics: wanted_labels.add("infographic") - # Per-row accumulators. - all_table: List[List[Dict[str, Any]]] = [] - all_chart: List[List[Dict[str, Any]]] = [] - all_infographic: List[List[Dict[str, Any]]] = [] - all_text: List[str] = [] - all_ocr_meta: List[Dict[str, Any]] = [] + # This bounds the outer crop list passed to the persistent model wrapper. + # Nemotron's internal detector_max_batch_size is a separate control. + local_invoke_batch_size = 0 + if not use_remote: + if inference_batch_size is None or inference_batch_size < 1: + raise ValueError(f"inference_batch_size must be set and greater than 0. Value: {inference_batch_size}") + local_invoke_batch_size = int(inference_batch_size) t0_total = time.perf_counter() + prepared_rows, row_results = _prepare_ocr_rows( + batch_df, + wanted_labels=wanted_labels, + extract_text=extract_text, + ) - for row in batch_df.itertuples(index=False): - table_items: List[Dict[str, Any]] = [] - chart_items: List[Dict[str, Any]] = [] - infographic_items: List[Dict[str, Any]] = [] - row_ocr_text_blocks: List[Dict[str, Any]] = [] - row_error: Any = None - - try: - # --- get page elements detections --- - pe = getattr(row, "page_elements_v3", None) - dets: List[Dict[str, Any]] = [] - if isinstance(pe, dict): - dets = pe.get("detections") or [] - if not isinstance(dets, list): - dets = [] - - # --- get page image --- - page_image = getattr(row, "page_image", None) or {} - page_image_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None - - if not isinstance(page_image_b64, str) or not page_image_b64: - meta = getattr(row, "metadata", None) or {} - upstream_err = meta.get("error") if isinstance(meta, dict) else None - page_num = getattr(row, "page_number", "?") - path = getattr(row, "path", "?") - if upstream_err: - _logger.warning( - "OCR skipping page %s of %s — no page image (upstream error: %s)", - page_num, - path, - upstream_err, - ) - else: - _logger.debug( - "OCR skipping page %s of %s — no page image (text-only or raster not requested)", - page_num, - path, - ) - all_table.append(table_items) - all_chart.append(chart_items) - all_infographic.append(infographic_items) - all_text.append(None) - all_ocr_meta.append({"timing": None, "error": upstream_err, "num_detections": 0, "counts_by_label": {}}) - continue - - # --- determine per-row labels (text/title only for pages needing OCR) --- - row_wanted = wanted_labels - if extract_text: - meta = getattr(row, "metadata", None) or {} - needs_ocr = meta.get("needs_ocr_for_text", False) if isinstance(meta, dict) else False - if needs_ocr: - row_wanted = wanted_labels | _TEXT_LABELS - - # --- decode page image once, crop all matching detections --- - if use_remote: - crops = _crop_all_from_page(page_image_b64, dets, row_wanted, as_b64=True) - crop_b64s: List[str] = [b64 for _label, _bbox, b64 in crops] - crop_meta: List[Tuple[str, List[float]]] = [(label, bbox) for label, bbox, _b64 in crops] - - if crop_b64s: - _invoke_kw = dict( - invoke_url=invoke_url, - image_b64_list=crop_b64s, - api_key=api_key, - timeout_s=float(request_timeout_s), - max_batch_size=int(kwargs.get("inference_batch_size", 8)), - max_retries=int(retry.remote_max_retries), - max_429_retries=int(retry.remote_max_429_retries), - ) - if nim_client is not None: - response_items = nim_client.invoke_image_inference_batches(**_invoke_kw) - else: - response_items = invoke_image_inference_batches( - **_invoke_kw, - max_pool_workers=int(retry.remote_max_pool_workers), - ) - if len(response_items) != len(crop_meta): - raise RuntimeError(f"Expected {len(crop_meta)} OCR responses, got {len(response_items)}") - - for i, (label_name, bbox) in enumerate(crop_meta): - preds = _extract_remote_ocr_item(response_items[i]) - - blocks = _parse_ocr_result(preds) - if label_name == "table": - crop_hw_table: Tuple[int, int] = (0, 0) - try: - _raw = base64.b64decode(crop_b64s[i]) - with Image.open(io.BytesIO(_raw)) as _cim: - _cw, _ch = _cim.size - crop_hw_table = (_ch, _cw) - except Exception: - pass - text = "" - if use_table_structure: - ts_match = _find_ts_detections_for_bbox(row, bbox) - if ts_match is not None: - ts_dets, ts_hw = ts_match - text = join_table_structure_and_ocr_output(ts_dets, preds, ts_hw or crop_hw_table) - if not text: - text = _blocks_to_pseudo_markdown(blocks, crop_hw=crop_hw_table) or _blocks_to_text( - blocks - ) - else: - text = _blocks_to_text(blocks) - entry = {"bbox_xyxy_norm": bbox, "text": text} - if label_name == "table": - table_items.append(entry) - elif label_name == "chart": - chart_items.append(entry) - elif label_name == "infographic": - infographic_items.append(entry) - elif label_name in _TEXT_LABELS: - row_ocr_text_blocks.extend(blocks) - else: - crops = _crop_all_from_page(page_image_b64, dets, row_wanted) - - if inference_batch_size is None or inference_batch_size < 1: - raise ValueError( - f"inference_batch_size must be set and greater than 0. Value: {inference_batch_size}" - ) - - local_batch_size = max(1, int(inference_batch_size)) - - # Tables require word-level merging; charts/infographics use paragraph-level. - # Group by merge level so each batched invoke uses one consistent setting. - local_jobs: Dict[str, List[Tuple[str, List[float], np.ndarray]]] = {"word": [], "paragraph": []} - for label_name, bbox, crop_array in crops: - ml = "word" if label_name == "table" else "paragraph" - local_jobs[ml].append((label_name, bbox, crop_array)) - - def _append_local_result( - label_name: str, bbox: List[float], preds: Any, crop_hw: Tuple[int, int] = (0, 0) - ) -> None: - blocks = _parse_ocr_result(preds) - if label_name == "table": - text = "" - if use_table_structure: - ts_match = _find_ts_detections_for_bbox(row, bbox) - if ts_match is not None: - ts_dets, ts_hw = ts_match - text = join_table_structure_and_ocr_output(ts_dets, preds, ts_hw or crop_hw) - if not text: - text = _blocks_to_pseudo_markdown(blocks, crop_hw=crop_hw) - if not text: - text = _blocks_to_text(blocks) - else: - text = _blocks_to_text(blocks) - entry = {"bbox_xyxy_norm": bbox, "text": text} - if label_name == "table": - table_items.append(entry) - elif label_name == "chart": - chart_items.append(entry) - elif label_name == "infographic": - infographic_items.append(entry) - elif label_name in _TEXT_LABELS: - row_ocr_text_blocks.extend(blocks) - - for ml, jobs in local_jobs.items(): - if not jobs: - continue - for start in range(0, len(jobs), local_batch_size): - batch_jobs = jobs[start : start + local_batch_size] - batch_crops = [crop_array for _, _, crop_array in batch_jobs] - - # Try batched invoke first; if backend does not return one response - # per input, fall back to per-item to preserve correctness. - try: - batch_preds = model.invoke(batch_crops, merge_level=ml) - except Exception: - batch_preds = None - - if isinstance(batch_preds, list) and len(batch_preds) == len(batch_jobs): - for (label_name, bbox, crop_array), preds in zip(batch_jobs, batch_preds): - _append_local_result( - label_name, bbox, preds, crop_hw=(crop_array.shape[0], crop_array.shape[1]) - ) - else: - for label_name, bbox, crop_array in batch_jobs: - preds = model.invoke(crop_array, merge_level=ml) - _append_local_result( - label_name, bbox, preds, crop_hw=(crop_array.shape[0], crop_array.shape[1]) - ) - - except BaseException as e: - print(f"Warning: OCR failed: {type(e).__name__}: {e}") - row_error = { - "stage": "ocr_page_elements", - "type": e.__class__.__name__, - "message": str(e), - "traceback": "".join(traceback.format_exception(type(e), e, e.__traceback__)), - } - - # Assemble OCR'd text from text/title detections for this row. - # Use None as sentinel for "keep existing native text". - if extract_text and row_ocr_text_blocks: - all_text.append(_blocks_to_text(row_ocr_text_blocks)) - else: - all_text.append(None) - - row_det_count = len(table_items) + len(chart_items) + len(infographic_items) + len(row_ocr_text_blocks) - row_counts: Dict[str, int] = {} - if table_items: - row_counts["table"] = len(table_items) - if chart_items: - row_counts["chart"] = len(chart_items) - if infographic_items: - row_counts["infographic"] = len(infographic_items) - if row_ocr_text_blocks: - row_counts["text"] = len(row_ocr_text_blocks) - - all_table.append(table_items) - all_chart.append(chart_items) - all_infographic.append(infographic_items) - all_ocr_meta.append( - {"timing": None, "error": row_error, "num_detections": row_det_count, "counts_by_label": row_counts} + if use_remote: + _run_remote_ocr( + prepared_rows, + row_results, + invoke_url=invoke_url, + api_key=api_key, + request_timeout_s=request_timeout_s, + # Preserve the existing remote behavior. The named + # inference_batch_size parameter is local policy in this path. + max_batch_size=int(kwargs.get("inference_batch_size", 8)), + retry=retry, + nim_client=nim_client, + use_table_structure=use_table_structure, + ) + else: + jobs_by_merge_level = _collect_local_crop_jobs(prepared_rows, row_results) + _run_local_ocr_batches( + model, + jobs_by_merge_level, + row_results, + batch_size=local_invoke_batch_size, + use_table_structure=use_table_structure, ) - elapsed = time.perf_counter() - t0_total - - for meta in all_ocr_meta: - meta["timing"] = {"seconds": float(elapsed)} - - out = batch_df.copy() - if extract_tables or "table" not in out.columns: - out["table"] = all_table - if extract_charts or "chart" not in out.columns: - out["chart"] = all_chart - if extract_infographics or "infographic" not in out.columns: - out["infographic"] = all_infographic - if extract_text and "text" in out.columns: - for i, ocr_text in enumerate(all_text): - if ocr_text is not None: - out.iat[i, out.columns.get_loc("text")] = ocr_text - elif extract_text: - out["text"] = [t if t is not None else "" for t in all_text] - out["ocr"] = all_ocr_meta - out["ocr_v1_num_detections"] = [m["num_detections"] for m in all_ocr_meta] - out["ocr_v1_counts_by_label"] = [m["counts_by_label"] for m in all_ocr_meta] - return out + return _build_ocr_page_elements_output( + batch_df, + row_results, + extract_text=extract_text, + extract_tables=extract_tables, + extract_charts=extract_charts, + extract_infographics=extract_infographics, + elapsed_s=time.perf_counter() - t0_total, + ) # --------------------------------------------------------------------------- diff --git a/nemo_retriever/src/nemo_retriever/common/modality/txt/split.py b/nemo_retriever/src/nemo_retriever/common/modality/txt/split.py index e1352bca18..165cd35862 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/txt/split.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/txt/split.py @@ -15,24 +15,30 @@ from typing import Any, Dict, List, Optional import pandas as pd +from nemo_retriever.common.inline_text import is_inline_text_source from nemo_retriever.common.params import TextChunkParams +from nemo_retriever.models import VL_EMBED_MODEL -DEFAULT_TOKENIZER_MODEL_ID = "nvidia/llama-nemotron-embed-1b-v2" +from .tokenizer_provider import ChunkTokenizer, load_chunk_tokenizer + +DEFAULT_TOKENIZER_MODEL_ID = VL_EMBED_MODEL DEFAULT_MAX_TOKENS = 1024 DEFAULT_OVERLAP_TOKENS = 0 -def _get_tokenizer(model_id: str, cache_dir: Optional[str] = None): # noqa: ANN201 - """Lazy-load HuggingFace tokenizer.""" - from transformers import AutoTokenizer +def empty_text_chunks_df() -> pd.DataFrame: + """Return the canonical empty result for raw text ingestion.""" + return pd.DataFrame(columns=["text", "content", "path", "page_number", "metadata"]).astype({"page_number": "int64"}) - from nemo_retriever.models.hf_model_registry import get_hf_revision - return AutoTokenizer.from_pretrained( +def _get_tokenizer( + model_id: str, + cache_dir: Optional[str] = None, +) -> ChunkTokenizer: + """Load the exact lightweight tokenizer configured for chunking.""" + return load_chunk_tokenizer( model_id, - revision=get_hf_revision(model_id), cache_dir=cache_dir, - trust_remote_code=True, ) @@ -55,7 +61,7 @@ def split_text_by_tokens( text : str Input text to split. tokenizer - HuggingFace tokenizer (e.g. AutoTokenizer) with encode/decode. + Lightweight :class:`ChunkTokenizer` instance with encode/decode. max_tokens : int Maximum tokens per chunk. overlap_tokens : int @@ -104,7 +110,7 @@ def split_df( Re-chunk a DataFrame's ``text`` column by token count. This is a **post-extraction** transform: it takes rows that already have a - ``text`` column (produced by ``extract`` / ``extract_txt`` / etc.) and + ``text`` column produced by extraction and splits long texts into multiple rows using :func:`split_text_by_tokens`. All other columns (``path``, ``page_number``, ``metadata``, …) are preserved on every output row. Each chunk row's ``metadata`` dict is @@ -179,57 +185,78 @@ def txt_file_to_chunks_df( path: str, params: TextChunkParams | None = None, ) -> pd.DataFrame: - """ - Read a .txt file and return a DataFrame of chunks (one row per chunk). + """Read a text file and return one row per token-bounded chunk. - Columns: text, path, page_number (chunk index, 1-based), metadata. - Shape is compatible with embed_text_from_primitives_df and LanceDB row build. + The file path is resolved before being copied to ``path`` and + ``metadata.source_path``. Empty or whitespace-only files return the same + typed schema without loading a tokenizer. Parameters ---------- path : str - Path to the .txt file. - max_tokens : int - Max tokens per chunk (default 512). - overlap_tokens : int - Overlap between consecutive chunks (default 0). - tokenizer_model_id : str, optional - HuggingFace model id for tokenizer (default: same as embedder). - encoding : str - File encoding (default utf-8). - tokenizer_cache_dir : str, optional - HuggingFace cache directory for tokenizer. + Path to the text file. + params : TextChunkParams, optional + File encoding, tokenizer, chunk-size, and overlap configuration. Returns ------- pd.DataFrame - Columns: text, path, page_number, metadata. + Columns: ``text``, ``content``, ``path``, ``page_number``, and + ``metadata``. """ chunk_params = params or TextChunkParams() - max_tokens = chunk_params.max_tokens - overlap_tokens = chunk_params.overlap_tokens - tokenizer_model_id = chunk_params.tokenizer_model_id - encoding = chunk_params.encoding - tokenizer_cache_dir = chunk_params.tokenizer_cache_dir path = str(Path(path).resolve()) - raw = Path(path).read_text(encoding=encoding, errors="replace") - if not raw or not raw.strip(): - return pd.DataFrame( - columns=["text", "path", "page_number", "metadata"], - ).astype({"page_number": "int64"}) - model_id = tokenizer_model_id or DEFAULT_TOKENIZER_MODEL_ID - tokenizer = _get_tokenizer(model_id, cache_dir=tokenizer_cache_dir) + raw = Path(path).read_text(encoding=chunk_params.encoding, errors="replace") + return text_to_chunks_df(raw, path, params=chunk_params) + + +def text_to_chunks_df( + text: str, + source_id: str, + params: TextChunkParams | None = None, +) -> pd.DataFrame: + """Split decoded text while preserving its logical source identifier. + + Unlike the file and byte adapters, this helper deliberately does not + resolve ``source_id`` as a filesystem path. This permits identifiers such + as ``inline://00000000`` to survive through embedding and vector storage. + + Parameters + ---------- + text : str + Decoded source document. Empty or whitespace-only text produces an + empty result without loading the tokenizer. + source_id : str + Logical source identity copied to ``path`` and + ``metadata.source_path`` on every chunk. + params : TextChunkParams, optional + Tokenizer, chunk-size, and overlap configuration. + + Returns + ------- + pd.DataFrame + Chunk rows with ``text``, ``content``, ``path``, ``page_number``, and + ``metadata`` columns. Empty input returns the same typed empty schema. + + Raises + ------ + ValueError + If the configured maximum token count is not positive. + """ + chunk_params = params or TextChunkParams() + if not text or not text.strip(): + return empty_text_chunks_df() + + model_id = chunk_params.tokenizer_model_id or DEFAULT_TOKENIZER_MODEL_ID + tokenizer = _get_tokenizer(model_id, cache_dir=chunk_params.tokenizer_cache_dir) chunk_texts = split_text_by_tokens( - raw, + text, tokenizer=tokenizer, - max_tokens=max_tokens, - overlap_tokens=overlap_tokens, + max_tokens=chunk_params.max_tokens, + overlap_tokens=chunk_params.overlap_tokens, ) - if not chunk_texts: - return pd.DataFrame( - columns=["text", "path", "page_number", "metadata"], - ).astype({"page_number": "int64"}) + return empty_text_chunks_df() rows: List[Dict[str, Any]] = [] for i, chunk in enumerate(chunk_texts): @@ -237,10 +264,10 @@ def txt_file_to_chunks_df( { "text": chunk, "content": chunk, - "path": path, + "path": source_id, "page_number": i + 1, "metadata": { - "source_path": path, + "source_path": source_id, "chunk_index": i, "content_metadata": {"type": "text"}, "content": chunk, @@ -259,43 +286,12 @@ def txt_bytes_to_chunks_df( Decode bytes to text and return a DataFrame of chunks (same shape as txt_file_to_chunks_df). Used by batch TxtSplitActor when input is bytes + path from read_binary_files. + Service-mode inline text also crosses HTTP as bytes; its ``inline://`` + identity is preserved and its transport encoding is always UTF-8. """ chunk_params = params or TextChunkParams() - max_tokens = chunk_params.max_tokens - overlap_tokens = chunk_params.overlap_tokens - tokenizer_model_id = chunk_params.tokenizer_model_id - encoding = chunk_params.encoding - tokenizer_cache_dir = chunk_params.tokenizer_cache_dir - path = str(Path(path).resolve()) + is_inline = is_inline_text_source(path) + source_id = path if is_inline else str(Path(path).resolve()) + encoding = "utf-8" if is_inline else chunk_params.encoding raw = content_bytes.decode(encoding, errors="replace") - model_id = tokenizer_model_id or DEFAULT_TOKENIZER_MODEL_ID - tokenizer = _get_tokenizer(model_id, cache_dir=tokenizer_cache_dir) - chunk_texts = split_text_by_tokens( - raw, - tokenizer=tokenizer, - max_tokens=max_tokens, - overlap_tokens=overlap_tokens, - ) - - if not chunk_texts: - return pd.DataFrame( - columns=["text", "path", "page_number", "metadata"], - ).astype({"page_number": "int64"}) - - rows: List[Dict[str, Any]] = [] - for i, chunk in enumerate(chunk_texts): - rows.append( - { - "text": chunk, - "content": chunk, - "path": path, - "page_number": i + 1, - "metadata": { - "source_path": path, - "chunk_index": i, - "content_metadata": {"type": "text"}, - "content": chunk, - }, - } - ) - return pd.DataFrame(rows) + return text_to_chunks_df(raw, source_id, params=chunk_params) diff --git a/nemo_retriever/src/nemo_retriever/common/modality/txt/tokenizer_provider.py b/nemo_retriever/src/nemo_retriever/common/modality/txt/tokenizer_provider.py new file mode 100644 index 0000000000..44e4ae29ec --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/common/modality/txt/tokenizer_provider.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dependency-light, revision-pinned tokenizers for text chunking.""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Any, Protocol + +from nemo_retriever.models.hf_model_registry import ( + get_hf_revision, + hf_hub_download_with_pinned_revision, +) + + +class TokenizerUnavailableError(RuntimeError): + """Raised when an exact configured tokenizer cannot be loaded.""" + + +class ChunkTokenizer(Protocol): + """Minimal tokenizer interface required by the text splitter.""" + + def encode(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + """Encode text into token IDs.""" + + def decode(self, token_ids: list[int], *, skip_special_tokens: bool = True) -> str: + """Decode token IDs into text.""" + + +class _FastTokenizer: + """Adapt ``tokenizers.Tokenizer`` to :class:`ChunkTokenizer`.""" + + def __init__(self, tokenizer: Any) -> None: + self._tokenizer = tokenizer + + def encode(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + """Encode text using the pinned tokenizer artifact.""" + encoding = self._tokenizer.encode(text, add_special_tokens=add_special_tokens) + return list(encoding.ids) + + def decode(self, token_ids: list[int], *, skip_special_tokens: bool = True) -> str: + """Decode token IDs using the pinned tokenizer artifact.""" + return str(self._tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)) + + +@lru_cache(maxsize=16) +def load_chunk_tokenizer( + model_id: str, + cache_dir: str | None = None, +) -> ChunkTokenizer: + """Load an immutable tokenizer artifact without model weights. + + Args: + model_id: Hugging Face model identifier registered with a pinned revision. + cache_dir: Optional Hugging Face cache directory. + + Returns: + A dependency-light tokenizer suitable for deterministic chunking. + + Raises: + TokenizerUnavailableError: If the pinned tokenizer cannot be resolved. + """ + revision: str | None = None + try: + revision = get_hf_revision(model_id) + tokenizer_path = hf_hub_download_with_pinned_revision( + repo_id=model_id, + filename="tokenizer.json", + revision=revision, + cache_dir=cache_dir, + ) + from tokenizers import Tokenizer + + return _FastTokenizer(Tokenizer.from_file(tokenizer_path)) + except Exception as exc: + raise TokenizerUnavailableError( + "Unable to load the exact tokenizer required for text chunking: " + f"model={model_id!r}, revision={revision!r}. Pre-cache tokenizer.json " + "for this revision (service image builds: " + "DOWNLOAD_DEFAULT_TOKENIZER=True) or allow Hugging Face Hub access " + "in this runtime." + ) from exc diff --git a/nemo_retriever/src/nemo_retriever/common/params/models.py b/nemo_retriever/src/nemo_retriever/common/params/models.py index 0ab4c3d420..0801c63a30 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/models.py +++ b/nemo_retriever/src/nemo_retriever/common/params/models.py @@ -568,6 +568,13 @@ def _auto_enable_features(self) -> "ExtractParams": self.table_output_format = "markdown" if self.use_table_structure else "pseudo_markdown" if self.ocr_version == "v1" and self.ocr_lang is not None: raise ValueError("ocr_lang is only supported when ocr_version='v2'.") + if self.method != "nemotron_parse" and ( + self.nemotron_parse_invoke_url is not None or self.nemotron_parse_model is not None + ): + raise ValueError( + "`nemotron_parse_invoke_url` and `nemotron_parse_model` require " + "`method='nemotron_parse'`; Parse-specific configuration is otherwise ignored." + ) if not self.use_page_elements: consumers = [("use_table_structure", self.use_table_structure and self.extract_tables)] enabled = [name for name, on in consumers if on] @@ -585,6 +592,7 @@ class EmbedParams(_ParamsModel): embedding_endpoint: Optional[str] = None embed_invoke_url: Optional[str] = None embed_model_name: Optional[str] = None + embed_model_revision: Optional[str] = None embed_model_provider_prefix: Optional[str] = None api_key: Optional[str] = None input_type: str = "passage" diff --git a/nemo_retriever/src/nemo_retriever/common/params/utils.py b/nemo_retriever/src/nemo_retriever/common/params/utils.py index bf1b9b2d36..09417387ba 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/utils.py +++ b/nemo_retriever/src/nemo_retriever/common/params/utils.py @@ -8,8 +8,6 @@ from typing import TYPE_CHECKING, Any, Dict -from nemo_retriever.common.api.util.string_processing import prepend_model_provider_prefix - if TYPE_CHECKING: from nemo_retriever.common.params.models import BatchTuningParams @@ -29,7 +27,7 @@ def coerce_params[T](params: T | None, model_cls: type[T], kwargs: dict[str, Any def normalize_embed_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: - """Normalize embedding endpoint aliases in an existing kwargs dict.""" + """Normalize embedding endpoint aliases without changing model identity.""" normalized = dict(kwargs) embed_invoke_url = ( str(normalized.get("embed_invoke_url") or "").strip() if "embed_invoke_url" in normalized else None @@ -52,12 +50,6 @@ def normalize_embed_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: if "embed_invoke_url" in normalized: normalized.setdefault("embedding_endpoint", normalized["embed_invoke_url"]) - endpoint = normalized.get("embedding_endpoint") or normalized.get("embed_invoke_url") - model_provider_prefix = normalized.pop("embed_model_provider_prefix", None) - if endpoint and model_provider_prefix: - for key in ("model_name", "embed_model_name"): - if key in normalized: - normalized[key] = prepend_model_provider_prefix(normalized[key], str(model_provider_prefix)) return normalized @@ -75,6 +67,7 @@ def build_embed_option_kwargs( embed_batch_size: int | None = None, embed_cpus_per_actor: float | None = None, embed_gpus_per_actor: float | None = None, + embed_model_revision: str | None = None, ) -> Dict[str, Any]: """Build ``EmbedParams`` kwargs from CLI/request option values.""" embed_kwargs: Dict[str, Any] = {} @@ -84,6 +77,8 @@ def build_embed_option_kwargs( # Remote HTTP embedding reads model_name; local/GPU paths read embed_model_name. embed_kwargs["model_name"] = embed_model_name embed_kwargs["embed_model_name"] = embed_model_name + if embed_model_revision is not None: + embed_kwargs["embed_model_revision"] = embed_model_revision if local_ingest_embed_backend is not None: embed_kwargs["local_ingest_embed_backend"] = local_ingest_embed_backend if embed_api_key is not None: diff --git a/nemo_retriever/src/nemo_retriever/common/policy.py b/nemo_retriever/src/nemo_retriever/common/policy.py index db5cf8051e..cfd799b576 100644 --- a/nemo_retriever/src/nemo_retriever/common/policy.py +++ b/nemo_retriever/src/nemo_retriever/common/policy.py @@ -57,6 +57,7 @@ "ocr_invoke_url", "table_structure_invoke_url", "nemotron_parse_invoke_url", + "nemotron_parse_model", "profile_name", ) diff --git a/nemo_retriever/src/nemo_retriever/common/schemas/collections.py b/nemo_retriever/src/nemo_retriever/common/schemas/collections.py new file mode 100644 index 0000000000..c61b55f146 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/common/schemas/collections.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public collection-management wire models. + +These models are shared by the service and Python SDK. Keeping the contract +here prevents agent adapters from depending on service implementation details +or LanceDB-specific names. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Annotated, Any, Literal + +from pydantic import Field, StringConstraints, field_validator + +from nemo_retriever.common.api.util.converters.datetools import ( + normalize_timezone_aware_iso8601_to_utc, +) +from nemo_retriever.common.schemas.base import RichModel + +CollectionStatus = Literal["active", "deleting"] +DeleteStatus = Literal["deleting", "deleted"] + + +class IngestOperation(str, Enum): + """How an accepted upload is applied to its target collection.""" + + APPEND = "append" + REPLACE = "replace" + + +_LogicalResourceId = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", + ), +] +CollectionName = _LogicalResourceId +DocumentId = _LogicalResourceId + + +def _normalize_expires_at(value: str | None) -> str | None: + if value is None: + return None + return normalize_timezone_aware_iso8601_to_utc(value) + + +class CollectionCreateRequest(RichModel): + """Properties accepted when creating a logical collection.""" + + name: CollectionName + description: str | None = Field(default=None, max_length=4096) + metadata: dict[str, Any] = Field(default_factory=dict) + expires_at: str | None = None + + @field_validator("expires_at", mode="before") + @classmethod + def _validate_expiry(cls, value: str | None) -> str | None: + return _normalize_expires_at(value) + + +class CollectionUpdateRequest(RichModel): + """Mutable properties of an existing logical collection.""" + + description: str | None = Field(default=None, max_length=4096) + metadata: dict[str, Any] | None = None + expires_at: str | None = None + + @field_validator("expires_at", mode="before") + @classmethod + def _validate_expiry(cls, value: str | None) -> str | None: + return _normalize_expires_at(value) + + +class CollectionInfo(RichModel): + """Public metadata and lifecycle state for a scoped collection.""" + + name: str + scope: str + status: CollectionStatus = "active" + description: str | None = None + metadata: dict[str, Any] = Field(default_factory=dict) + created_at: str + updated_at: str + expires_at: str | None = None + + +class CollectionPage(RichModel): + """One page of collections and its opaque continuation token.""" + + items: list[CollectionInfo] = Field(default_factory=list) + next_token: str | None = None + + +class DocumentInfo(RichModel): + """Public identity and ingestion state for a collection document.""" + + document_id: DocumentId + collection_name: str + scope: str + filename: str + content_sha256: str + document_version: str + status: str + chunk_count: int = 0 + job_id: str | None = None + created_at: str + updated_at: str + error: str | None = None + + +class DocumentPage(RichModel): + """One page of collection documents and its opaque continuation token.""" + + items: list[DocumentInfo] = Field(default_factory=list) + next_token: str | None = None + + +class DocumentDeleteResult(RichModel): + """Outcome of a document deletion request.""" + + document_id: DocumentId + collection_name: str + scope: str + existed: bool + deleted: bool + status: DeleteStatus + cleanup_pending: bool = False + + +class CollectionDeleteResult(RichModel): + """Outcome of a collection deletion request.""" + + name: str + scope: str + existed: bool + deleted: bool + status: DeleteStatus + cleanup_pending: bool = False + + +class QueryHit(RichModel): + """Citation-ready hit returned to agentic applications.""" + + chunk_id: str + document_id: DocumentId + text: str + distance: float = Field( + allow_inf_nan=False, + description="Native dense-vector distance; lower values are more similar.", + ) + filename: str + page_number: int | None = Field( + default=None, + ge=1, + description="One-based human-facing document page, or null when not applicable.", + ) + content_type: str | None = None + source: Any = None + source_id: str | None = None + stored_image_uri: str | None = None + bbox: Any = None + metadata: dict[str, Any] = Field(default_factory=dict) diff --git a/nemo_retriever/src/nemo_retriever/common/schemas/requests.py b/nemo_retriever/src/nemo_retriever/common/schemas/requests.py index e5a0277bd0..f21b84bf6b 100644 --- a/nemo_retriever/src/nemo_retriever/common/schemas/requests.py +++ b/nemo_retriever/src/nemo_retriever/common/schemas/requests.py @@ -6,9 +6,10 @@ from typing import Any -from pydantic import Field +from pydantic import Field, model_validator from nemo_retriever.common.schemas.base import RichModel +from nemo_retriever.common.schemas.collections import CollectionName, DocumentId, IngestOperation from nemo_retriever.common.schemas.pipeline_spec import PipelineSpec @@ -32,6 +33,14 @@ class IngestRequest(RichModel): pipeline: PipelineSpec | None = None +class DocumentManifestEntry(RichModel): + """One immutable file identity in a resumable ingestion job.""" + + manifest_entry_id: str = Field(min_length=64, max_length=64, pattern=r"^[0-9a-f]{64}$") + filename: str + content_sha256: str = Field(min_length=64, max_length=64, pattern=r"^[0-9a-f]{64}$") + + class JobCreateRequest(RichModel): """Body for ``POST /v1/ingest/job`` — open a new ingestion job. @@ -55,3 +64,25 @@ class JobCreateRequest(RichModel): "``result_data``." ), ) + collection_name: CollectionName | None = None + operation: IngestOperation = IngestOperation.APPEND + target_document_id: DocumentId | None = None + idempotency_key: str | None = Field(default=None, min_length=1, max_length=256) + document_manifest: list[DocumentManifestEntry] = Field(default_factory=list) + + @model_validator(mode="after") + def _validate_job_contract(self) -> "JobCreateRequest": + if self.document_manifest and len(self.document_manifest) != self.expected_documents: + raise ValueError("document_manifest length must match expected_documents") + if len({entry.manifest_entry_id for entry in self.document_manifest}) != len(self.document_manifest): + raise ValueError("document_manifest contains duplicate manifest_entry_id values") + if self.operation is IngestOperation.REPLACE: + if not self.collection_name: + raise ValueError("replace requires collection_name") + if self.expected_documents != 1: + raise ValueError("replace requires exactly one document") + if not self.target_document_id: + raise ValueError("replace requires target_document_id") + elif self.target_document_id is not None: + raise ValueError("append does not accept target_document_id") + return self diff --git a/nemo_retriever/src/nemo_retriever/common/schemas/responses.py b/nemo_retriever/src/nemo_retriever/common/schemas/responses.py index be143d22df..c5988e8b7a 100644 --- a/nemo_retriever/src/nemo_retriever/common/schemas/responses.py +++ b/nemo_retriever/src/nemo_retriever/common/schemas/responses.py @@ -9,12 +9,14 @@ from pydantic import Field from nemo_retriever.common.schemas.base import RichModel +from nemo_retriever.common.schemas.collections import IngestOperation class IngestAccepted(RichModel): """Response for the general ``POST /v1/ingest`` endpoint.""" document_id: str + attempt_id: str job_id: str | None = None content_sha256: str status: str @@ -36,6 +38,7 @@ class DocumentIngestAccepted(RichModel): """Response for ``POST /v1/ingest/document`` (whole document upload).""" document_id: str + attempt_id: str filename: str file_size_bytes: int content_sha256: str @@ -77,6 +80,8 @@ class JobCreatedResponse(RichModel): created_at: str label: str | None = None trace_id: str | None = None + collection_name: str | None = None + operation: IngestOperation = IngestOperation.APPEND class JobAggregateResponse(RichModel): @@ -100,6 +105,8 @@ class JobAggregateResponse(RichModel): counts: dict[str, int] = Field(default_factory=dict) document_ids: list[str] = Field(default_factory=list) documents: list[dict[str, Any]] | None = None + collection_name: str | None = None + operation: IngestOperation = IngestOperation.APPEND class DocumentStatusResponse(RichModel): @@ -112,6 +119,7 @@ class DocumentStatusResponse(RichModel): """ document_id: str + attempt_id: str job_id: str status: str submitted_at: str @@ -122,6 +130,8 @@ class DocumentStatusResponse(RichModel): result_rows: int | None = None result_data: list[dict[str, Any]] | None = None error: str | None = None + collection_name: str | None = None + content_sha256: str | None = None class JobDocumentsPage(RichModel): diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/README.md b/nemo_retriever/src/nemo_retriever/common/vdb/README.md index 8ae7e624ad..fbd543abf0 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/README.md +++ b/nemo_retriever/src/nemo_retriever/common/vdb/README.md @@ -11,17 +11,34 @@ The root CLI is intentionally LanceDB-first: `retriever ingest ...` writes Lance --- +## Collection capabilities + +`VDB` defines required collection and document capabilities for the service API. +Backends implement the CRUD methods plus `write_collection()` and +`retrieve_collection()`; callers never pass logical collection identity through +legacy `run()` or `retrieval(**kwargs)`. Maintenance and health retain safe empty +defaults for backends without recoverable lifecycle work or additional health +details. + +`CollectionWriteContext` carries immutable logical write identity. The service +and graph operators pass that context through unchanged; concrete backends own +physical names, schemas, native ranking fields, locks, and persistence. LanceDB +initializes its private collection catalog lazily, so ordinary fixed-table +construction and the existing CLI paths do not create collection metadata. + +--- + ## `IngestVdbOperator` (ingestion) ### Role -`IngestVdbOperator` adapts **flat graph / DataFrame rows** (the shape produced after extract → embed in NeMo Retriever) into the **nested ingestion-pipeline record batches** expected by client VDBs, then calls **`VDB.run(records)`** once per batch. +`IngestVdbOperator` adapts **flat graph / DataFrame rows** (the shape produced after extract → embed in NeMo Retriever) into the **nested ingestion-pipeline record batches** expected by client VDBs. Legacy calls use **`VDB.run(records)`** once per batch; an explicit `CollectionWriteContext` dispatches to **`VDB.write_collection(records, context=...)`**. -Flow (see `operators.py` and `records.py`): +Flow (see `operators/vdb.py` and `common/vdb/records.py`): 1. **`to_client_vdb_records(data)`** — converts rows to `list[list[dict]]` (one outer batch). Dense rows require an **embedding** plus either nonblank **text** or concrete image backing. Image-backed rows without text are stored as `type=image` and `text=""`; they are searchable through dense retrieval but add no FTS terms. The answer-oriented evidence formatter omits every hit without nonblank text and reports the omission in coverage. Sparse-only ingestion continues to require nonblank text. 2. Optional **sidecar metadata** — if `vdb_kwargs` contains `meta_dataframe` / `meta_source_field` / `meta_fields`, those keys are stripped for the concrete DB constructor and merged onto records via `sidecar_metadata.py`. -3. **`self._vdb.run(records)`** — delegates to the concrete backend (e.g. `LanceDB.run`). +3. **Explicit dispatch** — calls `VDB.run(records)` for fixed-table ingestion or `VDB.write_collection(records, context=...)` for a scoped collection. ### Ray batch pipelines (`RayDataExecutor`) @@ -94,10 +111,12 @@ Common constructor arguments include: ### Role -`RetrieveVdbOperator` wraps the same concrete **`VDB`** instance but calls **`retrieval(vectors, **kwargs)`** instead of `run`. It merges per-call kwargs with the operator’s stored `vdb_kwargs` and returns **`normalize_retrieval_results(...)`** output (see `operators.py`, `records.py`). +`RetrieveVdbOperator` wraps the same concrete **`VDB`** instance. Fixed-table calls use **`retrieval(vectors, **kwargs)`** and normalize legacy hit shapes; requests with both `scope` and `collection_name` use the explicit **`retrieve_collection(...)`** capability and validate/project its results into the canonical public hit contract. See `operators/vdb.py` and `common/vdb/records.py`. Important: retrieval here expects **`vectors`** — a list of query embedding vectors — as the primary input. String queries are embedded elsewhere (e.g. in `Retriever`). Hybrid backends that need raw text receive aligned `query_texts` as execution-only call context. +Before embedding, `Retriever` asks the operator for `get_index_metadata("embedding_model_name")`. The base `VDB` implementation returns `None`; a backend can override the method to expose metadata from its selected table or index. LanceDB exposes both `embedding_model_name` and `retrieval_mode` through this lookup. + ### LanceDB inside `RetrieveVdbOperator` For `vdb_op="lancedb"`, **`LanceDB.retrieval`**: diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/adt_vdb.py b/nemo_retriever/src/nemo_retriever/common/vdb/adt_vdb.py index faafa90307..10330206f5 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/adt_vdb.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/adt_vdb.py @@ -24,8 +24,68 @@ """ from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import Any +from nemo_retriever.common.schemas.collections import ( + CollectionCreateRequest, + CollectionDeleteResult, + CollectionInfo, + CollectionPage, + CollectionUpdateRequest, + DocumentDeleteResult, + DocumentInfo, + DocumentPage, + IngestOperation, +) + + +class UnsupportedVDBOperation(NotImplementedError): + """The selected VDB backend does not implement an optional operation.""" + + +class VDBResourceNotFound(LookupError): + """A logical resource does not exist in the selected VDB backend.""" + + +class VDBResourceConflict(RuntimeError): + """A VDB operation conflicts with the resource's current state.""" + + +class VDBInvalidRequest(ValueError): + """A request cannot be applied by the VDB backend.""" + + +@dataclass(frozen=True, slots=True) +class CollectionWriteContext: + """Logical identity and operation metadata for one collection write.""" + + scope: str + collection_name: str + document_id: str + document_version: str + content_sha256: str + filename: str + job_id: str | None = None + operation: IngestOperation = IngestOperation.APPEND + + def __post_init__(self) -> None: + """Coerce a wire-level operation string into its enum member. + + Callers outside the service layer construct this directly with a + plain string, so normalise here to keep identity comparisons sound. + """ + if not isinstance(self.operation, IngestOperation): + object.__setattr__(self, "operation", IngestOperation(self.operation)) + + +@dataclass(frozen=True, slots=True) +class CollectionWriteResult: + """Backend-neutral counts returned after a collection write.""" + + written: int + total_rows: int + class VDB(ABC): """Abstract base class for vector-database operators. @@ -148,6 +208,10 @@ def retrieval(self, queries: list, **kwargs): """ pass + def get_index_metadata(self, key: str, **kwargs: Any) -> str | None: + """Return one metadata value for the selected index, if available.""" + return None + def put(self, records: list, **kwargs: Any) -> dict[str, Any]: """Replace a batch of existing rows in the target table/index. @@ -231,6 +295,134 @@ def put(self, records: list, **kwargs: Any) -> dict[str, Any]: "in-place stable-key puts are not supported by this VDB backend." ) + @abstractmethod + def create_collection( + self, + *, + scope: str, + request: CollectionCreateRequest, + ) -> CollectionInfo: + """Create a logical collection. + + Backends must isolate the logical collection by both ``scope`` and + ``request.name``. + """ + pass + + @abstractmethod + def get_collection( + self, + *, + scope: str, + collection_name: str, + ) -> CollectionInfo: + """Return one logical collection visible within ``scope``.""" + pass + + @abstractmethod + def list_collections( + self, + *, + scope: str, + limit: int, + continuation_token: str | None, + ) -> CollectionPage: + """List logical collections visible within ``scope``.""" + pass + + @abstractmethod + def update_collection( + self, + *, + scope: str, + collection_name: str, + request: CollectionUpdateRequest, + ) -> CollectionInfo: + """Update one logical collection visible within ``scope``.""" + pass + + @abstractmethod + def delete_collection( + self, + *, + scope: str, + collection_name: str, + if_exists: bool, + ) -> CollectionDeleteResult: + """Delete one logical collection and its backend-owned vector data.""" + pass + + @abstractmethod + def get_document( + self, + *, + scope: str, + collection_name: str, + document_id: str, + ) -> DocumentInfo: + """Return one document from a logical collection.""" + pass + + @abstractmethod + def list_documents( + self, + *, + scope: str, + collection_name: str, + limit: int, + continuation_token: str | None, + ) -> DocumentPage: + """List documents from a logical collection.""" + pass + + @abstractmethod + def delete_document( + self, + *, + scope: str, + collection_name: str, + document_id: str, + if_exists: bool, + ) -> DocumentDeleteResult: + """Delete one document and its backend-owned vector data.""" + pass + + @abstractmethod + def write_collection( + self, + records: list, + *, + context: CollectionWriteContext, + ) -> CollectionWriteResult: + """Write canonical NRL records to an explicitly scoped collection.""" + pass + + @abstractmethod + def retrieve_collection( + self, + vectors: list, + *, + scope: str, + collection_name: str, + query_texts: list[str], + top_k: int, + **kwargs: Any, + ) -> tuple[list[list[dict[str, Any]]], list[str]]: + """Retrieve canonical hits and strategy names from one collection.""" + pass + + def reconcile_collections(self) -> dict[str, int]: + """Resume optional collection lifecycle work. + + Backends without durable collection lifecycle state have nothing to + reconcile and therefore return zero work rather than failing. + """ + return {"successes": 0, "failures": 0} + + def health(self) -> dict[str, Any]: + """Return optional backend-specific operational health details.""" + return {} + @abstractmethod def run(self, records): """Pipeline entry point: ensure the index exists, then ingest. diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py index 3ab48161d1..ed538c0492 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py @@ -5,18 +5,40 @@ import json import logging import os +import threading import time from collections.abc import Iterable, Sequence from datetime import timedelta +from types import SimpleNamespace from typing import Any, Final, FrozenSet import lancedb import pyarrow as pa import pyarrow.compute as pc -from nemo_retriever.common.vdb.adt_vdb import VDB - +from nemo_retriever.common.schemas.collections import ( + CollectionCreateRequest, + CollectionDeleteResult, + CollectionInfo, + CollectionPage, + CollectionUpdateRequest, + DocumentDeleteResult, + DocumentInfo, + DocumentPage, +) +from nemo_retriever.common.vdb.adt_vdb import ( + CollectionWriteContext, + CollectionWriteResult, + VDB, +) +from nemo_retriever.common.vdb.lancedb_capabilities import inspect_lancedb_table_object +from nemo_retriever.common.vdb.lancedb_schema import ( + build_lancedb_row, + infer_vector_dim, + lancedb_schema, + normalize_content_type, +) logger = logging.getLogger(__name__) @@ -25,6 +47,8 @@ _VALID_ON_BAD_VECTORS: Final[FrozenSet[str]] = frozenset({"drop", "fill", "null", "error"}) _RETRIEVAL_MODE_METADATA_KEY: Final[bytes] = b"retrieval_mode" _NEMO_RETRIEVER_RETRIEVAL_MODE_METADATA_KEY: Final[bytes] = b"nemo_retriever.retrieval_mode" +_EMBEDDING_MODEL_METADATA_KEY: Final[bytes] = b"nemo_retriever.embedding_model_name" +_EMBEDDING_MODEL_REVISION_METADATA_KEY: Final[bytes] = b"nemo_retriever.embedding_model_revision" def _normalize_on_bad_vectors(value: str) -> str: @@ -115,17 +139,32 @@ def _effective_ivf_num_partitions(num_rows: int, requested: int) -> int | None: return min(int(requested), max(1, cap)) -def _with_retrieval_mode_metadata(schema: pa.Schema, retrieval_mode: str | None) -> pa.Schema: +def _with_retrieval_mode_metadata( + schema: pa.Schema, + retrieval_mode: str | None, + embedding_model_name: str | None = None, + embedding_model_revision: str | None = None, +) -> pa.Schema: if retrieval_mode is None: return schema metadata = dict(schema.metadata or {}) encoded_mode = str(retrieval_mode).encode("utf-8") metadata[_RETRIEVAL_MODE_METADATA_KEY] = encoded_mode metadata[_NEMO_RETRIEVER_RETRIEVAL_MODE_METADATA_KEY] = encoded_mode + if embedding_model_name: + metadata[_EMBEDDING_MODEL_METADATA_KEY] = embedding_model_name.encode("utf-8") + if embedding_model_revision: + metadata[_EMBEDDING_MODEL_REVISION_METADATA_KEY] = embedding_model_revision.encode("utf-8") return schema.with_metadata(metadata) -def _lancedb_arrow_schema(vector_dim: int, *, retrieval_mode: str | None = None) -> pa.Schema: +def _lancedb_arrow_schema( + vector_dim: int, + *, + retrieval_mode: str | None = None, + embedding_model_name: str | None = None, + embedding_model_revision: str | None = None, +) -> pa.Schema: schema = pa.schema( [ pa.field("vector", pa.list_(pa.float32(), int(vector_dim))), @@ -135,7 +174,12 @@ def _lancedb_arrow_schema(vector_dim: int, *, retrieval_mode: str | None = None) pa.field("id", pa.string()), ] ) - return _with_retrieval_mode_metadata(schema, retrieval_mode) + return _with_retrieval_mode_metadata( + schema, + retrieval_mode, + embedding_model_name, + embedding_model_revision, + ) def _sparse_lancedb_arrow_schema(*, retrieval_mode: str | None = "sparse") -> pa.Schema: @@ -155,6 +199,17 @@ def _table_schema(table: Any) -> pa.Schema: return schema() if callable(schema) else schema +def _schema_vector_dim(schema: pa.Schema) -> int | None: + """Return a fixed vector width from a LanceDB table schema when present.""" + try: + vector_type = schema.field("vector").type + except KeyError: + return None + if pa.types.is_fixed_size_list(vector_type): + return int(vector_type.list_size) + return None + + def lancedb_row_count(uri: str, table_name: str) -> int: """Return the number of rows in a LanceDB table.""" table = lancedb.connect(uri).open_table(table_name) @@ -181,6 +236,47 @@ def _validate_append_schema(table: Any, expected_schema: pa.Schema, *, table_nam ) +def _validate_append_embedding_model( + table: Any, + embedding_model_name: str | None, + embedding_model_revision: str | None, + *, + table_name: str, + uri: str, +) -> None: + """Reject appends that would mix known embedding models in one table.""" + if not embedding_model_name: + return + + metadata = _table_schema(table).metadata or {} + stored_value = metadata.get(_EMBEDDING_MODEL_METADATA_KEY) + if stored_value is None: + return + + stored_model = stored_value.decode("utf-8", errors="replace").strip() + if stored_model and stored_model != embedding_model_name: + raise ValueError( + f"LanceDB table {table_name!r} at {uri!r} uses embedding model {stored_model!r}; " + f"cannot append vectors from {embedding_model_name!r}. Use the table model or overwrite the table." + ) + + stored_revision_value = metadata.get(_EMBEDDING_MODEL_REVISION_METADATA_KEY) + if stored_revision_value is None: + return + stored_revision = stored_revision_value.decode("utf-8", errors="replace").strip() + if stored_revision and not embedding_model_revision: + raise ValueError( + f"LanceDB table {table_name!r} at {uri!r} uses embedding model revision {stored_revision!r}; " + "cannot append vectors without a known revision. Use the table revision or overwrite the table." + ) + if stored_revision and stored_revision != embedding_model_revision: + raise ValueError( + f"LanceDB table {table_name!r} at {uri!r} uses embedding model revision {stored_revision!r}; " + f"cannot append vectors from revision {embedding_model_revision!r}. " + "Use the table revision or overwrite the table." + ) + + def _is_missing_lancedb_table_error(exc: ValueError) -> bool: return "was not found" in str(exc) @@ -355,6 +451,56 @@ def _create_lancedb_results( return lancedb_rows, counts +def _to_service_lancedb_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Adapt canonical dense rows to the established service table schema.""" + wide_rows: list[dict[str, Any]] = [] + for row in rows: + content_metadata = _maybe_parse_json(row.get("metadata")) + if not isinstance(content_metadata, dict): + content_metadata = {} + source_metadata = _maybe_parse_json(row.get("source")) + if not isinstance(source_metadata, dict): + source_metadata = {} + source_id = next( + ( + str(value).strip() + for value in ( + source_metadata.get("source_id"), + source_metadata.get("source_name"), + ) + if isinstance(value, str) and value.strip() + ), + "", + ) + content_type = normalize_content_type(content_metadata.get("type") or content_metadata.get("_content_type")) + if content_type: + content_metadata = dict(content_metadata) + content_metadata["type"] = content_type + content_metadata["_content_type"] = content_type + wide_row = build_lancedb_row( + SimpleNamespace( + metadata={ + "embedding": row.get("vector"), + "source_path": source_id, + "content_metadata": content_metadata, + }, + path=source_id, + page_number=content_metadata.get("page_number"), + text=row.get("text") or "", + _stored_image_uri=content_metadata.get("stored_image_uri"), + _content_type=content_type, + _bbox_xyxy_norm=content_metadata.get("bbox_xyxy_norm"), + ) + ) + if wide_row is None: + continue + wide_row["metadata"] = _json_str(content_metadata) + wide_row["source"] = _json_str(source_metadata) + wide_row["content_type"] = content_type or "" + wide_rows.append(wide_row) + return wide_rows + + def _create_sparse_lancedb_results(results) -> tuple[list, dict[str, int]]: """Transform NRL records into LanceDB rows for FTS-only sparse retrieval.""" lancedb_rows: list = [] @@ -394,7 +540,11 @@ def _create_sparse_lancedb_results(results) -> tuple[list, dict[str, int]]: "dropped_no_text": dropped_no_text, } if dropped_no_text: - logger.warning("_create_sparse_lancedb_results: accepted=%d dropped_no_text=%d", accepted, dropped_no_text) + logger.warning( + "_create_sparse_lancedb_results: accepted=%d dropped_no_text=%d", + accepted, + dropped_no_text, + ) return lancedb_rows, counts @@ -413,20 +563,24 @@ def __init__( hybrid: bool = False, sparse: bool = False, fts_language: str = "English", - vector_dim: int = _DEFAULT_VECTOR_DIM, + embedding_model_name: str | None = None, + vector_dim: int | None = _DEFAULT_VECTOR_DIM, on_bad_vectors: str = "drop", fill_value: float = 0.0, validate_vector_length: bool = True, build_index: bool | None = None, + expiration_cleanup_enabled: bool = True, + embedding_model_revision: str | None = None, **kwargs, ): create_index = kwargs.pop("create_index", None) + service_table_schema = bool(kwargs.pop("_service_table_schema", False)) if build_index is None: build_index = True if create_index is None else bool(create_index) elif create_index is not None and bool(create_index) != bool(build_index): raise ValueError("Pass only one index toggle: build_index or create_index.") - if int(vector_dim) <= 0: + if vector_dim is not None and int(vector_dim) <= 0: raise ValueError(f"vector_dim must be positive; got {vector_dim}") if sparse and hybrid: raise ValueError("LanceDB sparse ingest cannot also be hybrid; pass only one retrieval mode.") @@ -441,12 +595,251 @@ def __init__( self.hybrid = hybrid self.sparse = bool(sparse) self.fts_language = fts_language - self.vector_dim = int(vector_dim) + self.embedding_model_name = embedding_model_name + self.embedding_model_revision = embedding_model_revision + self.vector_dim = int(vector_dim) if vector_dim is not None else None self.on_bad_vectors = _normalize_on_bad_vectors(on_bad_vectors) self.fill_value = float(fill_value) self.validate_vector_length = bool(validate_vector_length) + self.expiration_cleanup_enabled = bool(expiration_cleanup_enabled) + self._service_table_schema = service_table_schema + self._collection_store: Any | None = None + self._collection_store_init_failed = False + self._collection_store_lock = threading.Lock() super().__init__(**kwargs) + def _get_collection_store(self) -> Any: + """Lazily initialize collection catalogs only when a collection API is used.""" + + store = self._collection_store + if store is None: + with self._collection_store_lock: + store = self._collection_store + if store is None: + from nemo_retriever.common.vdb.lancedb_collections import ( + LanceDBCollectionStore, + ) + + try: + store = LanceDBCollectionStore( + self, + expiration_cleanup_enabled=self.expiration_cleanup_enabled, + ) + except Exception: + self._collection_store_init_failed = True + raise + self._collection_store_init_failed = False + self._collection_store = store + return store + + def create_collection( + self, + *, + scope: str, + request: CollectionCreateRequest, + ) -> CollectionInfo: + """Create a logical collection through the LanceDB collection store.""" + + return self._get_collection_store().create_collection(scope, request) + + def get_collection( + self, + *, + scope: str, + collection_name: str, + ) -> CollectionInfo: + """Return a logical collection from the LanceDB collection store.""" + + return self._get_collection_store().get_collection(scope, collection_name) + + def list_collections( + self, + *, + scope: str, + limit: int, + continuation_token: str | None, + ) -> CollectionPage: + """List logical collections through the LanceDB collection store.""" + + return self._get_collection_store().list_collections( + scope, + limit, + continuation_token, + ) + + def update_collection( + self, + *, + scope: str, + collection_name: str, + request: CollectionUpdateRequest, + ) -> CollectionInfo: + """Update a logical collection through the LanceDB collection store.""" + + return self._get_collection_store().update_collection( + scope, + collection_name, + request, + ) + + def delete_collection( + self, + *, + scope: str, + collection_name: str, + if_exists: bool, + ) -> CollectionDeleteResult: + """Delete a logical collection through the LanceDB collection store.""" + + return self._get_collection_store().delete_collection( + scope, + collection_name, + if_exists, + ) + + def get_document( + self, + *, + scope: str, + collection_name: str, + document_id: str, + ) -> DocumentInfo: + """Return one collection document through the LanceDB collection store.""" + + return self._get_collection_store().get_document( + scope, + collection_name, + document_id, + ) + + def list_documents( + self, + *, + scope: str, + collection_name: str, + limit: int, + continuation_token: str | None, + ) -> DocumentPage: + """List collection documents through the LanceDB collection store.""" + + return self._get_collection_store().list_documents( + scope, + collection_name, + limit, + continuation_token, + ) + + def delete_document( + self, + *, + scope: str, + collection_name: str, + document_id: str, + if_exists: bool, + ) -> DocumentDeleteResult: + """Delete one collection document through the LanceDB collection store.""" + + return self._get_collection_store().delete_document( + scope, + collection_name, + document_id, + if_exists, + ) + + def write_collection( + self, + records: list, + *, + context: CollectionWriteContext, + ) -> CollectionWriteResult: + """Write canonical records using the collection lifecycle contract.""" + + return self._get_collection_store().write_collection(records, context=context) + + def retrieve_collection( + self, + vectors: list, + *, + scope: str, + collection_name: str, + query_texts: list[str], + top_k: int, + **kwargs: Any, + ) -> tuple[list[list[dict[str, Any]]], list[str]]: + """Retrieve scoped collection hits using LanceDB's collection contract.""" + + return self._get_collection_store().retrieve_collection( + vectors, + scope=scope, + collection_name=collection_name, + query_texts=query_texts, + top_k=top_k, + **kwargs, + ) + + def reconcile_collections(self) -> dict[str, int]: + """Resume interrupted collection and document lifecycle operations.""" + + return self._get_collection_store().reconcile_collections() + + def health(self) -> dict[str, Any]: + """Return legacy table and optional collection-store health.""" + + from nemo_retriever.common.vdb.lancedb_collections import LanceDBCollectionStore + + db = lancedb.connect(uri=self.uri) + table_exists = self.table_name in db.list_tables().tables + total_rows = 0 + effective_mode: str | None = None + retrieval_strategies: list[str] = [] + if table_exists: + try: + total_rows = int(db.open_table(self.table_name).count_rows()) + except Exception: + logger.warning( + "Failed to count rows in the default LanceDB table", + exc_info=True, + ) + try: + capabilities = inspect_lancedb_table_object(db.open_table(self.table_name)) + mode = capabilities.retrieval_mode + if mode in {"dense", "hybrid"}: + effective_mode = str(mode) + retrieval_strategies = [str(mode)] + else: + effective_mode = "unknown" + except Exception: + effective_mode = "unknown" + logger.warning( + "Failed to resolve the default LanceDB retrieval mode", + exc_info=True, + ) + + if self._collection_store_init_failed: + raise RuntimeError("Collection catalog initialization failed") + store = self._collection_store + collection_health = store.health() if store is not None else LanceDBCollectionStore.empty_health() + return { + **collection_health, + "total_rows": total_rows, + "table_exists": table_exists, + "effective_retrieval_mode": effective_mode, + "retrieval_strategies": retrieval_strategies, + } + + def get_index_metadata(self, key: str, **kwargs: Any) -> str | None: + """Read one NeMo Retriever metadata value from the selected table.""" + uri = str(kwargs.get("table_path") or kwargs.get("uri") or kwargs.get("lancedb_uri") or self.uri) + table_name = str(kwargs.get("table_name") or kwargs.get("lancedb_table") or self.table_name) + table = lancedb.connect(uri=uri).open_table(table_name) + metadata = _table_schema(table).metadata or {} + value = metadata.get(f"nemo_retriever.{key}".encode("utf-8")) + if value is None and key == "retrieval_mode": + value = metadata.get(_RETRIEVAL_MODE_METADATA_KEY) + if value is None: + return None + return value.decode("utf-8", errors="replace").strip() or None + def create_index(self, records=None, table_name: str = "nv-ingest", **kwargs): """Create or update a LanceDB table and populate it with transformed records. @@ -462,19 +855,51 @@ def create_index(self, records=None, table_name: str = "nv-ingest", **kwargs): connect_start = time.perf_counter() db = lancedb.connect(uri=self.uri) _record_timing("lancedb.connect", time.perf_counter() - connect_start) + record_batches = list(records or []) if self.sparse: - results, counts = _create_sparse_lancedb_results(records or []) + results, counts = _create_sparse_lancedb_results(record_batches) schema = _sparse_lancedb_arrow_schema() write_kwargs: dict[str, Any] = {} else: - if self.validate_vector_length and self.on_bad_vectors != "error": - expected_dim: int | None = self.vector_dim + enforce_dim = self.validate_vector_length and self.on_bad_vectors != "error" + vector_dim = self.vector_dim + if vector_dim is None and not self.overwrite: + try: + existing_table = db.open_table(self.table_name) + except ValueError as exc: + if not _is_missing_lancedb_table_error(exc): + raise + else: + vector_dim = _schema_vector_dim(_table_schema(existing_table)) + + if vector_dim is None: + results, counts = _create_lancedb_results(record_batches, expected_dim=None) + vector_dim = infer_vector_dim(results) + if vector_dim <= 0: + raise ValueError("Cannot infer LanceDB vector_dim because no non-empty embedding was produced.") + if enforce_dim: + results, counts = _create_lancedb_results(record_batches, expected_dim=vector_dim) else: - expected_dim = None + results, counts = _create_lancedb_results( + record_batches, expected_dim=vector_dim if enforce_dim else None + ) - results, counts = _create_lancedb_results(records or [], expected_dim=expected_dim) - schema = _lancedb_arrow_schema(self.vector_dim, retrieval_mode="hybrid" if self.hybrid else "dense") + if self._service_table_schema: + results = _to_service_lancedb_rows(results) + schema = _with_retrieval_mode_metadata( + lancedb_schema(vector_dim), + "hybrid" if self.hybrid else "dense", + embedding_model_name=self.embedding_model_name, + embedding_model_revision=self.embedding_model_revision, + ) + else: + schema = _lancedb_arrow_schema( + vector_dim, + retrieval_mode="hybrid" if self.hybrid else "dense", + embedding_model_name=self.embedding_model_name, + embedding_model_revision=self.embedding_model_revision, + ) write_kwargs = { "on_bad_vectors": self.on_bad_vectors, @@ -513,6 +938,13 @@ def create_index(self, records=None, table_name: str = "nv-ingest", **kwargs): else: _validate_append_schema(table, schema, table_name=table_name, uri=self.uri) if results: + _validate_append_embedding_model( + table, + self.embedding_model_name, + self.embedding_model_revision, + table_name=table_name, + uri=self.uri, + ) existing_rows = int(table.count_rows()) logger.warning( "Appending %d row(s) to existing LanceDB table %r at %s " @@ -638,7 +1070,10 @@ def run(self, records): fts_language=self.fts_language, ) else: - logger.info("Skipping LanceDB index creation for table %r because build_index=False.", self.table_name) + logger.info( + "Skipping LanceDB index creation for table %r because build_index=False.", + self.table_name, + ) return records def put( diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py new file mode 100644 index 0000000000..ddfa0ec533 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py @@ -0,0 +1,1269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Private LanceDB persistence for collection-managed service data.""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import logging +import math +import threading +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +import lancedb +import pyarrow as pa + +from nemo_retriever.common.schemas.collections import ( + CollectionCreateRequest, + CollectionDeleteResult, + CollectionInfo, + CollectionPage, + CollectionUpdateRequest, + DocumentDeleteResult, + DocumentInfo, + DocumentPage, + IngestOperation, +) +from nemo_retriever.common.vdb.adt_vdb import ( + CollectionWriteContext, + CollectionWriteResult, + UnsupportedVDBOperation, + VDBInvalidRequest, + VDBResourceConflict, + VDBResourceNotFound, +) +from nemo_retriever.common.vdb.lancedb_capabilities import ( + LanceRetrievalMode, + LanceTableCapabilities, + inspect_lancedb_table_object, +) +from nemo_retriever.common.vdb.lancedb_schema import ( + create_or_append_lancedb_table, + infer_vector_dim, + lancedb_schema, +) +from nemo_retriever.common.vdb.records import ( + RetrievalContractError, + normalize_content_type, + normalize_retrieval_results, +) + +logger = logging.getLogger(__name__) + +_COLLECTIONS_TABLE = "_nrl_collections" +_DOCUMENTS_TABLE = "_nrl_documents" +_CATALOG_SCHEMA_VERSION = 2 +_CATALOG_SCAN_LIMIT = 100_000 +_NATIVE_SCORE_FIELDS = frozenset({"_distance", "_score"}) + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _is_uncommitted_initial_append(row: Mapping[str, Any]) -> bool: + return row.get("recovery_state") == "appending" and not row.get("current_document_version") + + +def _physical_table(scope: str, collection_name: str) -> str: + digest = hashlib.sha256(f"{scope}\0{collection_name}".encode()).hexdigest() + return f"nrl_{digest[:40]}" + + +def _quoted(value: str) -> str: + return "'" + value.replace("'", "''") + "'" + + +def _encode_cursor(resource: str, scope: str, collection: str | None, last: list[str]) -> str: + """Encode a pagination position bound to its logical resource context.""" + + payload = { + "v": 1, + "resource": resource, + "scope": scope, + "collection": collection, + "last": last, + } + return ( + base64.urlsafe_b64encode(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()) + .decode() + .rstrip("=") + ) + + +def _decode_cursor( + token: str | None, + *, + resource: str, + scope: str, + collection: str | None, +) -> list[str] | None: + """Decode a cursor after validating its resource, scope, and collection context.""" + + if not token: + return None + try: + payload = json.loads(base64.urlsafe_b64decode(token + "=" * (-len(token) % 4)).decode()) + except Exception as exc: + raise VDBInvalidRequest("Invalid continuation token") from exc + if ( + not isinstance(payload, dict) + or payload.get("v") != 1 + or payload.get("resource") != resource + or payload.get("scope") != scope + or payload.get("collection") != collection + or not isinstance(payload.get("last"), list) + ): + raise VDBInvalidRequest("Continuation token does not match this resource context") + return [str(value) for value in payload["last"]] + + +def _json_string(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value or {}, ensure_ascii=False, separators=(",", ":"), default=str) + + +def _content_text(record: dict[str, Any], metadata: dict[str, Any]) -> str: + content = metadata.get("content") + if isinstance(content, str): + return content + document_type = str(record.get("document_type") or "") + if document_type == "structured": + table_metadata = metadata.get("table_metadata") + if isinstance(table_metadata, dict): + return str(table_metadata.get("table_content") or "") + if document_type == "image": + image_metadata = metadata.get("image_metadata") + if isinstance(image_metadata, dict): + content_metadata = metadata.get("content_metadata") + is_page = isinstance(content_metadata, dict) and content_metadata.get("subtype") == "page_image" + return str(image_metadata.get("text" if is_page else "caption") or "") + if document_type == "audio": + audio_metadata = metadata.get("audio_metadata") + if isinstance(audio_metadata, dict): + return str(audio_metadata.get("audio_transcript") or "") + return "" + + +def _positive_or_unknown_page(value: Any) -> int: + if isinstance(value, bool): + return -1 + try: + page = int(value) + except (TypeError, ValueError): + return -1 + return page if page > 0 else -1 + + +def _collection_rows( + records: list, + *, + context: CollectionWriteContext, +) -> list[dict[str, Any]]: + """Convert canonical NRL record batches into collection-managed LanceDB rows.""" + rows: list[dict[str, Any]] = [] + created_at = _now() + row_index = 0 + + for batch in records or []: + if not isinstance(batch, list): + continue + for record in batch: + if not isinstance(record, dict): + continue + metadata = record.get("metadata") + if not isinstance(metadata, dict): + continue + vector = metadata.get("embedding") + if not isinstance(vector, (list, tuple)) or not vector: + continue + content_metadata = metadata.get("content_metadata") + if not isinstance(content_metadata, dict): + content_metadata = {} + source_metadata = metadata.get("source_metadata") + if not isinstance(source_metadata, dict): + source_metadata = {} + + text = _content_text(record, metadata) + content_type = normalize_content_type(content_metadata.get("type") or record.get("document_type")) + content_type = content_type or "" + if content_type: + content_metadata = dict(content_metadata) + content_metadata["type"] = content_type + content_metadata["_content_type"] = content_type + if not text.strip() and content_type != "image": + continue + + source_id = str( + source_metadata.get("source_id") or source_metadata.get("source_name") or context.filename or "" + ) + source_path = Path(source_id) if source_id else None + filename = context.filename or (source_path.name if source_path else "") + pdf_basename = source_path.stem if source_path else Path(filename).stem + page_number = _positive_or_unknown_page(content_metadata.get("page_number")) + pdf_page = f"{pdf_basename}_{page_number}" if pdf_basename and page_number > 0 else "" + stored_image_uri = str(content_metadata.get("stored_image_uri") or "") + bbox = content_metadata.get("bbox_xyxy_norm") + + rows.append( + { + "vector": list(vector), + "pdf_page": pdf_page, + "filename": filename, + "pdf_basename": pdf_basename, + "page_number": page_number, + "source": _json_string(source_metadata), + "source_id": source_id, + "path": source_id, + "text": text, + "metadata": _json_string(content_metadata), + "stored_image_uri": stored_image_uri, + "content_type": content_type, + "bbox_xyxy_norm": _json_string(bbox) if bbox else "", + "chunk_id": hashlib.sha256( + f"{context.document_id}\0{context.document_version}\0{row_index}".encode() + ).hexdigest(), + "document_id": context.document_id, + "document_version": context.document_version, + "content_sha256": context.content_sha256, + "created_at": created_at, + } + ) + row_index += 1 + return rows + + +def _public_collection_hit(hit: dict[str, Any]) -> dict[str, Any]: + """Expose a finite native distance without leaking LanceDB score fields.""" + raw = hit.get("_distance") + if isinstance(raw, bool): + raise RetrievalContractError("Dense collection hit is missing a numeric _distance") + try: + distance = float(raw) + except (TypeError, ValueError) as exc: + raise RetrievalContractError("Dense collection hit is missing a numeric _distance") from exc + if not math.isfinite(distance): + raise RetrievalContractError("Dense collection hit has a non-finite _distance") + + public_hit = {key: value for key, value in hit.items() if key not in _NATIVE_SCORE_FIELDS} + content_type = str(public_hit.get("content_type") or "").lower() + page_number = _positive_or_unknown_page(public_hit.get("page_number")) + if content_type.startswith(("audio", "video")) or page_number < 0: + public_hit["page_number"] = None + public_hit["pdf_page"] = "" + else: + public_hit["page_number"] = page_number + public_hit["distance"] = distance + return public_hit + + +def _normalize_collection_results( + raw_results: Any, + *, + expected_queries: int, +) -> list[list[dict[str, Any]]]: + """Strictly validate collection result cardinality and hit shape.""" + if not isinstance(raw_results, list) or len(raw_results) != expected_queries: + raise RetrievalContractError("Collection retrieval returned an invalid query-result cardinality") + for query_index, hits in enumerate(raw_results): + if not isinstance(hits, list): + raise RetrievalContractError(f"Collection retrieval result {query_index} is not a hit list") + for hit_index, hit in enumerate(hits): + if not isinstance(hit, Mapping): + raise RetrievalContractError(f"Collection retrieval hit {query_index}:{hit_index} is not a mapping") + return normalize_retrieval_results(raw_results) + + +class LanceDBCollectionStore: + """Implement optional VDB collection capabilities for LanceDB. + + The store owns private catalogs and maps logical ``(scope, collection)`` + identities to physical tables. Table-user leases keep deletion and recovery + from racing with active reads or writes. + """ + + def __init__(self, backend: Any, *, expiration_cleanup_enabled: bool = True) -> None: + self._backend = backend + self._uri = backend.uri + self.expiration_cleanup_enabled = expiration_cleanup_enabled + self.reconciliation_successes = 0 + self.reconciliation_failures = 0 + self._write_lock = threading.Lock() + self._collection_write_lock = threading.Lock() + self._table_user_condition = threading.Condition(self._write_lock) + self._active_table_users: dict[str, int] = {} + self._db = lancedb.connect(uri=self._uri) + self._opened_tables: dict[str, Any] = {} + self._ensure_catalogs() + + def _ensure_catalogs(self) -> None: + collection_schema = pa.schema( + [ + pa.field("scope", pa.string()), + pa.field("name", pa.string()), + pa.field("physical_table", pa.string()), + pa.field("status", pa.string()), + pa.field("description", pa.string()), + pa.field("metadata_json", pa.string()), + pa.field("created_at", pa.string()), + pa.field("updated_at", pa.string()), + pa.field("expires_at", pa.string()), + pa.field("deletion_phase", pa.string()), + pa.field("retry_count", pa.int64()), + pa.field("next_retry_at", pa.string()), + pa.field("last_error", pa.string()), + pa.field("delete_started_at", pa.string()), + ] + ) + document_schema = pa.schema( + [ + pa.field("scope", pa.string()), + pa.field("collection_name", pa.string()), + pa.field("document_id", pa.string()), + pa.field("job_id", pa.string()), + pa.field("filename", pa.string()), + pa.field("content_sha256", pa.string()), + pa.field("document_version", pa.string()), + pa.field("status", pa.string()), + pa.field("chunk_count", pa.int64()), + pa.field("created_at", pa.string()), + pa.field("updated_at", pa.string()), + pa.field("error", pa.string()), + pa.field("current_document_version", pa.string()), + pa.field("pending_document_version", pa.string()), + pa.field("recovery_state", pa.string()), + ] + ) + for name, schema in ( + (_COLLECTIONS_TABLE, collection_schema), + (_DOCUMENTS_TABLE, document_schema), + ): + if self._has_table(name): + table = self._db.open_table(name) + else: + table = self._db.create_table(name, schema=schema, mode="create") + existing = {field.name: field for field in table.schema} + for field in schema: + if field.name in existing and existing[field.name].type != field.type: + raise RuntimeError( + f"Incompatible {name} catalog column {field.name!r}: " + f"expected {field.type}, found {existing[field.name].type}" + ) + missing = [field for field in schema if field.name not in existing] + if missing: + missing_names = ", ".join(sorted(field.name for field in missing)) + raise RuntimeError(f"Incompatible {name} catalog: missing required columns: {missing_names}") + index_columns = ( + ("scope", "name", "status", "expires_at") + if name == _COLLECTIONS_TABLE + else ("scope", "collection_name", "document_id", "status") + ) + for column in index_columns: + table.create_scalar_index(column, replace=True) + + def _rows( + self, + table_name: str, + where: str | None = None, + columns: list[str] | None = None, + ) -> list[dict[str, Any]]: + query = self._db.open_table(table_name).search() + if where: + query = query.where(where) + if columns: + query = query.select(columns) + return query.limit(_CATALOG_SCAN_LIMIT).to_list() + + def _has_table(self, table_name: str) -> bool: + return table_name in self._db.list_tables().tables + + def _open_table(self, table_name: str) -> Any: + table = self._opened_tables.get(table_name) + if table is None: + table = self._db.open_table(table_name) + self._opened_tables[table_name] = table + return table + + def _acquire_table_user_locked(self, table_name: str) -> None: + self._active_table_users[table_name] = self._active_table_users.get(table_name, 0) + 1 + + def _release_table_user(self, table_name: str) -> None: + with self._table_user_condition: + remaining = self._active_table_users.get(table_name, 0) - 1 + if remaining > 0: + self._active_table_users[table_name] = remaining + else: + self._active_table_users.pop(table_name, None) + self._table_user_condition.notify_all() + + def _wait_for_table_users_locked(self, table_name: str) -> None: + """Wait under the state lock before destructively mutating an active table.""" + + while self._active_table_users.get(table_name, 0): + self._table_user_condition.wait() + + def _collection_row(self, scope: str, name: str, *, active: bool = False) -> dict[str, Any] | None: + rows = self._rows( + _COLLECTIONS_TABLE, + f"scope = {_quoted(scope)} AND name = {_quoted(name)}", + ) + row = rows[0] if rows else None + if row and active and row["status"] != "active": + raise VDBInvalidRequest(f"Collection {name!r} is {row['status']}") + if row and active and row.get("expires_at"): + expires = datetime.fromisoformat(str(row["expires_at"])) + if expires <= datetime.now(timezone.utc): + raise VDBInvalidRequest(f"Collection {name!r} is expired") + return row + + @staticmethod + def _collection_info(row: dict[str, Any]) -> CollectionInfo: + return CollectionInfo( + name=row["name"], + scope=row["scope"], + status=row["status"], + description=row.get("description") or None, + metadata=json.loads(row.get("metadata_json") or "{}"), + created_at=row["created_at"], + updated_at=row["updated_at"], + expires_at=row.get("expires_at") or None, + ) + + @staticmethod + def _document_info(row: dict[str, Any]) -> DocumentInfo: + return DocumentInfo(**{key: row.get(key) for key in DocumentInfo.model_fields}) + + def create_collection(self, scope: str, request: CollectionCreateRequest) -> CollectionInfo: + """Create a scoped logical collection without exposing its physical table.""" + + with self._write_lock: + if self._collection_row(scope, request.name): + raise VDBResourceConflict(f"Collection {request.name!r} already exists") + now = _now() + row = { + "scope": scope, + "name": request.name, + "physical_table": _physical_table(scope, request.name), + "status": "active", + "description": request.description or "", + "metadata_json": json.dumps(request.metadata, sort_keys=True), + "created_at": now, + "updated_at": now, + "expires_at": request.expires_at or "", + "deletion_phase": "", + "retry_count": 0, + "next_retry_at": "", + "last_error": "", + "delete_started_at": "", + } + self._db.open_table(_COLLECTIONS_TABLE).add([row]) + return self._collection_info(row) + + def get_collection(self, scope: str, name: str) -> CollectionInfo: + """Return one scoped collection or raise when it does not exist.""" + + row = self._collection_row(scope, name) + if not row: + raise VDBResourceNotFound("Collection not found") + return self._collection_info(row) + + def list_collections( + self, + scope: str, + limit: int, + continuation_token: str | None, + ) -> CollectionPage: + """List collections in one scope using a context-bound cursor.""" + + rows = sorted( + self._rows(_COLLECTIONS_TABLE, f"scope = {_quoted(scope)}"), + key=lambda row: row["name"], + ) + last = _decode_cursor( + continuation_token, + resource="collections", + scope=scope, + collection=None, + ) + if last is not None: + if len(last) != 1: + raise VDBInvalidRequest("Invalid collection continuation token") + rows = [row for row in rows if row["name"] > last[0]] + page = rows[:limit] + next_token = ( + _encode_cursor("collections", scope, None, [page[-1]["name"]]) if len(rows) > limit and page else None + ) + return CollectionPage(items=[self._collection_info(row) for row in page], next_token=next_token) + + def update_collection( + self, + scope: str, + name: str, + request: CollectionUpdateRequest, + ) -> CollectionInfo: + """Update mutable metadata for an active scoped collection.""" + + with self._write_lock: + row = self._collection_row(scope, name, active=True) + if not row: + raise VDBResourceNotFound("Collection not found") + update = request.model_dump(exclude_unset=True) + row["description"] = update.get("description", row["description"]) or "" + if "metadata" in update: + row["metadata_json"] = json.dumps(update["metadata"] or {}, sort_keys=True) + now = _now() + if "expires_at" in update: + row["expires_at"] = update["expires_at"] or "" + row["updated_at"] = now + else: + self._refresh_collection_activity_row(row, activity_at=now) + ( + self._db.open_table(_COLLECTIONS_TABLE) + .merge_insert(["scope", "name"]) + .when_matched_update_all() + .execute([row]) + ) + return self._collection_info(row) + + def _persist_collection_row(self, row: dict[str, Any]) -> None: + ( + self._db.open_table(_COLLECTIONS_TABLE) + .merge_insert(["scope", "name"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute([row]) + ) + + @staticmethod + def _refresh_collection_activity_row(row: dict[str, Any], *, activity_at: str) -> None: + """Advance collection activity while preserving its expiration window.""" + + expires_at = str(row.get("expires_at") or "") + updated_at = str(row.get("updated_at") or "") + row["updated_at"] = activity_at + if not expires_at or not updated_at: + return + ttl = datetime.fromisoformat(expires_at) - datetime.fromisoformat(updated_at) + if ttl.total_seconds() > 0: + row["expires_at"] = (datetime.fromisoformat(activity_at) + ttl).isoformat() + + def _refresh_collection_activity_locked(self, scope: str, name: str, *, activity_at: str) -> None: + """Persist successful indexing activity for an active collection.""" + + row = self._collection_row(scope, name) + if not row: + raise VDBResourceNotFound("Collection not found") + if row["status"] != "active": + return + updated_at = str(row.get("updated_at") or "") + if updated_at and datetime.fromisoformat(updated_at) >= datetime.fromisoformat(activity_at): + return + self._refresh_collection_activity_row(row, activity_at=activity_at) + self._persist_collection_row(row) + + @staticmethod + def _retry_at(retry_count: int) -> str: + delay = min(3600, 2 ** min(max(retry_count, 1), 12)) + return (datetime.now(timezone.utc) + timedelta(seconds=delay)).isoformat() + + def _schedule_collection_retry(self, row: dict[str, Any], phase: str, exc: Exception) -> None: + retries = int(row.get("retry_count") or 0) + 1 + row.update( + { + "status": "deleting", + "deletion_phase": phase, + "retry_count": retries, + "next_retry_at": self._retry_at(retries), + "last_error": str(exc)[:2000], + "updated_at": _now(), + } + ) + self._persist_collection_row(row) + + def _mark_collection_deleting_locked(self, row: dict[str, Any]) -> None: + """Enter the first deletion phase with a fresh retry budget. + + Callers must already hold ``_write_lock``; the transition is persisted here + so an interrupted process resumes from a durable phase. + """ + now = _now() + row.update( + { + "status": "deleting", + "deletion_phase": "drop_table", + "retry_count": 0, + "next_retry_at": "", + "last_error": "", + "delete_started_at": now, + "updated_at": now, + } + ) + self._persist_collection_row(row) + + def _cleanup_collection_locked(self, row: dict[str, Any]) -> bool: + phase = str(row.get("deletion_phase") or "drop_table") + try: + if phase == "drop_table": + self._wait_for_table_users_locked(row["physical_table"]) + self._db.drop_table(row["physical_table"], ignore_missing=True) + self._opened_tables.pop(row["physical_table"], None) + row["deletion_phase"] = phase = "delete_catalog" + row["updated_at"] = _now() + self._persist_collection_row(row) + if phase == "delete_catalog": + self._db.open_table(_DOCUMENTS_TABLE).delete( + f"scope = {_quoted(row['scope'])} AND collection_name = {_quoted(row['name'])}" + ) + self._db.open_table(_COLLECTIONS_TABLE).delete( + f"scope = {_quoted(row['scope'])} AND name = {_quoted(row['name'])}" + ) + return True + except Exception as exc: + logger.exception("Collection cleanup paused at phase %s", phase) + self._schedule_collection_retry(row, phase, exc) + return False + + def delete_collection( + self, + scope: str, + name: str, + if_exists: bool, + ) -> CollectionDeleteResult: + """Delete a collection through the retryable table-and-catalog lifecycle.""" + + with self._write_lock: + row = self._collection_row(scope, name) + if not row: + if if_exists: + return CollectionDeleteResult( + name=name, + scope=scope, + existed=False, + deleted=False, + status="deleted", + cleanup_pending=False, + ) + raise VDBResourceNotFound("Collection not found") + if row["status"] != "deleting": + self._mark_collection_deleting_locked(row) + deleted = self._cleanup_collection_locked(row) + return CollectionDeleteResult( + name=name, + scope=scope, + existed=True, + deleted=deleted, + status="deleted" if deleted else "deleting", + cleanup_pending=not deleted, + ) + + def _resolved_table(self, scope: str, name: str) -> str: + row = self._collection_row(scope, name, active=True) + if not row: + raise VDBResourceNotFound("Collection not found") + table_name = row["physical_table"] + return table_name + + def _table_capabilities(self, table_name: str) -> LanceTableCapabilities: + """Inspect an existing table; callers must have confirmed it exists.""" + return inspect_lancedb_table_object(self._open_table(table_name)) + + def _resolve_effective_retrieval_mode( + self, + table_name: str, + capabilities: LanceTableCapabilities | None, + ) -> LanceRetrievalMode: + if capabilities is None: + raise RetrievalContractError(f"Unable to inspect collection table {table_name!r}") + mode: LanceRetrievalMode = capabilities.retrieval_mode + if mode == "unknown": + raise RetrievalContractError("Collection table has no supported vector or FTS search capability") + if mode != "dense": + raise UnsupportedVDBOperation( + f"{mode.capitalize()} collection retrieval is not supported; collection queries require dense vectors" + ) + return mode + + def _persist_document_row(self, row: dict[str, Any]) -> None: + ( + self._db.open_table(_DOCUMENTS_TABLE) + .merge_insert(["scope", "collection_name", "document_id"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute([row]) + ) + + def _document_rows( + self, + scope: str, + collection_name: str, + document_id: str, + ) -> list[dict[str, Any]]: + return self._rows( + _DOCUMENTS_TABLE, + f"scope = {_quoted(scope)} AND collection_name = {_quoted(collection_name)} " + f"AND document_id = {_quoted(document_id)}", + ) + + def write_collection( + self, + records: list, + *, + context: CollectionWriteContext, + ) -> CollectionWriteResult: + """Append or replace one document using stable, retry-safe chunk identities.""" + + with self._collection_write_lock: + return self._write_collection_serialized(records, context=context) + + def _write_collection_serialized( + self, + records: list, + *, + context: CollectionWriteContext, + ) -> CollectionWriteResult: + """Persist recovery state around LanceDB I/O without blocking unrelated queries.""" + + rows = _collection_rows(records, context=context) + completed_row: dict[str, Any] | None = None + with self._write_lock: + table_name = self._resolved_table(context.scope, context.collection_name) + if records and not rows: + raise VDBInvalidRequest("Collection records produced no writable vector rows") + existing = self._document_rows( + context.scope, + context.collection_name, + context.document_id, + ) + if context.operation is IngestOperation.REPLACE: + if not existing: + raise VDBResourceNotFound("Document not found") + elif existing: + document = existing[0] + known_versions = { + str(document.get(field) or "") + for field in ( + "document_version", + "current_document_version", + "pending_document_version", + ) + if document.get(field) + } + if document.get("recovery_state") not in { + "", + "appending", + } or known_versions != {context.document_version}: + raise VDBResourceConflict("append cannot change an existing document; use replace") + stored_hash = str(document.get("content_sha256") or "") + if stored_hash and stored_hash != context.content_sha256: + raise VDBResourceConflict("append content does not match the existing document; use replace") + + if rows: + now = _now() + created_at = existing[0]["created_at"] if existing else now + if context.operation is IngestOperation.APPEND: + marker = ( + dict(existing[0]) + if existing + else { + "scope": context.scope, + "collection_name": context.collection_name, + "document_id": context.document_id, + "job_id": context.job_id or "", + "filename": context.filename, + "content_sha256": context.content_sha256, + "document_version": "", + "status": "appending", + "chunk_count": 0, + "created_at": created_at, + "updated_at": now, + "error": "", + "current_document_version": "", + "pending_document_version": "", + "recovery_state": "", + } + ) + marker.update( + { + "job_id": context.job_id or "", + "pending_document_version": context.document_version, + "recovery_state": "appending", + "updated_at": now, + "error": "", + } + ) + self._persist_document_row(marker) + elif existing: + marker = dict(existing[0]) + marker.update( + { + "status": "replacing", + "pending_document_version": context.document_version, + "recovery_state": "replacing", + "updated_at": now, + "error": "", + } + ) + self._persist_document_row(marker) + + completed_row = { + "scope": context.scope, + "collection_name": context.collection_name, + "document_id": context.document_id, + "job_id": context.job_id or "", + "filename": context.filename, + "content_sha256": context.content_sha256, + "document_version": context.document_version, + "status": "completed", + "chunk_count": len(rows), + "created_at": created_at, + "updated_at": now, + "error": "", + "current_document_version": context.document_version, + "pending_document_version": "", + "recovery_state": "refreshing_collection_activity", + } + cached_table = self._opened_tables.get(table_name) + self._acquire_table_user_locked(table_name) + + try: + table_exists = self._has_table(table_name) + table = cached_table + if rows: + if not table_exists: + vector_dim = infer_vector_dim(rows) + if vector_dim == 0: + raise VDBInvalidRequest("Cannot infer vector dimension from collection records") + schema = lancedb_schema(vector_dim=vector_dim, collection_managed=True) + table = create_or_append_lancedb_table( + self._db, + table_name, + rows, + schema, + overwrite=True, + ) + logger.info( + "Created collection LanceDB table %r with %d rows (dim=%d)", + table_name, + len(rows), + vector_dim, + ) + else: + table = table or self._db.open_table(table_name) + if context.operation is IngestOperation.REPLACE: + predicate = f"document_id = {_quoted(context.document_id)}" + ( + table.merge_insert("chunk_id") + .when_matched_update_all() + .when_not_matched_insert_all() + .when_not_matched_by_source_delete(predicate) + .execute(rows) + ) + else: + ( + table.merge_insert("chunk_id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(rows) + ) + logger.info( + "Wrote %d rows to collection table %r operation=%s", + len(rows), + table_name, + context.operation, + ) + elif table_exists: + table = table or self._db.open_table(table_name) + + total_rows = int(table.count_rows()) if table is not None else 0 + with self._write_lock: + if table is not None: + self._opened_tables[table_name] = table + if completed_row is not None: + self._persist_document_row(completed_row) + self._refresh_collection_activity_locked( + context.scope, + context.collection_name, + activity_at=completed_row["updated_at"], + ) + completed_row["recovery_state"] = "" + self._persist_document_row(completed_row) + return CollectionWriteResult(written=len(rows), total_rows=total_rows) + finally: + self._release_table_user(table_name) + + def retrieve_collection( + self, + vectors: list, + *, + scope: str, + collection_name: str, + query_texts: list[str], + top_k: int, + **kwargs: Any, + ) -> tuple[list[list[dict[str, Any]]], list[str]]: + """Run scoped dense retrieval and expose finite native vector distances.""" + + if len(query_texts) != len(vectors): + raise RetrievalContractError("query_texts must contain one entry per query vector") + with self._write_lock: + table_name = self._resolved_table(scope, collection_name) + if not self._has_table(table_name): + return ([[] for _ in vectors], ["dense"]) + capabilities = self._table_capabilities(table_name) + self._resolve_effective_retrieval_mode(table_name, capabilities) + retrieval_kwargs: dict[str, Any] = { + **kwargs, + "table_name": table_name, + "top_k": top_k, + "hybrid": False, + } + pending_document_ids = sorted( + str(row["document_id"]) + for row in self._rows( + _DOCUMENTS_TABLE, + f"scope = {_quoted(scope)} AND collection_name = {_quoted(collection_name)}", + ) + if _is_uncommitted_initial_append(row) + ) + if pending_document_ids: + visibility_filter = " AND ".join( + f"document_id != {_quoted(document_id)}" for document_id in pending_document_ids + ) + requested_filter = retrieval_kwargs.get("where", retrieval_kwargs.get("_filter")) + if requested_filter is not None and str(requested_filter).strip(): + visibility_filter = f"({str(requested_filter).strip()}) AND ({visibility_filter})" + retrieval_kwargs["where"] = visibility_filter + + if capabilities is not None and capabilities.vector_column and capabilities.vector_column != "vector": + retrieval_kwargs["vector_column_name"] = capabilities.vector_column + self._acquire_table_user_locked(table_name) + + try: + raw_results = self._backend.retrieval(vectors, **retrieval_kwargs) + normalized_results = _normalize_collection_results(raw_results, expected_queries=len(vectors)) + public_results = [[_public_collection_hit(hit) for hit in hits] for hits in normalized_results] + return public_results, ["dense"] + finally: + self._release_table_user(table_name) + + def list_documents( + self, + scope: str, + collection_name: str, + limit: int, + continuation_token: str | None, + ) -> DocumentPage: + """List committed documents in a collection using a context-bound cursor.""" + + self._resolved_table(scope, collection_name) + rows = self._rows( + _DOCUMENTS_TABLE, + f"scope = {_quoted(scope)} AND collection_name = {_quoted(collection_name)}", + ) + rows = [row for row in rows if not _is_uncommitted_initial_append(row)] + rows.sort(key=lambda row: (row["created_at"], row["document_id"])) + last = _decode_cursor( + continuation_token, + resource="documents", + scope=scope, + collection=collection_name, + ) + if last is not None: + if len(last) != 2: + raise VDBInvalidRequest("Invalid document continuation token") + rows = [row for row in rows if (row["created_at"], row["document_id"]) > (last[0], last[1])] + page = rows[:limit] + return DocumentPage( + items=[self._document_info(row) for row in page], + next_token=( + _encode_cursor( + "documents", + scope, + collection_name, + [page[-1]["created_at"], page[-1]["document_id"]], + ) + if len(rows) > limit and page + else None + ), + ) + + def get_document( + self, + scope: str, + collection_name: str, + document_id: str, + ) -> DocumentInfo: + """Return one committed document from a scoped collection.""" + + self._resolved_table(scope, collection_name) + rows = self._document_rows(scope, collection_name, document_id) + if not rows or _is_uncommitted_initial_append(rows[0]): + raise VDBResourceNotFound("Document not found") + return self._document_info(rows[0]) + + def _reconcile_document_row_locked(self, row: dict[str, Any], table_name: str) -> bool: + """Complete or roll back one interrupted document lifecycle operation.""" + + activity_refresh = "refreshing_collection_activity" + state = str(row.get("recovery_state") or "") + scope = str(row["scope"]) + collection_name = str(row["collection_name"]) + document_id = str(row["document_id"]) + try: + if state in {"appending", "replacing"}: + pending = str(row.get("pending_document_version") or "") + pending_chunks: list[dict[str, Any]] = [] + if self._has_table(table_name): + chunks = self._rows( + table_name, + f"document_id = {_quoted(document_id)}", + ["document_version", "content_sha256", "filename"], + ) + pending_chunks = [chunk for chunk in chunks if str(chunk.get("document_version") or "") == pending] + if pending and pending_chunks: + pending_chunk = pending_chunks[0] + activity_at = _now() + row.update( + { + "document_version": pending, + "current_document_version": pending, + "content_sha256": str( + pending_chunk.get("content_sha256") or row.get("content_sha256") or "" + ), + "filename": str(pending_chunk.get("filename") or row.get("filename") or ""), + "chunk_count": len(pending_chunks), + "pending_document_version": "", + "status": "completed", + "recovery_state": activity_refresh, + "updated_at": activity_at, + "error": "", + } + ) + self._persist_document_row(row) + state = activity_refresh + elif state == "appending" and not row.get("current_document_version"): + self._db.open_table(_DOCUMENTS_TABLE).delete( + f"scope = {_quoted(scope)} " + f"AND collection_name = {_quoted(collection_name)} " + f"AND document_id = {_quoted(document_id)}" + ) + return True + else: + row.update( + { + "pending_document_version": "", + "status": "completed", + "recovery_state": "", + "updated_at": _now(), + "error": "", + } + ) + self._persist_document_row(row) + return True + if state == activity_refresh: + self._refresh_collection_activity_locked( + scope, + collection_name, + activity_at=row["updated_at"], + ) + row.update({"recovery_state": "", "error": ""}) + self._persist_document_row(row) + return True + if state == "deleting_chunks": + if self._has_table(table_name): + self._wait_for_table_users_locked(table_name) + self._open_table(table_name).delete(f"document_id = {_quoted(document_id)}") + self._db.open_table(_DOCUMENTS_TABLE).delete( + f"scope = {_quoted(scope)} AND collection_name = {_quoted(collection_name)} " + f"AND document_id = {_quoted(document_id)}" + ) + return True + return state == "" + except Exception as exc: + row["error"] = str(exc)[:2000] + if state == activity_refresh: + row["recovery_state"] = activity_refresh + else: + row["updated_at"] = _now() + self._persist_document_row(row) + logger.exception("Document reconciliation paused in state %s", state) + return False + + def delete_document( + self, + scope: str, + collection_name: str, + document_id: str, + if_exists: bool, + ) -> DocumentDeleteResult: + """Delete a document's chunks and catalog record with recovery state.""" + + with self._write_lock: + table_name = self._resolved_table(scope, collection_name) + rows = self._document_rows(scope, collection_name, document_id) + if not rows: + if if_exists: + return DocumentDeleteResult( + document_id=document_id, + collection_name=collection_name, + scope=scope, + existed=False, + deleted=False, + status="deleted", + cleanup_pending=False, + ) + raise VDBResourceNotFound("Document not found") + row = rows[0] + if row.get("recovery_state") != "deleting_chunks": + row.update( + { + "status": "deleting", + "recovery_state": "deleting_chunks", + "updated_at": _now(), + "error": "", + } + ) + self._persist_document_row(row) + deleted = self._reconcile_document_row_locked(row, table_name) + return DocumentDeleteResult( + document_id=document_id, + collection_name=collection_name, + scope=scope, + existed=True, + deleted=deleted, + status="deleted" if deleted else "deleting", + cleanup_pending=not deleted, + ) + + def reconcile_collections(self) -> dict[str, int]: + """Resume recoverable VDB lifecycle work and expire due collections.""" + successes = 0 + failures = 0 + now = datetime.now(timezone.utc) + now_quoted = _quoted(now.isoformat()) + + for candidate in self._rows(_DOCUMENTS_TABLE, "recovery_state != ''"): + with self._write_lock: + rows = self._document_rows( + candidate["scope"], + candidate["collection_name"], + candidate["document_id"], + ) + if not rows or not rows[0].get("recovery_state"): + continue + row = rows[0] + collection = self._collection_row(row["scope"], row["collection_name"]) + if not collection: + continue + table_name = collection["physical_table"] + self._wait_for_table_users_locked(table_name) + rows = self._document_rows(row["scope"], row["collection_name"], row["document_id"]) + if not rows or not rows[0].get("recovery_state"): + continue + if self._reconcile_document_row_locked(rows[0], table_name): + successes += 1 + else: + failures += 1 + + collection_filter = f"status = 'deleting' AND (next_retry_at = '' OR next_retry_at <= {now_quoted})" + if self.expiration_cleanup_enabled: + collection_filter += " OR (status = 'active' AND expires_at != '' " f"AND expires_at <= {now_quoted})" + for candidate in self._rows(_COLLECTIONS_TABLE, collection_filter): + with self._write_lock: + row = self._collection_row(candidate["scope"], candidate["name"]) + if not row: + continue + if ( + self.expiration_cleanup_enabled + and row.get("status") == "active" + and row.get("expires_at") + and datetime.fromisoformat(str(row["expires_at"])) <= now + ): + self._mark_collection_deleting_locked(row) + if row.get("status") != "deleting": + continue + retry_at = str(row.get("next_retry_at") or "") + if retry_at and datetime.fromisoformat(retry_at) > now: + continue + if self._cleanup_collection_locked(row): + successes += 1 + else: + failures += 1 + + with self._write_lock: + self.reconciliation_successes += successes + self.reconciliation_failures += failures + return {"successes": successes, "failures": failures} + + @staticmethod + def empty_health() -> dict[str, Any]: + """Return collection health before the lazy store is initialized.""" + return { + "catalog": { + "healthy": True, + "initialized": False, + "schema_version": _CATALOG_SCHEMA_VERSION, + }, + "collections": {"active": 0, "deleting": 0, "expired": 0}, + "cleanup": { + "pending": 0, + "oldest_age_seconds": 0.0, + }, + "reconciliation": { + "successes": 0, + "failures": 0, + }, + "open_table_cache_count": 0, + } + + def health(self) -> dict[str, Any]: + """Summarize catalog, cleanup, and reconciliation state without identifiers.""" + + now = datetime.now(timezone.utc) + collections = self._rows( + _COLLECTIONS_TABLE, + columns=["status", "expires_at", "delete_started_at"], + ) + documents = self._rows(_DOCUMENTS_TABLE, columns=["recovery_state", "updated_at"]) + active = sum(row.get("status") == "active" for row in collections) + deleting = sum(row.get("status") == "deleting" for row in collections) + expired = sum( + bool(row.get("expires_at")) and datetime.fromisoformat(str(row["expires_at"])) <= now for row in collections + ) + pending_times: list[datetime] = [] + for row in collections: + if row.get("status") == "deleting" and row.get("delete_started_at"): + pending_times.append(datetime.fromisoformat(str(row["delete_started_at"]))) + for row in documents: + if row.get("recovery_state") and row.get("updated_at"): + pending_times.append(datetime.fromisoformat(str(row["updated_at"]))) + oldest_age = max(((now - started).total_seconds() for started in pending_times), default=0.0) + return { + "catalog": { + "healthy": True, + "initialized": True, + "schema_version": _CATALOG_SCHEMA_VERSION, + }, + "collections": { + "active": active, + "deleting": deleting, + "expired": expired, + }, + "cleanup": { + "pending": len(pending_times), + "oldest_age_seconds": round(oldest_age, 3), + }, + "reconciliation": { + "successes": self.reconciliation_successes, + "failures": self.reconciliation_failures, + }, + "open_table_cache_count": len(self._opened_tables), + } diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py index 685d6df880..d684add889 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py @@ -10,18 +10,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Tuple -_CONTENT_TYPE_ALIASES: dict[str, str] = { - "chart": "chart", - "chart_caption": "chart", - "image": "image", - "image_caption": "image", - "images": "image", - "infographic": "infographic", - "infographic_caption": "infographic", - "table": "table", - "table_caption": "table", - "text": "text", -} +from nemo_retriever.common.vdb.records import normalize_content_type def extract_embedding_from_row( @@ -132,13 +121,6 @@ def _build_detection_metadata(row: Any) -> Dict[str, Any]: return out -def normalize_content_type(value: Any) -> str | None: - normalized = str(value or "").strip().lower() - if not normalized: - return None - return _CONTENT_TYPE_ALIASES.get(normalized, normalized) - - def update_metadata_with_content_type(metadata_obj: Dict[str, Any], *, content_type: Any) -> None: normalized = normalize_content_type(content_type) if normalized is None: @@ -242,27 +224,36 @@ def build_lancedb_rows( return rows -def lancedb_schema(vector_dim: int = 2048) -> Any: +def lancedb_schema(vector_dim: int = 2048, *, collection_managed: bool = False) -> Any: """Return a PyArrow schema for the standard LanceDB table layout.""" import pyarrow as pa # type: ignore - return pa.schema( - [ - pa.field("vector", pa.list_(pa.float32(), vector_dim)), - pa.field("pdf_page", pa.string()), - pa.field("filename", pa.string()), - pa.field("pdf_basename", pa.string()), - pa.field("page_number", pa.int32()), - pa.field("source", pa.string()), - pa.field("source_id", pa.string()), - pa.field("path", pa.string()), - pa.field("text", pa.string()), - pa.field("metadata", pa.string()), - pa.field("stored_image_uri", pa.string()), - pa.field("content_type", pa.string()), - pa.field("bbox_xyxy_norm", pa.string()), - ] - ) + fields = [ + pa.field("vector", pa.list_(pa.float32(), vector_dim)), + pa.field("pdf_page", pa.string()), + pa.field("filename", pa.string()), + pa.field("pdf_basename", pa.string()), + pa.field("page_number", pa.int32()), + pa.field("source", pa.string()), + pa.field("source_id", pa.string()), + pa.field("path", pa.string()), + pa.field("text", pa.string()), + pa.field("metadata", pa.string()), + pa.field("stored_image_uri", pa.string()), + pa.field("content_type", pa.string()), + pa.field("bbox_xyxy_norm", pa.string()), + ] + if collection_managed: + fields.extend( + [ + pa.field("chunk_id", pa.string()), + pa.field("document_id", pa.string()), + pa.field("document_version", pa.string()), + pa.field("content_sha256", pa.string()), + pa.field("created_at", pa.string()), + ] + ) + return pa.schema(fields) def infer_vector_dim(rows: List[Dict[str, Any]]) -> int: diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/records.py b/nemo_retriever/src/nemo_retriever/common/vdb/records.py index c72721479a..c492ede1c9 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/records.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/records.py @@ -11,6 +11,88 @@ from pathlib import Path from typing import Any, TypedDict +from pydantic import ValidationError + +from nemo_retriever.common.schemas.collections import QueryHit + +_CONTENT_TYPE_ALIASES: dict[str, str] = { + "chart_caption": "chart", + "image_caption": "image", + "images": "image", + "infographic_caption": "infographic", + "table_caption": "table", +} + + +def normalize_content_type(value: Any) -> str | None: + """Return the canonical public modality for an extracted content type.""" + normalized = str(value or "").strip().lower() + return _CONTENT_TYPE_ALIASES.get(normalized, normalized) or None + + +class RetrievalContractError(RuntimeError): + """A backend retrieval result cannot satisfy the canonical hit contract.""" + + +def validate_collection_retrieval_results( + result: Any, + *, + expected_queries: int, +) -> tuple[list[list[dict[str, Any]]], list[str]]: + """Validate the backend-neutral collection retrieval contract.""" + + if not isinstance(result, tuple) or len(result) != 2: + raise RetrievalContractError("collection retrieval must return hits and strategies") + + hits_by_query, strategies = result + if not isinstance(hits_by_query, list) or len(hits_by_query) != expected_queries: + raise RetrievalContractError("collection retrieval returned an unexpected number of result sets") + if ( + not isinstance(strategies, list) + or not strategies + or any(not isinstance(strategy, str) or not strategy.strip() for strategy in strategies) + ): + raise RetrievalContractError("collection retrieval returned invalid strategies") + + validated: list[list[dict[str, Any]]] = [] + for hits in hits_by_query: + if not isinstance(hits, list): + raise RetrievalContractError("each collection query result must be a list") + validated_hits: list[dict[str, Any]] = [] + for hit in hits: + if not isinstance(hit, Mapping): + raise RetrievalContractError("each collection query hit must be a mapping") + payload = dict(hit) + if "bbox" not in payload and "bbox_xyxy_norm" in payload: + payload["bbox"] = payload["bbox_xyxy_norm"] + try: + canonical = QueryHit.model_validate(payload) + except ValidationError as exc: + raise RetrievalContractError("collection query hit does not satisfy the public contract") from exc + public_hit = canonical.model_dump(exclude_none=True, exclude={"bbox"}) + if canonical.bbox is not None: + public_hit["bbox_xyxy_norm"] = canonical.bbox + validated_hits.append(public_hit) + validated.append(validated_hits) + return validated, list(strategies) + + +_LEGACY_ENTITY_FIELDS = frozenset( + { + "content", + "content_metadata", + "metadata", + "page_number", + "path", + "pdf_basename", + "pdf_page", + "source", + "source_id", + "source_metadata", + "text", + } +) + class RetrievalHit(TypedDict, total=False): """Shape of a single hit returned by ``Retriever.query`` / ``Retriever.queries``. @@ -20,6 +102,10 @@ class RetrievalHit(TypedDict, total=False): a re-encoded string leak back out here. See ``_normalize_hit`` for the contract enforcement point. + ``_distance`` is a backend-native vector distance and ``_score`` is a + backend-native FTS/BM25 score. Dense collection REST responses expose the finite + native distance without reinterpreting it as confidence or similarity. + ``total=False`` because optional fields (``stored_image_uri``, ``content_type``, ``bbox_xyxy_norm``, scores) are only set when present. """ @@ -37,6 +123,11 @@ class RetrievalHit(TypedDict, total=False): bbox_xyxy_norm: list[float] _distance: float _score: float + chunk_id: str + document_id: str + filename: str + document_version: str + content_sha256: str def _embedding_from_graph_row(row: dict[str, Any], metadata: dict[str, Any]) -> Any: @@ -70,13 +161,58 @@ def _optional_int(value: Any) -> int | None: return None +def _page_number_from_graph_row( + row: dict[str, Any], + content_metadata: dict[str, Any], +) -> int | None: + """Preserve the original document page carried by service pre-splitting.""" + candidates = ( + content_metadata.get("page_number"), + row.get("page_number"), + row.get("_page_number"), + ) + parsed = [page for value in candidates if (page := _optional_int(value)) is not None] + return max(parsed) if parsed else None + + +def _add_detection_metadata( + row: dict[str, Any], + content_metadata: dict[str, Any], +) -> None: + """Carry graph extraction diagnostics into the canonical record metadata.""" + count = _optional_int(row.get("page_elements_v3_num_detections")) + if count is not None: + content_metadata.setdefault("page_elements_v3_num_detections", count) + + counts_by_label = row.get("page_elements_v3_counts_by_label") + if isinstance(counts_by_label, dict): + normalized_counts = { + str(key): count + for key, value in counts_by_label.items() + if isinstance(key, str) and (count := _optional_int(value)) is not None + } + if normalized_counts: + content_metadata.setdefault("page_elements_v3_counts_by_label", normalized_counts) + + for content_type in ("table", "chart", "infographic"): + detections = row.get(content_type) + if isinstance(detections, list): + content_metadata.setdefault(f"ocr_{content_type}_detections", len(detections)) + + def _dict_or_empty(value: Any) -> dict[str, Any]: return dict(value) if isinstance(value, dict) else {} def _is_image_backed_row(row: dict[str, Any]) -> bool: """Return whether a post-embed graph row retains its image or stored URI.""" - return bool(_first_str(row.get("_image_b64"), row.get("_stored_image_uri"), row.get("stored_image_uri"))) + return bool( + _first_str( + row.get("_image_b64"), + row.get("_stored_image_uri"), + row.get("stored_image_uri"), + ) + ) def _derive_fidelity(content_type: Any, metadata: dict[str, Any], content_metadata: dict[str, Any]) -> str | None: @@ -110,18 +246,17 @@ def _client_record_from_graph_row(row: dict[str, Any], *, require_embedding: boo return None content_metadata = _dict_or_empty(metadata.get("content_metadata")) - page_number = _optional_int(content_metadata.get("page_number")) - if page_number is None: - page_number = _optional_int(row.get("page_number")) + page_number = _page_number_from_graph_row(row, content_metadata) if page_number is not None: - content_metadata.setdefault("page_number", page_number) + content_metadata["page_number"] = page_number + _add_detection_metadata(row, content_metadata) if image_only: content_type = "image" content_metadata["type"] = content_type content_metadata.pop("fidelity", None) else: - content_type = row.get("_content_type") or row.get("content_type") + content_type = normalize_content_type(row.get("_content_type") or row.get("content_type")) if content_type: content_metadata.setdefault("type", content_type) fidelity = _derive_fidelity(content_type, metadata, content_metadata) @@ -134,7 +269,13 @@ def _client_record_from_graph_row(row: dict[str, Any], *, require_embedding: boo if bbox: content_metadata.setdefault("bbox_xyxy_norm", bbox) - for key in ("segment_start_seconds", "segment_end_seconds", "frame_timestamp_seconds"): + for key in ( + "chunk_index", + "chunk_count", + "segment_start_seconds", + "segment_end_seconds", + "frame_timestamp_seconds", + ): if key in metadata: content_metadata.setdefault(key, metadata[key]) @@ -163,7 +304,7 @@ def _client_record_from_graph_row(row: dict[str, Any], *, require_embedding: boo return {"document_type": str(document_type), "metadata": record_metadata} -def to_client_vdb_records(rows: list[dict[str, Any]]) -> list[list[dict[str, Any]]]: +def to_client_vdb_records(rows: Any) -> list[list[dict[str, Any]]]: """Convert graph-ingest rows into the nested record shape expected by client VDBs. Dense rows require an embedding and either nonblank text or concrete image backing. @@ -172,6 +313,10 @@ def to_client_vdb_records(rows: list[dict[str, Any]]) -> list[list[dict[str, Any When at least one row converts, returns ``[batch]`` with a single non-empty inner list (never ``[[]]``, which would be truthy and could trip backends on an empty insert). """ + if isinstance(rows, list) and all(isinstance(batch, list) for batch in rows): + return rows + if hasattr(rows, "to_pandas"): + rows = rows.to_pandas() if hasattr(rows, "to_dict"): rows = rows.to_dict("records") # Walrus: bind conversion once per row — a plain ``if f(row)`` + ``f(row)`` list comp @@ -216,31 +361,47 @@ def _mapping(value: Any) -> dict[str, Any]: return parsed if isinstance(parsed, dict) else {} +def _flatten_legacy_entity_hit(hit: dict[str, Any]) -> dict[str, Any]: + """Flatten the legacy nested ``entity`` shape into canonical fields. + + Only fields that predate collection management are accepted from the + nested shape. Top-level values are authoritative when both shapes provide + the same field; collection identity and version fields must be top-level. + """ + + top_level = {key: value for key, value in hit.items() if key != "entity"} + entity = hit.get("entity") + if not isinstance(entity, dict): + return top_level + legacy = {key: entity[key] for key in _LEGACY_ENTITY_FIELDS if key in entity} + return {**legacy, **top_level} + + def _normalize_hit(hit: dict[str, Any]) -> RetrievalHit: """Adapt LanceDB client hit shapes to Retriever hits.""" - entity = hit.get("entity") if isinstance(hit.get("entity"), dict) else hit + hit = _flatten_legacy_entity_hit(hit) - source = _mapping(entity.get("source") or hit.get("source") or entity.get("source_metadata")) - if not source and isinstance(entity.get("source"), str): - source = {"source_id": entity["source"]} - content_metadata = _mapping(entity.get("content_metadata") or hit.get("content_metadata") or entity.get("metadata")) + source = _mapping(hit.get("source") or hit.get("source_metadata")) + if not source and isinstance(hit.get("source"), str): + source = {"source_id": hit["source"]} + content_metadata = _mapping(hit.get("content_metadata") or hit.get("metadata")) source_id = _first_str( source.get("source_id"), source.get("source_name"), - entity.get("source_id"), hit.get("source_id"), hit.get("path"), ) page_number = content_metadata.get("page_number") if isinstance(content_metadata, dict) else None if page_number is None: - page_number = entity.get("page_number", hit.get("page_number")) + page_number = hit.get("page_number") page_number = _optional_int(page_number) + content_type = normalize_content_type(_first_str(hit.get("content_type"), content_metadata.get("type"))) or "" path = Path(source_id) if source_id else None pdf_basename = path.stem if path is not None else "" normalized: RetrievalHit = { - "text": _first_str(entity.get("text"), entity.get("content"), hit.get("text")), + "text": _first_str(hit.get("text"), hit.get("content")), # Keep `metadata` as a native dict on the API boundary. The LanceDB # storage layer JSON-encodes it on write (see `_json_str` in # `vdb/lancedb.py`); we already parse it back on read in @@ -254,13 +415,32 @@ def _normalize_hit(hit: dict[str, Any]) -> RetrievalHit: "path": source_id, "page_number": page_number, "pdf_basename": pdf_basename, - "pdf_page": f"{pdf_basename}_{page_number}" if pdf_basename and page_number is not None else "", + "pdf_page": (f"{pdf_basename}_{page_number}" if pdf_basename and page_number is not None else ""), } - for key in ("stored_image_uri", "content_type", "bbox_xyxy_norm", "_distance", "_score"): + chunk_id = hit.get("chunk_id") + if chunk_id: + normalized.update( + { + "chunk_id": str(chunk_id), + "document_id": str(hit.get("document_id") or ""), + "filename": str(hit.get("filename") or (path.name if path else "")), + "document_version": str(hit.get("document_version") or ""), + "content_sha256": str(hit.get("content_sha256") or ""), + } + ) + stored_image_uri = _first_str(hit.get("stored_image_uri"), content_metadata.get("stored_image_uri")) + if stored_image_uri: + normalized["stored_image_uri"] = stored_image_uri + if content_type: + normalized["content_type"] = content_type + bbox = hit.get("bbox_xyxy_norm") + if bbox is None or (isinstance(bbox, str) and not bbox.strip()): + bbox = content_metadata.get("bbox_xyxy_norm") + if bbox is not None and not (isinstance(bbox, str) and not bbox.strip()): + normalized["bbox_xyxy_norm"] = bbox + for key in ("_distance", "_score"): if key in hit: normalized[key] = hit[key] - elif key in entity: - normalized[key] = entity[key] return normalized @@ -278,7 +458,11 @@ def _hit_to_dict(hit: Any) -> dict[str, Any] | None: return None -def normalize_retrieval_results(results: Any) -> list[list[RetrievalHit]]: +def normalize_retrieval_results( + results: Any, +) -> list[list[RetrievalHit]]: + """Canonicalize backend result shapes without changing native scores.""" + if results is None: return [] if isinstance(results, dict): diff --git a/nemo_retriever/src/nemo_retriever/graph/retriever.py b/nemo_retriever/src/nemo_retriever/graph/retriever.py index cf98f4f521..d342523a87 100644 --- a/nemo_retriever/src/nemo_retriever/graph/retriever.py +++ b/nemo_retriever/src/nemo_retriever/graph/retriever.py @@ -11,7 +11,7 @@ import pandas as pd -from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL +from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL, resolve_embed_model, resolve_embed_model_spec from nemo_retriever.graph.retriever_utils import ( filter_retrieval_kwargs, rerank_long_dataframe_to_hits, @@ -115,7 +115,14 @@ def _merge_embed_params(self, extra: Optional[dict[str, Any]] = None) -> Any: "embed_inference_batch_size": 32, "local_ingest_embed_backend": "hf", } - merged = {**base, **dict(self.embed_kwargs or {}), **dict(extra or {})} + overrides = {**dict(self.embed_kwargs or {}), **dict(extra or {})} + merged = {**base, **overrides} + endpoint = str(merged.get("embedding_endpoint") or merged.get("embed_invoke_url") or "").strip() + if self.run_mode == "local" and not endpoint and overrides.get("local_ingest_embed_backend") is None: + model_id = str(merged.get("embed_model_name") or merged.get("model_name") or "").strip() + spec = resolve_embed_model_spec(model_id, revision=merged.get("embed_model_revision")) + merged["local_ingest_embed_backend"] = "vllm" if spec.requires_vllm else "hf" + merged["embed_model_revision"] = spec.revision if "local_ingest_embed_backend" in merged and merged["local_ingest_embed_backend"] is not None: merged["local_ingest_embed_backend"] = normalize_backend( str(merged["local_ingest_embed_backend"]), @@ -302,6 +309,38 @@ def _resolve_lancedb_query_mode( return mode, caps, uri, table_name, mode_override != "auto" + @staticmethod + def _embedding_model_from_kwargs(kwargs: Optional[dict[str, Any]]) -> str | None: + values = dict(kwargs or {}) + for key in ("model_name", "embed_model_name"): + value = str(values.get(key) or "").strip() + if value: + return value + return None + + def _resolve_embed_kwargs( + self, + index_model: str | None, + runtime_embed_kwargs: Optional[dict[str, Any]], + index_revision: str | None = None, + ) -> dict[str, Any]: + """Choose the query model snapshot: explicit override, index metadata, or default.""" + resolved = dict(runtime_embed_kwargs or {}) + runtime_model = self._embedding_model_from_kwargs(runtime_embed_kwargs) + configured_model = self._embedding_model_from_kwargs(self.embed_kwargs) + explicit_model = runtime_model or configured_model + model_name = explicit_model or index_model + model_name = resolve_embed_model(model_name) + resolved["model_name"] = model_name + resolved["embed_model_name"] = model_name + if runtime_model and "embed_model_revision" not in resolved: + resolved_configured = resolve_embed_model(configured_model) if configured_model else None + if resolved_configured != model_name: + resolved["embed_model_revision"] = None + if explicit_model is None and index_revision: + resolved.setdefault("embed_model_revision", index_revision) + return resolved + def _execute_sparse_lancedb_queries( self, query_texts: list[str], @@ -399,6 +438,16 @@ def queries( retrieval_top_k = candidate_top_k * refine if self.rerank else candidate_top_k vdb_call_kwargs = dict(vdb_kwargs or {}) + index_model: str | None = None + index_revision: str | None = None + explicit_model = self._embedding_model_from_kwargs(embed_kwargs) or self._embedding_model_from_kwargs( + self.embed_kwargs + ) + if self.graph is None and explicit_model is None: + metadata_reader = RetrieveVdbOperator(**_coerce_vdb_init(self.vdb_kwargs)) + index_model = metadata_reader.get_index_metadata("embedding_model_name", **vdb_call_kwargs) + index_revision = metadata_reader.get_index_metadata("embedding_model_revision", **vdb_call_kwargs) + lancedb_mode = self._resolve_lancedb_query_mode(vdb_call_kwargs) for key in _QUERY_ROUTING_VDB_KWARGS: vdb_call_kwargs.pop(key, None) @@ -428,6 +477,8 @@ def queries( vdb_call_kwargs["hybrid"] = False if caps.vector_column and caps.vector_column != "vector": vdb_call_kwargs.setdefault("vector_column_name", caps.vector_column) + if self.graph is None: + embed_kwargs = self._resolve_embed_kwargs(index_model, embed_kwargs, index_revision) raw_hits = self._execute_queries_graph( query_texts, diff --git a/nemo_retriever/src/nemo_retriever/harness/beir_runner.py b/nemo_retriever/src/nemo_retriever/harness/beir_runner.py index 051863093b..beb26e96e7 100644 --- a/nemo_retriever/src/nemo_retriever/harness/beir_runner.py +++ b/nemo_retriever/src/nemo_retriever/harness/beir_runner.py @@ -222,8 +222,8 @@ def _agentic_retrieve( try: cfg = build_agentic_config(query_request, top_k=agentic_target_top_k("beir", list(ks))) except (ValueError, TypeError) as exc: - # Invalid agentic config (e.g. out-of-range temperature, backend_top_k < top_k) - # surfaces as a structured harness failure rather than an unhandled exception. + # Invalid agentic config (e.g. out-of-range temperature) surfaces as a + # structured harness failure rather than an unhandled exception. raise HarnessRunError( EXIT_INVALID, FailurePayload( diff --git a/nemo_retriever/src/nemo_retriever/harness/execution.py b/nemo_retriever/src/nemo_retriever/harness/execution.py index 99af0d0338..6e6fe0a4f2 100644 --- a/nemo_retriever/src/nemo_retriever/harness/execution.py +++ b/nemo_retriever/src/nemo_retriever/harness/execution.py @@ -443,7 +443,10 @@ def run_prepared_benchmark( writer.path("run.log"), label="service_ingest" if service_mode else "ingest" ): if service_mode: - ingest_summary = execute_service_ingest_request(ingest_request).to_summary_dict() + ingest_summary = execute_service_ingest_request( + ingest_request, + return_results=False, + ).to_summary_dict() else: ingest_summary = run_ingest_workflow(ingest_plan, dry_run=False) except Exception as exc: diff --git a/nemo_retriever/src/nemo_retriever/harness/resolution.py b/nemo_retriever/src/nemo_retriever/harness/resolution.py index 006530e42a..87c7287fb6 100644 --- a/nemo_retriever/src/nemo_retriever/harness/resolution.py +++ b/nemo_retriever/src/nemo_retriever/harness/resolution.py @@ -134,11 +134,11 @@ def _override_child_keys(prefix: str, paths: set[str]) -> set[str]: "query.agentic_local_max_model_len", "query.agentic_local_max_num_seqs", "query.agentic_reasoning_effort", - "query.agentic_backend_top_k", "query.agentic_react_max_steps", "query.agentic_text_truncation", "query.agentic_num_concurrent", "query.agentic_temperature", + "query.agentic_llm_client", } EVALUATION_OVERRIDE_PATHS = { "evaluation.mode", @@ -447,11 +447,11 @@ def build_query_request(resolved: dict[str, Any], query_text: str) -> QueryReque local_max_model_len=query.get("agentic_local_max_model_len"), local_max_num_seqs=query.get("agentic_local_max_num_seqs"), reasoning_effort=query.get("agentic_reasoning_effort"), - backend_top_k=int(query.get("agentic_backend_top_k") or 20), react_max_steps=int(query.get("agentic_react_max_steps") or 50), text_truncation=int(query.get("agentic_text_truncation") or 0), num_concurrent=int(query.get("agentic_num_concurrent") or 1), - temperature=float(query.get("agentic_temperature") or 0.0), + temperature=(float(query["agentic_temperature"]) if query.get("agentic_temperature") is not None else None), + llm_client=query.get("agentic_llm_client"), ), ) diff --git a/nemo_retriever/src/nemo_retriever/ingest/plan.py b/nemo_retriever/src/nemo_retriever/ingest/plan.py index b041e25d07..93456d8a7c 100644 --- a/nemo_retriever/src/nemo_retriever/ingest/plan.py +++ b/nemo_retriever/src/nemo_retriever/ingest/plan.py @@ -39,6 +39,8 @@ expand_input_file_patterns, resolve_input_files, ) +from nemo_retriever.models import resolve_embed_model +from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision IngestRunModeValue = Literal["inprocess", "batch"] IngestInputTypeValue = Literal["auto", "pdf", "doc", "txt", "html", "image", "audio", "video"] @@ -634,9 +636,19 @@ def resolve_ingest_plan(request: IngestPlanRequest) -> ResolvedIngestPlan: if extract_tuning is not None: extract_kwargs["batch_tuning"] = extract_tuning + embedding_model_name = None if validated_index_mode == "sparse" else resolve_embed_model(embed.embed_model_name) + embedding_model_revision = None + if embedding_model_name is not None and not str(embed.embed_invoke_url or "").strip(): + embedding_model_revision = resolve_embed_model_revision(embedding_model_name, None) + embed_runtime_model_name = ( + embedding_model_name + if embed.embed_model_name is not None or embed.embed_model_provider_prefix is not None + else None + ) embed_kwargs = build_embed_option_kwargs( embed.embed_invoke_url, - embed.embed_model_name, + embed_runtime_model_name, + embed_model_revision=embedding_model_revision if embed_runtime_model_name is not None else None, local_ingest_embed_backend=embed.local_ingest_embed_backend, embed_api_key=embed.embed_api_key, embed_model_provider_prefix=embed.embed_model_provider_prefix, @@ -661,6 +673,12 @@ def resolve_ingest_plan(request: IngestPlanRequest) -> ResolvedIngestPlan: vdb_upload_kwargs["sparse"] = True elif validated_index_mode == "hybrid": vdb_upload_kwargs["hybrid"] = True + if embedding_model_name is not None: + vdb_upload_kwargs["embedding_model_name"] = embedding_model_name + if embed.embed_model_name is not None: + vdb_upload_kwargs["vector_dim"] = None + if embedding_model_revision is not None: + vdb_upload_kwargs["embedding_model_revision"] = embedding_model_revision vdb_params = VdbUploadParams(vdb_kwargs=vdb_upload_kwargs) caption_params = build_caption_params( enabled=request.caption.enabled, diff --git a/nemo_retriever/src/nemo_retriever/ingest/service.py b/nemo_retriever/src/nemo_retriever/ingest/service.py index 606a2d058e..a4623160a0 100644 --- a/nemo_retriever/src/nemo_retriever/ingest/service.py +++ b/nemo_retriever/src/nemo_retriever/ingest/service.py @@ -132,9 +132,10 @@ class ServiceIngestExecutionResult: """Structured result from executing a resolved service ingest request. Service mode does not locally verify the remote vector database after - ingest. ``result_n_rows`` counts rows from the service ingest result when - available, and ``n_rows`` mirrors that value so root CLI summaries keep the - same top-level row-count contract as local ingest results. + ingest. ``result_n_rows`` sums the row counts reported by successful + document-completion events, and ``n_rows`` mirrors that value so root CLI + summaries keep the same top-level row-count contract as local ingest + results without downloading retained result payloads. """ request: ServiceIngestRequest @@ -268,10 +269,19 @@ def build_service_ingestor(request: ServiceIngestRequest) -> Any: return ingestor -def execute_service_ingest_request(request: ServiceIngestRequest) -> ServiceIngestExecutionResult: - """Execute a service ingest request and return its structured result.""" +def execute_service_ingest_request( + request: ServiceIngestRequest, + *, + return_results: bool = True, +) -> ServiceIngestExecutionResult: + """Execute a service ingest request and return its structured result. + + ``return_results`` defaults to the user-facing service client behavior. + Callers that only need completion metadata may disable retained-result + downloads explicitly. + """ - result = build_service_ingestor(request).ingest() + result = build_service_ingestor(request).ingest(return_results=return_results) failures = list(getattr(result, "failures", ()) or ()) if failures: document, detail = failures[0] @@ -463,9 +473,26 @@ def _sanitize_service_caption_params(caption_params: CaptionParams) -> CaptionPa def _count_service_result_rows(result: object) -> int | None: dataframe = getattr(result, "dataframe", None) - if dataframe is None: - return None + if dataframe is not None: + try: + return len(dataframe) + except TypeError: + return None + try: - return len(dataframe) + events = iter(result) except TypeError: return None + + total = 0 + for event in events: + if not isinstance(event, dict) or event.get("status") != "completed": + continue + result_rows = event.get("result_rows", 0) + if result_rows is None: + continue + try: + total += int(result_rows) + except (TypeError, ValueError): + return None + return total diff --git a/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py b/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py index a7c08c84e8..8f414049b8 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py @@ -41,6 +41,7 @@ class ExtractionBranchExecutor: branches: tuple[ExtractionBranchPlan, ...] documents: list[str] buffers: list[tuple[str, BytesIO]] + inline_rows: list[dict[str, str]] split_config: dict[str, Any] extract_params: Any | None text_params: Any | None @@ -77,7 +78,7 @@ def execute(self) -> Any: return self._execute_inprocess() def _execute_batch(self) -> Any: - _ray, cluster_resources = self.ensure_batch_runtime() + ray_module, cluster_resources = self.ensure_batch_runtime() effective_allow_no_gpu = self.allow_no_gpu or cluster_resources.available_gpu_count() == 0 branch_datasets: list[Any] = [] for branch in self.branches: @@ -99,7 +100,11 @@ def _execute_batch(self) -> Any: video_frame_params=effective_extraction.video_frame_params, ) executor = self._ray_executor(graph, derived_overrides) - branch_datasets.append(executor.build_dataset(list(branch.input_paths))) + file_paths, inline_rows = self._partition_branch_inputs(branch) + if file_paths: + branch_datasets.append(executor.build_dataset(file_paths)) + if inline_rows: + branch_datasets.append(executor.build_dataset(ray_module.data.from_items(inline_rows))) normalized = normalize_ray_branch_datasets(branch_datasets) combined = normalized[0] @@ -198,7 +203,8 @@ def _ray_executor(self, graph: Any, derived_overrides: dict[str, dict[str, Any]] ) def _inprocess_branch_input(self, branch: ExtractionBranchPlan) -> Any: - if not self.buffers: + inline_by_path = self._inline_rows_by_path() + if not self.buffers and not any(path in inline_by_path for path in branch.input_paths): return list(branch.input_paths) import pandas as pd @@ -206,8 +212,11 @@ def _inprocess_branch_input(self, branch: ExtractionBranchPlan) -> Any: buffer_by_name = {name: buf for name, buf in self.buffers} file_paths: list[str] = [] buffer_rows: list[dict[str, Any]] = [] + inline_rows: list[dict[str, str]] = [] for path in branch.input_paths: - if path in buffer_by_name: + if path in inline_by_path: + inline_rows.append(inline_by_path[path]) + elif path in buffer_by_name: buffer_rows.append({"bytes": buffer_by_name[path].getvalue(), "path": path}) else: file_paths.append(path) @@ -217,8 +226,25 @@ def _inprocess_branch_input(self, branch: ExtractionBranchPlan) -> Any: frames.append(InprocessExecutor._load_files(file_paths)) if buffer_rows: frames.append(pd.DataFrame(buffer_rows)) + if inline_rows: + frames.append(pd.DataFrame(inline_rows)) return concat_dataframes(frames) + def _inline_rows_by_path(self) -> dict[str, dict[str, str]]: + return {row["path"]: row for row in self.inline_rows} + + def _partition_branch_inputs(self, branch: ExtractionBranchPlan) -> tuple[list[str], list[dict[str, str]]]: + inline_by_path = self._inline_rows_by_path() + file_paths: list[str] = [] + inline_rows: list[dict[str, str]] = [] + for path in branch.input_paths: + row = inline_by_path.get(path) + if row is None: + file_paths.append(path) + else: + inline_rows.append(row) + return file_paths, inline_rows + def merge_node_overrides( derived_overrides: dict[str, dict[str, Any]], diff --git a/nemo_retriever/src/nemo_retriever/ingestor/core.py b/nemo_retriever/src/nemo_retriever/ingestor/core.py index 6c195621c5..c300089cab 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/core.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/core.py @@ -16,7 +16,7 @@ from __future__ import annotations from io import BytesIO -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Self, Sequence, Tuple, Union from nemo_retriever.common.params import CaptionParams from nemo_retriever.common.params import DedupParams @@ -106,6 +106,10 @@ def files(self, documents: Union[str, List[str]]) -> "ingestor": """Add document paths/URIs for processing.""" self._not_implemented("files") + def texts(self, texts: Union[str, Sequence[str]]) -> Self: + """Set raw inline text documents for processing.""" + self._not_implemented("texts") + def buffers(self, buffers: Union[Tuple[str, BytesIO], List[Tuple[str, BytesIO]]]) -> "ingestor": """Add in-memory buffers for processing.""" self._not_implemented("buffers") diff --git a/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py b/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py index 641843da78..e25286b99f 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py @@ -30,7 +30,7 @@ import os from dataclasses import dataclass from io import BytesIO -from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Self, Sequence, Tuple, Union from nemo_retriever.graph import InprocessExecutor, RayDataExecutor from nemo_retriever.ingestor.branch_extraction import ExtractionBranchExecutor, merge_node_overrides @@ -44,6 +44,12 @@ resolve_branch_extraction_inputs, ) from nemo_retriever.ingestor import ingestor +from nemo_retriever.common.inline_text import ( + inline_text_source_id, + is_blank_inline_corpus, + is_inline_text_source, + normalize_inline_texts, +) from nemo_retriever.common.params import ( ASRParams, AudioChunkParams, @@ -72,7 +78,7 @@ from nemo_retriever.common.remote_auth import resolve_remote_api_key from nemo_retriever.common.ray_runtime import ensure_local_ray_runtime from nemo_retriever.common.ray_resource_hueristics import gather_cluster_resources - +from nemo_retriever.common.modality.txt.split import empty_text_chunks_df _ERROR_FIELD_KEYS = ("error", "errors", "exception", "traceback", "failed") _REMOTE_EMBED_ENDPOINT_FIELDS = ("embedding_endpoint", "embed_invoke_url") @@ -462,9 +468,10 @@ def __init__( self._error_policy = error_policy self._rd_dataset: Any = None self._buffers: list[tuple[str, BytesIO]] = [] + self._inline_texts: list[str] | None = None # Pipeline configuration accumulated by fluent methods - self._extraction_mode: str | None = "pdf" + self._extraction_mode: str | None = None self._extract_params: Any = None self._text_params: Any = None self._html_params: Any = None @@ -493,6 +500,18 @@ def files(self, documents: Union[str, List[str]]) -> "GraphIngestor": self._documents = [documents] if isinstance(documents, str) else list(documents) return self + def texts(self, texts: Union[str, Sequence[str]]) -> Self: + """Set raw inline text documents as the graph input. + + Each string is one logical source document. It receives a deterministic + ``inline://`` identifier and flows through the normal text splitter, + embedding, and sink stages without being written to a temporary file. + Inline text may be combined with file or buffer inputs; the manifest + planner routes each source through its matching extraction branch. + """ + self._inline_texts = normalize_inline_texts(texts) + return self + def buffers( self, buffers: Union[Tuple[str, BytesIO], List[Tuple[str, BytesIO]]], @@ -588,13 +607,6 @@ def extract_image_files( self._record_stage("extract") return self - def extract_txt(self, params: Optional[TextChunkParams] = None, **kwargs: Any) -> "GraphIngestor": - """Configure plain-text extraction (extraction_mode='text').""" - self._extraction_mode = "text" - self._text_params = _coerce(params, kwargs, default_factory=TextChunkParams) - self._record_stage("extract") - return self - def extract_html(self, params: Optional[HtmlChunkParams] = None, **kwargs: Any) -> "GraphIngestor": """Configure HTML extraction (extraction_mode='html').""" self._extraction_mode = "html" @@ -736,19 +748,28 @@ def ingest(self, params: Any = None, **kwargs: Any) -> Any: Returns ------- - ``run_mode='batch'`` - A materialized ``ray.data.Dataset``. - ``run_mode='inprocess'`` + ``run_mode='batch'`` or ``run_mode='inprocess'`` A ``pandas.DataFrame``. ``return_failures=True`` ``(result, failures)`` where ``failures`` is a list of service-style ``(source, error)`` tuples. """ return_failures = self._resolve_return_failures(params, kwargs) + if not self._documents and not self._buffers and is_blank_inline_corpus(self._inline_texts): + result = empty_text_chunks_df() + if self._run_mode == "batch": + self._rd_dataset = result + else: + self._rd_dataset = None + return self._finalize_ingest_result(result, return_failures=return_failures) + default_branches = self._plan_default_extraction_branches() + execute_branches = default_branches is not None and ( + len(default_branches) > 1 or self._has_mixed_inline_sources() + ) if default_branches is None: single_effective = self._resolve_effective_extraction_inputs() - elif len(default_branches) == 1: + elif not execute_branches: single_effective = self._resolve_branch_extraction_inputs(default_branches[0]) else: single_effective = None @@ -768,7 +789,7 @@ def ingest(self, params: Any = None, **kwargs: Any) -> Any: post_extract_order = tuple(s for s in self._stage_order if s != "extract") - if default_branches is not None and len(default_branches) > 1: + if execute_branches: result = self._execute_extraction_branches(default_branches, post_extract_order=post_extract_order) else: if single_effective is None: @@ -793,7 +814,7 @@ def _execute_single_graph_batch( *, post_extract_order: tuple[str, ...], ) -> Any: - _ray, cluster_resources = self._ensure_batch_runtime() + ray, cluster_resources = self._ensure_batch_runtime() graph = build_graph( extraction_mode=effective_extraction.extraction_mode, extract_params=effective_extraction.extract_params, @@ -831,7 +852,8 @@ def _execute_single_graph_batch( num_gpus=self._num_gpus, node_overrides=merge_node_overrides(derived_overrides, self._node_overrides), ) - result = executor.ingest(self._documents) + executor_input = self._inline_text_dataset(ray.data) if self._inline_texts else self._documents + result = executor.ingest(executor_input) self._rd_dataset = result return result @@ -862,6 +884,8 @@ def _execute_single_graph_inprocess( ) executor = InprocessExecutor(graph, show_progress=self._show_progress) self._rd_dataset = None + if self._inline_texts: + return executor.ingest(self._inline_text_dataframe()) if self._buffers: import pandas as pd @@ -880,6 +904,7 @@ def _execute_extraction_branches( branches=branches, documents=self._documents, buffers=self._buffers, + inline_rows=self._inline_text_rows(), split_config=self._split_config, extract_params=self._extract_params, text_params=self._text_params, @@ -916,6 +941,22 @@ def _ensure_batch_runtime(self) -> tuple[Any, Any]: # Internal helpers # ------------------------------------------------------------------ + def _has_mixed_inline_sources(self) -> bool: + return bool(self._inline_texts) and bool(self._documents or self._buffers) + + def _inline_text_rows(self) -> list[dict[str, str]]: + return [ + {"text": text, "path": inline_text_source_id(index)} for index, text in enumerate(self._inline_texts or []) + ] + + def _inline_text_dataframe(self) -> Any: + import pandas as pd + + return pd.DataFrame(self._inline_text_rows(), columns=["text", "path"]) + + def _inline_text_dataset(self, ray_data: Any) -> Any: + return ray_data.from_items(self._inline_text_rows()) + def _configured_input_paths(self) -> list[str]: paths: list[str] = [] for document in self._documents: @@ -924,10 +965,16 @@ def _configured_input_paths(self) -> list[str]: except FileNotFoundError: paths.append(os.fspath(document)) paths.extend(name for name, _ in self._buffers) + paths.extend(inline_text_source_id(index) for index, _ in enumerate(self._inline_texts or [])) return paths def _classified_input_paths(self) -> list[tuple[str, str | None]]: - return [(path, input_type_for_path(path)) for path in self._configured_input_paths()] + # Service workers receive inline text as a named byte buffer. Keep the + # logical URI as its source path while classifying it as decoded text. + return [ + (path, "txt" if is_inline_text_source(path) else input_type_for_path(path)) + for path in self._configured_input_paths() + ] @staticmethod def _input_type_examples(paths: Iterable[str], *, limit: int = 3) -> str: @@ -953,7 +1000,7 @@ def _validate_explicit_extraction_mode_inputs( raise ValueError(f"Input file type(s) do not match extraction_mode={extraction_mode!r}: {examples}") def _plan_default_extraction_branches(self) -> tuple[ExtractionBranchPlan, ...] | None: - if self._extraction_mode is not None: + if self._extraction_mode is not None and not self._has_mixed_inline_sources(): return None manifest = build_input_manifest(self._configured_input_paths()) branches = plan_extraction_branches(manifest) @@ -1317,12 +1364,6 @@ def _record_stage(self, name: str) -> None: self._stage_order.append(name) def _apply_split_config(self, split_config: dict[str, Any] | None) -> None: - """Resolve split_config when the caller opts in. - - Typed shortcuts (extract_audio, extract_video, extract_image_files) - leave the constructor's all-None default in place when split_config is - omitted. Only the unified .extract() resolves None into the natural - default-on set. - """ + """Resolve an explicitly supplied split configuration.""" if split_config is not None: self._split_config = resolve_split_params(split_config) diff --git a/nemo_retriever/src/nemo_retriever/ingestor/manifest.py b/nemo_retriever/src/nemo_retriever/ingestor/manifest.py index 57546cd92c..22678d808b 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/manifest.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/manifest.py @@ -21,6 +21,7 @@ VideoFrameTextDedupParams, ) from nemo_retriever.common.input_files import _is_explicit_glob_path, input_type_for_path +from nemo_retriever.common.inline_text import is_inline_text_source _AUDIO_SPLIT_INTERVAL = 500000 @@ -118,7 +119,7 @@ def build_input_manifest(input_paths: Iterable[str]) -> InputManifest: unsupported: list[str] = [] for path in input_paths: is_glob = _is_explicit_glob_path(path) - input_type = None if is_glob else input_type_for_path(path) + input_type = None if is_glob else ("txt" if is_inline_text_source(path) else input_type_for_path(path)) entries.append(ManifestEntry(path=path, input_type=input_type, is_explicit_glob=is_glob)) if input_type is None and not is_glob: unsupported.append(path) diff --git a/nemo_retriever/src/nemo_retriever/models/__init__.py b/nemo_retriever/src/nemo_retriever/models/__init__.py index c83d3ac603..b376aa0550 100644 --- a/nemo_retriever/src/nemo_retriever/models/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/__init__.py @@ -6,21 +6,18 @@ from typing import TYPE_CHECKING, Any +from nemo_retriever.models.embed_model_spec import ( + EmbedModelSpec, + resolve_embed_model_spec, + validate_embed_model_backend, +) + if TYPE_CHECKING: from nemo_retriever.models.model import BaseModel VL_EMBED_MODEL = "nvidia/llama-nemotron-embed-vl-1b-v2" VL_RERANK_MODEL = "nvidia/llama-nemotron-rerank-vl-1b-v2" -_VL_EMBED_MODEL_IDS = frozenset( - { - VL_EMBED_MODEL, - "llama-nemotron-embed-vl-1b-v2", - "llama-3.2-nemoretriever-1b-vlm-embed-v1", - "nvidia/llama-3.2-nemoretriever-1b-vlm-embed-v1", - } -) - _VL_RERANK_MODEL_IDS = frozenset( { VL_RERANK_MODEL, @@ -50,8 +47,8 @@ def resolve_embed_model(model_name: str | None) -> str: def is_vl_embed_model(model_name: str | None) -> bool: - """Return True if *model_name* refers to the VL embedding model.""" - return resolve_embed_model(model_name) in _VL_EMBED_MODEL_IDS + """Return True when a legacy model ID or alias names the default VL embedder.""" + return resolve_embed_model(model_name) == VL_EMBED_MODEL def is_vl_rerank_model(model_name: str | None) -> bool: @@ -71,6 +68,7 @@ def create_local_embedder( normalize: bool = True, max_length: int = 8192, query_max_length: int = 128, + revision: str | None = None, ) -> Any: """Create the appropriate local embedding model (VL or non-VL). @@ -92,13 +90,21 @@ def create_local_embedder( Note: ``gpu_memory_utilization``, ``enforce_eager``, ``dimensions``, ``normalize``, and ``max_length`` apply to vLLM paths only; the HF VL path ignores them. + + Local checkpoints and compatible Hub fine-tunes are routed from their + immutable config. Compatibility requires a supported dense Nemotron + embedding architecture and average pooling. Output dimensions and text + prefixes are derived from checkpoint metadata. """ b = (backend or "vllm").strip().lower() if b not in ("vllm", "hf"): raise ValueError(f"backend must be 'vllm' or 'hf', got {backend!r}") + model_id = resolve_embed_model(model_name) + spec = resolve_embed_model_spec(model_id, revision=revision, hf_cache_dir=hf_cache_dir) + validate_embed_model_backend(spec, b) - if is_vl_embed_model(model_name): + if spec.family == "vl": if b == "hf": from nemo_retriever.models.local.llama_nemotron_embed_vl_1b_v2_embedder import ( LlamaNemotronEmbedVL1BV2Embedder, @@ -108,6 +114,8 @@ def create_local_embedder( device=device, hf_cache_dir=hf_cache_dir, model_id=model_id, + revision=spec.revision, + output_dimension=spec.output_dimension, ) from nemo_retriever.models.local.llama_nemotron_embed_vl_1b_v2_embedder import ( @@ -118,8 +126,12 @@ def create_local_embedder( model_id=model_id, device=device, hf_cache_dir=hf_cache_dir, + revision=spec.revision, gpu_memory_utilization=gpu_memory_utilization, enforce_eager=enforce_eager, + output_dimension=spec.output_dimension, + query_prefix=spec.query_prefix, + document_prefix=spec.document_prefix, ) if b == "hf": @@ -134,6 +146,9 @@ def create_local_embedder( max_length=int(max_length), query_max_length=int(query_max_length), model_id=model_id, + revision=spec.revision, + query_prefix=spec.query_prefix, + document_prefix=spec.document_prefix, ) from nemo_retriever.models.local.llama_nemotron_embed_1b_v2_embedder import ( @@ -149,6 +164,9 @@ def create_local_embedder( dimensions=dimensions, normalize=normalize, max_length=int(max_length), + revision=spec.revision, + query_prefix=spec.query_prefix, + document_prefix=spec.document_prefix, ) @@ -181,6 +199,7 @@ def create_local_query_embedder( normalize: bool = True, max_length: int = 8192, query_max_length: int = 128, + revision: str | None = None, ) -> Any: """Create a local embedder for *query* vectors in retrieval (Retriever / recall). @@ -188,6 +207,9 @@ def create_local_query_embedder( - ``backend="hf"``: HuggingFace for both VL and non-VL models. - ``backend="vllm"``: vLLM for both VL and non-VL models. + + Model architecture and quantization requirements are resolved from the + checkpoint config; see :func:`create_local_embedder`. """ b = normalize_backend(backend, _LOCAL_QUERY_BACKENDS, field_name="backend", default="hf") @@ -202,6 +224,7 @@ def create_local_query_embedder( normalize=normalize, max_length=int(max_length), query_max_length=int(query_max_length), + revision=revision, ) diff --git a/nemo_retriever/src/nemo_retriever/models/embed_model_spec.py b/nemo_retriever/src/nemo_retriever/models/embed_model_spec.py new file mode 100644 index 0000000000..cf6eb3a980 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/models/embed_model_spec.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolve dense Nemotron embedding checkpoints from immutable HF config.""" + +from __future__ import annotations + +import json +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +from nemo_retriever.models.hf_model_registry import HF_MODEL_REVISIONS + +EmbedModelFamily = Literal["text", "vl"] + +_MODEL_PROFILES: dict[str, tuple[EmbedModelFamily, str]] = { + "llama_bidirec": ("text", "LlamaBidirectionalModel"), + "llama_nemotron_vl": ("vl", "LlamaNemotronVLModel"), +} +_DEFAULT_QUERY_PREFIX = "query: " +_DEFAULT_DOCUMENT_PREFIX = "passage: " +_MODEL_CONFIG_FILENAME = "config.json" +_PROMPT_CONFIG_FILENAME = "config_sentence_transformers.json" +_COMMIT_SHA_RE = re.compile(r"^[0-9a-fA-F]{40}$") + + +@dataclass(frozen=True) +class EmbedModelSpec: + """Immutable loading and input-format information for a dense embedder.""" + + model_id: str + revision: str | None + family: EmbedModelFamily + output_dimension: int + query_prefix: str + document_prefix: str + quantization: str | None = None + requires_vllm: bool = False + + +def _read_config(path: Path, *, model_id: str) -> dict[str, Any]: + try: + config = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Embedding model {model_id!r} has an unreadable {path.name}: {exc}") from exc + if not isinstance(config, dict): + raise ValueError(f"Embedding model {model_id!r} {path.name} must contain a JSON object.") + return config + + +def _local_config_path(model_id: str) -> Path | None: + path = Path(model_id).expanduser() + if path.is_dir(): + config_path = path / "config.json" + if not config_path.is_file(): + raise ValueError(f"Local embedding model directory {model_id!r} does not contain config.json.") + return config_path + if path.is_absolute() or model_id.startswith((".", "~")) or os.path.exists(path): + raise ValueError(f"Local embedding model path {model_id!r} is not a model directory containing config.json.") + return None + + +def _hub_revision(model_id: str, revision: str | None) -> str: + candidate = revision or HF_MODEL_REVISIONS.get(model_id) + if candidate and _COMMIT_SHA_RE.fullmatch(candidate): + return candidate + + from huggingface_hub import HfApi + + try: + info = HfApi().model_info(model_id, revision=candidate) + except Exception as exc: + requested = candidate or "main" + raise ValueError( + f"Could not resolve Hugging Face embedding model {model_id!r} at revision {requested!r}: {exc}" + ) from exc + sha = str(info.sha or "").strip() + if not _COMMIT_SHA_RE.fullmatch(sha): + raise ValueError(f"Hugging Face did not return an immutable commit SHA for embedding model {model_id!r}.") + return sha + + +def _hub_json( + model_id: str, + filename: str, + revision: str | None, + *, + hf_cache_dir: str | None = None, + optional: bool = False, +) -> dict[str, Any] | None: + """Read a JSON file from a Hub repo; ``optional`` maps a missing file to None.""" + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError + + try: + path = hf_hub_download( + repo_id=model_id, + filename=filename, + revision=revision, + cache_dir=hf_cache_dir, + ) + except Exception as exc: + if optional and isinstance(exc, EntryNotFoundError): + return None + raise ValueError( + f"Could not load {filename} for Hugging Face embedding model {model_id!r} at {revision!r}: {exc}" + ) from exc + return _read_config(Path(path), model_id=model_id) + + +def _local_prompt_config(model_id: str, config_path: Path) -> dict[str, Any] | None: + """Read prompt metadata sitting beside a local checkpoint's config.json.""" + path = config_path.with_name(_PROMPT_CONFIG_FILENAME) + return _read_config(path, model_id=model_id) if path.is_file() else None + + +def _prompt_prefixes(config: dict[str, Any] | None) -> tuple[str, str]: + prompts = config.get("prompts") if isinstance(config, dict) else None + if not isinstance(prompts, dict): + return _DEFAULT_QUERY_PREFIX, _DEFAULT_DOCUMENT_PREFIX + query = prompts.get("query", _DEFAULT_QUERY_PREFIX) + document = prompts.get("document", _DEFAULT_DOCUMENT_PREFIX) + if not isinstance(query, str) or not isinstance(document, str): + raise ValueError("Sentence Transformers query and document prompts must be strings.") + return query, document + + +def _spec_from_config( + model_id: str, + revision: str | None, + config: dict[str, Any], + prompt_config: dict[str, Any] | None = None, +) -> EmbedModelSpec: + model_type = str(config.get("model_type") or "").strip() + profile = _MODEL_PROFILES.get(model_type) + if profile is None: + supported = ", ".join(sorted(_MODEL_PROFILES)) + raise ValueError( + f"Embedding model {model_id!r} uses unsupported model_type {model_type!r}; " + f"supported Nemotron embed model types are: {supported}." + ) + family, expected_architecture = profile + + architectures = config.get("architectures") + if architectures != [expected_architecture]: + raise ValueError( + f"Embedding model {model_id!r} uses unsupported architectures {architectures!r}; " + f"expected [{expected_architecture!r}] for the {family} dense embedding profile." + ) + + dimension_config = config.get("llm_config") if family == "vl" else config + dimension = dimension_config.get("hidden_size") if isinstance(dimension_config, dict) else None + if isinstance(dimension, bool) or not isinstance(dimension, int) or dimension <= 0: + raise ValueError( + f"Embedding model {model_id!r} has invalid embedding dimension {dimension!r}; " + "a positive hidden_size is required." + ) + + pooling = str(config.get("pooling") or "").strip().lower() + if pooling != "avg": + raise ValueError( + f"Embedding model {model_id!r} uses unsupported pooling {pooling!r}; " + "dense Nemotron embedding profiles require 'avg'." + ) + + quantization_config = config.get("quantization_config") + quantization = None + requires_vllm = False + if isinstance(quantization_config, dict): + quant_method = str(quantization_config.get("quant_method") or "").strip().lower() + quantization = str(quantization_config.get("quant_algo") or quant_method or "").strip() or None + if quant_method == "modelopt": + requires_vllm = True + + query_prefix, document_prefix = _prompt_prefixes(prompt_config) + + return EmbedModelSpec( + model_id=model_id, + revision=revision, + family=family, + output_dimension=dimension, + query_prefix=query_prefix, + document_prefix=document_prefix, + quantization=quantization, + requires_vllm=requires_vllm, + ) + + +def resolve_embed_model_spec( + model_id: str, + *, + revision: str | None = None, + hf_cache_dir: str | None = None, +) -> EmbedModelSpec: + """Resolve and validate a dense Nemotron embedding checkpoint. + + Registered Hub repositories retain their project-pinned revision. Other + repositories are resolved to an immutable Hub commit before their config is + inspected. Compatible checkpoints use a supported dense embedding + architecture, a positive declared output dimension, and average pooling. + Query and document prefixes come from Sentence Transformers metadata when + present. Local directories are identified by config.json and do not carry + a Hub revision. + """ + local_config = _local_config_path(model_id) + if local_config is not None: + if revision is not None: + raise ValueError("A Hugging Face revision cannot be used with a local embedding model directory.") + return _spec_from_config( + model_id, + None, + _read_config(local_config, model_id=model_id), + _local_prompt_config(model_id, local_config), + ) + + resolved_revision = _hub_revision(model_id, revision) + config = _hub_json(model_id, _MODEL_CONFIG_FILENAME, resolved_revision, hf_cache_dir=hf_cache_dir) or {} + prompt_config = _hub_json( + model_id, _PROMPT_CONFIG_FILENAME, resolved_revision, hf_cache_dir=hf_cache_dir, optional=True + ) + return _spec_from_config(model_id, resolved_revision, config, prompt_config) + + +def resolve_embed_model_revision(model_id: str, revision: str | None) -> str | None: + """Return an explicit revision or resolve one for a directly constructed embedder.""" + if revision is not None: + return revision + if _local_config_path(model_id) is not None: + return None + return _hub_revision(model_id, None) + + +def validate_embed_model_backend(spec: EmbedModelSpec, backend: str) -> None: + """Reject backends that cannot load the resolved checkpoint format.""" + if spec.requires_vllm and backend != "vllm": + quantization = f" ({spec.quantization})" if spec.quantization else "" + raise ValueError(f"Embedding model {spec.model_id!r}{quantization} requires backend='vllm'.") diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py index b539751adb..4ec2b3c832 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -158,18 +158,22 @@ def __call__( tools: Optional[list[dict[str, Any]]] = None, tool_choice: str | dict[str, Any] = "auto", timeout_s: float = 120.0, - temperature: float = 0.0, + temperature: Optional[float] = 0.0, max_tokens: Optional[int] = None, extra_body: Optional[dict[str, Any]] = None, max_retries: int = 10, max_429_retries: int = 5, ) -> dict[str, Any]: + """Run one chat completion on the in-process engine. + + ``temperature=None`` means *unset* and maps to ``0.0`` (greedy). + """ _ = (invoke_url, api_key, timeout_s, max_retries, max_429_retries) if model and model != self._model_path: logger.debug("Ignoring per-call model=%r for local agent LLM already loaded as %r", model, self._model_path) sampling_params = self._sampling_params_cls( - temperature=float(temperature), + temperature=0.0 if temperature is None else float(temperature), max_tokens=int(max_tokens) if max_tokens is not None else self._max_tokens, ) chat_kwargs = self._build_chat_kwargs(extra_body) diff --git a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py index ec9dcf7cbe..4ef170705c 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py +++ b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py @@ -11,7 +11,7 @@ import torch from nemo_retriever.models.hf_cache import configure_global_hf_cache_base -from nemo_retriever.models.hf_model_registry import get_hf_revision +from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor: @@ -23,7 +23,7 @@ def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor: @dataclass class LlamaNemotronEmbed1BV2Embedder: """ - Local text embedder for ``nvidia/llama-nemotron-embed-1b-v2`` via vLLM. + vLLM embedder for compatible dense LlamaBidirectional Nemotron checkpoints. Uses vLLM's pooling runner (``llm.embed()``) for throughput. No HTTP remote calls — load and inference stay in-process. @@ -33,6 +33,7 @@ class LlamaNemotronEmbed1BV2Embedder: ``device`` is deprecated and ignored; it remains only for backward compatibility with callers that constructed the former HuggingFace embedder with ``device=``. + The legacy class name is retained for compatibility. """ model_id: Optional[str] = None @@ -43,6 +44,9 @@ class LlamaNemotronEmbed1BV2Embedder: dimensions: Optional[int] = None normalize: bool = True max_length: int = 8192 + revision: Optional[str] = None + query_prefix: str = "query: " + document_prefix: str = "passage: " _llm: Any = field(default=None, init=False, repr=False) @@ -68,7 +72,7 @@ def _ensure_loaded(self) -> None: max_model_len = int(self.max_length) if int(self.max_length) > 0 else None self._llm = create_vllm_llm( str(model_id), - revision=get_hf_revision(model_id), + revision=resolve_embed_model_revision(model_id, self.revision), dimensions=self.dimensions, gpu_memory_utilization=self.gpu_memory_utilization, enforce_eager=self.enforce_eager, @@ -90,10 +94,10 @@ def _finalize_vectors(self, vectors: List[List[float]]) -> torch.Tensor: return _l2_normalize(t) return t - def embed(self, texts: Sequence[str], *, batch_size: int = 64, prefix: str = "passage: ") -> torch.Tensor: + def embed(self, texts: Sequence[str], *, batch_size: int = 64, prefix: str | None = None) -> torch.Tensor: """Embed texts. Returns CPU tensor ``[N, D]``. - ``prefix`` is prepended to every string before encoding (default ``passage: ``). + ``prefix`` overrides the checkpoint-declared document prefix when set. """ self._ensure_loaded() from nemo_retriever.models.inference.vllm import embed_with_vllm_llm @@ -105,7 +109,7 @@ def embed(self, texts: Sequence[str], *, batch_size: int = 64, prefix: str = "pa texts_list, self._llm, batch_size=max(1, int(batch_size)), - prefix=prefix, + prefix=self.document_prefix if prefix is None else prefix, normalize=self.normalize, ) return self._finalize_vectors(vectors) @@ -122,7 +126,7 @@ def embed_queries(self, texts: Sequence[str], *, batch_size: int = 64) -> torch. texts_list, self._llm, batch_size=max(1, int(batch_size)), - prefix="query: ", + prefix=self.query_prefix, normalize=self.normalize, ) return self._finalize_vectors(vectors) diff --git a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_hf_embedder.py b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_hf_embedder.py index 7f9f6a46c2..f037f65450 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_hf_embedder.py +++ b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_hf_embedder.py @@ -2,11 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""HuggingFace-only text embedder for ``nvidia/llama-nemotron-embed-1b-v2``. - -Used when local query embedding should match classic HF pooling (e.g. recall -evaluation) while document ingestion uses vLLM elsewhere. -""" +"""Hugging Face loader for compatible LlamaBidirectional embedding checkpoints.""" from __future__ import annotations @@ -18,7 +14,7 @@ import torch from nemo_retriever.models.hf_cache import configure_global_hf_cache_base -from nemo_retriever.models.hf_model_registry import get_hf_revision +from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision logger = logging.getLogger(__name__) @@ -31,7 +27,7 @@ def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor: @dataclass class LlamaNemotronEmbed1BV2HFEmbedder: - """Mean-pooled HF embeddings with Nemotron-style ``query:`` / ``passage:`` prefixes.""" + """Mean-pooled HF embeddings with checkpoint-declared text prefixes.""" device: Optional[str] = None hf_cache_dir: Optional[str] = None @@ -41,6 +37,9 @@ class LlamaNemotronEmbed1BV2HFEmbedder: # values preserve more text but increase query embedding cost. query_max_length: int = 128 model_id: Optional[str] = None + revision: Optional[str] = None + query_prefix: str = "query: " + document_prefix: str = "passage: " def __post_init__(self) -> None: self._tokenizer = None @@ -56,7 +55,7 @@ def _ensure_loaded(self) -> None: model_id = self.model_id or _DEFAULT_EMBED_MODEL dev = torch.device(self.device or ("cuda" if torch.cuda.is_available() else "cpu")) hf_cache_dir = configure_global_hf_cache_base(self.hf_cache_dir) - _revision = get_hf_revision(model_id) + _revision = resolve_embed_model_revision(model_id, self.revision) self._tokenizer = AutoTokenizer.from_pretrained( model_id, revision=_revision, @@ -125,16 +124,22 @@ def _embed_local( return torch.cat(outs, dim=0) if outs else torch.empty((0, 0), dtype=torch.float32) - def embed(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: - """Document strings; each line is prefixed with ``passage:`` for parity with vLLM.""" - texts_list = [] - for t in texts: - raw = str(t) + @staticmethod + def _prepare_texts(texts: Sequence[str], prefix: str) -> List[str]: + """Drop blank inputs and apply *prefix* to any line that lacks it.""" + prepared: List[str] = [] + for text in texts: + raw = str(text) if not raw.strip(): continue - if not raw.lower().startswith("passage:"): - raw = "passage: " + raw - texts_list.append(raw) + if prefix and not raw.lower().startswith(prefix.lower()): + raw = prefix + raw + prepared.append(raw) + return prepared + + def embed(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: + """Embed documents after applying the checkpoint-declared prefix.""" + texts_list = self._prepare_texts(texts, self.document_prefix) if not texts_list: return torch.empty((0, 0), dtype=torch.float32) return self._embed_local(texts_list, batch_size=batch_size) @@ -172,15 +177,8 @@ def _warn_if_queries_truncated(self, texts: Sequence[str], *, max_length: int) - ) def embed_queries(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: - """Query strings; each line is prefixed with ``query:`` (same rules as vLLM embed_queries).""" - texts_list = [] - for t in texts: - raw = str(t) - if not raw.strip(): - continue - if not raw.lower().startswith("query:"): - raw = "query: " + raw - texts_list.append(raw) + """Embed queries after applying the checkpoint-declared prefix.""" + texts_list = self._prepare_texts(texts, self.query_prefix) if not texts_list: return torch.empty((0, 0), dtype=torch.float32) # The Nemotron text embedder is sensitive to padding length. Use fixed diff --git a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py index 891ad25c43..ea6e2df94e 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py +++ b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py @@ -11,7 +11,7 @@ import torch from nemo_retriever.models.hf_cache import configure_global_hf_cache_base -from nemo_retriever.models.hf_model_registry import get_hf_revision +from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision from nemo_retriever.common.nvtx import gpu_inference_range @@ -21,6 +21,13 @@ def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor: return x / denom +def _to_normalized_cpu(out: Any) -> torch.Tensor: + """L2-normalize a model output on CPU, accepting tensors or array-likes.""" + if isinstance(out, torch.Tensor): + return _l2_normalize(out.detach().cpu()) + return _l2_normalize(torch.as_tensor(out, dtype=torch.float32).cpu()) + + def _b64_to_pil(b64_str: str): import base64 import io @@ -36,24 +43,24 @@ def _b64_to_pil(b64_str: str): @dataclass class LlamaNemotronEmbedVL1BV2Embedder: """ - Multimodal embedder wrapper for ``nvidia/llama-nemotron-embed-vl-1b-v2``. + HF wrapper for compatible ``LlamaNemotronVLModel`` checkpoints. The VL model exposes ``encode_queries()`` and ``encode_documents()`` instead of the standard tokenizer + forward pass used by the embedqa - model. This class supports text, image, and text+image modalities. + model. This class supports text, image, and text+image modalities; its + legacy class name is retained for compatibility. """ device: Optional[str] = None hf_cache_dir: Optional[str] = None model_id: Optional[str] = None + revision: Optional[str] = None + output_dimension: int = 2048 - # Populated in __post_init__ + # Populated lazily by _ensure_loaded. _model: Any = field(default=None, init=False, repr=False) _device: Any = field(default=None, init=False, repr=False) - def __post_init__(self) -> None: - pass - def _ensure_loaded(self) -> None: if self._model is not None: return @@ -70,7 +77,7 @@ def _ensure_loaded(self) -> None: # device_map when requesting it. Fall back to sdpa/eager on CPU or # when flash-attn is not installed. use_gpu = dev.type == "cuda" - _revision = get_hf_revision(model_id) + _revision = resolve_embed_model_revision(model_id, self.revision) for attn_impl in ("flash_attention_2", "sdpa", "eager"): try: kwargs: dict[str, Any] = { @@ -105,53 +112,47 @@ def _set_p_max_length(self, modality: str) -> None: self._model.processor.p_max_length = p def embed(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: - """Embed document texts. Returns CPU tensor ``[N, 2048]``.""" + """Embed document texts. Returns CPU tensor ``[N, D]``.""" self._ensure_loaded() texts_list = [str(t) for t in texts if str(t).strip()] if not texts_list: - return torch.empty((0, 2048), dtype=torch.float32) + return torch.empty((0, self.output_dimension), dtype=torch.float32) with torch.inference_mode(), warnings.catch_warnings(): warnings.filterwarnings("ignore", message="`input_embeds` is deprecated", category=FutureWarning) self._set_p_max_length("text") with gpu_inference_range("LlamaNemotronEmbedVL1B", batch_size=len(texts_list), mode="doc_text"): out = self._model.encode_documents(texts=texts_list) - if isinstance(out, torch.Tensor): - return _l2_normalize(out.detach().cpu()) - return _l2_normalize(torch.as_tensor(out, dtype=torch.float32).cpu()) + return _to_normalized_cpu(out) def embed_queries(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: - """Embed query strings. Returns CPU tensor ``[N, 2048]``.""" + """Embed query strings. Returns CPU tensor ``[N, D]``.""" self._ensure_loaded() texts_list = [str(t) for t in texts] if not texts_list: - return torch.empty((0, 2048), dtype=torch.float32) + return torch.empty((0, self.output_dimension), dtype=torch.float32) with torch.inference_mode(), warnings.catch_warnings(): warnings.filterwarnings("ignore", message="`input_embeds` is deprecated", category=FutureWarning) with gpu_inference_range("LlamaNemotronEmbedVL1B", batch_size=len(texts_list), mode="query"): out = self._model.encode_queries(texts_list) - if isinstance(out, torch.Tensor): - return _l2_normalize(out.detach().cpu()) - return _l2_normalize(torch.as_tensor(out, dtype=torch.float32).cpu()) + return _to_normalized_cpu(out) def embed_images(self, images_b64: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: - """Embed images (base64-encoded). Returns CPU tensor ``[N, 2048]``.""" + """Embed images (base64-encoded). Returns CPU tensor ``[N, D]``.""" self._ensure_loaded() image_dicts = [{"base64": b64} for b64 in images_b64 if b64] if not image_dicts: - return torch.empty((0, 2048), dtype=torch.float32) + return torch.empty((0, self.output_dimension), dtype=torch.float32) with torch.inference_mode(), warnings.catch_warnings(): warnings.filterwarnings("ignore", message="`input_embeds` is deprecated", category=FutureWarning) self._set_p_max_length("image") with gpu_inference_range("LlamaNemotronEmbedVL1B", batch_size=len(image_dicts), mode="doc_image"): out = self._model.encode_documents(images=image_dicts) - if isinstance(out, torch.Tensor): - return _l2_normalize(out.detach().cpu()) - return _l2_normalize(torch.as_tensor(out, dtype=torch.float32).cpu()) + return _to_normalized_cpu(out) def embed_text_image( self, texts: Sequence[str], images_b64: Sequence[str], *, batch_size: int = 64 ) -> torch.Tensor: - """Embed paired text+image inputs. Returns CPU tensor ``[N, 2048]``.""" + """Embed paired text+image inputs. Returns CPU tensor ``[N, D]``.""" self._ensure_loaded() paired_texts: list[str] = [] paired_images: list[dict[str, str]] = [] @@ -160,15 +161,13 @@ def embed_text_image( paired_texts.append(str(t)) paired_images.append({"base64": b64}) if not paired_images: - return torch.empty((0, 2048), dtype=torch.float32) + return torch.empty((0, self.output_dimension), dtype=torch.float32) with torch.inference_mode(), warnings.catch_warnings(): warnings.filterwarnings("ignore", message="`input_embeds` is deprecated", category=FutureWarning) self._set_p_max_length("text_image") with gpu_inference_range("LlamaNemotronEmbedVL1B", batch_size=len(paired_images), mode="doc_text_image"): out = self._model.encode_documents(texts=paired_texts, images=paired_images) - if isinstance(out, torch.Tensor): - return _l2_normalize(out.detach().cpu()) - return _l2_normalize(torch.as_tensor(out, dtype=torch.float32).cpu()) + return _to_normalized_cpu(out) def unload(self) -> None: """Release GPU memory held by the HF model.""" @@ -182,11 +181,11 @@ def unload(self) -> None: @dataclass class LlamaNemotronEmbedVL1BV2VLLMEmbedder: """ - vLLM-backed embedder for ``nvidia/llama-nemotron-embed-vl-1b-v2``. + vLLM embedder for compatible ``LlamaNemotronVLModel`` checkpoints. Supports text, image, and text+image modalities via vLLM's Python API (bfloat16 + FLASH_ATTN, pooling runner). Requires vLLM >= 0.17.0. - + The legacy class name is retained for compatibility. """ model_id: Optional[str] = None @@ -194,6 +193,10 @@ class LlamaNemotronEmbedVL1BV2VLLMEmbedder: hf_cache_dir: Optional[str] = None gpu_memory_utilization: float = 0.45 enforce_eager: bool = False + revision: Optional[str] = None + output_dimension: int = 2048 + query_prefix: str = "query: " + document_prefix: str = "passage: " _llm: Any = field(default=None, init=False, repr=False) @@ -234,7 +237,7 @@ def _ensure_loaded(self) -> None: model_id = self.model_id or "nvidia/llama-nemotron-embed-vl-1b-v2" self._llm = create_vllm_llm( str(model_id), - revision=get_hf_revision(model_id), + revision=resolve_embed_model_revision(model_id, self.revision), gpu_memory_utilization=self.gpu_memory_utilization, enforce_eager=self.enforce_eager, limit_mm_per_prompt={"image": 1}, @@ -244,62 +247,67 @@ def _ensure_loaded(self) -> None: def is_remote(self) -> bool: return False + def _finalize_vectors(self, vectors: Sequence[Sequence[float]]) -> torch.Tensor: + """Zero-pad rows vLLM failed to embed to the returned width, then normalize.""" + valid = [v for v in vectors if v] + if not valid: + return torch.empty((0, self.output_dimension), dtype=torch.float32) + dim = len(valid[0]) + padded = [v if v else [0.0] * dim for v in vectors] + return _l2_normalize(torch.tensor(padded, dtype=torch.float32)) + def embed(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: - """Embed document texts. Returns CPU tensor ``[N, 2048]``.""" + """Embed document texts. Returns CPU tensor ``[N, D]``.""" self._ensure_loaded() from nemo_retriever.models.inference.vllm import embed_with_vllm_llm texts_list = [str(t) for t in texts if str(t).strip()] if not texts_list: - return torch.empty((0, 2048), dtype=torch.float32) - vectors = embed_with_vllm_llm(texts_list, self._llm, batch_size=max(1, int(batch_size)), prefix="passage: ") - valid = [v for v in vectors if v] - if not valid: - return torch.empty((0, 2048), dtype=torch.float32) - dim = len(valid[0]) - padded = [v if v else [0.0] * dim for v in vectors] - return _l2_normalize(torch.tensor(padded, dtype=torch.float32)) + return torch.empty((0, self.output_dimension), dtype=torch.float32) + vectors = embed_with_vllm_llm( + texts_list, + self._llm, + batch_size=max(1, int(batch_size)), + prefix=self.document_prefix, + ) + return self._finalize_vectors(vectors) def embed_queries(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: - """Embed query strings. Returns CPU tensor ``[N, 2048]``.""" + """Embed query strings. Returns CPU tensor ``[N, D]``.""" self._ensure_loaded() from nemo_retriever.models.inference.vllm import embed_with_vllm_llm texts_list = [str(t) for t in texts if str(t).strip()] if not texts_list: - return torch.empty((0, 2048), dtype=torch.float32) - vectors = embed_with_vllm_llm(texts_list, self._llm, batch_size=max(1, int(batch_size)), prefix="query: ") - valid = [v for v in vectors if v] - if not valid: - return torch.empty((0, 2048), dtype=torch.float32) - dim = len(valid[0]) - padded = [v if v else [0.0] * dim for v in vectors] - return _l2_normalize(torch.tensor(padded, dtype=torch.float32)) + return torch.empty((0, self.output_dimension), dtype=torch.float32) + vectors = embed_with_vllm_llm( + texts_list, + self._llm, + batch_size=max(1, int(batch_size)), + prefix=self.query_prefix, + ) + return self._finalize_vectors(vectors) def embed_images(self, images_b64: Sequence[str], *, batch_size: int = 64) -> torch.Tensor: - """Embed images (base64-encoded). Returns CPU tensor ``[N, 2048]``.""" + """Embed images (base64-encoded). Returns CPU tensor ``[N, D]``.""" self._ensure_loaded() from nemo_retriever.models.inference.vllm import embed_multimodal_with_vllm_llm valid_b64 = [b64 for b64 in images_b64 if b64 and str(b64).strip()] if not valid_b64: - return torch.empty((0, 2048), dtype=torch.float32) + return torch.empty((0, self.output_dimension), dtype=torch.float32) prompt_dicts = [ - {"prompt": "passage: ", "multi_modal_data": {"image": _b64_to_pil(b64)}} for b64 in valid_b64 + {"prompt": f"{self.document_prefix} ", "multi_modal_data": {"image": _b64_to_pil(b64)}} + for b64 in valid_b64 ] vectors = embed_multimodal_with_vllm_llm(prompt_dicts, self._llm, batch_size=max(1, int(batch_size))) - valid = [v for v in vectors if v] - if not valid: - return torch.empty((0, 2048), dtype=torch.float32) - dim = len(valid[0]) - padded = [v if v else [0.0] * dim for v in vectors] - return _l2_normalize(torch.tensor(padded, dtype=torch.float32)) + return self._finalize_vectors(vectors) def embed_text_image( self, texts: Sequence[str], images_b64: Sequence[str], *, batch_size: int = 64 ) -> torch.Tensor: - """Embed paired text+image inputs. Returns CPU tensor ``[N, 2048]``.""" + """Embed paired text+image inputs. Returns CPU tensor ``[N, D]``.""" self._ensure_loaded() from nemo_retriever.models.inference.vllm import embed_multimodal_with_vllm_llm @@ -311,19 +319,14 @@ def embed_text_image( paired_b64.append(b64) if not paired_b64: - return torch.empty((0, 2048), dtype=torch.float32) + return torch.empty((0, self.output_dimension), dtype=torch.float32) prompt_dicts = [ - {"prompt": f"passage: {text}", "multi_modal_data": {"image": _b64_to_pil(b64)}} + {"prompt": f"{self.document_prefix} {text}", "multi_modal_data": {"image": _b64_to_pil(b64)}} for text, b64 in zip(paired_texts, paired_b64) ] vectors = embed_multimodal_with_vllm_llm(prompt_dicts, self._llm, batch_size=max(1, int(batch_size))) - valid = [v for v in vectors if v] - if not valid: - return torch.empty((0, 2048), dtype=torch.float32) - dim = len(valid[0]) - padded = [v if v else [0.0] * dim for v in vectors] - return _l2_normalize(torch.tensor(padded, dtype=torch.float32)) + return self._finalize_vectors(vectors) def unload(self) -> None: """Release GPU memory held by the vLLM engine.""" diff --git a/nemo_retriever/src/nemo_retriever/models/nim/chat_completions.py b/nemo_retriever/src/nemo_retriever/models/nim/chat_completions.py index dddb790dde..1ab6465da2 100644 --- a/nemo_retriever/src/nemo_retriever/models/nim/chat_completions.py +++ b/nemo_retriever/src/nemo_retriever/models/nim/chat_completions.py @@ -97,7 +97,7 @@ def invoke_chat_completion_step( tools: Optional[List[Dict[str, Any]]] = None, tool_choice: str = "auto", timeout_s: float = 120.0, - temperature: float = 0.0, + temperature: Optional[float] = 0.0, max_tokens: Optional[int] = None, extra_body: Optional[Dict[str, Any]] = None, max_retries: int = 10, @@ -117,6 +117,9 @@ def invoke_chat_completion_step( tool_choice ``"auto"`` (default) lets the model decide; ``"none"`` suppresses tool use; or a specific tool name dict. + temperature + Sampling temperature. ``None`` omits the field from the payload entirely + so the endpoint/model default applies. """ token = (api_key or "").strip() headers: Dict[str, str] = {"Accept": "application/json", "Content-Type": "application/json"} @@ -126,10 +129,10 @@ def invoke_chat_completion_step( invoke_urls = _parse_invoke_urls(invoke_url) endpoint_url = invoke_urls[0] - payload: Dict[str, Any] = { - "messages": messages, - "temperature": temperature, - } + payload: Dict[str, Any] = {"messages": messages} + # Unset (None) => omit, so the endpoint/model default applies. + if temperature is not None: + payload["temperature"] = temperature if model: payload["model"] = model if max_tokens is not None: @@ -159,7 +162,7 @@ def invoke_chat_completions_images( timeout_s: float = 120.0, task_prompt: Optional[str] = None, temperature: float = 0.0, - repetition_penalty: float = 1.1, + repetition_penalty: Optional[float] = 1.1, extra_body: Optional[Dict[str, Any]] = None, max_pool_workers: int = 16, max_retries: int = 10, diff --git a/nemo_retriever/src/nemo_retriever/models/nim/nim.py b/nemo_retriever/src/nemo_retriever/models/nim/nim.py index 56a38851cb..6613b50a60 100644 --- a/nemo_retriever/src/nemo_retriever/models/nim/nim.py +++ b/nemo_retriever/src/nemo_retriever/models/nim/nim.py @@ -214,7 +214,7 @@ def _post_with_retries( if 400 <= status_code < 500: raise requests.HTTPError( - f"HTTP {status_code} from {invoke_url}: {response.text}", + f"HTTP {status_code} from {_safe_endpoint_attribute(invoke_url)}: {response.text}", response=response, ) response.raise_for_status() @@ -450,7 +450,7 @@ def invoke_chat_completions_images( timeout_s: float = 120.0, task_prompt: Optional[str] = None, temperature: float = 0.0, - repetition_penalty: float = 1.1, + repetition_penalty: Optional[float] = 1.1, extra_body: Optional[Dict[str, Any]] = None, max_retries: int = 10, max_429_retries: int = 5, @@ -473,7 +473,9 @@ def invoke_chat_completions_images( ) messages_list.append([{"role": "user", "content": content}]) - merged_extra: Dict[str, Any] = {"repetition_penalty": repetition_penalty} + merged_extra: Dict[str, Any] = {} + if repetition_penalty is not None: + merged_extra["repetition_penalty"] = repetition_penalty if extra_body: merged_extra.update(extra_body) diff --git a/nemo_retriever/src/nemo_retriever/models/nim/primitives/model_interface/yolox.py b/nemo_retriever/src/nemo_retriever/models/nim/primitives/model_interface/yolox.py index 347f203120..f3414b6c10 100644 --- a/nemo_retriever/src/nemo_retriever/models/nim/primitives/model_interface/yolox.py +++ b/nemo_retriever/src/nemo_retriever/models/nim/primitives/model_interface/yolox.py @@ -669,7 +669,11 @@ def postprocess_annotations(self, annotation_dicts, final_score=None, **kwargs): if annotation_dicts and running_v3: annotation_dicts = [ - postprocess_page_elements_v3(annotation_dict, labels=YOLOX_PAGE_V3_CLASS_LABELS) + postprocess_page_elements_v3( + annotation_dict, + labels=YOLOX_PAGE_V3_CLASS_LABELS, + final_score=final_score, + ) for annotation_dict in annotation_dicts ] else: @@ -857,12 +861,18 @@ def expand_chart_bboxes(annotation_dict, labels=None): return annotation_dict -def postprocess_page_elements_v3(annotation_dict, labels=None): +def postprocess_page_elements_v3( + annotation_dict: Dict[str, List[List[float]]], + labels: Optional[List[str]] = None, + final_score: Optional[Dict[str, float]] = None, +) -> Dict[str, List[List[float]]]: """ Expand bounding boxes of tables/charts/infographics and titles based on the bounding boxes of the other class. Args: annotation_dict: output of postprocess_results, a dictionary with keys: "table", "chart", "infographics", "title", "paragraph", "header_footer". + labels: ordered class labels corresponding to the annotation dictionary. + final_score: per-class thresholds used by the final output filter. Returns: annotation_dict: same as input, with expanded bboxes for page elements. @@ -870,6 +880,8 @@ def postprocess_page_elements_v3(annotation_dict, labels=None): """ if not labels: labels = list(annotation_dict.keys()) + if final_score is None: + final_score = YOLOX_PAGE_V3_FINAL_SCORE if not annotation_dict: return annotation_dict @@ -897,6 +909,20 @@ def postprocess_page_elements_v3(annotation_dict, labels=None): label_idxs = np.concatenate(label_idxs) bboxes, confidences, label_idxs = remove_overlapping_boxes_using_wbf(bboxes, confidences, label_idxs) + + # Apply final thresholds here only to determine matching eligibility; the + # caller still owns final output filtering. A low-confidence page-sized table + # can otherwise consume a valid title before both detections are removed. + matching_thresholds = np.array([final_score.get(labels[int(label_idx)], 0.0) for label_idx in label_idxs]) + eligible_for_matching = confidences >= matching_thresholds + bboxes, confidences, label_idxs = ( + bboxes[eligible_for_matching], + confidences[eligible_for_matching], + label_idxs[eligible_for_matching], + ) + if not len(bboxes): + return {label: [] for label in labels} + bboxes, confidences, label_idxs, found_title = match_structured_boxes_with_title( bboxes, confidences, label_idxs, labels ) diff --git a/nemo_retriever/src/nemo_retriever/operators/embed/cpu_operator.py b/nemo_retriever/src/nemo_retriever/operators/embed/cpu_operator.py index 935fafb86b..95e6da2413 100644 --- a/nemo_retriever/src/nemo_retriever/operators/embed/cpu_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/embed/cpu_operator.py @@ -12,6 +12,7 @@ from nemo_retriever.operators.cpu_operator import CPUOperator from nemo_retriever.models.nim.probe import probe_endpoint from nemo_retriever.common.params import EmbedParams +from nemo_retriever.common.api.util.string_processing import prepend_model_provider_prefix from nemo_retriever.models.inference.runtime import embed_text_main_text_embed from nemo_retriever.models.inference.shared import build_embed_kwargs @@ -45,7 +46,9 @@ def __init__(self, params: EmbedParams) -> None: # Probe the /embeddings path with a model-name-only body — auth is # checked before body validation so a bad key returns 401 without # triggering inference. A valid key with an empty input returns 400. - model_name = self._kwargs.get("model_name", "") + model_name = prepend_model_provider_prefix( + self._kwargs.get("model_name"), self._kwargs.get("embed_model_provider_prefix") + ) probe_url = ( endpoint if endpoint.rstrip("/").endswith("/embeddings") else endpoint.rstrip("/") + "/embeddings" ) diff --git a/nemo_retriever/src/nemo_retriever/operators/embed/gpu_operator.py b/nemo_retriever/src/nemo_retriever/operators/embed/gpu_operator.py index b4756b61d3..44c2804aef 100644 --- a/nemo_retriever/src/nemo_retriever/operators/embed/gpu_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/embed/gpu_operator.py @@ -64,6 +64,7 @@ def __init__(self, params: EmbedParams) -> None: normalize=bool(self._kwargs.get("normalize", True)), max_length=int(self._kwargs.get("max_length", 8192)), query_max_length=int(self._kwargs.get("query_max_length", 128)), + revision=self._kwargs.get("embed_model_revision"), ) def preprocess(self, data: Any, **kwargs: Any) -> Any: diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/html/ray_data.py b/nemo_retriever/src/nemo_retriever/operators/extract/html/ray_data.py index 79218dba16..6fbf48661a 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/html/ray_data.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/html/ray_data.py @@ -12,6 +12,9 @@ import pandas as pd +from nemo_retriever.common.modality.txt.tokenizer_provider import ( + TokenizerUnavailableError, +) from nemo_retriever.common.params import HtmlChunkParams from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.operators.cpu_operator import CPUOperator @@ -62,6 +65,8 @@ def process(self, data: Any, **kwargs: Any) -> Any: chunk_df = html_bytes_to_chunks_df(payload, path_str, params=params) if not chunk_df.empty: out_dfs.append(chunk_df) + except TokenizerUnavailableError: + raise except Exception: continue if not out_dfs: diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/parse/nemotron_parse.py b/nemo_retriever/src/nemo_retriever/operators/extract/parse/nemotron_parse.py index cdf3253104..46098e39ea 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/parse/nemotron_parse.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/parse/nemotron_parse.py @@ -12,7 +12,10 @@ from __future__ import annotations +from dataclasses import dataclass +from enum import Enum from typing import Any, Dict, List, Optional, Tuple +from urllib.parse import urlsplit import base64 import io @@ -47,6 +50,7 @@ # --------------------------------------------------------------------------- NEMOTRON_PARSE_REMOTE_DEFAULT_MODEL = "nvidia/nemotron-parse-v1.2" +NEMOTRON_PARSE_HOSTED_MODEL = "nvidia/nemotron-parse" NEMOTRON_PARSE_LOCAL_DEFAULT_MODEL = "nvidia/NVIDIA-Nemotron-Parse-v1.2" NEMOTRON_PARSE_DEFAULT_TASK_PROMPT = "" @@ -147,19 +151,79 @@ def _route_parsed_elements( return table_items, chart_items, infographic_items, page_text +class _NemotronParseContractProfile(str, Enum): + HOSTED_TOOL_CALL = "hosted_tool_call" + LEGACY_TOOL_CALL = "legacy_tool_call" + V1_2_TAGGED = "v1_2_tagged" + + +@dataclass(frozen=True) +class _ResolvedNemotronParseContract: + model: str + profile: _NemotronParseContractProfile + has_build_endpoint: bool + + @property + def uses_tool_call_routing(self) -> bool: + return self.profile in { + _NemotronParseContractProfile.HOSTED_TOOL_CALL, + _NemotronParseContractProfile.LEGACY_TOOL_CALL, + } + + def _is_legacy_nemotron_parse_model(model_name: str) -> bool: normalized = model_name.lower() return bool(re.search(r"v1[._][01](?!\d)", normalized)) -def _route_parsed_elements_v1( +def _is_nvidia_build_endpoint(invoke_url: str) -> bool: + return (urlsplit(invoke_url).hostname or "").lower() == "integrate.api.nvidia.com" + + +def _resolve_nemotron_parse_contract( + invoke_url: str, + model_name: Optional[str], +) -> _ResolvedNemotronParseContract: + """Resolve the internal request/response contract for a chat endpoint.""" + + invoke_urls = [part.strip() for part in str(invoke_url or "").split(",") if part.strip()] + if not invoke_urls: + raise ValueError("Nemotron Parse invoke_url is required.") + + build_endpoints = [_is_nvidia_build_endpoint(url) for url in invoke_urls] + explicit_model = str(model_name or "").strip() + if not explicit_model and any(build_endpoints) and not all(build_endpoints): + raise ValueError( + "Nemotron Parse endpoint lists cannot mix NVIDIA Build and self-hosted endpoints " + "unless `nemotron_parse_model` is set explicitly." + ) + + resolved_model = explicit_model or ( + NEMOTRON_PARSE_HOSTED_MODEL if all(build_endpoints) else NEMOTRON_PARSE_REMOTE_DEFAULT_MODEL + ) + normalized_model = resolved_model.lower() + if normalized_model == NEMOTRON_PARSE_HOSTED_MODEL: + profile = _NemotronParseContractProfile.HOSTED_TOOL_CALL + elif _is_legacy_nemotron_parse_model(normalized_model): + profile = _NemotronParseContractProfile.LEGACY_TOOL_CALL + else: + profile = _NemotronParseContractProfile.V1_2_TAGGED + + return _ResolvedNemotronParseContract( + model=resolved_model, + profile=profile, + has_build_endpoint=any(build_endpoints), + ) + + +def _route_tool_call_elements( raw_json_text: str, *, extract_tables: bool, extract_charts: bool, extract_infographics: bool, ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]], Optional[str]]: - """Route v1.0/v1.1 tool-call JSON into pipeline content channels.""" + """Route hosted or legacy tool-call JSON into pipeline content channels.""" try: parsed = json.loads(raw_json_text) @@ -167,12 +231,19 @@ def _route_parsed_elements_v1( return [], [], [], None elements: List[Dict[str, Any]] = [] - if isinstance(parsed, list): - for item in parsed: - if isinstance(item, list): - elements.extend(elem for elem in item if isinstance(elem, dict)) - elif isinstance(item, dict): - elements.append(item) + + def _collect(value: Any) -> None: + if isinstance(value, dict): + if "type" in value and ("text" in value or "bbox" in value): + elements.append(value) + return + for nested in value.values(): + _collect(nested) + elif isinstance(value, list): + for nested in value: + _collect(nested) + + _collect(parsed) table_items: List[Dict[str, Any]] = [] chart_items: List[Dict[str, Any]] = [] @@ -297,23 +368,27 @@ def nemotron_parse_pages( # -- Phase 2: run model inference in a single batch ------------------ raw_texts: List[str] = [""] * len(batch_indices) - used_v1_api = False + uses_tool_call_routing = False + contract: _ResolvedNemotronParseContract | None = None if batch_images: try: if use_remote: if "/v1/chat/completions" in invoke_url: - _model_name = nemotron_parse_model or NEMOTRON_PARSE_REMOTE_DEFAULT_MODEL - used_v1_api = _is_legacy_nemotron_parse_model(_model_name) + contract = _resolve_nemotron_parse_contract(invoke_url, nemotron_parse_model) + uses_tool_call_routing = contract.uses_tool_call_routing extra_body: Dict[str, Any] = {"max_tokens": 8192} - if used_v1_api: + if contract.profile == _NemotronParseContractProfile.LEGACY_TOOL_CALL: extra_body["tools"] = [{"type": "function", "function": {"name": "markdown_bbox"}}] _chat_kw = dict( invoke_url=invoke_url, image_b64_list=batch_images, - model=_model_name, + model=contract.model, api_key=api_key, timeout_s=float(request_timeout_s), - task_prompt=None if used_v1_api else task_prompt, + task_prompt=None if uses_tool_call_routing else task_prompt, + repetition_penalty=( + None if contract.profile == _NemotronParseContractProfile.HOSTED_TOOL_CALL else 1.1 + ), extra_body=extra_body, max_retries=int(retry.remote_max_retries), max_429_retries=int(retry.remote_max_429_retries), @@ -351,6 +426,22 @@ def nemotron_parse_pages( else: raw_texts = [str(model.invoke(img, task_prompt=task_prompt) or "").strip() for img in batch_images] except BaseException as e: + if ( + contract is not None + and nemotron_parse_model + and contract.has_build_endpoint + and contract.profile == _NemotronParseContractProfile.V1_2_TAGGED + and "text input" in str(e).lower() + ): + hint = ValueError( + "Nemotron Parse model/contract mismatch: NVIDIA Build model " + "`nvidia/nemotron-parse` uses an image-only tool-call contract, but " + f"`{contract.model}` selected the v1.2 text-control-token contract. " + "Use `nemotron_parse_model='nvidia/nemotron-parse'` with Build, or send " + "the versioned v1.2 model to a compatible self-hosted endpoint." + ) + hint.__cause__ = e + e = hint print(f"Warning: Nemotron Parse batch failed: {type(e).__name__}: {e}") err = { "stage": "nemotron_parse_pages", @@ -363,7 +454,7 @@ def nemotron_parse_pages( raw_texts = [] # -- Phase 3: route parsed elements into content channels ------------ - route_fn = _route_parsed_elements_v1 if used_v1_api else _route_parsed_elements + route_fn = _route_tool_call_elements if uses_tool_call_routing else _route_parsed_elements for pos, raw_text in enumerate(raw_texts): idx = batch_indices[pos] try: diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/txt/ray_data.py b/nemo_retriever/src/nemo_retriever/operators/extract/txt/ray_data.py index 2676ba921b..6713421f30 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/txt/ray_data.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/txt/ray_data.py @@ -8,17 +8,23 @@ from __future__ import annotations +import logging from typing import Any, Dict, List # noqa: F401 import pandas as pd +from nemo_retriever.common.modality.txt.tokenizer_provider import ( + TokenizerUnavailableError, +) from nemo_retriever.common.params import TextChunkParams from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.operators.cpu_operator import CPUOperator from nemo_retriever.graph.designer import designer_component from nemo_retriever.operators.operator_archetype import ArchetypeOperator -from nemo_retriever.common.modality.txt.split import txt_bytes_to_chunks_df +from nemo_retriever.common.modality.txt.split import empty_text_chunks_df, text_to_chunks_df, txt_bytes_to_chunks_df + +logger = logging.getLogger(__name__) @designer_component( @@ -70,9 +76,9 @@ def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: ) class TxtSplitCPUActor(AbstractOperator, CPUOperator): """ - Ray Data map_batches callable: DataFrame with bytes, path -> DataFrame of chunks. + Ray Data map_batches callable: DataFrame with bytes/text, path -> DataFrame of chunks. - Each output row has: text, path, page_number, metadata (same shape as txt_file_to_chunks_df). + Each output row has: text, content, path, page_number, metadata (same shape as txt_file_to_chunks_df). """ def __init__(self, params: TextChunkParams | None = None) -> None: @@ -81,7 +87,7 @@ def __init__(self, params: TextChunkParams | None = None) -> None: def preprocess(self, data: Any, **kwargs: Any) -> Any: if not isinstance(data, pd.DataFrame) or data.empty: - return pd.DataFrame(columns=["text", "path", "page_number", "metadata"]) + return empty_text_chunks_df() return data def process(self, data: Any, **kwargs: Any) -> Any: @@ -94,18 +100,25 @@ def process(self, data: Any, **kwargs: Any) -> Any: raw = row.get("bytes") text = row.get("text") path = row.get("path") - if (raw is None and text is None) or path is None: + if (not isinstance(raw, (bytes, bytearray)) and not isinstance(text, str)) or path is None: continue path_str = str(path) if path is not None else "" try: - payload = raw or text.encode("utf-8") - chunk_df = txt_bytes_to_chunks_df(payload, path_str, params=params) + if isinstance(raw, (bytes, bytearray)): + chunk_df = txt_bytes_to_chunks_df(bytes(raw), path_str, params=params) + elif isinstance(text, str): + chunk_df = text_to_chunks_df(text, path_str, params=params) + else: + continue if not chunk_df.empty: out_dfs.append(chunk_df) + except TokenizerUnavailableError: + raise except Exception: + logger.warning("Failed to split text source %r", path_str, exc_info=True) continue if not out_dfs: - return pd.DataFrame(columns=["text", "path", "page_number", "metadata"]) + return empty_text_chunks_df() return pd.concat(out_dfs, ignore_index=True) def postprocess(self, data: Any, **kwargs: Any) -> Any: diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py index 4750706432..54a283a1f7 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py @@ -2,29 +2,46 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Operator that runs a ReAct agentic retrieval loop per query.""" +"""Operator that runs a ReAct agentic retrieval loop per query. + +The agent logic itself lives in the private :mod:`nemo_retriever._agentic.nemo_agent` +library. This operator is a thin adapter: it builds an +:class:`~nemo_retriever._agentic.nemo_agent.Agent` (select mode) from a subset of the +library's configuration, runs it once per query, and flattens the resulting +retrieval log and final doc-id list into the exploded DataFrame that +:class:`RRFAggregatorOperator` and :class:`SelectionAgentOperator` consume. +""" from __future__ import annotations -import json import logging -import os - -import requests from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Callable, Dict, List, Literal, Optional +from typing import Any, Callable, Dict, List, Optional import pandas as pd +from nemo_retriever._agentic.nemo_agent import Agent, AgentConfig, create_retrieve_tool +from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.operators.cpu_operator import CPUOperator -from nemo_retriever.models.nim.chat_completions import invoke_chat_completion_step logger = logging.getLogger(__name__) _LOG_PREVIEW_CHARS = 300 _LOG_DOC_ID_LIMIT = 20 +#: Output DataFrame columns emitted by :func:`_build_output_rows` / this operator. +_OUTPUT_COLUMNS = [ + "query_id", + "query_text", + "step_idx", + "doc_id", + "text", + "rank", + "has_valid_final_results", + "is_final_result", +] + def _preview_text(value: Any, *, limit: int = _LOG_PREVIEW_CHARS) -> str: text = " ".join(str(value or "").split()) @@ -33,270 +50,20 @@ def _preview_text(value: Any, *, limit: int = _LOG_PREVIEW_CHARS) -> str: return text[:limit].rstrip() + "..." -def _preview_doc_ids(docs: List[Dict[str, Any]], *, limit: int = _LOG_DOC_ID_LIMIT) -> List[str]: - return [str(doc.get("doc_id", "")) for doc in docs[:limit]] - - -# --------------------------------------------------------------------------- -# Prompt rendering (verbatim content of 02_v1.j2, rendered via Python) -# --------------------------------------------------------------------------- - -_GOAL = """\ -You are a retrieval agent that finds all documents related to a given query. - - -You are given a search query and a list of documents retrieved for that query. Your task is to write new \ -queries and use the given search tool to find *ALL* the related and somewhat related documents to the given \ -query (i.e., maximize recall). -If the user's query is a question, you should not answer the question yourself. Instead, you should find \ -the related documents for the given query. -""" - -_RELEVANCE_DEFINITION = """ - - -- You should be careful, in the context of this task, what it means to be a "query", "document", and \ -"relevant" can sometimes be very complex and might not follow the traditional definition of these terms \ -in standard information retrieval. -- In standard retrieval, a query is usually a user question (like a web search query), the document is \ -some sort of content that provides information (e.g., a web page), and these two are considered relevant if \ -the document provides information that answers the user's query. -- However, in our setting, this could be different. Here are some examples: - * the query is a programming problem and documents are programming language syntax references. A document \ -is relevant if it contains the reference for the programming syntax used for solving the problem. - * both query and documents are descriptions programming problems and a query and document are relevant if \ -the same approach is used to solve them. - * the query is a math problem and documents are theorems. Relevant documents (theorems) are the ones \ -that are useful for solving the math problem. - * the query and document are both math problems. A query and a document are relevant if the same theorem \ -is used for solving them. - * the query is a task description (e.g., for an API programmer) and documents are descriptions of \ -available APIs. Relevant documents (e.g., APIs) are the ones needed for completing the task. -- This is not an exhaustive list. These are just some examples to show you the complexity of queries, \ -documents, and the concept of relevance in this task. -- Note that even here, the relevant documents are still the ones that are useful for a user who is \ -searching for the given query. But the relation is more nuanced. -- You should analyze the query and some of the available documents. And then reason about what could be a \ -meaningful definition of relevance in this case, and what the user could be looking for. -- Moreover, sometimes, the query could be even a prompt that is given to a Large Language Model (LLM) and \ -the user wants to find the useful documents for the LLM that help answering/solving this prompt. -""" - -_WORKFLOW_TEMPLATE = """ - -- You are given a retrieval tool, powered by a dense embedding model, that takes a text query and returns \ -the most similar documents. -{extended_relevance_line}\ -- You can call the search tool multiple times. -- Search for related documents to the user's query from different angles. -- If needed, revise your search queries based on the documents you find in previous steps. -- Once you are confident that you have found all the related and somewhat related documents and there are \ -no more related documents in the corpus, call the "final_results" tool to finish the task. -{final_results_count_line}\ -- When calling the "final_results" tool, the list of documents must be sorted in the decreasing level of \ -relevance to the query. I.e., the first document is the most relevant to the query, the second document is \ -the second most relevant to the query, and so on. -""" - -_BEST_PRACTICES_TEMPLATE = """ - - -- You should be thorough and find all related and somewhat related documents. -- The goal is to increase the **Recall** of your search attempt. So, if multiple documents are relevant \ -to the given query, you should find and report all of them even if only a subset of them is enough \ -for answering the query. -{with_init_docs_line}\ -""" - - -def _render_react_agent_prompt( - top_k: int, - *, - with_init_docs: bool = True, - extended_relevance: bool = False, -) -> str: - """Render the ReAct agent system prompt (verbatim 02_v1.j2 logic).""" - parts = [_GOAL] - if extended_relevance: - parts.append(_RELEVANCE_DEFINITION) - - ext_line = ( - "- As explained above, reason and figure out what the meaning of relevance is in this case, " - "and what could be relevant and useful information for the given query.\n" - if extended_relevance - else "" - ) - final_results_count_line = ( - f'- When calling "final_results", you must select exactly the {top_k} most relevant documents ' - "among all documents you have retrieved.\n" - ) - parts.append( - _WORKFLOW_TEMPLATE.format( - extended_relevance_line=ext_line, - final_results_count_line=final_results_count_line, - ) - ) - - init_docs_line = ( - "- **TIP**: you can look at the list of documents retrieved using the original query and think " - "what other queries you can use to find the potentially related documents that are missing in these results.\n" - if with_init_docs - else "" - ) - parts.append(_BEST_PRACTICES_TEMPLATE.format(with_init_docs_line=init_docs_line)) - return "".join(parts) - - -# --------------------------------------------------------------------------- -# Tool specs (verbatim from retrieval_bench/nemo_agentic/tool_helpers.py) -# --------------------------------------------------------------------------- - - -def _make_think_tool_spec(extended_relevance: bool = False) -> Dict[str, Any]: - ext = "" - if extended_relevance: - ext = ( - "- When it is difficult to understand what is the intent of the user and what they are trying " - "to find with this query, use this tool to think about potential definitions of relevance that " - "could be meaningful/useful to the user for this task.\n" - "- If the intention of the user is vague especially given the available documents, use this tool " - "to think how you should decide what documents are relevant and what the metric of relevance is.\n" - ) - description = ( - "Use the tool to think about something. It will not obtain new information or make any changes, " - "but just log the thought. Use it when complex reasoning or brainstorming is needed.\n\n" - "Common use cases:\n" - f"{ext}" - "- When processing a complex query, use this tool to organize your thoughts and think about " - "the sub queries that you need to search for to find the relevant information\n" - "- If a query is vague is very difficult to find information for it, you can use this tool to think " - "about clues in the query that you can use to narrow down the search and spot relevant pieces of information.\n" - "- When finding related documents that help you create better search queries in the next step, use this " - "tool to think about what pieces of information from these documents are helpful to search for.\n" - "- When you fail to find any related information to the query, use this tool to think about other " - "search strategies that you can take to retrieve the related documents\n\n" - "The tool simply logs your thought process for better transparency and does not make any changes." - ) - return { - "type": "function", - "function": { - "name": "think", - "description": description, - "parameters": { - "type": "object", - "properties": {"thought": {"type": "string", "description": "The thought to log."}}, - "required": ["thought"], - }, - }, - } - - -def _make_retrieve_tool_spec(top_k: int) -> Dict[str, Any]: - return { - "type": "function", - "function": { - "name": "retrieve", - "description": ( - "Search for documents relevant to the given query using a dense embedding retrieval system. " - "Returns the most semantically similar documents from the corpus." - ), - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query to retrieve documents for.", - }, - }, - "required": ["query"], - }, - }, - } - - -def _make_final_results_tool_spec(top_k: int) -> Dict[str, Any]: - tk_ins = f"- You must choose exactly {top_k} document IDs when calling this function.\n" - - description = ( - "Signals the completion of the search process for the current query.\n\n" - "Use this tool when:\n" - "- You have found all the relevant documents to the query.\n" - "- Despite several attempts, you cannot find good documents for the given query.\n\n" - "The message should include:\n" - "- A brief summary of your exploration and the results\n" - "- Explanation if the search was unsuccessful\n\n" - "When reporting the selected document IDs, make sure:\n" - "- the list of document IDs is sorted in the decreasing level of relevance to the query. " - "I.e., the first document in the list is the most relevant to the query, the second is the " - "second most relevant to the query, and so on.\n" - f"{tk_ins}" - "\nThe successful_search field should be set to true if you believed you have found the most " - "relevant documents to the user's query, and false otherwise. And partial if it is in between." - ) - return { - "type": "function", - "function": { - "name": "final_results", - "description": description, - "parameters": { - "type": "object", - "required": ["doc_ids", "message", "search_successful"], - "properties": { - "message": { - "type": "string", - "description": ( - "A message for the user to explain why you think you found all the related " - "documents and there is no related document is missing. Also, include a short " - "description of your exploration process. If your attempts to find related " - "documents were unsuccessful, explain why." - ), - }, - "doc_ids": { - "type": "array", - "items": {"type": "string"}, - "minItems": 1, - "description": ( - "List of document IDs that are relevant to the user's query sorted descending " - "by their level of relevance to the user's query. I.e., the first document is " - "the most relevant to the query, the second is the second most relevant to the " - "query, and so on." - ), - }, - "search_successful": { - "type": "string", - "enum": ["true", "false", "partial"], - "description": "Whether you managed to find all the related documents to the query.", - }, - }, - }, - }, - } - - -# --------------------------------------------------------------------------- -# Operator -# --------------------------------------------------------------------------- - -#: Message sent when the LLM produces a stop without calling any tool. -_AUTO_USER_MSG = ( - "continue with the task. Do not re-read the query. Do not summarize your progress. " - "If you believe you have done all the required steps, call the `final_results` tool" -) - - class ReActAgentOperator(AbstractOperator, CPUOperator): """Run an iterative ReAct retrieval loop per query and emit the full retrieval log. - Each query row is processed independently by an LLM-driven ReAct loop - (Reason + Act) that has access to three tools: ``think``, ``retrieve``, - and ``final_results``. The operator emits one output row per retrieved - document per retrieval step, enabling downstream + Each query row is processed independently by an + :class:`~nemo_retriever._agentic.nemo_agent.Agent` running in ``select`` mode. The + agent owns the ReAct loop, prompt rendering, tool schemas, and the + over-fetch/dedup retrieval bookkeeping; this operator only supplies the + retrieval callback and translates results into the library's exploded + DataFrame convention. The operator emits one output row per retrieved + document per retrieval step (plus a synthetic final step for the agent's + ``final_results`` selection), enabling downstream :class:`RRFAggregatorOperator` to fuse the ranked lists with Reciprocal Rank Fusion. - The system prompt is a verbatim Python rendering of the retrieval-bench - ``02_v1.j2`` template, including the optional ``extended_relevance`` block. - Input DataFrame schema ---------------------- query_id : str — unique query identifier @@ -305,82 +72,67 @@ class ReActAgentOperator(AbstractOperator, CPUOperator): Output DataFrame schema ----------------------- - query_id : str — same ``query_id`` as the input - query_text : str — same ``query_text`` (passed through for downstream) - step_idx : int — 0 = initial seed retrieval; 1 … N = per-loop retrieve calls - doc_id : str — retrieved document identifier - text : str — document text - rank : int — 1-indexed rank within this step (1 = most relevant) + query_id : str — same ``query_id`` as the input + query_text : str — same ``query_text`` (passed through) + step_idx : int — 0 = seed retrieval; 1 … N = per-loop retrieve + calls; ``len(retrieval_log)`` = synthetic final step + doc_id : str — retrieved document identifier + text : str — document text + rank : int — 1-indexed rank within this step + has_valid_final_results : bool — True iff the agent returned a non-empty final list + is_final_result : bool — True only on the synthetic final-step rows Parameters ---------- invoke_url : str - Full ``/v1/chat/completions`` endpoint URL. + LLM endpoint. Forwarded as the LLM config's ``base_url``. llm_model : str - Model identifier forwarded to the endpoint. + Model identifier forwarded verbatim to the backend (litellm + provider-prefix transform is deferred). retriever_fn : Callable[[str, int], list[dict]] - ``(query_text, top_k) → [{doc_id: str, text: str, ...}]``. - The callable is invoked for every retrieve tool call the agent makes. - Each returned dict must contain ``doc_id`` and ``text`` keys. + ``(query_text, top_k) → [{doc_id: str, text: str, score: float}]``. + Wrapped by ``create_retrieve_tool`` after renaming ``doc_id`` → ``id``. retriever_top_k : int - Number of documents fetched per retrieve call. Defaults to ``500``. + Default number of documents requested per retrieve call (the tool's + ``default_top_k``). Defaults to ``500``. target_top_k : int - Number of final documents to select, communicated to the LLM via the - system prompt and ``final_results`` tool spec. Defaults to ``10``. - user_msg_type : {"with_results", "simple"} - ``"with_results"`` (default): make one upfront retrieval call with the - original query and include those documents in the first user message, - mirroring the retrieval-bench ``with_results`` mode. - ``"simple"``: start the loop with just the query text. - extended_relevance : bool - Include the ```` block in the system prompt for - tasks with non-standard relevance definitions. Defaults to ``False``. + Number of final documents the agent targets. Defaults to ``10``. max_steps : int - Maximum ReAct loop iterations per query before forced exit. - Defaults to ``10``. + Maximum agent LLM steps per query. Defaults to ``200``. num_concurrent : int Number of queries processed concurrently via ``ThreadPoolExecutor``. - Defaults to ``8``. api_key : str, optional - Literal API key **or** an ``"os.environ/VAR_NAME"`` reference. + Literal API key **or** an ``"os.environ/VAR_NAME"`` reference (resolved + by the LLM backend). max_tokens : int, optional - Upper bound on tokens in each LLM response. + Per-request completion budget (the LLM config's ``max_completion_tokens``). + parallel_tool_calls : bool, optional + Forwarded as the LLM config's ``parallel_tool_calls`` (sent to the + provider only when set). + reasoning_effort : str, optional + Forwarded as the LLM config's ``reasoning_effort``. + temperature : float, optional + Forwarded as the LLM config's ``temperature`` (sent to the provider + only when set). + backend : {"callable", "litellm"} + LLM backend to build. ``"callable"`` (default) drives an OpenAI-compatible + completion callable: ``chat_completion_fn`` when supplied (the in-process + vLLM adapter), otherwise the shared ``invoke_chat_completion_step`` HTTP + client against ``invoke_url``. + chat_completion_fn : callable, optional + OpenAI-compatible completion callable (e.g. the local in-process vLLM + adapter). When set, forwarded to the ``"callable"`` LLM backend. Notes ----- - ``retriever_fn`` must be serialisable when used with ``RayDataExecutor``. - Prefer module-level functions or picklable callable objects over lambdas. - - Examples - -------- - :: - - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - from nemo_retriever.operators.graph_ops.rrf_aggregator_operator import RRFAggregatorOperator - from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator - from nemo_retriever.graph.executor import InprocessExecutor - - def my_retriever(query_text: str, top_k: int) -> list[dict]: - # Returns [{doc_id, text, score?}, ...] - ... - - pipeline = ( - ReActAgentOperator( - invoke_url="https://integrate.api.nvidia.com/v1/chat/completions", - llm_model="nvidia/llama-3.3-nemotron-super-49b-v1", - retriever_fn=my_retriever, - retriever_top_k=500, - target_top_k=10, - ) - >> RRFAggregatorOperator(k=60) - >> SelectionAgentOperator( - invoke_url="https://integrate.api.nvidia.com/v1/chat/completions", - llm_model="nvidia/llama-3.3-nemotron-super-49b-v1", - top_k=10, - ) - ) - - result_df = InprocessExecutor(pipeline).ingest(query_df) + ``retriever_fn`` must be picklable when used with ``RayDataExecutor``. The + :class:`~nemo_retriever._agentic.nemo_agent.Agent` and LLM backend are built lazily on + first ``process`` call (not stored as constructor kwargs) so the operator + stays reconstructable via ``get_constructor_kwargs``. + + ``run_sync`` performs one ``asyncio.run`` per query; it must not be called + from inside a running event loop. This is safe because the graph is executed + synchronously (single query in the calling thread, batches in worker threads). """ _NVIDIA_BUILD_ENDPOINT = "https://integrate.api.nvidia.com/v1/chat/completions" @@ -393,16 +145,14 @@ def __init__( retriever_fn: Callable[[str, int], List[Dict[str, Any]]], retriever_top_k: int = 500, target_top_k: int = 10, - user_msg_type: Literal["with_results", "simple"] = "with_results", - extended_relevance: bool = False, - max_steps: int = 10, + max_steps: int = 200, num_concurrent: int = 8, api_key: Optional[str] = None, max_tokens: Optional[int] = None, - parallel_tool_calls: bool = True, + parallel_tool_calls: Optional[bool] = None, reasoning_effort: Optional[str] = None, - backend_top_k: Optional[int] = None, - temperature: float = 0.0, + temperature: Optional[float] = None, + backend: str = "callable", chat_completion_fn: Optional[Callable[..., Dict[str, Any]]] = None, ) -> None: super().__init__() @@ -411,26 +161,89 @@ def __init__( self._retriever_fn = retriever_fn self._retriever_top_k = retriever_top_k self._target_top_k = target_top_k - self._user_msg_type = user_msg_type - self._extended_relevance = extended_relevance self._max_steps = max_steps self._num_concurrent = num_concurrent self._api_key = api_key self._max_tokens = max_tokens self._parallel_tool_calls = parallel_tool_calls - self._reasoning_effort = reasoning_effort - self._backend_top_k = backend_top_k self._temperature = temperature + self._reasoning_effort = reasoning_effort + self._backend = backend + # When set, an OpenAI-compatible completion callable (e.g. the in-process + # local vLLM adapter). Forwarded to the "callable" LLM backend by _build_llm. self._chat_completion_fn = chat_completion_fn + # Built lazily on first process() so the live Agent/LLM (which hold a + # litellm client + lock) are never part of the picklable ctor state. + self._agent: Optional[Agent] = None - def _build_extra_body(self) -> Optional[Dict[str, Any]]: - """Assemble per-call extra payload fields (parallel_tool_calls, reasoning_effort).""" - extra: Dict[str, Any] = {} - if not self._parallel_tool_calls: - extra["parallel_tool_calls"] = False - if self._reasoning_effort: - extra["reasoning_effort"] = self._reasoning_effort - return extra or None + # ------------------------------------------------------------------ + # private agent construction (lazy, memoized) + # ------------------------------------------------------------------ + + def _build_llm(self) -> Any: + config = create_llm_config( + self._backend, + model=str(self._llm_model), + base_url=self._invoke_url, + api_key=self._api_key, + reasoning_effort=self._reasoning_effort or None, + temperature=self._temperature, + parallel_tool_calls=self._parallel_tool_calls, + max_completion_tokens=self._max_tokens, + ) + completion_fn = self._chat_completion_fn + if self._backend == "callable" and completion_fn is None: + # Remote run on the default backend: supply the shared chat-completions + # client. Imported HERE and injected, never imported by the agent + # library, which must not depend on the rest of nemo_retriever. + from nemo_retriever.models.nim.chat_completions import invoke_chat_completion_step + + completion_fn = invoke_chat_completion_step + kwargs = {"completion_fn": completion_fn} if completion_fn is not None else {} + return create_llm(config, **kwargs) + + def _retrieve_adapter(self, query: str, top_k: int) -> List[Dict[str, Any]]: + """Adapt ``retriever_fn`` output to the private agent's ``id``/``score``/``text`` contract.""" + top_k = min(top_k, 1_000) + out: List[Dict[str, Any]] = [] + for doc in self._retriever_fn(query, top_k): + doc_id = str(doc.get("doc_id", doc.get("id", ""))) + if not doc_id: + continue + out.append( + { + "id": doc_id, + "score": float(doc.get("score", 0.0)), + "text": str(doc.get("text", "")), + } + ) + return out + + def _ensure_agent(self) -> Agent: + if self._agent is None: + retrieve_tool = create_retrieve_tool( + "default", + self._retrieve_adapter, + name="retrieve", + default_top_k=int(self._retriever_top_k), + ) + self._agent = Agent( + config=AgentConfig( + mode="select", + target_top_k=int(self._target_top_k), + enforce_top_k=True, + user_msg_type="with_results", + extended_relevance=True, + enable_think=False, + ensure_new_docs=True, + end_tool_with_msg=False, + max_steps=int(self._max_steps), + on_error="never_raise", + ), + llm=self._build_llm(), + retrieve_tool=retrieve_tool, + ) + return self._agent # ------------------------------------------------------------------ # AbstractOperator interface @@ -448,46 +261,35 @@ def preprocess(self, data: Any, **kwargs: Any) -> pd.DataFrame: return data[["query_id", "query_text"]].copy() def process(self, data: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: - """Run the ReAct loop for each query, concurrently up to num_concurrent.""" - api_key = self._resolve_api_key() + """Run the agent for each query, concurrently up to num_concurrent.""" + self._ensure_agent() rows: List[Dict[str, Any]] = [] query_rows = [(str(r["query_id"]), str(r["query_text"])) for _, r in data.iterrows()] + if not query_rows: + return pd.DataFrame(columns=_OUTPUT_COLUMNS) if len(query_rows) == 1: - # Fast path: single query, no threading overhead - qid, qtxt = query_rows[0] - rows.extend(self._run_single_query(qid, qtxt, api_key)) + # Fast path: single query, no threading overhead. + rows.extend(self._run_single_query(*query_rows[0])) else: - # Collect per-query results keyed by query_id, then re-emit in the ORIGINAL - # input order. as_completed() yields futures in nondeterministic completion - # order; emitting in that order would make downstream groupby(sort=False) - # output order depend on which query finished first. Re-ordering here keeps - # the operator output deterministic regardless of concurrency. + # Collect per-query results keyed by query_id, then re-emit in the + # ORIGINAL input order so downstream groupby(sort=False) output is + # deterministic regardless of thread completion order. results_by_qid: Dict[str, List[Dict[str, Any]]] = {} with ThreadPoolExecutor(max_workers=min(self._num_concurrent, len(query_rows))) as executor: - futures = { - executor.submit(self._run_single_query, qid, qtxt, api_key): (qid, qtxt) for qid, qtxt in query_rows - } + futures = {executor.submit(self._run_single_query, qid, qtxt): qid for qid, qtxt in query_rows} for future in as_completed(futures): - qid, qtxt = futures[future] + qid = futures[future] try: results_by_qid[qid] = future.result() - except TimeoutError as exc: - logger.warning("ReActAgentOperator: query %r timed out: %s", qid, exc, exc_info=True) - except RuntimeError as exc: - logger.warning("ReActAgentOperator: query %r retries exhausted: %s", qid, exc, exc_info=True) - except requests.RequestException as exc: - logger.warning("ReActAgentOperator: query %r HTTP error: %s", qid, exc, exc_info=True) - except (json.JSONDecodeError, ValueError) as exc: - logger.warning("ReActAgentOperator: query %r data error: %s", qid, exc, exc_info=True) - except Exception as exc: # catches unexpected worker errors not covered above + except Exception as exc: # production: one bad query must not kill the batch logger.warning("ReActAgentOperator: query %r failed: %s", qid, exc, exc_info=True) for qid, _qtxt in query_rows: rows.extend(results_by_qid.get(qid, [])) if not rows: - return pd.DataFrame(columns=["query_id", "query_text", "step_idx", "doc_id", "text", "rank"]) + return pd.DataFrame(columns=_OUTPUT_COLUMNS) return pd.DataFrame(rows) @@ -495,409 +297,39 @@ def postprocess(self, data: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: return data # ------------------------------------------------------------------ - # Internal: single query ReAct loop + # Internal: single query # ------------------------------------------------------------------ - def _run_single_query( - self, - query_id: str, - query_text: str, - api_key: Optional[str], - ) -> List[Dict[str, Any]]: - """Run the full ReAct loop for one query; return a list of row dicts.""" - with_init_docs = self._user_msg_type == "with_results" - - system_prompt = _render_react_agent_prompt( - self._target_top_k, - with_init_docs=with_init_docs, - extended_relevance=self._extended_relevance, - ) - tools = [ - _make_think_tool_spec(self._extended_relevance), - _make_retrieve_tool_spec(self._retriever_top_k), - _make_final_results_tool_spec(self._target_top_k), - ] - - messages: List[Dict[str, Any]] = [{"role": "system", "content": system_prompt}] - - # Retrieval log: one list per step, each item is {doc_id, text, score?} - retrieval_log: List[List[Dict[str, Any]]] = [] - seen_doc_ids: set[str] = set() - step_counter = 0 - + def _run_single_query(self, query_id: str, query_text: str) -> List[Dict[str, Any]]: + """Run the agent for one query and translate its result into output rows.""" + agent = self._ensure_agent() logger.info( - "ReActAgentOperator: query=%s start max_steps=%d retriever_top_k=%d target_top_k=%d query=%r", + "ReActAgentOperator: query=%s start max_steps=%d target_top_k=%d query=%r", query_id, - self._max_steps, - self._retriever_top_k, - self._target_top_k, + int(self._max_steps), + int(self._target_top_k), _preview_text(query_text), ) - - # ------ optional initial retrieval (with_results mode) ------ - if with_init_docs: - init_docs = self._call_retriever(query_text, seen_doc_ids, api_key) - retrieval_log.append(init_docs) - step_counter += 1 - for d in init_docs: - seen_doc_ids.add(d["doc_id"]) - logger.info( - "ReActAgentOperator: query=%s initial_retrieve docs=%d seen=%d doc_ids=%s", - query_id, - len(init_docs), - len(seen_doc_ids), - _preview_doc_ids(init_docs), - ) - - doc_content = _docs_to_message_content(init_docs) - user_msg_content: List[Dict[str, Any]] = [ - {"type": "text", "text": f"Query:\n{query_text}\n\nRetrieved Documents:"} - ] + doc_content - messages.append({"role": "user", "content": user_msg_content}) - else: - messages.append({"role": "user", "content": f"Query:\n{query_text}"}) - - final_doc_ids: Optional[List[str]] = None - - # ------ main ReAct loop ------ - for _step in range(self._max_steps): - logger.info("ReActAgentOperator: query=%s step=%d begin seen_docs=%d", query_id, _step, len(seen_doc_ids)) - try: - chat_completion_fn = self._chat_completion_fn or invoke_chat_completion_step - response = chat_completion_fn( - invoke_url=self._invoke_url, - messages=messages, - model=self._llm_model, - api_key=api_key, - tools=tools, - tool_choice="auto", - temperature=self._temperature, - max_tokens=self._max_tokens, - extra_body=self._build_extra_body(), - ) - except TimeoutError as exc: - logger.warning( - "ReActAgentOperator: LLM call timed out on step %d for query %r: %s", - _step, - query_id, - exc, - exc_info=True, - ) - break - except RuntimeError as exc: - logger.warning( - "ReActAgentOperator: LLM retries exhausted on step %d for query %r: %s", - _step, - query_id, - exc, - exc_info=True, - ) - break - except requests.RequestException as exc: - logger.warning( - "ReActAgentOperator: LLM HTTP error on step %d for query %r: %s", - _step, - query_id, - exc, - exc_info=True, - ) - break - except json.JSONDecodeError as exc: - logger.warning( - "ReActAgentOperator: LLM returned invalid JSON on step %d for query %r: %s", - _step, - query_id, - exc, - exc_info=True, - ) - break - - if not response.get("choices"): - logger.warning( - "ReActAgentOperator: empty choices in API response on step %d for query %r", _step, query_id - ) - break - choice = response["choices"][0] - msg = choice["message"] - finish_reason = choice.get("finish_reason") - tool_calls = msg.get("tool_calls") or [] - if msg.get("content"): - # Agent reasoning can quote document text/PII; keep content at DEBUG. - logger.debug( - "ReActAgentOperator: query=%s step=%d assistant content=%r", - query_id, - _step, - _preview_text(msg.get("content")), - ) - usage = response.get("usage") or {} - logger.info( - "ReActAgentOperator: query=%s step=%d finish_reason=%s tool_calls=%s " - "prompt_tokens=%s completion_tokens=%s total_tokens=%s", - query_id, - _step, - finish_reason, - [((tc.get("function") or {}).get("name") or "") for tc in tool_calls], - usage.get("prompt_tokens"), - usage.get("completion_tokens"), - usage.get("total_tokens"), - ) - - # Append assistant turn - assistant_turn: Dict[str, Any] = {"role": "assistant"} - if msg.get("content"): - assistant_turn["content"] = msg["content"] - if tool_calls: - assistant_turn["tool_calls"] = tool_calls - messages.append(assistant_turn) - - if finish_reason == "stop" or not tool_calls: - logger.info( - "ReActAgentOperator: query=%s step=%d no tool call; requesting continuation", - query_id, - _step, - ) - messages.append({"role": "user", "content": _AUTO_USER_MSG}) - continue - - tool_messages: List[Dict[str, Any]] = [] - loop_done = False - - for tc in tool_calls: - tc_id = tc.get("id", "") - fn = tc.get("function", {}) - fn_name = fn.get("name", "") - try: - fn_args = json.loads(fn.get("arguments", "{}")) - except json.JSONDecodeError: - tool_messages.append( - {"role": "tool", "tool_call_id": tc_id, "content": "Error: could not parse tool arguments."} - ) - continue - if not isinstance(fn_args, dict): - tool_messages.append( - { - "role": "tool", - "tool_call_id": tc_id, - "content": "Error: tool arguments must be a JSON object.", - } - ) - continue - - if fn_name == "think": - # Agent thoughts can quote document text/PII; keep content at DEBUG. - logger.debug( - "ReActAgentOperator: query=%s step=%d think=%r", - query_id, - _step, - _preview_text(fn_args.get("thought")), - ) - tool_messages.append( - {"role": "tool", "tool_call_id": tc_id, "content": "Your thought has been logged."} - ) - - elif fn_name == "retrieve": - subquery = str(fn_args.get("query", query_text)) - logger.info( - "ReActAgentOperator: query=%s step=%d retrieve subquery=%r seen_before=%d", - query_id, - _step, - _preview_text(subquery), - len(seen_doc_ids), - ) - retrieved = self._call_retriever(subquery, seen_doc_ids, api_key) - retrieval_log.append(retrieved) - step_counter += 1 - for d in retrieved: - seen_doc_ids.add(d["doc_id"]) - logger.info( - "ReActAgentOperator: query=%s step=%d retrieve docs=%d seen_after=%d doc_ids=%s", - query_id, - _step, - len(retrieved), - len(seen_doc_ids), - _preview_doc_ids(retrieved), - ) - doc_content = _docs_to_message_content(retrieved) - tool_content: List[Dict[str, Any]] = [ - {"type": "text", "text": f"Retrieved {len(retrieved)} documents:"} - ] + doc_content - tool_messages.append({"role": "tool", "tool_call_id": tc_id, "content": tool_content}) - - elif fn_name == "final_results": - raw_ids: List[str] = fn_args.get("doc_ids", []) - logger.info( - "ReActAgentOperator: query=%s step=%d final_results search_successful=%s doc_ids=%s", - query_id, - _step, - fn_args.get("search_successful"), - raw_ids[:_LOG_DOC_ID_LIMIT] if isinstance(raw_ids, list) else raw_ids, - ) - # Message can quote document text/PII; keep content at DEBUG. - logger.debug( - "ReActAgentOperator: query=%s step=%d final_results message=%r", - query_id, - _step, - _preview_text(fn_args.get("message")), - ) - validation_error = self._validate_final_results_args(fn_args, valid_doc_ids=seen_doc_ids) - if validation_error is None: - final_doc_ids = list(raw_ids) - tool_messages.append( - { - "role": "tool", - "tool_call_id": tc_id, - "content": "The results have been successfully logged and the interaction ended.", - } - ) - loop_done = True - else: - tool_messages.append( - { - "role": "tool", - "tool_call_id": tc_id, - "content": f"Error: {validation_error}", - } - ) - - else: - tool_messages.append( - {"role": "tool", "tool_call_id": tc_id, "content": f"Error: unknown tool '{fn_name}'."} - ) - - messages.extend(tool_messages) - if loop_done: - break + result = agent.run_sync(str(query_text), query_id=str(query_id), raw_log_dir=None) + + # Private agent retrieval_log entries are {"input", "tool_name", + # "query_type", "output": [ {id, score, text|note, ...} ]}. The exploded + # schema needs one score-desc ranked list per step (the private agent already + # sorts each output by score descending). + step_lists = [entry.get("output", []) or [] for entry in result.retrieval_log] + # Empty final_doc_ids (failure / no valid final_results) maps to None so + # has_valid_final_results and the synthetic final step stay consistent + # (both driven by the same truthiness). Never index end_payload here. + final_doc_ids = result.final_doc_ids or None logger.info( - "ReActAgentOperator: query=%s done retrieval_steps=%d seen_docs=%d final_doc_ids=%s", + "ReActAgentOperator: query=%s done retrieval_steps=%d succeeded=%s final_doc_ids=%s", query_id, - len(retrieval_log), - len(seen_doc_ids), - final_doc_ids[:_LOG_DOC_ID_LIMIT] if final_doc_ids else [], + len(step_lists), + result.succeeded, + (final_doc_ids or [])[:_LOG_DOC_ID_LIMIT], ) - return _build_output_rows(query_id, query_text, retrieval_log, final_doc_ids) - - def _call_retriever( - self, - query_text: str, - seen_doc_ids: set[str], - api_key: Optional[str], - ) -> List[Dict[str, Any]]: - """Call retriever_fn, over-fetching to ensure new results after dedup.""" - fetch_k = self._retriever_top_k + len(seen_doc_ids) - # Optional fixed ceiling on backend depth, matching Path A's --retriever-top-k - # cap. Once the agent has seen the whole capped pool, retrieves return no new - # docs, so the prompt stops growing (prevents context-window overflow). - if self._backend_top_k: - fetch_k = min(fetch_k, int(self._backend_top_k)) - try: - raw = self._retriever_fn(query_text, fetch_k) - except TimeoutError as exc: - logger.warning( - "ReActAgentOperator: retriever_fn timed out for query %r: %s", query_text, exc, exc_info=True - ) - return [] - except (TypeError, ValueError) as exc: - logger.warning( - "ReActAgentOperator: retriever_fn bad call/return for query %r: %s", query_text, exc, exc_info=True - ) - return [] - except Exception as exc: # retriever_fn is user-supplied; catches remaining unexpected errors. - logger.warning("ReActAgentOperator: retriever_fn failed for query %r: %s", query_text, exc, exc_info=True) - return [] - - # Walk the ranked results, normalising keys and de-duplicating within the - # batch. Already-seen docs are always re-presented as short stubs, matching - # Path A's retrieve_with_guarantees, and do not count toward top_k. - results: List[Dict[str, Any]] = [] - batch_doc_ids: set[str] = set() - new_count = 0 - for item in raw: - doc_id = str(item.get("doc_id", item.get("id", ""))) - if not doc_id or doc_id in batch_doc_ids: - continue - score = float(item.get("score", 0.0)) - if doc_id in seen_doc_ids: - batch_doc_ids.add(doc_id) - results.append( - { - "doc_id": doc_id, - "text": ( - "This document was retrieved before. See the earlier retrieval " - f"results for its content (id: {doc_id})." - ), - "score": score, - } - ) - continue - batch_doc_ids.add(doc_id) - results.append( - { - "doc_id": doc_id, - "text": str(item.get("text", "")), - "score": score, - } - ) - new_count += 1 - if new_count >= self._retriever_top_k: - break - - return results - - def _validate_final_results_args( - self, - fn_args: Dict[str, Any], - *, - valid_doc_ids: Optional[set[str]] = None, - ) -> Optional[str]: - """Validate final_results tool args outside the prompt/schema.""" - message = fn_args.get("message") - if not isinstance(message, str): - return f"`message` must be a string. Got `{type(message)}` type." - - doc_ids = fn_args.get("doc_ids") - if not isinstance(doc_ids, list): - return f"`doc_ids` must be a list. Got `{type(doc_ids)}` type." - if len(doc_ids) == 0: - return "`doc_ids` cannot be empty. You must choose at least one relevant document." - if not all(isinstance(doc_id, str) for doc_id in doc_ids): - return "Items in `doc_ids` must be of type string (i.e., python's `str` type)." - if not all(doc_id.strip() for doc_id in doc_ids): - return "Items in `doc_ids` must be non-empty (no blank or whitespace-only IDs)." - - search_successful = fn_args.get("search_successful") - if not isinstance(search_successful, str): - return f"`search_successful` must be a string. Got `{type(search_successful)}` type." - if search_successful not in {"true", "false", "partial"}: - return ( - f"`search_successful` must be one of `true`, `false`, or `partial`. Got `{search_successful}` instead." - ) - - if valid_doc_ids is not None: - invalid_doc_ids = [doc_id for doc_id in doc_ids if doc_id not in valid_doc_ids] - if invalid_doc_ids: - preview = invalid_doc_ids[:_LOG_DOC_ID_LIMIT] - return f"`doc_ids` contains IDs that were not retrieved: {preview}." - - if len(doc_ids) != self._target_top_k: - return ( - f"`doc_ids` must contain exactly {self._target_top_k} documents. " - f"But got {len(doc_ids)} document IDs instead." - ) - - return None - - def _resolve_api_key(self) -> Optional[str]: - api_key = self._api_key - if api_key is not None and api_key.strip().startswith("os.environ/"): - var = api_key.strip().removeprefix("os.environ/") - value = os.environ.get(var) - if value is None: - raise ValueError( - f"Environment variable '{var}' is not set. " f"Set it with: export {var}=" - ) - return value - return api_key + return _build_output_rows(str(query_id), str(query_text), step_lists, final_doc_ids) # --------------------------------------------------------------------------- @@ -905,29 +337,18 @@ def _resolve_api_key(self) -> Optional[str]: # --------------------------------------------------------------------------- -def _docs_to_message_content(docs: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Convert a list of doc dicts to LLM message content blocks.""" - content: List[Dict[str, Any]] = [] - for doc in docs: - doc_id = doc.get("doc_id", "") - text = doc.get("text", "").strip() - entry: Dict[str, Any] = {"id": doc_id} - if text: - entry["text"] = text - score = doc.get("score") - if score is not None: - entry["score"] = score - content.append({"type": "text", "text": json.dumps(entry)}) - return content - - def _build_output_rows( query_id: str, query_text: str, retrieval_log: List[List[Dict[str, Any]]], final_doc_ids: Optional[List[str]], ) -> List[Dict[str, Any]]: - """Convert the retrieval log to one row per (step_idx, rank, doc_id).""" + """Convert the retrieval log to one row per (step_idx, rank, doc_id). + + ``retrieval_log`` is a list of per-step document lists (the private agent's + ``retrieval_log[*]["output"]``); each document carries an ``id`` (private-agent + key) — ``doc_id`` is also accepted for robustness. + """ rows: List[Dict[str, Any]] = [] for step_idx, step_docs in enumerate(retrieval_log): for rank, doc in enumerate(step_docs, 1): @@ -936,22 +357,23 @@ def _build_output_rows( "query_id": query_id, "query_text": query_text, "step_idx": step_idx, - "doc_id": doc.get("doc_id", ""), - "text": doc.get("text", ""), + "doc_id": str(doc.get("id", doc.get("doc_id", ""))), + "text": str(doc.get("text", "")), "has_valid_final_results": final_doc_ids is not None, "is_final_result": False, "rank": rank, } ) - # If final_results was called, also emit those as a synthetic final step - # (step_idx = len(retrieval_log)) so RRF can weight the agent's final - # judgment in addition to the raw retrieval history. + # If final_results was returned, also emit those as a synthetic final step + # (step_idx = len(retrieval_log)) so RRF/selection can recover the agent's + # final ranking. Gated on final_doc_ids truthiness — the same predicate that + # drives has_valid_final_results above. if final_doc_ids: first_doc_by_id: Dict[str, Dict[str, Any]] = {} for step_docs in retrieval_log: for doc in step_docs: - doc_id = str(doc.get("doc_id", "")) + doc_id = str(doc.get("id", doc.get("doc_id", ""))) if doc_id and doc_id not in first_doc_by_id: first_doc_by_id[doc_id] = doc @@ -969,7 +391,7 @@ def _build_output_rows( "query_text": query_text, "step_idx": final_step_idx, "doc_id": doc_id, - "text": doc.get("text", ""), + "text": str(doc.get("text", "")), "has_valid_final_results": True, "is_final_result": True, "rank": rank, diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/rrf_aggregator_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/rrf_aggregator_operator.py index 94056ffb93..c13371dab3 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/rrf_aggregator_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/rrf_aggregator_operator.py @@ -109,27 +109,38 @@ def process(self, data: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: else False ) - # Process each step's ranked list + # Process each step's ranked list. The synthetic final-results step + # (is_final_result=True) is recorded for react_final_rank but must + # NOT contribute to the RRF score — the reference fuses only the + # retrieve outputs (see retrieval_bench/common/rrf.py). for _step_idx, sgroup in qgroup.groupby("step_idx", sort=True): # Sort by rank ascending so rank=1 is processed first for _, row in sgroup.sort_values("rank").iterrows(): doc_id = str(row["doc_id"]) rank = int(row["rank"]) - rrf_scores[doc_id] += 1.0 / (rank + k) + is_final = bool(row.get("is_final_result", False)) + if not is_final: + rrf_scores[doc_id] += 1.0 / (rank + k) if doc_id not in first_text: first_text[doc_id] = str(row["text"]) - if bool(row.get("is_final_result", False)): + if is_final: previous = react_final_rank.get(doc_id) if previous is None or rank < previous: react_final_rank[doc_id] = rank - for doc_id, score in sorted(rrf_scores.items(), key=lambda kv: kv[1], reverse=True): + # Emit every candidate that appears in a retrieve step OR only in the + # synthetic final step, so a final-only doc keeps its react_final_rank + # (needed for the ReAct pass-through) even though it scored 0 in RRF. + all_doc_ids = set(rrf_scores) | set(react_final_rank) + # Tie-break by doc_id so equal RRF scores order deterministically across + # runs (set iteration is hash-randomized); score stays descending. + for doc_id in sorted(all_doc_ids, key=lambda d: (-rrf_scores.get(d, 0.0), d)): rows.append( { "query_id": query_id, "query_text": query_text, "doc_id": doc_id, - "rrf_score": score, + "rrf_score": rrf_scores.get(doc_id, 0.0), "text": first_text.get(doc_id, ""), "has_valid_final_results": has_valid_final_results, "react_final_rank": react_final_rank.get(doc_id), diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py index eb8ee037f2..24efb1259d 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py @@ -2,28 +2,36 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Operator that re-ranks retrieved documents using an LLM-based selection agent.""" +"""Operator that re-ranks retrieved documents using an LLM-based selection agent. + +The selection logic lives in the private +:class:`~nemo_retriever._agentic.nemo_agent.SelectionAgent`. This operator adapts the +RRF-stage DataFrame into that library's inputs and applies a three-tier gate per +query: pass through the ReAct agent's ``final_results`` when present, otherwise +run the selection agent, otherwise fall back to the RRF ranking. +""" from __future__ import annotations -import json import logging -import os - -import requests from typing import Any, Callable, Dict, List, Optional import pandas as pd +from nemo_retriever._agentic.nemo_agent import SelectionAgent, SelectionAgentConfig +from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config +from nemo_retriever._agentic.nemo_agent.results import AgentRunResult from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.operators.cpu_operator import CPUOperator -from nemo_retriever.models.nim.chat_completions import invoke_chat_completion_step logger = logging.getLogger(__name__) _LOG_PREVIEW_CHARS = 300 _LOG_DOC_ID_LIMIT = 20 +#: Max LLM steps for the selection sub-agent (mirrors the reference workflow). +_SELECTION_MAX_STEPS = 10 + def _preview_text(value: Any, *, limit: int = _LOG_PREVIEW_CHARS) -> str: text = " ".join(str(value or "").split()) @@ -32,155 +40,84 @@ def _preview_text(value: Any, *, limit: int = _LOG_PREVIEW_CHARS) -> str: return text[:limit].rstrip() + "..." -# --------------------------------------------------------------------------- -# Prompt rendering (verbatim content of 01_v0.j2, rendered via Python) -# --------------------------------------------------------------------------- - -_ROLE = """\ -You are a document re-ranker agent, which is the final stage in an information retrieval pipeline. - - -You are given a search query and a list of retrieved candidate documents that are potentially relevant to \ -the given query. Your goal is to help the users identify the most relevant documents to the given query \ -from the list of candidate documents. -""" - -_RELEVANCE_DEFINITION = """\ - - -- You should be careful, in the context of this task, what it means to be a "query", "document", and \ -"relevant" can sometimes be very complex and might not follow the traditional definition of these terms \ -in standard re-ranking and retrieval. -- In standard re-ranking/retrieval, a query is usually a user question (like a web search query), the \ -document is some sort of content that provides information (e.g., a web page), and these two are considered \ -relevant if the document provides information that answers the user's query. -- However, in our setting, this could be different. Here are some examples: - * the query is a programming problem and documents are programming language syntax references. A document \ -is relevant if it contains the reference for the programming syntax used for solving the problem. - * both query and documents are descriptions programming problems and a query and document are relevant if \ -the same approach is used to solve them. - * the query is a math problem and documents are theorems. Relevant documents (theorems) are the ones \ -that are useful for solving the math problem. - * the query and document are both math problems. A query and a document are relevant if the same theorem \ -is used for solving them. - * the query is a task description (e.g., for an API programmer) and documents are descriptions of \ -available APIs. Relevant documents (e.g., APIs) are the ones needed for completing the task. -- This is not an exhaustive list. These are just some examples to show you the complexity of queries, \ -documents, and the concept of relevance in this task. -- Note that even here, the relevant documents are still the ones that are useful for a user who is \ -searching for the given query. But the relation is more nuanced. -- You should analyze the query and the available documents. And then reason about what could be a meaningful \ -definition of relevance in this case, and what the user could be looking for. -- Moreover, sometimes, the query could be even a prompt that is given to a Large Language Model (LLM) and \ -the user wants to find the useful documents for the LLM that help answering/solving this prompt. -""" - -_WORKFLOW_TEMPLATE = """\ - - -* You are given a search query and a list of candidate documents. You have access to the ID and content of \ -each candidate document. -* You should read the query carefully and understand it. -{extended_relevance_line}\ -* Then you should compare the query with each one of the candidate documents. In this comparison, you want \ -to identify if the document is relevant/useful for the given query and to what extent. -* Select the {top_k} most relevant candidate documents for the given query. -* Note that just selecting the most relevant documents is not enough. You should identify the relative level \ -of relevance between the query and selected documents. This helps you sort the selected documents later \ -based on how relevant they are to the query. -* Once you have this information, you should call the "log_selected_documents" function to report the final \ -results and signal the completion of the task. -* Note that the selected document IDs must be reported in the decreasing level of relevance. I.e., The \ -first document in the list is the most relevant, the second is the second most relevant, and so on. This \ -is similar to what a search engine (e.g., Google Search) does (it shows you the relevant results in a \ -sorted order, where the most relevant results appear on top of the list). -""" - -_THINKING_TIPS = """ - - -* you have access to a "think" tool that you can use for complex thinking and analysis. Here are examples \ -of cases where the think tool might be useful: - - complex analysis and thinking to understand the meaning and intent of the query. E.g., what is the \ -user trying to find with this query? what kind of information is helpful for the user? - - extended thinking to analyze how each candidate document could or could not be relevant to the given query. - - reasoning to identify the relative level of relevance between the query and selected documents. It \ -helps you sort the documents correctly when reporting the final answer. -""" - - -def _render_selection_prompt(top_k: int, *, extended_relevance: bool = False) -> str: - """Render the selection agent system prompt (verbatim 01_v0.j2 logic).""" - parts = [_ROLE] - if extended_relevance: - parts.append(_RELEVANCE_DEFINITION) - ext_line = ( - "* As explained above, reason and figure out what the meaning of relevance is in this case, " - "and what could be relevant and useful information for the given query.\n" - if extended_relevance - else "" - ) - parts.append(_WORKFLOW_TEMPLATE.format(top_k=top_k, extended_relevance_line=ext_line)) - parts.append(_THINKING_TIPS) - return "".join(parts) - - -# --------------------------------------------------------------------------- -# Operator -# --------------------------------------------------------------------------- - - class SelectionAgentOperator(AbstractOperator, CPUOperator): - """Re-rank a set of retrieved documents using an LLM-based selection agent. + """Re-rank retrieved documents using an LLM-based selection agent. - For each ``query_id`` group in the input DataFrame, the operator runs an - agentic LLM loop that reads the query and all candidate documents, then - calls a ``log_selected_documents`` tool to report the final ranked list. - The loop also has access to a ``think`` tool for extended reasoning. + For each ``query_id`` group produced by :class:`RRFAggregatorOperator`, the + operator applies a three-tier gate: - The system prompt matches the retrieval-bench ``01_v0.j2`` template verbatim, - with an optional ``extended_relevance`` mode for complex retrieval tasks. + 1. **final_results** — if the ReAct agent produced a non-empty final list + (recovered from ``react_final_rank``), pass those doc ids through. + 2. **selection_agent** — otherwise run + :class:`~nemo_retriever._agentic.nemo_agent.SelectionAgent` over the RRF-ranked + candidates (with the RRF scores arming its context-overflow shrink retry). + 3. **rrf** — if selection produced nothing (failure / empty), fall back to + the top RRF-ranked candidates. Input DataFrame schema ---------------------- query_id : str — unique query identifier query_text : str — original query text shown to the LLM - doc_id : str — unique document identifier + doc_id : str — candidate document identifier text : str — document text content shown to the LLM + rrf_score : float, optional — used to order candidates and as the + selection shrink-retry priority / RRF fallback + react_final_rank : int, optional — ReAct final ordering (drives tier 1) (any additional columns are ignored) Output DataFrame schema ----------------------- - query_id : str — same ``query_id`` as the input - doc_id : str — selected document ID - rank : int — 1-indexed rank (1 = most relevant) - message : str — LLM explanation of the selection + query_id : str — same ``query_id`` as the input + doc_id : str — selected document ID + rank : int — 1-indexed rank (1 = most relevant) + message : str — always empty (retained for schema compatibility) + result_source : str — one of ``{"final_results", "selection_agent", "rrf"}`` Parameters ---------- llm_model : str - Model identifier forwarded to the endpoint. + Model identifier forwarded verbatim to the backend. invoke_url : str - Full ``/v1/chat/completions`` endpoint URL. + LLM endpoint. Forwarded as the LLM config's ``base_url``. top_k : int - Number of documents to select per query. Defaults to ``5``. + Number of documents to select per query. Defaults to ``10``. api_key : str, optional Literal API key **or** an ``"os.environ/VAR_NAME"`` reference. max_tokens : int, optional - Upper bound on tokens in each LLM response. + Per-request completion budget (the LLM config's ``max_completion_tokens``). max_steps : int - Maximum agentic loop iterations per query. Defaults to ``10``. - extended_relevance : bool - When ``True``, include the ```` block in the - system prompt for tasks with non-standard relevance definitions. - Defaults to ``False``. + Maximum selection-agent LLM steps per query. Defaults to ``10``. system_prompt_override : str, optional - Fully custom system prompt. Use ``{top_k}`` as a placeholder. + Forwarded as the selection agent's ``system_prompt`` (``None`` selects + the packaged default selection prompt). text_truncation : int - Maximum characters of each document's text shown to the LLM. - Defaults to ``2000``. + Maximum characters of each document's text passed to the agent. + ``0`` disables truncation. Defaults to ``0``. + parallel_tool_calls : bool, optional + Forwarded as the LLM config's ``parallel_tool_calls`` (sent to the + provider only when set). base_url : str, optional - Deprecated alias for ``invoke_url``. Prefer ``invoke_url``. + Deprecated alias for ``invoke_url``. + reasoning_effort : str, optional + Forwarded as the LLM config's ``reasoning_effort``. + temperature : float, optional + Forwarded as the LLM config's ``temperature`` (sent to the provider + only when set). + backend : {"callable", "litellm"} + LLM backend to build. ``"callable"`` (default) drives an OpenAI-compatible + completion callable: ``chat_completion_fn`` when supplied (the in-process + vLLM adapter), otherwise the shared ``invoke_chat_completion_step`` HTTP + client against ``invoke_url``. + chat_completion_fn : callable, optional + OpenAI-compatible completion callable (e.g. the local in-process vLLM + adapter). When set, forwarded to the ``"callable"`` LLM backend. + + Notes + ----- + The :class:`~nemo_retriever._agentic.nemo_agent.SelectionAgent` and LLM backend are + built lazily on first ``process`` call so the operator stays reconstructable + via ``get_constructor_kwargs``. ``select_sync`` performs one ``asyncio.run`` + per query and must not be called from inside a running event loop. """ _NVIDIA_BUILD_ENDPOINT = "https://integrate.api.nvidia.com/v1/chat/completions" @@ -190,31 +127,33 @@ def __init__( *, llm_model: str, invoke_url: Optional[str] = None, - top_k: int = 5, + top_k: int = 10, api_key: Optional[str] = None, max_tokens: Optional[int] = None, - max_steps: int = 10, - extended_relevance: bool = False, + max_steps: int = _SELECTION_MAX_STEPS, system_prompt_override: Optional[str] = None, - text_truncation: int = 2000, - parallel_tool_calls: bool = True, + text_truncation: int = 0, + parallel_tool_calls: Optional[bool] = None, base_url: Optional[str] = None, reasoning_effort: Optional[str] = None, - temperature: float = 0.0, + temperature: Optional[float] = None, + backend: str = "callable", chat_completion_fn: Optional[Callable[..., Dict[str, Any]]] = None, ) -> None: super().__init__() - self._reasoning_effort = reasoning_effort - self._temperature = temperature self._llm_model = llm_model self._top_k = top_k self._api_key = api_key self._max_tokens = max_tokens self._max_steps = max_steps - self._extended_relevance = extended_relevance self._system_prompt_override = system_prompt_override self._text_truncation = text_truncation self._parallel_tool_calls = parallel_tool_calls + self._temperature = temperature + self._reasoning_effort = reasoning_effort + self._backend = backend + # When set, an OpenAI-compatible completion callable (e.g. the in-process + # local vLLM adapter). Forwarded to the "callable" LLM backend by _build_llm. self._chat_completion_fn = chat_completion_fn if invoke_url is not None: @@ -231,6 +170,52 @@ def __init__( else: self._invoke_url = self._NVIDIA_BUILD_ENDPOINT + # Built lazily on first process(); never part of the picklable ctor state. + self._sel: Optional[SelectionAgent] = None + + # ------------------------------------------------------------------ + # private agent construction (lazy, memoized) + # ------------------------------------------------------------------ + + def _build_llm(self) -> Any: + config = create_llm_config( + self._backend, + model=str(self._llm_model), + base_url=self._invoke_url, + api_key=self._api_key, + reasoning_effort=self._reasoning_effort or None, + temperature=self._temperature, + parallel_tool_calls=self._parallel_tool_calls, + max_completion_tokens=self._max_tokens, + ) + completion_fn = self._chat_completion_fn + if self._backend == "callable" and completion_fn is None: + # Remote run on the default backend: supply the shared chat-completions + # client. Imported HERE and injected, never imported by the agent + # library, which must not depend on the rest of nemo_retriever. + from nemo_retriever.models.nim.chat_completions import invoke_chat_completion_step + + completion_fn = invoke_chat_completion_step + kwargs = {"completion_fn": completion_fn} if completion_fn is not None else {} + return create_llm(config, **kwargs) + + def _ensure_agent(self) -> SelectionAgent: + if self._sel is None: + self._sel = SelectionAgent( + config=SelectionAgentConfig( + target_top_k=int(self._top_k), + extended_relevance=True, + enable_think=False, + end_tool_with_msg=False, + shrink_attempts=2, + max_steps=int(self._max_steps), + on_error="never_raise", + system_prompt=self._system_prompt_override, + ), + llm=self._build_llm(), + ) + return self._sel + # ------------------------------------------------------------------ # AbstractOperator interface # ------------------------------------------------------------------ @@ -248,70 +233,70 @@ def preprocess(self, data: Any, **kwargs: Any) -> pd.DataFrame: return data.copy() def process(self, data: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: - """Run the selection agent loop for each query group.""" + """Apply the three-tier selection gate for each query group.""" + self._ensure_agent() rows: List[Dict[str, Any]] = [] for query_id, group in data.groupby("query_id", sort=False): - query_text = str(group["query_text"].iloc[0]) - ordered_group = group - if "rrf_score" in group.columns: - ordered_group = group.sort_values("rrf_score", ascending=False) - docs = [ - { - "id": str(row["doc_id"]), - "text": str(row["text"]), - } - for _, row in ordered_group.iterrows() - ] + query_text = str(group["query_text"].iloc[0]) if "query_text" in group.columns else "" + ordered = group.sort_values("rrf_score", ascending=False) if "rrf_score" in group.columns else group + + # Candidate documents in RRF-descending order, deduplicated (first wins), + # plus the {doc_id: rrf_score} priority side-table (covers every candidate). + documents: List[Dict[str, Any]] = [] + seen: set[str] = set() + for _, row in ordered.iterrows(): + doc_id = str(row["doc_id"]) + if doc_id in seen: + continue + seen.add(doc_id) + text = str(row["text"]) + if self._text_truncation and int(self._text_truncation) > 0: + text = text[: int(self._text_truncation)] + documents.append({"id": doc_id, "text": text}) + scores: Optional[Dict[str, float]] = None + if "rrf_score" in ordered.columns: + scores = {str(row["doc_id"]): float(row["rrf_score"]) for _, row in ordered.iterrows()} + logger.info( - "SelectionAgentOperator: query=%s start candidates=%d unique_candidates=%d query=%r", + "SelectionAgentOperator: query=%s candidates=%d query=%r", query_id, - len(docs), - len({doc["id"] for doc in docs}), + len(documents), _preview_text(query_text), ) - preferred_doc_ids, message, result_source = self._preferred_doc_ids(ordered_group) - if preferred_doc_ids is None: - result = self._select_documents(query_text, docs) - message = result.get("message", "") - doc_ids = list(result.get("doc_ids", [])) - result_source = "selection_agent" + + # Tier 1: ReAct produced a final list (success or salvage) -> pass through. + react_final = self._react_final_doc_ids(ordered) + if react_final: + doc_ids = list(react_final) + result_source = "final_results" else: - doc_ids = preferred_doc_ids - if not doc_ids: - if preferred_doc_ids is None: - doc_ids = ordered_group["doc_id"].astype(str).drop_duplicates().head(int(self._top_k)).tolist() - message = ( - f"{message} Falling back to top {len(doc_ids)} RRF-ranked candidates." - if message - else f"Falling back to top {len(doc_ids)} RRF-ranked candidates." - ) - result_source = "candidate_ranking" - logger.warning( - "SelectionAgentOperator: query=%s selection failed; " - "falling back to candidate ranking doc_ids=%s", - query_id, - doc_ids[:_LOG_DOC_ID_LIMIT], - ) + doc_ids = [] + result_source = "" + # Tier 2: run the selection agent over the RRF candidates. + if documents: + result = self._run_selection(query_text, documents, scores, str(query_id)) + if result is not None and result.succeeded and result.final_doc_ids: + doc_ids = [str(d) for d in result.final_doc_ids][: int(self._top_k)] + result_source = "selection_agent" + # Tier 3: fall back to the RRF ranking. + if not doc_ids: + doc_ids = ordered["doc_id"].astype(str).drop_duplicates().head(int(self._top_k)).tolist() + result_source = "rrf" + logger.info( "SelectionAgentOperator: query=%s result_source=%s selected=%s", query_id, result_source, doc_ids[:_LOG_DOC_ID_LIMIT], ) - # Message can quote document text/PII; keep content at DEBUG. - logger.debug( - "SelectionAgentOperator: query=%s message=%r", - query_id, - _preview_text(message), - ) for rank, doc_id in enumerate(doc_ids, 1): rows.append( { "query_id": query_id, - "doc_id": doc_id, + "doc_id": str(doc_id), "rank": rank, - "message": message, + "message": "", "result_source": result_source, } ) @@ -328,102 +313,19 @@ def postprocess(self, data: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: # Internal helpers # ------------------------------------------------------------------ - def _resolve_api_key(self) -> Optional[str]: - api_key = self._api_key - if api_key is not None and api_key.strip().startswith("os.environ/"): - var = api_key.strip().removeprefix("os.environ/") - value = os.environ.get(var) - if value is None: - raise ValueError( - f"Environment variable '{var}' is not set. " f"Set it with: export {var}=" - ) - return value - return api_key - - def _build_system_prompt(self, top_k: int) -> str: - if self._system_prompt_override: - return self._system_prompt_override.format(top_k=top_k) - return _render_selection_prompt(top_k, extended_relevance=self._extended_relevance) - - def _build_tools(self, top_k: int, valid_doc_ids: List[str]) -> List[Dict[str, Any]]: - """Return the two tool specs for the selection agent loop.""" - return [ - { - "type": "function", - "function": { - "name": "think", - "description": ( - "Use this tool to think through complex analysis before making a decision. " - "It logs your reasoning without making any changes. Use it to compare " - "documents against the query or to reason about relevance." - ), - "parameters": { - "type": "object", - "properties": { - "thought": { - "type": "string", - "description": "Your reasoning or analysis.", - } - }, - "required": ["thought"], - }, - }, - }, - { - "type": "function", - "function": { - "name": "log_selected_documents", - "description": ( - f"Records the {top_k} most relevant documents and ends the task. " - f"Call this when you have finished evaluating all candidate documents. " - f"The doc_ids list must be sorted from most to least relevant. " - f"Valid document IDs are: {valid_doc_ids}." - ), - "parameters": { - "type": "object", - "required": ["doc_ids", "message"], - "properties": { - "doc_ids": { - "type": "array", - "items": {"type": "string"}, - "description": ( - f"The IDs of the {top_k} most relevant documents, sorted from " - "most to least relevant. Must be valid document IDs from the candidates." - ), - }, - "message": { - "type": "string", - "description": "A brief explanation of your selection and the relevance ordering.", - }, - }, - }, - }, - }, - ] - - def _preferred_doc_ids(self, ordered_group: pd.DataFrame) -> tuple[List[str] | None, str, str]: - """Apply retrieval-bench-style source priority before invoking selection.""" - doc_ids = self._react_final_doc_ids(ordered_group) - if doc_ids is not None: - return doc_ids, "Using ReAct final_results.", "final_results" - - if "rrf_score" in ordered_group.columns: - doc_ids = ordered_group["doc_id"].astype(str).drop_duplicates().head(int(self._top_k)).tolist() - if doc_ids: - return doc_ids, "Using RRF ranking.", "rrf" - - return None, "", "" - - def _react_final_doc_ids(self, ordered_group: pd.DataFrame) -> List[str] | None: - if "has_valid_final_results" in ordered_group.columns and not bool( - ordered_group["has_valid_final_results"].astype(bool).any() - ): - return None + def _react_final_doc_ids(self, ordered_group: pd.DataFrame) -> List[str]: + """Recover the ReAct agent's final ranked doc ids from ``react_final_rank``. + + Returns an empty list when the ReAct agent produced no final results — the + gate branches on non-emptiness, so the None-vs-empty distinction that the + old code relied on is irrelevant here. + """ if "react_final_rank" not in ordered_group.columns: - return None - final_rows = ordered_group[ordered_group["react_final_rank"].notna()].copy() + return [] + final_rows = ordered_group[ordered_group["react_final_rank"].notna()] if final_rows.empty: - return [] if "has_valid_final_results" in ordered_group.columns else None + return [] + final_rows = final_rows.copy() final_rows["react_final_rank"] = final_rows["react_final_rank"].astype(int) doc_ids: List[str] = [] for doc_id in final_rows.sort_values("react_final_rank")["doc_id"].astype(str): @@ -433,260 +335,22 @@ def _react_final_doc_ids(self, ordered_group: pd.DataFrame) -> List[str] | None: break return doc_ids - def _build_user_message(self, query_text: str, docs: List[Dict[str, Any]]) -> Dict[str, Any]: - """Format query + candidate documents as a multi-part user message.""" - content: List[Dict[str, Any]] = [ - {"type": "text", "text": f"Query:\n{query_text}"}, - {"type": "text", "text": "Candidate Documents:"}, - ] - seen: set[str] = set() - for doc in docs: - doc_id = doc["id"] - if doc_id in seen: - continue - seen.add(doc_id) - content.append({"type": "text", "text": f"Doc ID: {doc_id}"}) - text = doc.get("text", "").strip() - if text: - if self._text_truncation > 0: - truncated = text[: self._text_truncation] - else: - truncated = text - if self._text_truncation > 0 and len(text) > self._text_truncation: - truncated += "..." - content.append({"type": "text", "text": f"Doc Text: {truncated}"}) - return {"role": "user", "content": content} - - def _select_documents( + def _run_selection( self, query_text: str, - docs: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """Run the agentic selection loop for a single query.""" - valid_ids = list(dict.fromkeys(d["id"] for d in docs)) - feasible_k = min(self._top_k, len(valid_ids)) - logger.info( - "SelectionAgentOperator: selecting top_k=%d feasible_k=%d valid_doc_ids=%s", - self._top_k, - feasible_k, - valid_ids[:_LOG_DOC_ID_LIMIT], - ) - - system_prompt = self._build_system_prompt(feasible_k) - tools = self._build_tools(feasible_k, valid_ids) - valid_id_set = set(valid_ids) - api_key = self._resolve_api_key() - - messages: List[Dict[str, Any]] = [ - {"role": "system", "content": system_prompt}, - self._build_user_message(query_text, docs), - ] - - extra_body: Dict[str, Any] = {} - if not self._parallel_tool_calls: - extra_body["parallel_tool_calls"] = False - if self._reasoning_effort: - extra_body["reasoning_effort"] = self._reasoning_effort - - for _step in range(self._max_steps): - logger.info( - "SelectionAgentOperator: step=%d begin candidates=%d feasible_k=%d", - _step, - len(valid_ids), - feasible_k, + documents: List[Dict[str, Any]], + scores: Optional[Dict[str, float]], + query_id: str, + ) -> Optional[AgentRunResult]: + """Run the selection agent, returning None on an unexpected failure.""" + try: + return self._ensure_agent().select_sync( + query_text, + documents, + scores=scores, + query_id=query_id, + raw_log_dir=None, ) - try: - chat_completion_fn = self._chat_completion_fn or invoke_chat_completion_step - response = chat_completion_fn( - invoke_url=self._invoke_url, - messages=messages, - model=self._llm_model, - api_key=api_key, - tools=tools, - tool_choice="auto", - temperature=self._temperature, - max_tokens=self._max_tokens, - extra_body=extra_body or None, - ) - except TimeoutError as exc: - logger.warning( - "SelectionAgentOperator: LLM call timed out on step %d for query %r: %s", - _step, - query_text, - exc, - exc_info=True, - ) - break - except RuntimeError as exc: - logger.warning( - "SelectionAgentOperator: LLM retries exhausted on step %d for query %r: %s", - _step, - query_text, - exc, - exc_info=True, - ) - break - except requests.RequestException as exc: - logger.warning( - "SelectionAgentOperator: LLM HTTP error on step %d for query %r: %s", - _step, - query_text, - exc, - exc_info=True, - ) - break - except json.JSONDecodeError as exc: - logger.warning( - "SelectionAgentOperator: LLM returned invalid JSON on step %d for query %r: %s", - _step, - query_text, - exc, - exc_info=True, - ) - break - - if not response.get("choices"): - logger.warning("SelectionAgentOperator: empty choices in API response on step %d", _step) - break - choice = response["choices"][0] - msg = choice["message"] - finish_reason = choice.get("finish_reason") - - # Append the assistant turn to history - assistant_turn: Dict[str, Any] = {"role": "assistant"} - if msg.get("content"): - assistant_turn["content"] = msg["content"] - # Agent reasoning can quote document text/PII; keep content at DEBUG. - logger.debug( - "SelectionAgentOperator: step=%d assistant content=%r", - _step, - _preview_text(msg.get("content")), - ) - tool_calls = msg.get("tool_calls") or [] - logger.info( - "SelectionAgentOperator: step=%d finish_reason=%s tool_calls=%s", - _step, - finish_reason, - [((tc.get("function") or {}).get("name") or "") for tc in tool_calls], - ) - if tool_calls: - assistant_turn["tool_calls"] = tool_calls - messages.append(assistant_turn) - - if finish_reason == "stop" or not tool_calls: - logger.info("SelectionAgentOperator: step=%d no tool call; asking for final selection", _step) - messages.append( - { - "role": "user", - "content": "Please call log_selected_documents to report your final selection.", - } - ) - continue - - tool_messages: List[Dict[str, Any]] = [] - should_end = False - end_kwargs: Dict[str, Any] = {} - - for tc in tool_calls: - tc_id = tc.get("id", "") - fn = tc.get("function", {}) - try: - fn_args = json.loads(fn.get("arguments", "{}")) - except json.JSONDecodeError: - tool_messages.append( - {"role": "tool", "tool_call_id": tc_id, "content": "Error: could not parse tool arguments."} - ) - continue - if not isinstance(fn_args, dict): - tool_messages.append( - { - "role": "tool", - "tool_call_id": tc_id, - "content": "Error: tool arguments must be a JSON object.", - } - ) - continue - - if fn.get("name") == "think": - # Agent thoughts can quote document text/PII; keep content at DEBUG - # (matches ReActAgentOperator's think logging). - logger.debug( - "SelectionAgentOperator: step=%d think=%r", - _step, - _preview_text(fn_args.get("thought")), - ) - tool_messages.append( - {"role": "tool", "tool_call_id": tc_id, "content": "Your thought has been logged."} - ) - - elif fn.get("name") == "log_selected_documents": - raw_doc_ids = fn_args.get("doc_ids", []) - if isinstance(raw_doc_ids, str): - try: - raw_doc_ids = json.loads(raw_doc_ids) - except json.JSONDecodeError: - raw_doc_ids = [] - if not isinstance(raw_doc_ids, list): - tool_messages.append( - { - "role": "tool", - "tool_call_id": tc_id, - "content": "Error: `doc_ids` must be a list of candidate document IDs.", - } - ) - continue - - invalid_doc_ids = [doc_id for doc_id in raw_doc_ids if doc_id not in valid_id_set] - if invalid_doc_ids: - logger.warning( - "SelectionAgentOperator: LLM returned doc_id(s) outside the candidate set " - "for query %r: %s", - query_text, - invalid_doc_ids[:_LOG_DOC_ID_LIMIT], - ) - tool_messages.append( - { - "role": "tool", - "tool_call_id": tc_id, - "content": ( - "Error: `doc_ids` contains IDs that are not candidate documents: " - f"{invalid_doc_ids[:_LOG_DOC_ID_LIMIT]}. Use only valid candidate IDs." - ), - } - ) - continue - - doc_ids = raw_doc_ids[:feasible_k] - logger.info( - "SelectionAgentOperator: step=%d log_selected_documents raw=%s accepted=%s", - _step, - raw_doc_ids[:_LOG_DOC_ID_LIMIT], - doc_ids[:_LOG_DOC_ID_LIMIT], - ) - # Message can quote document text/PII; keep content at DEBUG. - logger.debug( - "SelectionAgentOperator: step=%d log_selected_documents message=%r", - _step, - _preview_text(fn_args.get("message")), - ) - end_kwargs = {"doc_ids": doc_ids, "message": fn_args.get("message", "")} - should_end = True - - else: - tool_messages.append( - { - "role": "tool", - "tool_call_id": tc_id, - "content": f"Error: unknown tool '{fn.get('name')}'.", - } - ) - - if should_end: - return end_kwargs - - messages.extend(tool_messages) - - return { - "doc_ids": [], - "message": "Selection agent reached max steps without completing.", - } + except Exception as exc: # production: fall back to RRF rather than crash + logger.warning("SelectionAgentOperator: selection failed for query %r: %s", query_id, exc, exc_info=True) + return None diff --git a/nemo_retriever/src/nemo_retriever/operators/vdb.py b/nemo_retriever/src/nemo_retriever/operators/vdb.py index b0c8efd2d3..35bc6002ee 100644 --- a/nemo_retriever/src/nemo_retriever/operators/vdb.py +++ b/nemo_retriever/src/nemo_retriever/operators/vdb.py @@ -10,11 +10,15 @@ import pandas as pd -from nemo_retriever.common.vdb.adt_vdb import VDB +from nemo_retriever.common.vdb.adt_vdb import CollectionWriteContext, VDB from nemo_retriever.common.vdb.factory import get_vdb_op_cls from nemo_retriever.operators.abstract_operator import AbstractOperator -from nemo_retriever.common.vdb.records import normalize_retrieval_results, to_client_vdb_records +from nemo_retriever.common.vdb.records import ( + normalize_retrieval_results, + to_client_vdb_records, + validate_collection_retrieval_results, +) from nemo_retriever.common.vdb.sidecar_metadata import ( apply_sidecar_metadata_to_client_batches, build_sidecar_lookup, @@ -140,6 +144,13 @@ def process(self, data: Any, **kwargs: Any) -> Any: meta_fields=self._sidecar_spec["meta_fields"], join_key=self._sidecar_spec["meta_join_key"], ) + collection_context = kwargs.get("collection_context") + if collection_context is not None: + if not isinstance(collection_context, CollectionWriteContext): + raise TypeError("collection_context must be a CollectionWriteContext") + if not records or not any(records): + raise ValueError("Collection writes require at least one canonical VDB record") + return self._vdb.write_collection(records, context=collection_context) if records and any(batch for batch in records): self._vdb.run(records) return data @@ -216,25 +227,58 @@ def __init__( merged = dict(vdb_kwargs or {}) clean_kwargs, _sidecar = split_sidecar_from_vdb_kwargs(merged) clean_kwargs.pop("query_texts", None) - super().__init__(vdb=vdb, vdb_op=vdb_op, vdb_kwargs=clean_kwargs, explode_for_rerank=explode_for_rerank) + super().__init__( + vdb=vdb, + vdb_op=vdb_op, + vdb_kwargs=clean_kwargs, + explode_for_rerank=explode_for_rerank, + ) self._vdb_kwargs = clean_kwargs self._retrieval_vdb_kwargs = clean_kwargs self._vdb = _construct_vdb(vdb=vdb, vdb_op=vdb_op, vdb_kwargs=clean_kwargs) self._explode_for_rerank = bool(explode_for_rerank) + def get_index_metadata(self, key: str, **kwargs: Any) -> str | None: + """Read one index metadata value through the configured VDB.""" + return self._vdb.get_index_metadata(key, **{**self._vdb_kwargs, **kwargs}) + def preprocess(self, data: Any, **kwargs: Any) -> Any: if isinstance(data, pd.DataFrame): return query_vectors_from_embedded_dataframe(data) return data - def process(self, data: Any, **kwargs: Any) -> list[list[dict[str, Any]]]: + def process( + self, data: Any, **kwargs: Any + ) -> list[list[dict[str, Any]]] | tuple[list[list[dict[str, Any]]], list[str]]: from nemo_retriever.graph.retriever_utils import filter_retrieval_kwargs - retrieval_kwargs = {**self._retrieval_vdb_kwargs, **filter_retrieval_kwargs(kwargs)} + runtime_kwargs = dict(kwargs) + scope = runtime_kwargs.pop("scope", None) + collection_name = runtime_kwargs.pop("collection_name", None) + if (scope is None) != (collection_name is None): + raise ValueError("Collection retrieval requires both scope and collection_name") + + retrieval_kwargs = { + **self._retrieval_vdb_kwargs, + **filter_retrieval_kwargs(runtime_kwargs), + } if "hybrid" in retrieval_kwargs: effective_hybrid = bool(retrieval_kwargs["hybrid"]) else: effective_hybrid = bool(getattr(self._vdb, "hybrid", False)) + if collection_name is not None: + retrieval_kwargs.pop("collection_name", None) + top_k = int(retrieval_kwargs.pop("top_k", 10)) + result = self._vdb.retrieve_collection( + data, + scope=str(scope), + collection_name=str(collection_name), + query_texts=list(kwargs.get("query_texts") or []), + top_k=top_k, + **retrieval_kwargs, + ) + return validate_collection_retrieval_results(result, expected_queries=len(data)) + if effective_hybrid and "query_texts" in kwargs: retrieval_kwargs["query_texts"] = kwargs["query_texts"] return normalize_retrieval_results(self._vdb.retrieval(data, **retrieval_kwargs)) diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index dfd29ee635..428131b7c5 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -19,10 +19,12 @@ import pandas as pd +from nemo_retriever.common.params import build_embed_option_kwargs from nemo_retriever.operators.abstract_operator import AbstractOperator -from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL +from nemo_retriever.models import VL_RERANK_MODEL from nemo_retriever.query.agentic_options import ( - agentic_backend_top_k_error, + AGENTIC_DEFAULT_CLIENT, + agentic_llm_client_error, agentic_float_range_value, agentic_int_min_error, agentic_int_value, @@ -48,14 +50,12 @@ AGENTIC_RETRIEVER_TOP_K = 10 AGENTIC_TARGET_TOP_K = 10 -AGENTIC_BACKEND_TOP_K = 20 # backend retrieve-pool depth. show-count stays AGENTIC_TARGET_TOP_K=10 -AGENTIC_SELECTION_TOP_K = 10 AGENTIC_NUM_CONCURRENT = 1 AGENTIC_TEXT_TRUNCATION = 0 -AGENTIC_PARALLEL_TOOL_CALLS = False +AGENTIC_PARALLEL_TOOL_CALLS = None AGENTIC_RRF_K = 60 AGENTIC_REACT_MAX_STEPS = 50 -AGENTIC_TEMPERATURE = 0.0 # agent LLM sampling temperature (0.0 = greedy) +AGENTIC_TEMPERATURE = None # agent LLM sampling temperature; None leaves it unset AGENTIC_MAX_TOKENS: Optional[int] = None AGENTIC_LLM_BACKEND = "in_process" AGENTIC_LLM_BACKENDS = frozenset({"openai_compatible", "in_process"}) @@ -133,12 +133,12 @@ class AgenticRetrievalConfig: vdb_op: str = "lancedb" vdb_kwargs: dict[str, Any] = field(default_factory=dict) - query_embedder: str = VL_EMBED_MODEL + query_embedder: Optional[str] = None query_embedder_provider_prefix: Optional[str] = None embedding_endpoint: Optional[str] = None embedding_api_key: str = "" local_hf_batch_size: int = 32 - local_query_embed_backend: str = "hf" + local_query_embed_backend: Optional[str] = None reranker: Optional[str] = None reranker_endpoint: Optional[str] = None reranker_api_key: str = "" @@ -160,10 +160,16 @@ class AgenticRetrievalConfig: # Forwarded verbatim as the OpenAI `reasoning_effort` field on every LLM # call when explicitly configured. reasoning_effort: Optional[str] = None - # Backend retrieve-pool depth, distinct from the final selected top_k. - backend_top_k: int = AGENTIC_BACKEND_TOP_K - # Sampling temperature sent on every agent LLM call (0.0 = greedy). - temperature: float = AGENTIC_TEMPERATURE + # Sampling temperature sent on every agent LLM call. ``None`` leaves it unset + # so the endpoint/model default applies. + temperature: Optional[float] = AGENTIC_TEMPERATURE + # LLM client used to build the ReAct and selection agent LLMs. Optional: + # when unset it defaults to ``"callable"`` for both in-process (local vLLM) + # runs and remote (openai_compatible) runs -- the difference is which + # completion callable the operator injects. A remote-only client (anything + # other than ``"callable"``) may be named explicitly to override the default; + # the valid set is the ``nemo_agent`` LLM backend registry. + llm_client: Optional[str] = None # Optional upper bound on tokens in each agent LLM response. max_tokens: Optional[int] = AGENTIC_MAX_TOKENS # Final number of documents the agent targets/selects and the pipeline returns. @@ -191,6 +197,28 @@ def __post_init__(self) -> None: ) object.__setattr__(self, "llm_backend", llm_backend) + # Two independent axes: ``llm_backend`` is WHERE compute runs (in-process + # vs a remote endpoint); ``llm_client`` is WHICH nemo_agent adapter drives + # it. ``callable`` spans both, because it wraps a completion callable that + # may be either an in-process engine or the shared HTTP client. + # Validity of an explicit name is delegated to the nemo_agent registry + # via ``agentic_llm_client_error``. + explicit_client = str(self.llm_client or "").strip().lower() or None + if explicit_client is not None: + client_error = agentic_llm_client_error(explicit_client, field_name="llm_client") + if client_error: + raise ValueError(client_error) + if llm_backend == "in_process": + if explicit_client is not None and explicit_client != "callable": + raise ValueError( + "in-process agentic runs use the 'callable' LLM client; " + "provide invoke_url to use a remote client." + ) + llm_client = "callable" + else: + llm_client = explicit_client or AGENTIC_DEFAULT_CLIENT + object.__setattr__(self, "llm_client", llm_client) + local_llm_backend = _normalize_agentic_choice( self.local_llm_backend, AGENTIC_LOCAL_LLM_BACKENDS, @@ -229,15 +257,6 @@ def __post_init__(self) -> None: raise ValueError(integer_error) object.__setattr__(self, field_name, agentic_int_value(value, field_name=field_name)) - backend_error = agentic_backend_top_k_error( - self.backend_top_k, - target_top_k=int(self.top_k), - field_name="backend_top_k", - ) - if backend_error: - raise ValueError(backend_error) - object.__setattr__(self, "backend_top_k", agentic_int_value(self.backend_top_k, field_name="backend_top_k")) - local_tp_error = agentic_int_min_error( self.local_tensor_parallel_size, field_name="local_tensor_parallel_size", min_value=1 ) @@ -267,15 +286,21 @@ def __post_init__(self) -> None: ) object.__setattr__(self, "local_gpu_memory_utilization", local_gpu_memory_utilization) - temperature_invoke_url = self.invoke_url if self.llm_backend == "openai_compatible" else "local://in-process" - temperature_error = agentic_temperature_error( - self.temperature, - invoke_url=temperature_invoke_url, - field_name="temperature", - ) - if temperature_error: - raise ValueError(temperature_error) - object.__setattr__(self, "temperature", float(self.temperature)) + # Temperature: ``None`` stays unset (endpoint/model default). The range + # check uses a local sentinel so in-process runs (invoke_url=None) get the + # in-process bound instead of the hosted-endpoint bound. + if self.temperature is not None: + temperature_invoke_url = ( + self.invoke_url if self.llm_backend == "openai_compatible" else "local://in-process" + ) + temperature_error = agentic_temperature_error( + self.temperature, + invoke_url=temperature_invoke_url, + field_name="temperature", + ) + if temperature_error: + raise ValueError(temperature_error) + object.__setattr__(self, "temperature", float(self.temperature)) def _normalize_agentic_choice(value: object, valid: frozenset[str], *, field_name: str, default: str) -> str: @@ -320,22 +345,28 @@ def __init__( self._doc_id_field = str(doc_id_field) if doc_id_field else None if self._doc_id_field is not None and self._doc_id_field not in VALID_BEIR_DOC_ID_FIELDS: raise ValueError(f"Unsupported doc_id_field: {self._doc_id_field}") + embed_kwargs = build_embed_option_kwargs( + cfg.embedding_endpoint, + cfg.query_embedder, + embed_api_key=cfg.embedding_api_key, + embed_model_provider_prefix=cfg.query_embedder_provider_prefix, + ) + if cfg.local_query_embed_backend is not None: + embed_kwargs["local_ingest_embed_backend"] = str(cfg.local_query_embed_backend) + embed_kwargs.update( + { + "input_type": "query", + "inference_batch_size": int(cfg.local_hf_batch_size), + "embed_inference_batch_size": int(cfg.local_hf_batch_size), + } + ) + self._retriever = Retriever( vdb_kwargs={ "vdb_op": str(cfg.vdb_op), "vdb_kwargs": dict(cfg.vdb_kwargs or {}), }, - embed_kwargs={ - "model_name": str(cfg.query_embedder or VL_EMBED_MODEL), - "embed_model_name": str(cfg.query_embedder or VL_EMBED_MODEL), - "embed_model_provider_prefix": cfg.query_embedder_provider_prefix, - "embedding_endpoint": cfg.embedding_endpoint, - "api_key": cfg.embedding_api_key, - "input_type": "query", - "local_ingest_embed_backend": str(cfg.local_query_embed_backend), - "inference_batch_size": int(cfg.local_hf_batch_size), - "embed_inference_batch_size": int(cfg.local_hf_batch_size), - }, + embed_kwargs=embed_kwargs, top_k=AGENTIC_RETRIEVER_TOP_K, rerank=bool(cfg.reranker), rerank_kwargs={ @@ -406,15 +437,13 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D retriever_fn=self._retrieve_for_agent, retriever_top_k=per_hop_top_k, target_top_k=target_top_k, - user_msg_type="with_results", max_steps=int(self._cfg.react_max_steps), - extended_relevance=True, api_key=_none_if_empty(self._cfg.api_key), parallel_tool_calls=AGENTIC_PARALLEL_TOOL_CALLS, num_concurrent=int(self._cfg.num_concurrent), reasoning_effort=self._cfg.reasoning_effort, - backend_top_k=self._cfg.backend_top_k, - temperature=float(self._cfg.temperature), + temperature=self._cfg.temperature, + backend=self._cfg.llm_client, max_tokens=self._cfg.max_tokens, chat_completion_fn=chat_completion_fn, ) @@ -425,10 +454,10 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D top_k=target_top_k, api_key=_none_if_empty(self._cfg.api_key), parallel_tool_calls=AGENTIC_PARALLEL_TOOL_CALLS, - extended_relevance=True, # match Path A text_truncation=int(self._cfg.text_truncation), reasoning_effort=self._cfg.reasoning_effort, - temperature=float(self._cfg.temperature), + temperature=self._cfg.temperature, + backend=self._cfg.llm_client, max_tokens=self._cfg.max_tokens, chat_completion_fn=chat_completion_fn, ) diff --git a/nemo_retriever/src/nemo_retriever/query/agentic_options.py b/nemo_retriever/src/nemo_retriever/query/agentic_options.py index 5bcf6823af..7d68a9d712 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic_options.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic_options.py @@ -15,6 +15,11 @@ AGENTIC_OPENAI_COMPATIBLE_TEMPERATURE_MAX = 2.0 AGENTIC_NVIDIA_TEMPERATURE_MAX = 1.0 +#: Default LLM client for remote (openai_compatible) agentic retrieval. The same +#: client serves in-process runs; the difference is which completion callable the +#: operator injects. +AGENTIC_DEFAULT_CLIENT = "callable" + def _parse_integer(value: object, *, field_name: str) -> tuple[int | None, str | None]: if isinstance(value, bool): @@ -117,6 +122,23 @@ def agentic_temperature_error( return None +def agentic_llm_client_error(client: object, *, field_name: str = "agentic_llm_client") -> str | None: + """Return an error string if *client* is not a registered LLM client. + + The valid set is sourced from ``nemo_agent.llm.get_available_backends`` (the + ``nemo_agent`` library still calls these "backends" internally; the + user-facing term is "client"). Reading the registry here keeps the check + future-proof: a newly registered backend is accepted with no edits. + """ + from nemo_retriever._agentic.nemo_agent.llm import get_available_backends + + valid = get_available_backends() + if str(client).strip() not in valid: + choices = ", ".join(valid) + return f"{field_name} must be one of: {choices}" + return None + + def agentic_target_top_k(evaluation_mode: str, beir_k: list[int] | tuple[int, ...] | None = None) -> int: """Resolve the final document count required by the selected evaluation.""" @@ -137,21 +159,3 @@ def agentic_target_top_k(evaluation_mode: str, beir_k: list[int] | tuple[int, .. if not positive_ks: raise ValueError("agentic evaluation requires at least one positive k") return max(positive_ks) - - -def agentic_backend_top_k_error( - backend_top_k: object, - *, - target_top_k: int, - field_name: str = "agentic_backend_top_k", -) -> str | None: - parsed, error = _parse_integer(backend_top_k, field_name=field_name) - if error or parsed is None: - return error or f"{field_name} must be an integer" - value = parsed - - if value < 1: - return f"{field_name} must be >= 1" - if value < int(target_top_k): - return f"{field_name} must be >= target top_k ({int(target_top_k)})" - return None diff --git a/nemo_retriever/src/nemo_retriever/query/evidence.py b/nemo_retriever/src/nemo_retriever/query/evidence.py index 55386da75c..8dd968b9be 100644 --- a/nemo_retriever/src/nemo_retriever/query/evidence.py +++ b/nemo_retriever/src/nemo_retriever/query/evidence.py @@ -67,12 +67,12 @@ def _evidence_item(hit: dict[str, Any]) -> dict[str, Any]: fidelity = meta.get("fidelity") or _derive_fidelity(raw_modality, meta, meta) or "verbatim" - if "_score" in hit and hit["_score"] is not None: - score: float = hit["_score"] - elif "_distance" in hit and hit["_distance"] is not None: - score = hit["_distance"] - else: - score = 0.0 + raw_score = hit.get("distance") + if raw_score is None: + raw_score = hit.get("_score") + if raw_score is None: + raw_score = hit.get("_distance") + score = float(raw_score) if raw_score is not None else 0.0 return { "text": hit.get("text", ""), @@ -129,5 +129,9 @@ def build_evidence_result(hits: list, strategies_used: list[str]) -> dict[str, A ) return { "evidence": evidence, - "coverage": {"strategies_used": strategies_used, "n_docs_seen": len(sources), "thin_spots": thin}, + "coverage": { + "strategies_used": strategies_used, + "n_docs_seen": len(sources), + "thin_spots": thin, + }, } diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index f9a0c4696a..74def4e52d 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/nemo_retriever/src/nemo_retriever/query/options.py @@ -26,6 +26,7 @@ class QueryEmbedOptions: embed_invoke_url: str | None = None embed_model_name: str | None = None embed_model_provider_prefix: str | None = None + embed_api_key: str | None = None @dataclass(frozen=True) @@ -64,11 +65,16 @@ class QueryAgenticOptions: local_max_model_len: int | None = None local_max_num_seqs: int | None = None reasoning_effort: str | None = None - backend_top_k: int = 20 react_max_steps: int = 50 text_truncation: int = 0 num_concurrent: int = 1 - temperature: float = 0.0 + temperature: float | None = None + # LLM client (see AgenticRetrievalConfig.llm_client). Optional: defaults to + # ``callable`` for both in-process and remote runs when unset. + llm_client: str | None = None + # Accepted for service-layer compatibility only. The agent derives its own + # per-hop retrieval depth from ``top_k``, so this value is never read. + backend_top_k: int | None = None @dataclass(frozen=True) diff --git a/nemo_retriever/src/nemo_retriever/query/shaping.py b/nemo_retriever/src/nemo_retriever/query/shaping.py index fa81a8f822..c55b164550 100644 --- a/nemo_retriever/src/nemo_retriever/query/shaping.py +++ b/nemo_retriever/src/nemo_retriever/query/shaping.py @@ -6,12 +6,13 @@ from typing import Any, Sequence, cast -from nemo_retriever.common.vdb.lancedb_schema import normalize_content_type -from nemo_retriever.common.vdb.records import RetrievalHit +from nemo_retriever.common.vdb.records import RetrievalHit, normalize_content_type from nemo_retriever.common.vdb.sidecar_metadata import parse_hit_content_metadata -def normalize_query_content_type_allowlist(content_types: str | Sequence[str] | None) -> set[str] | None: +def normalize_query_content_type_allowlist( + content_types: str | Sequence[str] | None, +) -> set[str] | None: """Normalize query-time content type filters to stored hit metadata values.""" if content_types is None: return None diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index e7d9bb21db..c9ddf88789 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -163,7 +163,7 @@ def build_agentic_config(request: QueryRequest, *, top_k: int | None = None) -> """ from nemo_retriever.query.agentic import AgenticRetrievalConfig - api_key = resolve_remote_api_key() + api_key = resolve_remote_api_key(request.embed.embed_api_key) vdb_kwargs: dict[str, Any] = {"uri": request.storage.lancedb_uri, "table_name": request.storage.table_name} if request.retrieval.retrieval_mode != "auto": vdb_kwargs["retrieval_mode"] = request.retrieval.retrieval_mode @@ -183,11 +183,11 @@ def build_agentic_config(request: QueryRequest, *, top_k: int | None = None) -> "local_max_num_seqs": request.agentic.local_max_num_seqs, "api_key": api_key, "reasoning_effort": request.agentic.reasoning_effort, - "backend_top_k": int(request.agentic.backend_top_k), "react_max_steps": int(request.agentic.react_max_steps), "text_truncation": int(request.agentic.text_truncation), "num_concurrent": int(request.agentic.num_concurrent), - "temperature": float(request.agentic.temperature), + "temperature": request.agentic.temperature, + "llm_client": request.agentic.llm_client, } if request.agentic.llm_backend: cfg_kwargs["llm_backend"] = request.agentic.llm_backend @@ -234,10 +234,13 @@ def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]: result = result.sort_values("rank") ranked: list[dict[str, Any]] = [] for _, row in result.iterrows(): + doc_id = str(row.get("doc_id", "")).strip() + if not doc_id: + continue ranked.append( { "rank": int(row.get("rank", len(ranked) + 1)), - "doc_id": str(row.get("doc_id", "")), + "doc_id": doc_id, "result_source": str(row.get("result_source", "")), } ) diff --git a/nemo_retriever/src/nemo_retriever/service/agentic_query.py b/nemo_retriever/src/nemo_retriever/service/agentic_query.py new file mode 100644 index 0000000000..906d4beac0 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/service/agentic_query.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Service boundary for in-process agentic retrieval over the VectorDB query API.""" + +from __future__ import annotations + +from typing import Any + +from nemo_retriever.query.options import ( + QueryAgenticOptions, + QueryEmbedOptions, + QueryRequest as WorkflowQueryRequest, + QueryRetrievalOptions, + QueryStorageOptions, +) +from nemo_retriever.query.workflow import agentic_query_documents +from nemo_retriever.service.config import AgenticConfig +from nemo_retriever.service.query_schema import QueryResponse, QueryResult + + +def agentic_ranked_to_hits(ranked: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Map agentic document ranks onto the ``/v1/query`` hits envelope. + + Agentic retrieval selects documents (``doc_id``), not chunks. Each hit places + ``doc_id`` in ``source``, records ``result_source`` / ``rank`` under + ``metadata``, and leaves chunk-level fields (``text``, ``page_number``, + scores, …) unset. + + Rows without a non-empty ``doc_id`` are a contract violation: the agentic + workflow already skips blank ids on retrieve hops, and a hit with + ``source=null`` is useless to clients. Raise rather than emit a null source. + """ + hits: list[dict[str, Any]] = [] + for item in ranked: + doc_id = str(item.get("doc_id") or "").strip() + if not doc_id: + raise ValueError( + "agentic ranked result is missing a non-empty doc_id " + f"(rank={item.get('rank')!r}, result_source={item.get('result_source')!r})" + ) + hits.append( + { + "text": None, + "metadata": { + "result_source": item.get("result_source"), + "rank": item.get("rank"), + }, + "source": doc_id, + "source_id": None, + "path": None, + "page_number": None, + "pdf_basename": None, + "pdf_page": None, + } + ) + return hits + + +def build_agentic_query_request( + *, + query: str, + top_k: int, + config: AgenticConfig, + lancedb_uri: str, + table_name: str, + embed_endpoint: str, + embed_model: str, + embed_model_provider_prefix: str | None, + embed_api_key: str, +) -> WorkflowQueryRequest: + """Map server-owned service settings onto the shared agentic query request.""" + return WorkflowQueryRequest( + query=query, + retrieval=QueryRetrievalOptions(top_k=top_k), + embed=QueryEmbedOptions( + embed_invoke_url=embed_endpoint or None, + embed_model_name=embed_model or None, + embed_model_provider_prefix=embed_model_provider_prefix, + embed_api_key=embed_api_key or None, + ), + storage=QueryStorageOptions( + lancedb_uri=lancedb_uri, + table_name=table_name, + ), + agentic=QueryAgenticOptions( + enabled=True, + llm_model=config.llm_model, + invoke_url=config.invoke_url, + reasoning_effort=config.reasoning_effort, + backend_top_k=config.backend_top_k, + react_max_steps=config.react_max_steps, + text_truncation=config.text_truncation, + temperature=config.temperature, + ), + ) + + +def run_agentic_query( + *, + query: str, + top_k: int, + config: AgenticConfig, + lancedb_uri: str, + table_name: str, + embed_endpoint: str, + embed_model: str, + embed_model_provider_prefix: str | None, + embed_api_key: str, +) -> QueryResponse: + """Execute one agentic retrieval query and return ``QueryResponse`` hits.""" + query_request = build_agentic_query_request( + query=query, + top_k=top_k, + config=config, + lancedb_uri=lancedb_uri, + table_name=table_name, + embed_endpoint=embed_endpoint, + embed_model=embed_model, + embed_model_provider_prefix=embed_model_provider_prefix, + embed_api_key=embed_api_key, + ) + ranked = agentic_query_documents(query_request) + return QueryResponse( + results=[QueryResult(hits=agentic_ranked_to_hits(ranked))], + query_mode="agentic", + ) diff --git a/nemo_retriever/src/nemo_retriever/service/app.py b/nemo_retriever/src/nemo_retriever/service/app.py index 10ad647f9a..e232f9dc5c 100644 --- a/nemo_retriever/src/nemo_retriever/service/app.py +++ b/nemo_retriever/src/nemo_retriever/service/app.py @@ -42,7 +42,11 @@ def _configure_logging(config: ServiceConfig) -> None: file_handler.setFormatter(fmt) root.addHandler(file_handler) - logger.info("Logging configured: level=%s file=%s", config.logging.level, config.logging.file) + logger.info( + "Logging configured: level=%s file=%s", + config.logging.level, + config.logging.file, + ) def _apply_resource_limits(config: ServiceConfig) -> None: @@ -95,7 +99,10 @@ def _check_media_dependencies(mode: str) -> None: ) if is_media_available(): - logger.info("Media dependencies (ffmpeg, ffprobe) detected — audio/video ingestion enabled (mode=%s)", mode) + logger.info( + "Media dependencies (ffmpeg, ffprobe) detected — audio/video ingestion enabled (mode=%s)", + mode, + ) return missing = ", ".join(missing_media_dependencies()) or "ffmpeg, ffprobe" @@ -125,10 +132,19 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]: config: ServiceConfig = app.state.config mode = config.mode - from nemo_retriever.service.services.event_bus import init_event_bus, shutdown_event_bus - from nemo_retriever.service.services.job_tracker import init_job_tracker, shutdown_job_tracker + from nemo_retriever.service.services.event_bus import ( + init_event_bus, + shutdown_event_bus, + ) + from nemo_retriever.service.services.job_tracker import ( + init_job_tracker, + shutdown_job_tracker, + ) from nemo_retriever.service.services.metrics import init_metrics, shutdown_metrics - from nemo_retriever.service.services.pipeline_pool import init_pipeline_pool, shutdown_pipeline_pool + from nemo_retriever.service.services.pipeline_pool import ( + init_pipeline_pool, + shutdown_pipeline_pool, + ) from nemo_retriever.service.services.proxy import init_proxy, shutdown_proxy from nemo_retriever.service.services.sidecar_store import init_sidecar_store, shutdown_sidecar_store from nemo_retriever.service.services.worker_result_store import validate_result_store @@ -142,6 +158,8 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]: app.state.metrics = None tracker = init_job_tracker() + if app.state.metrics is not None: + tracker.add_terminal_observer(app.state.metrics.record_terminal_transition) event_bus = init_event_bus() tracker.set_event_bus(event_bus) app.state.sidecar_store = init_sidecar_store() @@ -167,6 +185,7 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]: batch_work_fn=bt_fn, work_queue_config=config.work_queue, auth_config=config.auth, + internal_api_token=config.vectordb.internal_api_token, ) if ( @@ -176,7 +195,9 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]: ): import asyncio - from nemo_retriever.service.services.pipeline_executor import warmup_process_pool_workers + from nemo_retriever.service.services.pipeline_executor import ( + warmup_process_pool_workers, + ) warmup_status = await asyncio.to_thread(warmup_process_pool_workers) logger.info("Local model warmup status: %s", warmup_status) @@ -192,7 +213,9 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]: yield - from nemo_retriever.service.services.pipeline_executor import shutdown_process_executors + from nemo_retriever.service.services.pipeline_executor import ( + shutdown_process_executors, + ) shutdown_process_executors() await shutdown_work_broker() @@ -225,7 +248,10 @@ def create_app(config: ServiceConfig) -> FastAPI: try: from fastmcp.utilities.lifespan import combine_lifespans - from nemo_retriever.service.mcp_server import build_mcp_app, settings_from_service_config + from nemo_retriever.service.mcp_server import ( + build_mcp_app, + settings_from_service_config, + ) mcp_asgi_app = build_mcp_app(settings_from_service_config(config)) lifespan = combine_lifespans(_lifespan, mcp_asgi_app.lifespan) @@ -251,26 +277,30 @@ def create_app(config: ServiceConfig) -> FastAPI: app.add_middleware(_RequestIdMiddleware) - if config.auth.api_token: - from nemo_retriever.service.auth import BearerAuthMiddleware + from nemo_retriever.service.auth import BearerAuthMiddleware - app.add_middleware(BearerAuthMiddleware, config=config.auth) - logger.info( - "Bearer-token authentication ENABLED (header=%s, bypass=%s)", - config.auth.header_name, - config.auth.bypass_paths, - ) - else: - logger.info("Bearer-token authentication DISABLED (no api_token configured)") + app.add_middleware( + BearerAuthMiddleware, + config=config.auth, + internal_api_token=config.vectordb.internal_api_token, + ) + logger.info( + "Scope authorization configured (enabled=%s, header=%s, secret_file=%s, allow_unscoped_dev=%s)", + config.auth.enabled, + config.auth.header_name, + bool(config.auth.scope_token_file), + config.auth.allow_unscoped_dev, + ) if mcp_asgi_app is not None: app.mount(config.mcp.path, mcp_asgi_app) logger.info("FastMCP service endpoint mounted at %s", config.mcp.path) - from nemo_retriever.service.routers import admin, ingest, metrics, work + from nemo_retriever.service.routers import admin, collections, ingest, metrics, work from nemo_retriever.service.services.prometheus import instrument_app app.include_router(ingest.router, prefix="/v1") + app.include_router(collections.router, prefix="/v1") app.include_router(metrics.router, prefix="/v1") # Admin/internal endpoints — pool_stats etc. Registered on every # role; the handler self-reports an empty pool dict on gateway pods. @@ -302,7 +332,9 @@ async def health() -> dict: and config.local_models.enabled and config.local_models.warmup_on_startup ): - from nemo_retriever.service.services.pipeline_executor import get_service_warmup_status + from nemo_retriever.service.services.pipeline_executor import ( + get_service_warmup_status, + ) warmup = get_service_warmup_status() base["models_warm"] = bool(warmup.get("complete")) diff --git a/nemo_retriever/src/nemo_retriever/service/auth.py b/nemo_retriever/src/nemo_retriever/service/auth.py index 3744ce1c55..aadbb56ac5 100644 --- a/nemo_retriever/src/nemo_retriever/service/auth.py +++ b/nemo_retriever/src/nemo_retriever/service/auth.py @@ -2,20 +2,15 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Optional bearer-token authentication middleware. - -Activated only when ``ServiceConfig.auth.api_token`` is non-empty. The -middleware compares the incoming header value to the configured token in -constant time to avoid timing-attack leakage. - -Bypass paths are matched as **prefixes** so OpenAPI / docs / health all -work with just the four defaults in ``AuthConfig``. -""" +"""Bearer authentication and token-to-workspace scope authorization.""" from __future__ import annotations import hmac +import json import logging +from pathlib import Path +from typing import Any from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request @@ -27,6 +22,7 @@ _BEARER_PREFIX = "bearer " +_INTERNAL_TOKEN_HEADER = "X-NRL-Internal-Token" def _strip_bearer(value: str) -> str: @@ -44,34 +40,118 @@ def auth_headers(config: AuthConfig) -> dict[str, str]: return {config.header_name: value} -class BearerAuthMiddleware(BaseHTTPMiddleware): - """Reject requests that don't carry the configured token. +def internal_auth_headers(token: str | None) -> dict[str, str]: + """Build headers for gateway/worker calls to the VectorDB service.""" + token = (token or "").strip() + if not token: + return {} + return {"X-NRL-Internal-Token": token} + + +def authorized_scope(request: Request) -> str: + """Return the middleware-authorized scope; never trust a raw header here.""" + return str(getattr(request.state, "authorized_scope", "default")) + + +class ScopeAuthorizer: + """Resolve credentials to allowed logical scopes without logging secrets.""" + + def __init__(self, config: AuthConfig) -> None: + self.default_scope = config.default_scope.strip() or "default" + self.enabled = config.enabled + self.allow_unscoped_dev = config.allow_unscoped_dev + self._records: list[tuple[str, frozenset[str]]] = [] + if token := (config.api_token or "").strip(): + self._records.append((token, frozenset({self.default_scope}))) + if config.scope_token_file: + self._load_file(config.scope_token_file) + + def _load_file(self, path: str) -> None: + try: + payload: Any = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise ValueError(f"Unable to load scope-token secret file {path!r}: {exc}") from exc + records = payload.get("tokens") if isinstance(payload, dict) else None + if not isinstance(records, list): + raise ValueError("scope-token secret file must contain a 'tokens' list") + for record in records: + if not isinstance(record, dict): + raise ValueError("each scope-token record must be an object") + token = str(record.get("token") or "").strip() + scopes = record.get("scopes") + if not token or not isinstance(scopes, list) or not scopes: + raise ValueError("scope-token records require a non-empty token and scopes list") + self._records.append( + ( + token, + frozenset(str(scope).strip() for scope in scopes if str(scope).strip()), + ) + ) - A no-op when ``config.api_token`` is None or empty so deployments that - don't enable auth pay no overhead beyond the middleware dispatch. - """ + def authorize(self, provided_token: str, requested_scope: str | None) -> tuple[str | None, int | None]: + """Resolve an authorized scope without revealing whether a token was recognized.""" + + requested = (requested_scope or self.default_scope).strip() or self.default_scope + if not self.enabled: + return requested, None + if not self._records: + if self.allow_unscoped_dev: + return requested, None + return None, 401 + allowed: frozenset[str] | None = None + for token, scopes in self._records: + if hmac.compare_digest(provided_token, token): + allowed = scopes + if allowed is None: + return None, 401 + if requested not in allowed: + return None, 401 + return requested, None - def __init__(self, app, *, config: AuthConfig) -> None: + +class BearerAuthMiddleware(BaseHTTPMiddleware): + """Authenticate public requests and isolate pod-only callback routes.""" + + def __init__( + self, + app, + *, + config: AuthConfig, + internal_api_token: str | None = None, + ) -> None: super().__init__(app) - self._token = (config.api_token or "").strip() self._header = config.header_name self._bypass = tuple(config.bypass_paths) + self._authorizer = ScopeAuthorizer(config) + self._internal_api_token = (internal_api_token or "").strip() async def dispatch(self, request: Request, call_next): - if not self._token: - return await call_next(request) + """Authenticate the request and attach its authorized logical scope.""" path = request.url.path if any(path == p or path.startswith(p.rstrip("/") + "/") for p in self._bypass): + request.state.authorized_scope = self._authorizer.default_scope + return await call_next(request) + + if path.startswith("/v1/internal/") and self._internal_api_token: + supplied = request.headers.get(_INTERNAL_TOKEN_HEADER, "").strip() + if not supplied or not hmac.compare_digest(supplied, self._internal_api_token): + return JSONResponse( + status_code=401, + content={"detail": "Missing or invalid internal credential."}, + ) + request.state.authorized_scope = self._authorizer.default_scope return await call_next(request) provided = request.headers.get(self._header, "") provided_token = _strip_bearer(provided) - if not provided_token or not hmac.compare_digest(provided_token, self._token): + scope, failure = self._authorizer.authorize(provided_token, request.headers.get("X-NRL-Scope")) + if failure is not None: return JSONResponse( status_code=401, content={"detail": "Missing or invalid bearer token."}, headers={"WWW-Authenticate": "Bearer"}, ) + request.state.authorized_scope = scope return await call_next(request) diff --git a/nemo_retriever/src/nemo_retriever/service/cli.py b/nemo_retriever/src/nemo_retriever/service/cli.py index b8cf1d09c7..ed6676cff9 100644 --- a/nemo_retriever/src/nemo_retriever/service/cli.py +++ b/nemo_retriever/src/nemo_retriever/service/cli.py @@ -157,6 +157,17 @@ def mcp_stdio( ), concurrency: int = typer.Option(8, "--concurrency", min=1, help="Max concurrent MCP document uploads."), request_timeout_s: float = typer.Option(60.0, "--request-timeout", min=0.1, help="HTTP request timeout."), + query_methods: str = typer.Option( + "classic", + "--query-methods", + help="Retrieval MCP tools to expose: classic, agentic, or all.", + ), + agentic_request_timeout_s: float = typer.Option( + 1800.0, + "--agentic-request-timeout", + min=1.0, + help="HTTP timeout for agentic_query.", + ), ingest_timeout_s: float = typer.Option(1800.0, "--ingest-timeout", min=1.0, help="Document ingest timeout."), poll_interval_s: float = typer.Option(2.0, "--poll-interval", min=0.1, help="Status polling interval."), enable_write_tools: bool = typer.Option( @@ -168,14 +179,20 @@ def mcp_stdio( """Run the retriever service MCP server over stdio for local agents.""" from nemo_retriever.service.mcp_server import ServiceMCPSettings, build_mcp + normalized = query_methods.strip().lower() + if normalized not in {"classic", "agentic", "all"}: + raise typer.BadParameter("query-methods must be one of: classic, agentic, all") + settings = ServiceMCPSettings( base_url=service_url, api_token=api_token, auth_header_name=auth_header_name, max_concurrency=concurrency, request_timeout_s=request_timeout_s, + agentic_request_timeout_s=agentic_request_timeout_s, ingest_timeout_s=ingest_timeout_s, poll_interval_s=poll_interval_s, enable_write_tools=enable_write_tools, + query_methods=normalized, # type: ignore[arg-type] ) build_mcp(settings).run(transport="stdio") diff --git a/nemo_retriever/src/nemo_retriever/service/client.py b/nemo_retriever/src/nemo_retriever/service/client.py index 52bb763149..86d84698d3 100644 --- a/nemo_retriever/src/nemo_retriever/service/client.py +++ b/nemo_retriever/src/nemo_retriever/service/client.py @@ -33,11 +33,13 @@ from __future__ import annotations import asyncio +import hashlib import json import logging import time +from concurrent.futures import ThreadPoolExecutor from pathlib import Path -from typing import Any, AsyncIterator, Callable, NamedTuple +from typing import Any, AsyncIterator, Callable, Coroutine, NamedTuple, TypeVar import httpx from pydantic import ValidationError @@ -50,6 +52,25 @@ TimeRemainingColumn, ) +from nemo_retriever.common.schemas.collections import ( + CollectionDeleteResult, + CollectionInfo, + CollectionPage, + DocumentDeleteResult, + DocumentInfo, + DocumentPage, + QueryHit, +) +from nemo_retriever.common.schemas.responses import ( + JobAggregateResponse, + JobDocumentsPage, +) +from nemo_retriever.service.errors import ( + RetrieverServiceConflictError, + RetrieverServiceError, + RetrieverServiceNotFoundError, + RetrieverServiceValidationError, +) from nemo_retriever.service.query_schema import QueryResponse logger = logging.getLogger(__name__) @@ -59,6 +80,8 @@ _MAX_UPLOAD_RETRIES = 10 _DEFAULT_RETRY_AFTER = 2.0 +_T = TypeVar("_T") + _TRANSIENT_ERRORS: tuple[type[Exception], ...] = ( httpx.ReadError, httpx.WriteError, @@ -74,6 +97,29 @@ class _CreatedJob(NamedTuple): trace_id: str | None = None +class _UploadOutcome(NamedTuple): + filename: str + document_id: str = "" + attempt_id: str = "" + error: str | None = None + + +class InMemoryUpload(NamedTuple): + """Document bytes that should be uploaded without a temporary file.""" + + filename: str + content: bytes + content_type: str = "application/octet-stream" + classification_filename: str | None = None + + +UploadInput = Path | InMemoryUpload + + +def _upload_filename(source: UploadInput) -> str: + return source.name if isinstance(source, Path) else source.filename + + # ------------------------------------------------------------------ # Errors # ------------------------------------------------------------------ @@ -185,7 +231,7 @@ def mark_failed(self, doc_id: str, error: str | None, event: dict[str, Any]) -> class RetrieverServiceClient: - """Submits documents to a running retriever service and queries its VectorDB endpoint. + """Manage scoped collections, ingest documents, and query a retriever service. Ingest opens a job aggregate with ``POST /v1/ingest/job`` (sized to the number of files), then uses @@ -203,47 +249,353 @@ def __init__( max_concurrency: int = 8, *, api_token: str | None = None, + scope: str | None = None, ) -> None: self._base_url = base_url.rstrip("/") self._max_concurrency = max_concurrency self._api_token = (api_token or "").strip() or None + self._scope = (scope or "").strip() or None @property def _auth_headers(self) -> dict[str, str]: - return {"Authorization": f"Bearer {self._api_token}"} if self._api_token else {} - - # ------------------------------------------------------------------ - # Query - # ------------------------------------------------------------------ + headers: dict[str, str] = {} + if self._scope: + headers["X-NRL-Scope"] = self._scope + if self._api_token: + headers["Authorization"] = f"Bearer {self._api_token}" + return headers + + @staticmethod + def _raise_for_response(resp: httpx.Response, operation: str) -> None: + if resp.status_code < 400: + return + detail = resp.text[:1000] if resp.text else "(empty)" + error_type: type[RetrieverServiceError] + if resp.status_code == 404: + error_type = RetrieverServiceNotFoundError + elif resp.status_code == 409: + error_type = RetrieverServiceConflictError + elif resp.status_code in (400, 422): + error_type = RetrieverServiceValidationError + else: + error_type = RetrieverServiceError + raise error_type( + f"{operation} failed: HTTP {resp.status_code}: {detail}", + status_code=resp.status_code, + ) - def query(self, query: str | list[str], *, top_k: int) -> list[list[dict[str, Any]]]: - """Search ingested documents through ``POST /v1/query``.""" - url = f"{self._base_url}/v1/query" - expected_results = len(query) if isinstance(query, list) else 1 - payload: dict[str, Any] = {"query": query, "top_k": int(top_k)} + async def _arequest(self, method: str, path: str, **kwargs: Any) -> Any: + # Construct the client inside the coroutine. ``_run`` may drive this on + # a worker thread's event loop, and a client bound to a different loop + # fails there — so do not hoist it to ``__init__`` to pool connections. try: - with httpx.Client( - timeout=httpx.Timeout(300.0, connect=30.0), - headers=self._auth_headers, + async with httpx.AsyncClient( + timeout=httpx.Timeout(300.0, connect=30.0), headers=self._auth_headers ) as client: - resp = client.post(url, json=payload) + resp = await client.request(method, f"{self._base_url}{path}", **kwargs) except httpx.HTTPError as exc: - raise RuntimeError(f"Service query failed: {type(exc).__name__}: {exc}") from exc + raise RetrieverServiceError(f"{method} {path} transport failure: {exc}") from exc + self._raise_for_response(resp, f"{method} {path}") + try: + return resp.json() if resp.content else None + except ValueError as exc: + raise RetrieverServiceError(f"{method} {path} returned malformed JSON") from exc - if resp.status_code >= 400: - detail = resp.text[:500] if resp.text else "(empty)" - raise RuntimeError(f"Service query failed: HTTP {resp.status_code}: {detail}") + @staticmethod + def _run(coro: Coroutine[Any, Any, _T]) -> _T: + """Drive an async operation to completion from synchronous code. + Each operation is implemented once, asynchronously; the synchronous + methods are thin facades over it. When the caller already runs inside + an event loop we hand the coroutine to a worker thread, because + ``asyncio.run`` refuses to nest. + """ try: - body = resp.json() - except ValueError as exc: - raise RuntimeError("Service query returned invalid JSON.") from exc + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(coro) + with ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(asyncio.run, coro).result() + + @staticmethod + def _model(model: Any, payload: Any, operation: str) -> Any: + try: + return model.model_validate(payload) + except (ValidationError, ValueError) as exc: + raise RetrieverServiceError(f"{operation} returned an invalid response: {exc}") from exc + + # ------------------------------------------------------------------ + # Collection and document lifecycle + # ------------------------------------------------------------------ + + def create_collection( + self, + name: str, + *, + description: str | None = None, + metadata: dict[str, Any] | None = None, + expires_at: str | None = None, + ) -> CollectionInfo: + """Create a logical collection in the client's configured scope.""" + + return self._run( + self.acreate_collection( + name, + description=description, + metadata=metadata, + expires_at=expires_at, + ) + ) + + async def acreate_collection( + self, + name: str, + *, + description: str | None = None, + metadata: dict[str, Any] | None = None, + expires_at: str | None = None, + ) -> CollectionInfo: + """Asynchronously create a logical collection.""" + + body = { + "name": name, + "description": description, + "metadata": metadata or {}, + "expires_at": expires_at, + } + return self._model( + CollectionInfo, + await self._arequest("POST", "/v1/collections", json=body), + "create collection", + ) + + def get_collection(self, name: str) -> CollectionInfo: + """Return one logical collection from the configured scope.""" + + return self._run(self.aget_collection(name)) + + async def aget_collection(self, name: str) -> CollectionInfo: + """Asynchronously return one logical collection.""" + + return self._model( + CollectionInfo, + await self._arequest("GET", f"/v1/collections/{name}"), + "get collection", + ) + + def list_collections(self, *, limit: int = 100, continuation_token: str | None = None) -> CollectionPage: + """List logical collections in the configured scope.""" + + return self._run(self.alist_collections(limit=limit, continuation_token=continuation_token)) + + async def alist_collections( + self, + *, + limit: int = 100, + continuation_token: str | None = None, + ) -> CollectionPage: + """Asynchronously list logical collections.""" + + params = { + "limit": limit, + "continuation_token": continuation_token, + } + return self._model( + CollectionPage, + await self._arequest("GET", "/v1/collections", params=params), + "list collections", + ) + + def update_collection(self, name: str, **changes: Any) -> CollectionInfo: + """Update mutable properties of a logical collection.""" + + return self._run(self.aupdate_collection(name, **changes)) + + async def aupdate_collection(self, name: str, **changes: Any) -> CollectionInfo: + """Asynchronously update a logical collection.""" + + return self._model( + CollectionInfo, + await self._arequest("PATCH", f"/v1/collections/{name}", json=changes), + "update collection", + ) + + def delete_collection(self, name: str, *, if_exists: bool = False) -> CollectionDeleteResult: + """Request deletion of a logical collection and its VectorDB-owned data.""" + + return self._run(self.adelete_collection(name, if_exists=if_exists)) + + async def adelete_collection(self, name: str, *, if_exists: bool = False) -> CollectionDeleteResult: + """Asynchronously delete a logical collection.""" + + body = await self._arequest("DELETE", f"/v1/collections/{name}", params={"if_exists": if_exists}) + return self._model(CollectionDeleteResult, body, "delete collection") + + def list_documents( + self, + collection_name: str, + *, + limit: int = 100, + continuation_token: str | None = None, + ) -> DocumentPage: + """List committed documents in a logical collection.""" + + return self._run( + self.alist_documents( + collection_name, + limit=limit, + continuation_token=continuation_token, + ) + ) + + async def alist_documents( + self, + collection_name: str, + *, + limit: int = 100, + continuation_token: str | None = None, + ) -> DocumentPage: + """Asynchronously list committed collection documents.""" + + params = { + "limit": limit, + "continuation_token": continuation_token, + } + return self._model( + DocumentPage, + await self._arequest("GET", f"/v1/collections/{collection_name}/documents", params=params), + "list documents", + ) + + def get_document(self, collection_name: str, document_id: str) -> DocumentInfo: + """Return one committed collection document.""" + + return self._run(self.aget_document(collection_name, document_id)) + + async def aget_document(self, collection_name: str, document_id: str) -> DocumentInfo: + """Asynchronously return one committed collection document.""" + + return self._model( + DocumentInfo, + await self._arequest("GET", f"/v1/collections/{collection_name}/documents/{document_id}"), + "get document", + ) + + def delete_document( + self, + collection_name: str, + document_id: str, + *, + if_exists: bool = False, + ) -> DocumentDeleteResult: + """Request deletion of one document and its collection chunks.""" + + return self._run(self.adelete_document(collection_name, document_id, if_exists=if_exists)) + async def adelete_document( + self, + collection_name: str, + document_id: str, + *, + if_exists: bool = False, + ) -> DocumentDeleteResult: + """Asynchronously delete one collection document.""" + + return self._model( + DocumentDeleteResult, + await self._arequest( + "DELETE", + f"/v1/collections/{collection_name}/documents/{document_id}", + params={"if_exists": if_exists}, + ), + "delete document", + ) + + def get_job(self, job_id: str) -> JobAggregateResponse: + """Return aggregate ingestion status for a job.""" + + return self._run(self.aget_job(job_id)) + + async def aget_job(self, job_id: str) -> JobAggregateResponse: + """Asynchronously return aggregate ingestion status.""" + + return self._model( + JobAggregateResponse, + await self._arequest("GET", f"/v1/ingest/job/{job_id}"), + "get job", + ) + + def list_job_documents(self, job_id: str, *, offset: int = 0, limit: int = 100) -> JobDocumentsPage: + """List per-document ingestion status for a job.""" + + return self._run(self.alist_job_documents(job_id, offset=offset, limit=limit)) + + async def alist_job_documents(self, job_id: str, *, offset: int = 0, limit: int = 100) -> JobDocumentsPage: + """Asynchronously list per-document ingestion status.""" + + return self._model( + JobDocumentsPage, + await self._arequest( + "GET", + f"/v1/ingest/job/{job_id}/documents", + params={"offset": offset, "limit": limit}, + ), + "list job documents", + ) + + # ------------------------------------------------------------------ + # Query + # ------------------------------------------------------------------ + + def query( + self, + query: str | list[str], + *, + top_k: int, + collection_name: str | None = None, + ) -> list[list[dict[str, Any]]] | list[QueryHit]: + """Search ingested documents through ``POST /v1/query``. + + Note: + ``top_k`` is required here but defaults to 10 on :meth:`aquery`. + That asymmetry is part of the released signature; do not unify it. + """ + return self._run(self.aquery(query, top_k=top_k, collection_name=collection_name)) + + def _query_hit(self, hit: dict[str, Any]) -> QueryHit: + return self._model( + QueryHit, + { + **hit, + "bbox": hit.get("bbox_xyxy_norm"), + "source": hit.get("source"), + "metadata": hit.get("metadata") or {}, + }, + "collection query", + ) + + async def aquery( + self, + query: str | list[str], + *, + top_k: int = 10, + collection_name: str | None = None, + ) -> list[list[dict[str, Any]]] | list[QueryHit]: + """Asynchronously search through ``POST /v1/query``.""" + + payload: dict[str, Any] = {"query": query, "top_k": int(top_k)} + if collection_name: + payload["collection_name"] = collection_name + body = await self._arequest("POST", "/v1/query", json=payload) try: - query_response = QueryResponse.model_validate(body) - return query_response.hits_by_query(expected_results=expected_results) + parsed = QueryResponse.model_validate(body).hits_by_query( + expected_results=len(query) if isinstance(query, list) else 1 + ) except (ValidationError, ValueError) as exc: - raise RuntimeError(f"Service query returned invalid response: {exc}") from exc + raise RetrieverServiceError(f"Service query returned invalid response: {exc}") from exc + if collection_name and isinstance(query, str): + return [self._query_hit(hit) for hit in parsed[0]] + return parsed # ------------------------------------------------------------------ # Job lifecycle @@ -256,6 +608,11 @@ async def _create_job( expected_documents: int, label: str | None = None, retain_results: bool = False, + collection_name: str | None = None, + operation: str = "append", + target_document_id: str | None = None, + idempotency_key: str | None = None, + document_manifest: list[dict[str, str]] | None = None, ) -> _CreatedJob: """Open a server-side job aggregate and return its client-visible metadata. @@ -270,7 +627,19 @@ async def _create_job( } if label is not None: payload["label"] = label - resp = await client.post(url, json=payload) + if collection_name is not None: + payload["collection_name"] = collection_name + payload["operation"] = operation + if target_document_id is not None: + payload["target_document_id"] = target_document_id + if idempotency_key is not None: + payload["idempotency_key"] = idempotency_key + if document_manifest: + payload["document_manifest"] = document_manifest + try: + resp = await client.post(url, json=payload) + except httpx.HTTPError as exc: + raise RetrieverServiceError(f"Job creation transport failure: {exc}") from exc # A 404/410 here means the deployed service does not advertise # the job-scoped ingest API. Surface a dedicated compatibility # error instead of a generic HTTPStatusError so callers see one @@ -284,14 +653,11 @@ async def _create_job( body=resp.text if resp.text else "", ) ) - if resp.status_code >= 400: - detail = resp.text[:500] if resp.text else "(empty)" - raise httpx.HTTPStatusError( - f"Job creation failed: HTTP {resp.status_code}: {detail}", - request=resp.request, - response=resp, - ) - body = resp.json() + self._raise_for_response(resp, "create ingestion job") + try: + body = resp.json() + except ValueError as exc: + raise RetrieverServiceError("Job creation returned malformed JSON") from exc job_id = body.get("job_id") if not job_id: raise RuntimeError(f"Job creation returned no job_id: {body!r}") @@ -301,6 +667,123 @@ async def _create_job( trace_id=trace_id if isinstance(trace_id, str) and trace_id else None, ) + async def asubmit_documents( + self, + collection_name: str, + files: list[str | Path], + *, + idempotency_key: str | None = None, + operation: str = "append", + target_document_id: str | None = None, + metadata: dict[str, Any] | None = None, + pipeline_spec: dict[str, Any] | None = None, + ) -> JobAggregateResponse: + """Create a job and upload its files, returning after acceptance. + + Processing continues on the service. Poll :meth:`get_job` or + :meth:`aget_job` until the aggregate reaches a terminal state. + """ + paths = [Path(file) for file in files] + if not paths: + raise RetrieverServiceValidationError("At least one file is required") + timeout = httpx.Timeout(timeout=None, connect=30.0) + limits = httpx.Limits(max_connections=200, max_keepalive_connections=100) + manifest = [] + for position, path in enumerate(paths): + content_sha256 = hashlib.sha256(path.read_bytes()).hexdigest() + manifest_entry_id = hashlib.sha256(f"{position}\0{path.name}\0{content_sha256}".encode("utf-8")).hexdigest() + manifest.append( + { + "manifest_entry_id": manifest_entry_id, + "filename": path.name, + "content_sha256": content_sha256, + } + ) + async with httpx.AsyncClient(timeout=timeout, limits=limits, headers=self._auth_headers) as client: + created = await self._create_job( + client, + expected_documents=len(paths), + collection_name=collection_name, + operation=operation, + target_document_id=target_document_id, + idempotency_key=idempotency_key, + document_manifest=manifest, + ) + sem = asyncio.Semaphore(self._max_concurrency) + + async def upload(path: Path, entry: dict[str, str]) -> None: + async with sem: + await self._upload_one( + client, + path, + job_id=created.job_id, + metadata=metadata, + pipeline_spec=pipeline_spec, + manifest_entry_id=entry["manifest_entry_id"], + ) + + await asyncio.gather(*(upload(path, entry) for path, entry in zip(paths, manifest, strict=True))) + resp = await client.get(f"{self._base_url}/v1/ingest/job/{created.job_id}") + self._raise_for_response(resp, "get accepted job") + try: + payload = resp.json() + except ValueError as exc: + raise RetrieverServiceError("Accepted job response was malformed JSON") from exc + return self._model(JobAggregateResponse, payload, "get accepted job") + + def submit_documents( + self, + collection_name: str, + files: list[str | Path], + **kwargs: Any, + ) -> JobAggregateResponse: + """Synchronous accepted-not-completed collection ingestion.""" + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(self.asubmit_documents(collection_name, files, **kwargs)) + raise RuntimeError("submit_documents cannot run inside an event loop; use asubmit_documents") + + def replace_document( + self, + collection_name: str, + document_id: str, + file: str | Path, + *, + idempotency_key: str | None = None, + **kwargs: Any, + ) -> JobAggregateResponse: + """Replace one stable document with a newly ingested file.""" + + return self.submit_documents( + collection_name, + [file], + operation="replace", + target_document_id=document_id, + idempotency_key=idempotency_key, + **kwargs, + ) + + async def areplace_document( + self, + collection_name: str, + document_id: str, + file: str | Path, + *, + idempotency_key: str | None = None, + **kwargs: Any, + ) -> JobAggregateResponse: + """Asynchronously replace one stable document.""" + + return await self.asubmit_documents( + collection_name, + [file], + operation="replace", + target_document_id=document_id, + idempotency_key=idempotency_key, + **kwargs, + ) + # ------------------------------------------------------------------ # Upload # ------------------------------------------------------------------ @@ -308,11 +791,12 @@ async def _create_job( async def _upload_one( self, client: httpx.AsyncClient, - file_path: Path, + source: UploadInput, *, job_id: str, metadata: dict[str, Any] | None = None, pipeline_spec: dict[str, Any] | None = None, + manifest_entry_id: str | None = None, ) -> dict[str, Any]: """Upload a file under an existing job, with retry on 429 + transient errors. @@ -321,9 +805,19 @@ async def _upload_one( key so the server can validate and apply it. Returns the parsed JSON response (contains ``document_id`` and ``job_id``). """ - file_bytes = file_path.read_bytes() - filename = file_path.name + if isinstance(source, Path): + file_bytes = source.read_bytes() + filename = source.name + content_type = "application/octet-stream" + classification_filename = None + else: + file_bytes = source.content + filename = source.filename + content_type = source.content_type + classification_filename = source.classification_filename meta_payload: dict[str, Any] = dict(metadata or {}) + if classification_filename is not None: + meta_payload.setdefault("filename", classification_filename) if pipeline_spec is not None: meta_payload["pipeline"] = pipeline_spec meta_json = json.dumps(meta_payload) @@ -334,15 +828,23 @@ async def _upload_one( try: resp = await client.post( url, - files={"file": (filename, file_bytes, "application/octet-stream")}, - data={"metadata": meta_json}, + files={"file": (filename, file_bytes, content_type)}, + data={ + "metadata": meta_json, + **({"manifest_entry_id": manifest_entry_id} if manifest_entry_id else {}), + }, ) except _TRANSIENT_ERRORS as exc: transport_attempts += 1 if transport_attempts > 5: - raise + raise RetrieverServiceError(f"Upload of {filename} transport failure after retries: {exc}") from exc delay = min(_DEFAULT_RETRY_AFTER * (2 ** (transport_attempts - 1)), 60.0) - logger.debug("Transient %s uploading %s, retry in %.1fs", type(exc).__name__, filename, delay) + logger.debug( + "Transient %s uploading %s, retry in %.1fs", + type(exc).__name__, + filename, + delay, + ) await asyncio.sleep(delay) continue @@ -367,17 +869,36 @@ async def _upload_one( ) ) - if resp.status_code >= 400: - detail = resp.text[:500] if resp.text else "(empty)" - raise httpx.HTTPStatusError( - f"Upload of {filename} returned HTTP {resp.status_code}: {detail}", - request=resp.request, - response=resp, - ) + self._raise_for_response(resp, f"upload {filename}") + try: + return resp.json() + except ValueError as exc: + raise RetrieverServiceError(f"Upload of {filename} returned malformed JSON") from exc - return resp.json() + raise RetrieverServiceError(f"Upload of {filename} failed after {_MAX_UPLOAD_RETRIES} retries") - raise RuntimeError(f"Upload of {filename} failed after {_MAX_UPLOAD_RETRIES} retries") + async def _upload_source( + self, + client: httpx.AsyncClient, + source: UploadInput, + *, + job_id: str, + pipeline_spec: dict[str, Any] | None, + semaphore: asyncio.Semaphore, + ) -> _UploadOutcome: + """Upload one source and capture the common success or failure outcome.""" + filename = _upload_filename(source) + async with semaphore: + try: + response = await self._upload_one(client, source, job_id=job_id, pipeline_spec=pipeline_spec) + except Exception as exc: + logger.error("Upload failed for %s: %s", filename, exc) + return _UploadOutcome(filename=filename, error=str(exc)) + return _UploadOutcome( + filename=filename, + document_id=response.get("document_id", ""), + attempt_id=response.get("attempt_id", ""), + ) # ------------------------------------------------------------------ # SSE consumer @@ -498,14 +1019,24 @@ def _is_done() -> bool: if _is_done(): break elif line.startswith(":"): - if _is_done(): + if uploads_done.is_set(): + _reconcile() + if pending: + logger.info( + "SSE keepalive arrived with %d items pending " + "after uploads completed; switching to bulk poll", + len(pending), + ) break except Exception as exc: logger.warning("SSE stream error: %s: %s", type(exc).__name__, exc) if pending: - logger.info("SSE closed with %d items pending — falling back to bulk poll", len(pending)) + logger.info( + "SSE closed with %d items pending — falling back to bulk poll", + len(pending), + ) await self._bulk_poll_fallback(client, pending, tracker, on_event) # ------------------------------------------------------------------ @@ -552,7 +1083,11 @@ async def _bulk_poll_fallback( status = info.get("status", "") if status in ("completed", "failed"): pending.discard(doc_id) - event = {"id": doc_id, "status": status, "result_rows": info.get("result_rows", 0)} + event = { + "id": doc_id, + "status": status, + "result_rows": info.get("result_rows", 0), + } error_msg = info.get("error") if error_msg: event["error"] = error_msg @@ -572,7 +1107,7 @@ async def _bulk_poll_fallback( async def ingest_documents( self, - files: list[Path], + files: list[UploadInput], *, on_file_submitted: Callable[[str, str], Any] | None = None, show_progress: bool = True, @@ -615,19 +1150,21 @@ async def ingest_documents( upload_sem = asyncio.Semaphore(self._max_concurrency) upload_failures: list[tuple[str, str]] = [] - async def _upload_one_file(fpath: Path) -> None: - async with upload_sem: - try: - resp_json = await self._upload_one(client, fpath, job_id=job_id, pipeline_spec=pipeline_spec) - doc_id = resp_json.get("document_id", "") - if doc_id: - pending.add(doc_id) - document_ids.append(doc_id) - if on_file_submitted: - on_file_submitted(fpath.name, doc_id) - except Exception as exc: - upload_failures.append((fpath.name, str(exc))) - logger.error("Upload failed for %s: %s", fpath.name, exc) + async def _upload_one_file(source: UploadInput) -> None: + outcome = await self._upload_source( + client, + source, + job_id=job_id, + pipeline_spec=pipeline_spec, + semaphore=upload_sem, + ) + if outcome.error is not None: + upload_failures.append((outcome.filename, outcome.error)) + elif outcome.document_id: + pending.add(outcome.attempt_id or outcome.document_id) + document_ids.append(outcome.document_id) + if on_file_submitted: + on_file_submitted(outcome.filename, outcome.document_id) progress_ctx = _make_progress() if show_progress else None @@ -684,7 +1221,7 @@ async def _upload_all() -> None: async def aingest_documents_stream( self, - files: list[Path], + files: list[UploadInput], *, pipeline_spec: dict[str, Any] | None = None, retain_results: bool = False, @@ -730,31 +1267,33 @@ async def aingest_documents_stream( yield event upload_sem = asyncio.Semaphore(self._max_concurrency) - async def _upload_one_file(fpath: Path) -> None: - async with upload_sem: - try: - resp_json = await self._upload_one(client, fpath, job_id=job_id, pipeline_spec=pipeline_spec) - doc_id = resp_json.get("document_id", "") - if doc_id: - pending.add(doc_id) - await event_queue.put( - { - "event": "upload_complete", - "filename": fpath.name, - "document_id": doc_id, - "job_id": job_id, - } - ) - except Exception as exc: - logger.error("Upload failed for %s: %s", fpath.name, exc) - await event_queue.put( - { - "event": "upload_failed", - "filename": fpath.name, - "error": str(exc), - "job_id": job_id, - } - ) + async def _upload_one_file(source: UploadInput) -> None: + outcome = await self._upload_source( + client, + source, + job_id=job_id, + pipeline_spec=pipeline_spec, + semaphore=upload_sem, + ) + if outcome.error is not None: + await event_queue.put( + { + "event": "upload_failed", + "filename": outcome.filename, + "error": outcome.error, + "job_id": job_id, + } + ) + elif outcome.document_id: + pending.add(outcome.attempt_id or outcome.document_id) + await event_queue.put( + { + "event": "upload_complete", + "filename": outcome.filename, + "document_id": outcome.document_id, + "job_id": job_id, + } + ) async def _upload_all() -> None: tasks = [asyncio.create_task(_upload_one_file(f)) for f in files] diff --git a/nemo_retriever/src/nemo_retriever/service/config.py b/nemo_retriever/src/nemo_retriever/service/config.py index 4d729de66d..e68a72c265 100644 --- a/nemo_retriever/src/nemo_retriever/service/config.py +++ b/nemo_retriever/src/nemo_retriever/service/config.py @@ -7,6 +7,7 @@ from __future__ import annotations from importlib import resources as importlib_resources +import os from pathlib import Path from typing import Any, Literal @@ -16,6 +17,7 @@ from nemo_retriever.common.schemas.base import RichModel ServiceMode = Literal["standalone", "gateway", "realtime", "batch"] +MCPQueryMethods = Literal["classic", "agentic", "all"] class ServerConfig(RichModel): @@ -78,6 +80,24 @@ class LocalAsrConfig(RichModel): enabled: bool = True +class LocalRerankConfig(RichModel): + """In-pod reranker used by the main service query API. + + This is deliberately separate from the ingestion process-pool settings: + query-time reranking lives in the main service process and is loaded lazily + on the first ``rerank=true`` request. + """ + + model_config = ConfigDict(extra="forbid") + + enabled: bool = False + model_name: str = "nvidia/llama-nemotron-rerank-1b-v2" + backend: Literal["hf", "vllm"] = "vllm" + gpu_memory_utilization: float = Field(default=0.5, gt=0, le=1) + max_length: int = Field(default=512, ge=1, le=8192) + batch_size: int = Field(default=32, ge=1) + + class LocalModelsConfig(RichModel): """Load Nemotron Hugging Face weights inside the service worker pod. @@ -121,6 +141,7 @@ class LocalModelsConfig(RichModel): extract: LocalExtractConfig = Field(default_factory=LocalExtractConfig) embed: LocalEmbedConfig = Field(default_factory=LocalEmbedConfig) asr: LocalAsrConfig = Field(default_factory=LocalAsrConfig) + rerank: LocalRerankConfig = Field(default_factory=LocalRerankConfig) class NimEndpointsConfig(RichModel): @@ -131,6 +152,23 @@ class NimEndpointsConfig(RichModel): page_elements_invoke_url: str | None = None ocr_invoke_url: str | None = None table_structure_invoke_url: str | None = None + nemotron_parse_invoke_url: str | None = Field( + default=None, + description=( + "Remote Nemotron Parse chat-completions endpoint. When set, " + "service-mode requests using method='nemotron_parse' call this " + "endpoint instead of loading the local Parse model." + ), + ) + nemotron_parse_model: str | None = Field( + default=None, + description=( + "Model identifier passed to the remote Nemotron Parse endpoint. " + "Use nvidia/nemotron-parse for NVIDIA-hosted inference and " + "nvidia/nemotron-parse-v1.2 for a self-hosted NIM. " + "Server-owned — clients cannot override the deployed Parse SKU." + ), + ) embed_invoke_url: str | None = None embed_model_name: str | None = Field( default=None, @@ -146,7 +184,20 @@ class NimEndpointsConfig(RichModel): "remote embedding endpoints that require namespaced model IDs." ), ) - rerank_invoke_url: str | None = None + rerank_invoke_url: str | None = Field( + default=None, + description=( + "Remote reranking endpoint used by the main service for /v1/query " + "requests with rerank=true. The endpoint, model, and API key are " + "server-owned." + ), + ) + rerank_model_name: str | None = Field( + default=None, + description=( + "Model identifier passed to rerank_invoke_url. Defaults to the " "Nemotron text reranker when omitted." + ), + ) audio_grpc_endpoint: str | None = Field( default=None, description=( @@ -173,6 +224,22 @@ class NimEndpointsConfig(RichModel): ) api_key: str | None = None + @model_validator(mode="after") + def _validate_nemotron_parse_config(self) -> "NimEndpointsConfig": + endpoint = (self.nemotron_parse_invoke_url or "").strip() + model = (self.nemotron_parse_model or "").strip() + self.nemotron_parse_invoke_url = endpoint or None + self.nemotron_parse_model = model or None + if model and not endpoint: + raise ValueError("nim_endpoints.nemotron_parse_model requires " "nim_endpoints.nemotron_parse_invoke_url") + rerank_endpoint = (self.rerank_invoke_url or "").strip() + rerank_model = (self.rerank_model_name or "").strip() + self.rerank_invoke_url = rerank_endpoint or None + self.rerank_model_name = rerank_model or None + if rerank_model and not rerank_endpoint: + raise ValueError("nim_endpoints.rerank_model_name requires " "nim_endpoints.rerank_invoke_url") + return self + class LLMConfig(RichModel): """Remote LLM configuration for service-mode RAG answer generation.""" @@ -200,6 +267,30 @@ def _validate_enabled_model(self) -> "LLMConfig": return self +class AgenticConfig(RichModel): + """Server-owned configuration for agentic (ReAct) retrieval queries.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool = False + llm_model: str | None = None + invoke_url: str | None = None + reasoning_effort: str | None = "high" + backend_top_k: int = Field(default=20, ge=1) + react_max_steps: int = Field(default=50, ge=1) + text_truncation: int = Field(default=0, ge=0) + temperature: float = Field(default=0.0, ge=0.0) + request_timeout_s: float = Field(default=1800.0, gt=0) + + @model_validator(mode="after") + def _validate_remote_model(self) -> "AgenticConfig": + if self.enabled and not (self.invoke_url or "").strip(): + raise ValueError("agentic.invoke_url must be set when agentic.enabled is true") + if self.enabled and not (self.llm_model or "").strip(): + raise ValueError("agentic.llm_model must be set when agentic.enabled is true") + return self + + class ResourceLimitsConfig(RichModel): model_config = ConfigDict(extra="forbid") @@ -214,11 +305,15 @@ class ResourceLimitsConfig(RichModel): class AuthConfig(RichModel): - """Optional bearer-token authentication.""" + """Bearer authentication and authorization for logical workspace scopes.""" model_config = ConfigDict(extra="forbid") + enabled: bool = False api_token: str | None = None + default_scope: str = "default" + scope_token_file: str | None = None + allow_unscoped_dev: bool = False header_name: str = "Authorization" bypass_paths: list[str] = Field(default_factory=lambda: ["/v1/health", "/docs", "/openapi.json", "/redoc"]) @@ -238,6 +333,14 @@ class MCPConfig(RichModel): ), ) enable_write_tools: bool = True + query_methods: MCPQueryMethods = Field( + default="classic", + description=( + "Which retrieval MCP tools to register: 'classic' (query only), " + "'agentic' (agentic_query only), or 'all' (both). Agentic tools are " + "still omitted when agentic.enabled is false." + ), + ) max_concurrency: int = Field(default=8, ge=1) request_timeout_s: float = Field(default=60.0, gt=0) ingest_timeout_s: float = Field(default=1800.0, gt=0) @@ -328,6 +431,16 @@ class VectorDbConfig(RichModel): default="http://nemo-retriever-vectordb:7671", description="URL of the vectordb service (for workers to POST embeddings to)", ) + internal_api_token: str | None = Field( + default=None, + description="Dedicated gateway/worker credential for the VectorDB service.", + ) + reconciliation_interval_seconds: int = Field( + default=60, + ge=0, + description="Local lifecycle reconciliation interval; zero disables the loop.", + ) + expiration_cleanup_enabled: bool = True class SinksConfig(RichModel): @@ -453,6 +566,7 @@ class ServiceConfig(RichModel): nim_endpoints: NimEndpointsConfig = Field(default_factory=NimEndpointsConfig) local_models: LocalModelsConfig = Field(default_factory=LocalModelsConfig) llm: LLMConfig = Field(default_factory=LLMConfig) + agentic: AgenticConfig = Field(default_factory=AgenticConfig) resources: ResourceLimitsConfig = Field(default_factory=ResourceLimitsConfig) auth: AuthConfig = Field(default_factory=AuthConfig) mcp: MCPConfig = Field(default_factory=MCPConfig) @@ -525,7 +639,7 @@ def load_config( """Load a :class:`ServiceConfig` from YAML with optional CLI overrides.""" path = _discover_config_path(config_path) if path is not None: - raw: dict[str, Any] = yaml.safe_load(path.read_text()) or {} + raw: dict[str, Any] = yaml.safe_load(os.path.expandvars(path.read_text())) or {} else: raw = {} @@ -539,9 +653,21 @@ def load_config( target = target.setdefault(part, {}) target[parts[-1]] = value + # Secret-backed runtime values intentionally bypass ConfigMaps and the + # rendered configuration tree. + if scope_file := os.environ.get("NRL_SCOPE_TOKEN_FILE"): + raw.setdefault("auth", {})["scope_token_file"] = scope_file + internal_token = os.environ.get("NRL_INTERNAL_VDB_TOKEN") + if not internal_token and (internal_token_file := os.environ.get("NRL_INTERNAL_VDB_TOKEN_FILE")): + internal_token = Path(internal_token_file).read_text(encoding="utf-8").strip() + if internal_token: + internal_token = internal_token.strip() + if internal_token: + raw.setdefault("vectordb", {})["internal_api_token"] = internal_token + config = ServiceConfig(**raw) - _REDACTED_FIELDS = frozenset({"api_key", "api_token", "password", "secret"}) + _REDACTED_FIELDS = frozenset({"api_key", "api_token", "internal_api_token", "password", "secret"}) from rich.console import Console from rich.tree import Tree diff --git a/nemo_retriever/src/nemo_retriever/service/errors.py b/nemo_retriever/src/nemo_retriever/service/errors.py new file mode 100644 index 0000000000..e5406db917 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/service/errors.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed errors raised by the public retriever service client.""" + + +class RetrieverServiceError(RuntimeError): + """Base error raised by the SDK for retriever service failures.""" + + def __init__(self, message: str, *, status_code: int | None = None) -> None: + super().__init__(message) + self.status_code = status_code + + +class RetrieverServiceNotFoundError(RetrieverServiceError): + """The requested scoped resource does not exist.""" + + +class RetrieverServiceConflictError(RetrieverServiceError): + """The request conflicts with current state or idempotency history.""" + + +class RetrieverServiceValidationError(RetrieverServiceError): + """The service rejected invalid input.""" diff --git a/nemo_retriever/src/nemo_retriever/service/mcp_server.py b/nemo_retriever/src/nemo_retriever/service/mcp_server.py index 4f3f6498ea..60b05c1a6c 100644 --- a/nemo_retriever/src/nemo_retriever/service/mcp_server.py +++ b/nemo_retriever/src/nemo_retriever/service/mcp_server.py @@ -25,7 +25,7 @@ from fastmcp import FastMCP from pydantic import BaseModel, ConfigDict, Field, model_validator -from nemo_retriever.service.config import ServiceConfig +from nemo_retriever.service.config import MCPQueryMethods, ServiceConfig from nemo_retriever.service.query_schema import QueryFormat logger = logging.getLogger(__name__) @@ -77,9 +77,11 @@ class ServiceMCPSettings: auth_header_name: str = "Authorization" max_concurrency: int = 8 request_timeout_s: float = 60.0 + agentic_request_timeout_s: float = 1800.0 ingest_timeout_s: float = 1800.0 poll_interval_s: float = 2.0 enable_write_tools: bool = True + query_methods: MCPQueryMethods = "classic" @property def normalized_base_url(self) -> str: @@ -92,6 +94,29 @@ def auth_headers(self) -> dict[str, str]: return {} return {self.auth_header_name: f"Bearer {token}"} + @property + def enable_classic_query(self) -> bool: + return self.query_methods in ("classic", "all") + + @property + def enable_agentic_query(self) -> bool: + return self.query_methods in ("agentic", "all") + + +def _effective_query_methods( + configured: MCPQueryMethods, + *, + agentic_enabled: bool, +) -> MCPQueryMethods: + """Drop agentic MCP tools when agentic retrieval is not configured.""" + if agentic_enabled: + return configured + if configured == "agentic": + return "classic" + if configured == "all": + return "classic" + return configured + def settings_from_service_config(config: ServiceConfig) -> ServiceMCPSettings: """Build MCP settings for the MCP app mounted inside ``retriever service start``.""" @@ -111,9 +136,14 @@ def settings_from_service_config(config: ServiceConfig) -> ServiceMCPSettings: auth_header_name=config.auth.header_name, max_concurrency=mcp_cfg.max_concurrency, request_timeout_s=mcp_cfg.request_timeout_s, + agentic_request_timeout_s=config.agentic.request_timeout_s, ingest_timeout_s=mcp_cfg.ingest_timeout_s, poll_interval_s=mcp_cfg.poll_interval_s, enable_write_tools=mcp_cfg.enable_write_tools, + query_methods=_effective_query_methods( + mcp_cfg.query_methods, + agentic_enabled=config.agentic.enabled, + ), ) @@ -201,17 +231,46 @@ async def query( *, top_k: int = 5, format: QueryFormat = "hits", + rerank: bool = False, + rerank_top_k: int | None = None, payload: dict[str, Any] | None = None, ) -> dict[str, Any]: body = dict(payload or {}) - body.setdefault("query", query) - body.setdefault("top_k", top_k) - body.setdefault("format", format) + # Typed tool arguments are authoritative; payload is only for additional + # service options such as filters. In particular, it must not enable + # reranking when the explicit rerank argument is false. + for key in ("query", "top_k", "format", "agentic", "rerank", "rerank_top_k"): + body.pop(key, None) + body.update({"query": query, "top_k": top_k, "format": format}) + if rerank: + body["rerank"] = True + if rerank_top_k is not None: + body["rerank_top_k"] = rerank_top_k async with self._client() as client: resp = await client.post("/v1/query", json=body) self._raise_for_status(resp) return dict(self._json_or_text(resp)) + async def agentic_query( + self, + query: str, + *, + top_k: int = 5, + ) -> dict[str, Any]: + """Call ``POST /v1/query`` with ``agentic=true`` (long timeout).""" + async with self._client(timeout_s=self._settings.agentic_request_timeout_s) as client: + resp = await client.post( + "/v1/query", + json={ + "query": query, + "top_k": top_k, + "format": "hits", + "agentic": True, + }, + ) + self._raise_for_status(resp) + return dict(self._json_or_text(resp)) + async def answer( self, query: str, @@ -448,7 +507,8 @@ def build_mcp(settings: ServiceMCPSettings | None = None) -> FastMCP: instructions=( "Use these tools to interact with a running NVIDIA NeMo Retriever " "service. Ingest documents, check job status, query the configured " - "VectorDB, and ask the configured answer-generation endpoint." + "VectorDB, run agentic retrieval when configured, and ask the " + "configured answer-generation endpoint." ), ) @@ -486,22 +546,48 @@ async def list_job_documents( async def get_document(job_id: str, document_id: str) -> dict[str, Any]: return await service.get_document(job_id, document_id) - @mcp.tool( - name="query", - description=( - "Search ingested documents through the service VectorDB endpoint. " - "format='hits' (default) returns raw retrieval hits; format='evidence' " - "returns the fidelity-tagged, citation-ready {evidence, coverage} shape. " - "Retrieval (dense vs hybrid) is auto-detected from the table's own indexes." - ), - ) - async def query( - query: str, - top_k: int = 5, - format: QueryFormat = "hits", - payload: dict[str, Any] | None = None, - ) -> dict[str, Any]: - return await service.query(query, top_k=top_k, format=format, payload=payload) + if settings.enable_classic_query: + + @mcp.tool( + name="query", + description=( + "Search ingested documents through the service VectorDB endpoint. " + "format='hits' (default) returns raw retrieval hits; format='evidence' " + "returns the fidelity-tagged, citation-ready {evidence, coverage} shape. " + "Retrieval (dense vs hybrid) is auto-detected from the table's own indexes. " + "Set rerank=true to rerank candidates through the configured service endpoint." + ), + ) + async def query( + query: str, + top_k: int = 5, + format: QueryFormat = "hits", + rerank: bool = False, + rerank_top_k: int | None = None, + payload: dict[str, Any] | None = None, + ) -> dict[str, Any]: + return await service.query( + query, + top_k=top_k, + format=format, + rerank=rerank, + rerank_top_k=rerank_top_k, + payload=payload, + ) + + if settings.enable_agentic_query: + + @mcp.tool( + name="agentic_query", + description=( + "Run the configured agentic (ReAct) retrieval workflow over ingested " + "documents via POST /v1/query with agentic=true. Returns the standard " + "hits envelope: source holds doc_id; result_source and rank are under " + "metadata; chunk-level fields are unset for document-level results." + ), + ) + async def agentic_query(query: str, top_k: int = 5) -> dict[str, Any]: + return await service.agentic_query(query, top_k=top_k) @mcp.tool(name="answer", description="Search ingested documents and generate an answer.") async def answer( diff --git a/nemo_retriever/src/nemo_retriever/service/query_schema.py b/nemo_retriever/src/nemo_retriever/service/query_schema.py index b921ab814c..29c850c116 100644 --- a/nemo_retriever/src/nemo_retriever/service/query_schema.py +++ b/nemo_retriever/src/nemo_retriever/service/query_schema.py @@ -6,21 +6,97 @@ from typing import Any, Literal -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator QueryFormat = Literal["hits", "evidence"] +QueryMode = Literal["classic", "agentic"] + +# Agentic queries are replayed into every step of a multi-step LLM loop, so an +# oversized query multiplies prompt cost and latency. Roughly 1k tokens of +# natural-language question is far above any realistic retrieval query. +MAX_AGENTIC_QUERY_CHARS = 4096 class QueryRequest(BaseModel): query: str | list[str] top_k: int = Field(default=10, ge=1, le=1000) + collection_name: str | None = Field(default=None, min_length=1, max_length=128) format: QueryFormat = Field( default="hits", description=( "Output shape: 'hits' (default) returns raw retrieval hits; 'evidence' " - "returns the fidelity-tagged, citation-ready {evidence, coverage} shape." + "returns the fidelity-tagged, citation-ready {evidence, coverage} shape. " + "Agentic queries require format='hits'." + ), + ) + agentic: bool = Field( + default=False, + description=( + "When true, run the server-configured agentic (ReAct) retrieval workflow. " + "Requires agentic.enabled in service configuration. Response uses the same " + "hits envelope as dense/hybrid query; document-level agentic results map " + "doc_id onto source, keep result_source/rank in metadata, and leave " + "chunk-level fields unset (null)." + ), + ) + + rerank: bool = Field( + default=False, + description=( + "When true, retrieve a larger candidate set from VectorDB, then rerank " + "it through the server-configured reranker before returning top_k hits." ), ) + rerank_top_k: int | None = Field( + default=None, + ge=1, + le=1000, + description=( + "Number of VectorDB candidates to retrieve before reranking. Defaults " + "to max(top_k, 50) when rerank is enabled." + ), + ) + + @model_validator(mode="after") + def _validate_agentic_request(self) -> "QueryRequest": + if self.rerank: + if self.agentic: + raise ValueError("rerank cannot be combined with agentic queries") + if self.format != "hits": + raise ValueError("rerank queries require format='hits'") + if self.rerank_top_k is not None and self.rerank_top_k < self.top_k: + raise ValueError("rerank_top_k must be greater than or equal to top_k") + if not self.agentic: + return self + if not isinstance(self.query, str): + raise ValueError("agentic queries require a single query string, not a list") + if not self.query.strip(): + raise ValueError("agentic query must be a non-empty string") + if len(self.query) > MAX_AGENTIC_QUERY_CHARS: + raise ValueError(f"agentic query exceeds max length of {MAX_AGENTIC_QUERY_CHARS} characters") + if self.format != "hits": + raise ValueError("agentic queries require format='hits'") + return self + + @model_validator(mode="before") + @classmethod + def _reject_raw_storage_keys(cls, value: Any) -> Any: + if isinstance(value, dict): + raw_keys = { + "table_name", + "table", + "physical_table", + "lancedb_uri", + "lance_uri", + "uri", + "table_path", + "database_uri", + "vdb_uri", + } + supplied = sorted(raw_keys.intersection(value)) + if supplied: + raise ValueError(f"client-selected storage is not supported: {', '.join(supplied)}") + return value class QueryResult(BaseModel): @@ -29,6 +105,13 @@ class QueryResult(BaseModel): class QueryResponse(BaseModel): results: list[QueryResult] + query_mode: QueryMode = Field( + default="classic", + description=( + "Which /v1/query workflow produced this response: 'classic' (dense/hybrid) " + "or 'agentic' (ReAct document ranking)." + ), + ) def hits_by_query(self, *, expected_results: int | None = None) -> list[list[dict[str, Any]]]: if expected_results is not None and len(self.results) != expected_results: @@ -72,3 +155,7 @@ class EvidenceResult(BaseModel): class EvidenceQueryResponse(BaseModel): results: list[EvidenceResult] + query_mode: QueryMode = Field( + default="classic", + description="Evidence format is classic retrieval only; always 'classic'.", + ) diff --git a/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml b/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml index c464c6aff6..d8b258b332 100644 --- a/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml +++ b/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml @@ -33,12 +33,23 @@ nim_endpoints: page_elements_invoke_url: null table_structure_invoke_url: null ocr_invoke_url: null + # Remote Nemotron Parse chat-completions endpoint and model. Both are + # server-owned; service clients select method="nemotron_parse" but cannot + # redirect the endpoint or change the deployed model. + nemotron_parse_invoke_url: null + nemotron_parse_model: null embed_invoke_url: null # Model name for the remote embed NIM (server-owned; must match the SKU). embed_model_name: null # Optional LiteLLM provider prefix prepended to embed_model_name for # proxies that require provider/model IDs. embed_model_provider_prefix: null + # Optional remote reranker used by POST /v1/query with rerank=true. The + # endpoint can be an in-cluster NIM, a vLLM server, or another compatible + # ranking API. This service config does not deploy or manage that endpoint. + rerank_invoke_url: null + # Optional server-owned model ID sent to the reranking endpoint. + rerank_model_name: null # gRPC endpoint for the Parakeet ASR NIM (e.g. parakeet-nim:50051). # When set, audio/video pipelines use remote ASR instead of loading # the local Parakeet model (which requires torch + GPU). @@ -74,6 +85,17 @@ local_models: gpu_memory_utilization: 0.45 asr: enabled: true + # Query-time reranking is loaded lazily in the main service process when a + # client sends ``rerank: true``. It is independent of ``local_models.enabled`` + # because it does not run in the ingestion process pools. Leave this disabled + # when using a remote/NIM reranker configured under ``nim_endpoints``. + rerank: + enabled: false + model_name: nvidia/llama-nemotron-rerank-1b-v2 + backend: vllm + gpu_memory_utilization: 0.5 + max_length: 512 + batch_size: 32 # Remote LLM endpoint used by POST /v1/answer. Helm auto-wires these # fields when the Super-49B NIM is enabled. @@ -92,6 +114,21 @@ llm: rag_system_prompt_prefix: null reasoning_enabled: true +# Agentic (ReAct) retrieval exposed through POST /v1/query with agentic=true. +# The agentic_query MCP tool is registered when mcp.query_methods is agentic or all +# and agentic.enabled is true. Service mode requires a remote OpenAI-compatible +# invoke_url and llm_model when enabled. +agentic: + enabled: false + llm_model: null + invoke_url: null + reasoning_effort: high + backend_top_k: 20 + react_max_steps: 50 + text_truncation: 0 + temperature: 0.0 + request_timeout_s: 1800.0 + # Pipeline worker pools. Workers are abstract dispatchers — sizing # depends on whether they do local GPU work or fan out to remote NIMs. # For CPU-only NIM-forwarding nodes, higher worker counts are fine. @@ -143,9 +180,10 @@ gateway: timeout_s: 300.0 max_connections: 100 -# Optional bearer-token authentication. When api_token is set, every -# request must carry "Authorization: Bearer ". +# Optional bearer-token authentication. When enabled, every request must carry +# "Authorization: Bearer " unless it is a bypass path. auth: + enabled: false api_token: null header_name: "Authorization" bypass_paths: @@ -163,6 +201,9 @@ mcp: # Defaults to loopback on server.port when null. base_url: null enable_write_tools: true + # Which retrieval tools to register: classic | agentic | all. + # Agentic tools are omitted when agentic.enabled is false. + query_methods: classic max_concurrency: 8 request_timeout_s: 60.0 ingest_timeout_s: 1800.0 diff --git a/nemo_retriever/src/nemo_retriever/service/routers/collections.py b/nemo_retriever/src/nemo_retriever/service/routers/collections.py new file mode 100644 index 0000000000..7831a2ac74 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/service/routers/collections.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public collection/document lifecycle routes. + +The gateway owns authentication and forwards only logical resource names and +the authenticated scope. LanceDB locations and physical table names never +cross this boundary. +""" + +from __future__ import annotations + +import logging + +import httpx +from fastapi import APIRouter, HTTPException, Request, Response + +from nemo_retriever.common.schemas.collections import DocumentId + +router = APIRouter(tags=["collections"]) +logger = logging.getLogger(__name__) + + +async def _forward(request: Request, suffix: str) -> Response: + """Forward the authorized scope and logical request to the internal VectorDB service.""" + + from nemo_retriever.service.auth import authorized_scope, internal_auth_headers + + config = request.app.state.config + if not config.vectordb.enabled: + raise HTTPException(404, "VectorDB is not enabled in the service configuration.") + if config.mode in ("realtime", "batch"): + raise HTTPException(404, "Collection management is available through the gateway.") + + target = f"{config.vectordb.vectordb_url.rstrip('/')}/v1/{suffix}" + headers = {"Content-Type": request.headers.get("content-type", "application/json")} + headers["X-NRL-Scope"] = authorized_scope(request) + headers.update(internal_auth_headers(config.vectordb.internal_api_token)) + try: + async with httpx.AsyncClient(timeout=60.0) as client: + response = await client.request( + request.method, + target, + content=await request.body(), + params=request.query_params, + headers=headers, + ) + except httpx.HTTPError as exc: + logger.exception("Failed to proxy collection request to VectorDB at %s", target) + raise HTTPException(502, "VectorDB service is unavailable.") from exc + return Response( + content=response.content, + status_code=response.status_code, + media_type=response.headers.get("content-type", "application/json"), + ) + + +@router.get("/collections") +async def list_collections(request: Request) -> Response: + """Forward a collection list request to VectorDB.""" + return await _forward(request, "collections") + + +@router.post("/collections") +async def create_collection(request: Request) -> Response: + """Forward a collection creation request to VectorDB.""" + return await _forward(request, "collections") + + +@router.get("/collections/{collection_name}") +async def get_collection(request: Request, collection_name: str) -> Response: + """Forward a collection lookup request to VectorDB.""" + return await _forward(request, f"collections/{collection_name}") + + +@router.patch("/collections/{collection_name}") +async def update_collection(request: Request, collection_name: str) -> Response: + """Forward a collection update request to VectorDB.""" + return await _forward(request, f"collections/{collection_name}") + + +@router.delete("/collections/{collection_name}") +async def delete_collection(request: Request, collection_name: str) -> Response: + """Forward a collection deletion request to VectorDB.""" + return await _forward(request, f"collections/{collection_name}") + + +@router.get("/collections/{collection_name}/documents") +async def collection_documents(request: Request, collection_name: str) -> Response: + """Forward a collection document list request to VectorDB.""" + return await _forward(request, f"collections/{collection_name}/documents") + + +@router.get("/collections/{collection_name}/documents/{document_id}") +async def get_collection_document(request: Request, collection_name: str, document_id: DocumentId) -> Response: + """Forward a collection document lookup request to VectorDB.""" + return await _forward(request, f"collections/{collection_name}/documents/{document_id}") + + +@router.delete("/collections/{collection_name}/documents/{document_id}") +async def delete_collection_document(request: Request, collection_name: str, document_id: DocumentId) -> Response: + """Forward a collection document deletion request to VectorDB.""" + return await _forward(request, f"collections/{collection_name}/documents/{document_id}") diff --git a/nemo_retriever/src/nemo_retriever/service/routers/dashboard.py b/nemo_retriever/src/nemo_retriever/service/routers/dashboard.py index c4cf1b0014..6d0f52bc2a 100644 --- a/nemo_retriever/src/nemo_retriever/service/routers/dashboard.py +++ b/nemo_retriever/src/nemo_retriever/service/routers/dashboard.py @@ -19,6 +19,8 @@ from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from pydantic import BaseModel, Field +from nemo_retriever.service.auth import internal_auth_headers + logger = logging.getLogger(__name__) router = APIRouter(tags=["dashboard"]) @@ -515,6 +517,7 @@ async def vdb_query(req: VdbQueryRequest, request: Request) -> JSONResponse: resp = await client.post( f"{vdb_cfg.vectordb_url}/v1/query", json={"query": req.query, "top_k": req.top_k}, + headers=internal_auth_headers(vdb_cfg.internal_api_token), ) resp.raise_for_status() return JSONResponse(resp.json()) diff --git a/nemo_retriever/src/nemo_retriever/service/routers/ingest.py b/nemo_retriever/src/nemo_retriever/service/routers/ingest.py index 347414bc0b..21d5cced22 100644 --- a/nemo_retriever/src/nemo_retriever/service/routers/ingest.py +++ b/nemo_retriever/src/nemo_retriever/service/routers/ingest.py @@ -32,6 +32,7 @@ from pydantic import BaseModel, Field, model_validator from starlette.responses import StreamingResponse +from nemo_retriever.common.schemas.collections import IngestOperation from nemo_retriever.common.schemas.pipeline_spec import PipelineSpec from nemo_retriever.common.schemas.requests import IngestRequest, JobCreateRequest from nemo_retriever.common.schemas.responses import ( @@ -45,6 +46,7 @@ PageIngestAccepted, SidecarUploadResponse, ) +from nemo_retriever.service.query_schema import QueryRequest, QueryResponse from nemo_retriever.common.policy import PolicyError, validate_pipeline_spec from nemo_retriever.models.llm.types import ( AnswerRequest as CoreAnswerRequest, @@ -53,9 +55,15 @@ build_answer_result, ) from nemo_retriever.service.services.event_bus import get_event_bus -from nemo_retriever.service.services.job_tracker import MarkOutcome, get_job_tracker +from nemo_retriever.service.services.job_tracker import ( + DocumentRecord, + JobAggregate, + MarkOutcome, + get_job_tracker, +) from nemo_retriever.service.services.metrics import get_metrics from nemo_retriever.service.services.pipeline_pool import ( + DocumentWriteContext, PoolType, WorkItem, get_pipeline_pool, @@ -74,6 +82,7 @@ ) from nemo_retriever.service.utils.file_type import ( FileCategory, + FileClassification, FileClassifier, enforce_media_dependencies, ) @@ -81,11 +90,6 @@ _RETRY_AFTER_SECONDS = "5" _RESULT_RETRY_AFTER_SECONDS = 60 _DRY_RUN_HEADER = "X-Nemo-Dry-Run" -_GATEWAY_DOC_ID_HEADER = "X-Gateway-Document-Id" -_GATEWAY_CALLBACK_HEADER = "X-Gateway-Callback-Url" -_GATEWAY_PIPELINE_SPEC_HEADER = "X-Gateway-Pipeline-Spec" -_GATEWAY_JOB_ID_HEADER = "X-Gateway-Job-Id" -_GATEWAY_RETAIN_RESULTS_HEADER = "X-Gateway-Retain-Results" _PAGE_THRESHOLD_FOR_BATCH = 5 # SSE keepalive cadence; tests monkey-patch this to a short value so @@ -146,18 +150,12 @@ def _is_worker(request: Request) -> bool: """Return True for split-mode worker pods (``realtime`` or ``batch``). Workers don't own the ``JobTracker`` aggregate — the gateway does. - When the gateway forwards an upload to a worker, the URL still - contains the ``job_id``, but the worker must trust it (and not - re-validate via ``_require_job``). + They receive work by claiming it from the gateway broker rather than + over these routes. """ return _mode(request) in ("realtime", "batch") -def _retain_results_from_request(request: Request) -> bool: - val = request.headers.get(_GATEWAY_RETAIN_RESULTS_HEADER, "").strip().lower() - return val in ("1", "true", "yes") - - def _job_retain_results(job_id: str | None) -> bool: if not job_id: return False @@ -167,18 +165,55 @@ def _job_retain_results(job_id: str | None) -> bool: return tracker.should_retain_results(job_id) -def _work_item_retain_results(request: Request, *, job_id: str | None) -> bool: - """Whether the worker pool should cache row payloads for this upload.""" - if request.headers.get(_GATEWAY_DOC_ID_HEADER): - return _retain_results_from_request(request) - return _job_retain_results(job_id) - - def _internal_auth_headers(request: Request) -> dict[str, str]: """Return service credentials for pod-to-pod callback traffic.""" - from nemo_retriever.service.auth import auth_headers + from nemo_retriever.service.auth import internal_auth_headers + + return internal_auth_headers(request.app.state.config.vectordb.internal_api_token) + + +def _proxied_response(response: httpx.Response) -> Response: + """Relay an upstream VectorDB response to the caller unchanged.""" + return Response( + content=response.content, + status_code=response.status_code, + media_type=response.headers.get("content-type", "application/json"), + ) + + +async def _vectordb_get(request: Request, url: str, *, scope: str, failure_detail: str) -> httpx.Response: + """Read one scoped resource from the internal VectorDB service.""" + try: + async with httpx.AsyncClient(timeout=30.0) as client: + return await client.get( + url, + headers={"X-NRL-Scope": scope, **_internal_auth_headers(request)}, + ) + except httpx.HTTPError as exc: + raise HTTPException(502, failure_detail) from exc + - return auth_headers(request.app.state.config.auth) +def _job_idempotency_fingerprint(body: JobCreateRequest) -> str: + """Hash the request fields that an idempotent replay must match. + + The field list is explicit rather than a full model dump so that adding a + request field cannot silently invalidate previously issued fingerprints. + """ + return hashlib.sha256( + json.dumps( + { + "expected_documents": body.expected_documents, + "collection_name": body.collection_name, + "operation": body.operation, + "target_document_id": body.target_document_id, + "metadata": body.metadata, + "retain_results": body.retain_results, + "document_manifest": [entry.model_dump(mode="json") for entry in body.document_manifest], + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + ).hexdigest() def _record_prometheus( @@ -204,7 +239,10 @@ def _register_document_under_job( document_id: str, job_id: str, filename: str | None = None, -) -> None: + content_sha256: str | None = None, + stable_document_id: str | None = None, + manifest_entry_id: str | None = None, +): """Register a per-document tracker entry inside an existing job. Maps :class:`JobTrackerError` subclasses to HTTP responses so the @@ -220,9 +258,16 @@ def _register_document_under_job( tracker = get_job_tracker() if tracker is None: - return + raise HTTPException(status_code=503, detail="Job tracker not available") try: - tracker.register_document(document_id, job_id=job_id, filename=filename) + return tracker.register_document_idempotent( + document_id, + job_id=job_id, + filename=filename, + content_sha256=content_sha256, + stable_document_id=stable_document_id, + manifest_entry_id=manifest_entry_id, + ) except JobNotFoundError as exc: raise HTTPException(status_code=404, detail=str(exc)) from exc except JobFullError as exc: @@ -233,7 +278,46 @@ def _register_document_under_job( raise HTTPException(status_code=getattr(exc, "status_code", 500), detail=str(exc)) from exc -def _require_job(job_id: str): +def _resolve_stable_document_id( + attempt_id: str, + *, + collection_name: str | None, + target_document_id: str | None, +) -> str: + """Keep legacy IDs pollable while separating collection document identity.""" + if target_document_id: + return target_document_id + if collection_name: + return uuid.uuid4().hex + return attempt_id + + +def _validate_manifest_entry(job, manifest_entry_id: str | None, filename: str, content_sha256: str) -> None: + """Bind an upload to exactly one immutable entry in its job manifest.""" + if not job.document_manifest: + return + if not manifest_entry_id: + raise HTTPException(409, "manifest_entry_id is required for this idempotent job") + entry = next( + (item for item in job.document_manifest if item.get("manifest_entry_id") == manifest_entry_id), + None, + ) + if not entry or entry.get("filename") != filename or entry.get("content_sha256") != content_sha256: + raise HTTPException(409, "Uploaded document does not match its job manifest entry") + + +def _validate_collection_pipeline_spec(job, spec: PipelineSpec | None) -> None: + """Prevent collection writes from bypassing the configured VDB boundary.""" + if not job.collection_name or spec is None: + return + if spec.vdb_upload_params is not None: + raise HTTPException( + 422, + "collection-aware ingestion cannot override VectorDB upload configuration", + ) + + +def _require_job(job_id: str, request: Request | None = None): """Look up an existing :class:`JobAggregate` or raise HTTP 404.""" tracker = get_job_tracker() if tracker is None: @@ -241,6 +325,11 @@ def _require_job(job_id: str): agg = tracker.get_job(job_id) if agg is None: raise HTTPException(status_code=404, detail=f"Job {job_id!r} not found") + if request is not None and not _is_worker(request): + from nemo_retriever.service.auth import authorized_scope + + if agg.scope != authorized_scope(request): + raise HTTPException(status_code=404, detail=f"Job {job_id!r} not found") return agg @@ -259,7 +348,9 @@ async def _enqueue_or_reject(pool_type: PoolType, item: WorkItem) -> None: ) -async def _fetch_result_data_from_workers(document_id: str) -> list[dict[str, Any]] | None: +async def _fetch_result_data_from_workers( + document_id: str, +) -> list[dict[str, Any]] | None: """Read rows already handed off to this gateway's retained store.""" try: rows = await asyncio.to_thread(get_result_data, document_id) @@ -286,19 +377,31 @@ def _worker_result_url( ) -> str: """Build a fixed-path owner URL from a validated worker pod IP.""" if not isinstance(worker_ip_value, str): - raise HTTPException(status_code=503, detail="Completion callback is missing result worker identity") + raise HTTPException( + status_code=503, + detail="Completion callback is missing result worker identity", + ) try: worker_ip = ipaddress.ip_address(worker_ip_value) except ValueError as exc: - raise HTTPException(status_code=400, detail="Completion callback has an invalid result worker IP") from exc + raise HTTPException( + status_code=400, + detail="Completion callback has an invalid result worker IP", + ) from exc if worker_ip.is_unspecified or worker_ip.is_multicast or worker_ip.is_loopback or worker_ip.is_link_local: - raise HTTPException(status_code=400, detail="Completion callback has an unroutable result worker IP") + raise HTTPException( + status_code=400, + detail="Completion callback has an unroutable result worker IP", + ) if callback_worker_ip is not None: try: advertised_ip = ipaddress.ip_address(callback_worker_ip) except ValueError as exc: - raise HTTPException(status_code=400, detail="Completion callback has an invalid result worker IP") from exc + raise HTTPException( + status_code=400, + detail="Completion callback has an invalid result worker IP", + ) from exc if advertised_ip != worker_ip: raise HTTPException(status_code=409, detail="Result worker IP does not match lease owner") @@ -342,9 +445,15 @@ async def _pull_and_store_worker_result( payload = response.json() rows = payload.get("result_data") if isinstance(payload, dict) else None except ValueError as exc: - raise HTTPException(status_code=503, detail=f"Worker returned invalid result data for {document_id!r}") from exc + raise HTTPException( + status_code=503, + detail=f"Worker returned invalid result data for {document_id!r}", + ) from exc if not rows or not isinstance(rows, list) or not all(isinstance(row, dict) for row in rows): - raise HTTPException(status_code=503, detail=f"Worker returned invalid result data for {document_id!r}") + raise HTTPException( + status_code=503, + detail=f"Worker returned invalid result data for {document_id!r}", + ) try: await asyncio.to_thread(store_result_data, document_id, rows) except (OSError, ValueError, TypeError) as exc: @@ -363,11 +472,19 @@ async def _gateway_enqueue( job_id: str, payload: bytes, filename: str | None, - pipeline_spec: PipelineSpec | None = None, - extra: dict[str, Any] | None = None, + pipeline_spec: dict[str, Any] | None = None, + write: DocumentWriteContext | None = None, ) -> None: - """Admit split-mode work to the gateway broker after atomic spooling.""" - from nemo_retriever.service.services.work_queue import WorkQueueFull, get_work_broker + """Admit split-mode work to the gateway broker after atomic spooling. + + The write context is nested under a single ``write`` key so the claiming + worker rebuilds it as one typed object rather than a splat of loose + fields that :class:`WorkItem` would silently discard. + """ + from nemo_retriever.service.services.work_queue import ( + WorkQueueFull, + get_work_broker, + ) broker = get_work_broker() if broker is None: @@ -383,9 +500,9 @@ async def _gateway_enqueue( payload=payload, filename=filename, retain_results=_job_retain_results(job_id), - pipeline_spec=pipeline_spec.model_dump(mode="json") if pipeline_spec is not None else None, + pipeline_spec=pipeline_spec, trace_context=_safe_inject_trace_context(), - extra=extra, + extra={"write": write.model_dump(mode="json")} if write is not None else None, ) except WorkQueueFull as exc: tracker = get_job_tracker() @@ -442,11 +559,10 @@ def _route_by_page_count( * Audio / video files are always routed to **batch** — they involve heavyweight ASR / frame-extraction pipelines. - * Image files are always routed to **realtime** — they are single-page - and latency-sensitive. - * Documents (PDF, DOCX, PPTX) and other types use the original - page-count heuristic: small docs ( PipelineSpe raise HTTPException(status_code=exc.status_code, detail=exc.detail) from exc -def _spec_from_gateway_header(request: Request) -> PipelineSpec | None: - """Recover and re-validate the spec forwarded by the gateway pod. +async def _prepare_job_work_item( + request: Request, + *, + job_id: str, + file: UploadFile, + meta: IngestRequest, + validated_spec: PipelineSpec | None, + manifest_entry_id: str | None, + job: JobAggregate, + pool_type: PoolType | None = None, +) -> tuple[WorkItem, PoolType, FileClassification]: + """Build the execution envelope shared by the document upload routes. + + Reads the upload exactly once and resolves every non-payload value the + gateway and standalone paths need: pool routing, content digest, + manifest binding, attempt identity and durable storage identity. + + Args: + request: The inbound upload request. + job_id: Job aggregate the upload belongs to, taken from the URL path. + file: The uploaded file. + meta: Parsed ``IngestRequest`` metadata accompanying the upload. + validated_spec: Policy-validated per-request pipeline overrides. + manifest_entry_id: Immutable manifest entry the upload claims, if any. + job: The job aggregate owning this upload. + pool_type: Fixed pool for callers that do not auto-route. + + Returns: + The work item, its target pool, and the file classification. + """ + classification = FileClassifier.classify(file, filename_override=meta.filename or "") + enforce_media_dependencies(classification) + + file_bytes = await file.read() + route = pool_type or _route_by_page_count(file_bytes, meta, file_category=classification.category) + content_sha256 = hashlib.sha256(file_bytes).hexdigest() + + _validate_manifest_entry(job, manifest_entry_id, file.filename or "", content_sha256) + + attempt_id = uuid.uuid4().hex + storage_document_id = _resolve_stable_document_id( + attempt_id, + collection_name=job.collection_name, + target_document_id=job.target_document_id, + ) + + item = WorkItem( + id=attempt_id, + payload=file_bytes, + filename=file.filename, + callback_headers=_internal_auth_headers(request), + job_id=job_id, + pipeline_spec=validated_spec.model_dump(mode="json") if validated_spec is not None else None, + retain_results=_job_retain_results(job_id), + write=DocumentWriteContext( + scope=job.scope, + collection_name=job.collection_name, + operation=job.operation, + content_sha256=content_sha256, + storage_document_id=storage_document_id, + document_metadata=meta.metadata, + ), + ) + return item, route, classification + - The gateway has already validated against its own copy of the policy, - but we re-validate on the worker as defense-in-depth: a misconfigured - gateway or a pod with a different ``pipeline_overrides`` config will - still see consistent enforcement. +async def _submit_job_work_item( + request: Request, + pool_type: PoolType, + item: WorkItem, + *, + manifest_entry_id: str | None, +) -> DocumentRecord | None: + """Register *item* under its job and admit it for execution. + + Args: + request: The inbound upload request. + pool_type: Pool the item is admitted to. + item: Envelope produced by :func:`_prepare_job_work_item`. + manifest_entry_id: Immutable manifest entry the upload claims, if any. + + Returns: + The already-registered record when an idempotent retry matched an + earlier attempt, otherwise ``None`` once the item is admitted. """ - raw = request.headers.get(_GATEWAY_PIPELINE_SPEC_HEADER) - if not raw: - return None - try: - spec = PipelineSpec.model_validate_json(raw) - except ValueError as exc: - raise HTTPException( - status_code=400, - detail=f"Malformed {_GATEWAY_PIPELINE_SPEC_HEADER!r} from gateway: {exc}", - ) from exc - policy = _build_policy(request) - try: - return validate_pipeline_spec(spec, policy) - except PolicyError as exc: - raise HTTPException(status_code=exc.status_code, detail=exc.detail) from exc + if item.job_id is None: + raise RuntimeError("job upload work item is missing job_id") + + record, created = _register_document_under_job( + document_id=item.id, + job_id=item.job_id, + filename=item.filename, + content_sha256=item.write.content_sha256, + stable_document_id=item.write.storage_document_id, + manifest_entry_id=manifest_entry_id, + ) + if not created: + return record + + if _is_gateway(request): + await _gateway_enqueue( + request, + pool_type, + work_id=item.id, + job_id=item.job_id, + payload=item.payload, + filename=item.filename, + pipeline_spec=item.pipeline_spec, + write=item.write, + ) + else: + await _enqueue_or_reject(pool_type, item) + + return None def _parse_backend_json(resp: Response) -> dict: @@ -534,7 +741,10 @@ def _safe_inject_trace_context() -> dict[str, str]: try: return dict(tracing.inject_trace_context()) except Exception as exc: - logger.warning("Trace context injection failed; continuing without propagated context: %s", exc) + logger.warning( + "Trace context injection failed; continuing without propagated context: %s", + exc, + ) return {} @@ -545,7 +755,10 @@ def _safe_extract_trace_context(carrier: dict[str, str] | None) -> Any | None: try: return tracing.extract_trace_context(carrier) except Exception as exc: - logger.warning("Trace context extraction failed; continuing without parent context: %s", exc) + logger.warning( + "Trace context extraction failed; continuing without parent context: %s", + exc, + ) return None @@ -607,6 +820,8 @@ def _aggregate_to_response(agg, *, documents: list[dict[str, Any]] | None = None counts=dict(agg.counts), document_ids=list(agg.document_ids), documents=documents, + collection_name=agg.collection_name, + operation=agg.operation, ) @@ -616,7 +831,7 @@ def _aggregate_to_response(agg, *, documents: list[dict[str, Any]] | None = None status_code=201, 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: """Open a job that will receive ``expected_documents`` uploads. The server returns an opaque ``job_id`` the client uses for every @@ -632,6 +847,34 @@ async def create_job(request: Request, response: Response, body: JobCreateReques tracker = get_job_tracker() if tracker is None: raise HTTPException(status_code=503, detail="Job tracker not available") + from nemo_retriever.service.auth import authorized_scope + + scope = authorized_scope(request) + if body.collection_name: + config = request.app.state.config + if not config.vectordb.enabled: + raise HTTPException(404, "VectorDB is not enabled in the service configuration") + target = f"{config.vectordb.vectordb_url.rstrip('/')}/v1/collections/{body.collection_name}" + collection_response = await _vectordb_get( + request, + target, + scope=scope, + failure_detail="Failed to validate collection with VectorDB service", + ) + if collection_response.status_code != 200: + return _proxied_response(collection_response) + if collection_response.json().get("status") != "active": + raise HTTPException(409, "Collection is not active") + if body.operation is IngestOperation.REPLACE and body.target_document_id: + document_response = await _vectordb_get( + request, + f"{target}/documents/{body.target_document_id}", + scope=scope, + failure_detail="Failed to validate replacement document", + ) + if document_response.status_code != 200: + return _proxied_response(document_response) + fingerprint = _job_idempotency_fingerprint(body) job_id = uuid.uuid4().hex trace_id: str | None = None inbound_trace_context = _trace_context_from_request_or_job(request, None) @@ -655,13 +898,23 @@ async def create_job(request: Request, response: Response, body: JobCreateReques 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], ) except JobTrackerError as exc: raise HTTPException(status_code=getattr(exc, "status_code", 500), detail=str(exc)) from exc - if trace_id: + created = agg.job_id == job_id + if not created: + response.status_code = 200 + if created and trace_id: response.headers[tracing.TRACE_ID_HEADER] = trace_id - if (m := get_metrics()) is not None: + if created and (m := get_metrics()) is not None: m.record_request("/v1/ingest/job") m.record_job_created(job_id) return JobCreatedResponse( @@ -671,6 +924,8 @@ async def create_job(request: Request, response: Response, body: JobCreateReques created_at=agg.created_at, label=agg.label, trace_id=agg.trace_id, + collection_name=agg.collection_name, + operation=agg.operation, ) @@ -689,7 +944,7 @@ async def get_job( Pass ``?include_documents=true`` to also return the per-document records (capped to the first 10k entries to keep payloads bounded). """ - agg = _require_job(job_id) + agg = _require_job(job_id, request) documents: list[dict[str, Any]] | None = None if include_documents: tracker = get_job_tracker() @@ -713,7 +968,8 @@ async def get_job( def _document_to_response(rec, *, result_data=None) -> DocumentStatusResponse: """Project a :class:`DocumentRecord` to the wire response shape.""" return DocumentStatusResponse( - document_id=rec.id, + document_id=rec.stable_document_id, + attempt_id=rec.id, job_id=rec.job_id, status=rec.status.value, submitted_at=rec.submitted_at, @@ -724,6 +980,8 @@ def _document_to_response(rec, *, result_data=None) -> DocumentStatusResponse: result_rows=rec.result_rows, result_data=result_data, error=rec.error, + collection_name=rec.collection_name, + content_sha256=rec.content_sha256, ) @@ -760,7 +1018,7 @@ async def get_job_documents( if limit < 1 or limit > 1000: raise HTTPException(status_code=400, detail="limit must be in [1, 1000]") - agg = _require_job(job_id) + agg = _require_job(job_id, request) tracker = get_job_tracker() docs = tracker.job_documents(job_id) if tracker is not None else [] @@ -812,7 +1070,7 @@ async def get_job_document( """ from nemo_retriever.service.services.job_tracker import DocumentStatus - _require_job(job_id) + _require_job(job_id, request) tracker = get_job_tracker() if tracker is None: raise HTTPException(status_code=503, detail="Job tracker is not available.") @@ -854,6 +1112,7 @@ async def submit_document_to_job( job_id: str, file: UploadFile = File(..., description="The file to ingest"), metadata: str = Form(default="{}", description="JSON-encoded IngestRequest metadata"), + manifest_entry_id: str | None = Form(default=None, description="Immutable entry ID from the job manifest"), ) -> IngestAccepted | Response: """General-purpose upload into a job. @@ -867,109 +1126,53 @@ async def submit_document_to_job( except (json.JSONDecodeError, ValueError) as exc: raise HTTPException(status_code=400, detail=f"Invalid metadata JSON: {exc}") - # Job lookup is gateway/standalone only — worker pods don't own the - # JobTracker, so we must trust the gateway-forwarded URL. - if not _is_worker(request): - _require_job(job_id) + job = _require_job(job_id, request) _check_upload_size(file, request) validated_spec = _resolve_pipeline_spec(request, meta) + _validate_collection_pipeline_spec(job, validated_spec) with _start_accept_span(request, job_id, "ingest.document.accept"): - if _is_gateway(request): - classification = FileClassifier.classify(file, filename_override=meta.filename or "") - enforce_media_dependencies(classification) - file_size = _file_size_from_upload(file) - - file_bytes = await file.read() - route = _route_by_page_count(file_bytes, meta, file_category=classification.category) - - document_id = uuid.uuid4().hex - content_sha256 = hashlib.sha256(file_bytes).hexdigest() - now = datetime.now(timezone.utc).isoformat() - - _register_document_under_job(document_id=document_id, job_id=job_id, filename=file.filename) - await _gateway_enqueue( - request, - route, - work_id=document_id, - job_id=job_id, - payload=file_bytes, - filename=file.filename, - pipeline_spec=validated_spec, - ) - - _record_prometheus(request, "/v1/ingest/job/document", "2xx", file_size=file_size) - if (m := get_metrics()) is not None: - m.record_request("/v1/ingest/job/document") - m.record_document_accepted( - document_id=document_id, - job_id=job_id, - filename=classification.filename, - file_category=classification.category.value, - content_type=classification.content_type, - file_size_bytes=file_size, - endpoint="/v1/ingest/job/document", - ) + item, route, classification = await _prepare_job_work_item( + request, + job_id=job_id, + file=file, + meta=meta, + validated_spec=validated_spec, + manifest_entry_id=manifest_entry_id, + job=job, + ) + now = datetime.now(timezone.utc).isoformat() + record = await _submit_job_work_item(request, route, item, manifest_entry_id=manifest_entry_id) + if record is not None: return IngestAccepted( - document_id=document_id, - job_id=job_id, - content_sha256=content_sha256, + document_id=record.stable_document_id, + attempt_id=record.id, + job_id=item.job_id, + content_sha256=record.content_sha256 or item.write.content_sha256, status="accepted", - created_at=now, + created_at=record.submitted_at, ) - # ── worker / standalone ────────────────────────────────────── - classification = FileClassifier.classify(file, filename_override=meta.filename or "") - enforce_media_dependencies(classification) - - file_bytes = await file.read() - route = _route_by_page_count(file_bytes, meta, file_category=classification.category) - content_sha256 = hashlib.sha256(file_bytes).hexdigest() - now = datetime.now(timezone.utc).isoformat() - - gw_doc_id = request.headers.get(_GATEWAY_DOC_ID_HEADER) - gw_callback_url = request.headers.get(_GATEWAY_CALLBACK_HEADER) - gw_job_id = request.headers.get(_GATEWAY_JOB_ID_HEADER) or job_id - document_id = gw_doc_id or uuid.uuid4().hex - - worker_spec = _spec_from_gateway_header(request) if gw_doc_id else validated_spec - - if not gw_callback_url: - _register_document_under_job(document_id=document_id, job_id=job_id, filename=file.filename) - - await _enqueue_or_reject( - route, - WorkItem( - id=document_id, - payload=file_bytes, - filename=file.filename, - callback_url=gw_callback_url, - callback_headers=_internal_auth_headers(request), - job_id=gw_job_id, - pipeline_spec=worker_spec.model_dump(mode="json") if worker_spec is not None else None, - retain_results=_work_item_retain_results(request, job_id=gw_job_id), - ), - ) - - _record_prometheus(request, "/v1/ingest/job/document", "2xx", file_size=len(file_bytes)) - + file_size = _file_size_from_upload(file) if _is_gateway(request) else len(item.payload) + _record_prometheus(request, "/v1/ingest/job/document", "2xx", file_size=file_size) if (m := get_metrics()) is not None: m.record_request("/v1/ingest/job/document") m.record_document_accepted( - document_id=document_id, - job_id=gw_job_id, + document_id=item.id, + job_id=item.job_id, filename=classification.filename, file_category=classification.category.value, content_type=classification.content_type, - file_size_bytes=len(file_bytes), + file_size_bytes=file_size, endpoint="/v1/ingest/job/document", ) return IngestAccepted( - document_id=document_id, - job_id=gw_job_id, - content_sha256=content_sha256, + document_id=item.write.storage_document_id, + attempt_id=item.id, + job_id=item.job_id, + content_sha256=item.write.content_sha256, status="accepted", created_at=now, ) @@ -985,14 +1188,18 @@ async def submit_page_to_job( request: Request, job_id: str, file: UploadFile = File(..., description="A single-page PDF or image"), - document_id: str = Form(..., description="Client-assigned ID grouping pages from the same source document"), + document_id: str = Form( + ..., + description="Client-assigned ID grouping pages from the same source document", + ), page_number: int = Form(..., description="1-based page number within the source document"), filename: str = Form(default="", description="Original source document filename"), ) -> PageIngestAccepted | Response: - # Job lookup is gateway/standalone only (workers don't own the - # JobTracker — they trust the gateway-forwarded URL). - if not _is_worker(request): - _require_job(job_id) + if _require_job(job_id, request).collection_name: + raise HTTPException( + 422, + "collection-aware ingestion does not support /page; use /document or /whole", + ) _check_upload_size(file, request) with _start_accept_span(request, job_id, "ingest.page.accept"): @@ -1008,7 +1215,11 @@ async def submit_page_to_job( now = datetime.now(timezone.utc).isoformat() if not dry_run: - _register_document_under_job(document_id=page_id, job_id=job_id, filename=filename or file.filename) + _register_document_under_job( + document_id=page_id, + job_id=job_id, + filename=filename or file.filename, + ) await _gateway_enqueue( request, PoolType.REALTIME, @@ -1016,10 +1227,6 @@ async def submit_page_to_job( job_id=job_id, payload=file_bytes, filename=file.filename, - extra={ - "source_document_id": document_id, - "page_number": page_number, - }, ) _record_prometheus( @@ -1034,6 +1241,7 @@ async def submit_page_to_job( m.record_page_accepted( page_id=page_id, document_id=document_id, + job_id=job_id, endpoint="/v1/ingest/job/page", page_number=page_number, file_size_bytes=file_size, @@ -1059,34 +1267,40 @@ async def submit_page_to_job( content_sha256 = hashlib.sha256(file_bytes).hexdigest() now = datetime.now(timezone.utc).isoformat() - gw_doc_id = request.headers.get(_GATEWAY_DOC_ID_HEADER) - gw_callback_url = request.headers.get(_GATEWAY_CALLBACK_HEADER) - gw_job_id = request.headers.get(_GATEWAY_JOB_ID_HEADER) or job_id - page_id = gw_doc_id or uuid.uuid4().hex + page_id = uuid.uuid4().hex if not dry_run: - if not gw_callback_url: - _register_document_under_job(document_id=page_id, job_id=job_id, filename=filename or file.filename) + _register_document_under_job( + document_id=page_id, + job_id=job_id, + filename=filename or file.filename, + ) await _enqueue_or_reject( PoolType.REALTIME, WorkItem( id=page_id, payload=file_bytes, filename=file.filename, - callback_url=gw_callback_url, callback_headers=_internal_auth_headers(request), - job_id=gw_job_id, - retain_results=_work_item_retain_results(request, job_id=gw_job_id), + job_id=job_id, + retain_results=_job_retain_results(job_id), ), ) - _record_prometheus(request, "/v1/ingest/job/page", "2xx", file_size=len(file_bytes), is_page=True) + _record_prometheus( + request, + "/v1/ingest/job/page", + "2xx", + file_size=len(file_bytes), + is_page=True, + ) if (m := get_metrics()) is not None: m.record_request("/v1/ingest/job/page") m.record_page_accepted( page_id=page_id, document_id=document_id, + job_id=job_id, endpoint="/v1/ingest/job/page", page_number=page_number, file_size_bytes=len(file_bytes), @@ -1115,117 +1329,64 @@ async def submit_whole_document_to_job( job_id: str, file: UploadFile = File(..., description="The full document to ingest"), metadata: str = Form(default="{}", description="JSON-encoded IngestRequest metadata"), + manifest_entry_id: str | None = Form(default=None, description="Immutable entry ID from the job manifest"), ) -> DocumentIngestAccepted | Response: try: meta = IngestRequest(**json.loads(metadata)) except (json.JSONDecodeError, ValueError) as exc: raise HTTPException(status_code=400, detail=f"Invalid metadata JSON: {exc}") - # Job lookup is gateway/standalone only (workers don't own the - # JobTracker — they trust the gateway-forwarded URL). - if not _is_worker(request): - _require_job(job_id) + job = _require_job(job_id, request) _check_upload_size(file, request) validated_spec = _resolve_pipeline_spec(request, meta) + _validate_collection_pipeline_spec(job, validated_spec) with _start_accept_span(request, job_id, "ingest.whole.accept"): - if _is_gateway(request): - dry_run = _is_dry_run(request) - classification = FileClassifier.classify(file, filename_override=meta.filename or "") - enforce_media_dependencies(classification) - file_size = _file_size_from_upload(file) - - document_id = uuid.uuid4().hex - file_bytes = await file.read() - content_sha256 = hashlib.sha256(file_bytes).hexdigest() - now = datetime.now(timezone.utc).isoformat() - - if not dry_run: - _register_document_under_job(document_id=document_id, job_id=job_id, filename=file.filename) - await _gateway_enqueue( - request, - PoolType.BATCH, - work_id=document_id, - job_id=job_id, - payload=file_bytes, - filename=file.filename, - pipeline_spec=validated_spec, - ) + item, route, classification = await _prepare_job_work_item( + request, + job_id=job_id, + file=file, + meta=meta, + validated_spec=validated_spec, + manifest_entry_id=manifest_entry_id, + job=job, + pool_type=PoolType.BATCH, + ) + now = datetime.now(timezone.utc).isoformat() - _record_prometheus(request, "/v1/ingest/job/whole", "2xx", file_size=file_size) - if (m := get_metrics()) is not None: - m.record_request("/v1/ingest/job/whole") - m.record_document_accepted( - document_id=document_id, - job_id=job_id, + if not _is_dry_run(request): + record = await _submit_job_work_item(request, route, item, manifest_entry_id=manifest_entry_id) + if record is not None: + return DocumentIngestAccepted( + document_id=record.stable_document_id, + attempt_id=record.id, filename=classification.filename, - file_category=classification.category.value, - content_type=classification.content_type, - file_size_bytes=file_size, - endpoint="/v1/ingest/job/whole", + file_size_bytes=len(item.payload), + content_sha256=record.content_sha256 or item.write.content_sha256, + status="accepted", + created_at=record.submitted_at, ) - return DocumentIngestAccepted( - document_id=document_id, - filename=classification.filename, - file_size_bytes=len(file_bytes), - content_sha256=content_sha256, - status="accepted", - created_at=now, - ) - - # ── worker / standalone ────────────────────────────────────── - dry_run = _is_dry_run(request) - classification = FileClassifier.classify(file, filename_override=meta.filename or "") - enforce_media_dependencies(classification) - - file_bytes = await file.read() - content_sha256 = hashlib.sha256(file_bytes).hexdigest() - now = datetime.now(timezone.utc).isoformat() - - gw_doc_id = request.headers.get(_GATEWAY_DOC_ID_HEADER) - gw_callback_url = request.headers.get(_GATEWAY_CALLBACK_HEADER) - gw_job_id = request.headers.get(_GATEWAY_JOB_ID_HEADER) or job_id - document_id = gw_doc_id or uuid.uuid4().hex - - worker_spec = _spec_from_gateway_header(request) if gw_doc_id else validated_spec - - if not dry_run: - if not gw_callback_url: - _register_document_under_job(document_id=document_id, job_id=job_id, filename=file.filename) - await _enqueue_or_reject( - PoolType.BATCH, - WorkItem( - id=document_id, - payload=file_bytes, - filename=file.filename, - callback_url=gw_callback_url, - callback_headers=_internal_auth_headers(request), - job_id=gw_job_id, - pipeline_spec=worker_spec.model_dump(mode="json") if worker_spec is not None else None, - retain_results=_work_item_retain_results(request, job_id=gw_job_id), - ), - ) - - _record_prometheus(request, "/v1/ingest/job/whole", "2xx", file_size=len(file_bytes)) - + file_size = _file_size_from_upload(file) if _is_gateway(request) else len(item.payload) + _record_prometheus(request, "/v1/ingest/job/whole", "2xx", file_size=file_size) if (m := get_metrics()) is not None: m.record_request("/v1/ingest/job/whole") m.record_document_accepted( - document_id=document_id, - job_id=gw_job_id, + document_id=item.id, + job_id=item.job_id, filename=classification.filename, file_category=classification.category.value, content_type=classification.content_type, - file_size_bytes=len(file_bytes), + file_size_bytes=file_size, endpoint="/v1/ingest/job/whole", ) return DocumentIngestAccepted( - document_id=document_id, + document_id=item.write.storage_document_id, + attempt_id=item.id, filename=classification.filename, - file_size_bytes=len(file_bytes), - content_sha256=content_sha256, + file_size_bytes=len(item.payload), + content_sha256=item.write.content_sha256, status="accepted", created_at=now, ) @@ -1256,6 +1417,7 @@ async def _status_response(request: Request, item_id: str) -> JSONResponse: rec = tracker.get_document(item_id) if rec is None: raise HTTPException(status_code=404, detail=f"No tracked document with id={item_id!r}") + _require_job(rec.job_id, request) is_terminal = rec.status in (DocumentStatus.COMPLETED, DocumentStatus.FAILED) result_data = tracker.get_result_data(item_id) if is_terminal else None @@ -1550,21 +1712,26 @@ async def answer(req: ServiceAnswerRequest, request: Request) -> Response | Answ target = f"{vectordb_url}/v1/query" try: + from nemo_retriever.service.auth import authorized_scope, internal_auth_headers + async with httpx.AsyncClient(timeout=60.0) as client: - resp = await client.post(target, json={"query": answer_req.query, "top_k": answer_req.top_k}) - except Exception as exc: + resp = await client.post( + target, + json={"query": answer_req.query, "top_k": answer_req.top_k}, + headers={ + "X-NRL-Scope": authorized_scope(request), + **internal_auth_headers(config.vectordb.internal_api_token), + }, + ) + except httpx.HTTPError: logger.exception("Failed to query vectordb at %s for answer generation", target) raise HTTPException( status_code=502, - detail=f"Failed to reach VectorDB service: {type(exc).__name__}: {exc}", + detail="VectorDB service is unavailable.", ) if resp.status_code != 200: - return Response( - content=resp.content, - status_code=resp.status_code, - media_type=resp.headers.get("content-type", "application/json"), - ) + return _proxied_response(resp) payload = resp.json() result_sets = payload.get("results") or [] @@ -1643,58 +1810,173 @@ async def answer(req: ServiceAnswerRequest, request: Request) -> Response | Answ # ------------------------------------------------------------------ -# POST /v1/query — vector search (proxied to vectordb pod) +# POST /v1/query — vector search / agentic retrieval (proxied to vectordb) # ------------------------------------------------------------------ @router.post( "/query", - summary="Search ingested documents by semantic similarity or hybrid retrieval", + summary="Search ingested documents by semantic similarity, hybrid, agentic, or reranked retrieval", ) async def query(request: Request) -> Response: - """Proxy a query request to the VectorDB service. + """Run the public query API through VectorDB and optional reranking. - * **gateway / standalone** — forwards the JSON body to the vectordb pod. - * **worker** — returns 404 (workers don't handle queries). + ``vectordb_app`` remains responsible only for nearest-neighbor retrieval. + With ``rerank=true``, the main service obtains a larger candidate set from + VectorDB, then uses the server-configured remote endpoint or local model + in the main service process, and returns the requested ``top_k`` hits. """ import httpx config = request.app.state.config - if not config.vectordb.enabled: raise HTTPException( status_code=404, detail="VectorDB is not enabled in the service configuration.", ) - - mode = _mode(request) - if mode in ("realtime", "batch"): + if _mode(request) in ("realtime", "batch"): raise HTTPException( status_code=404, detail="Query endpoint is not available on worker pods. Use the gateway.", ) - vectordb_url = config.vectordb.vectordb_url.rstrip("/") - target = f"{vectordb_url}/v1/query" - body = await request.body() + agentic = False + rerank_request: QueryRequest | None = None + query_body = body + try: + parsed = json.loads(body.decode("utf-8") if isinstance(body, (bytes, bytearray)) else body) + agentic = bool(parsed.get("agentic")) if isinstance(parsed, dict) else False + if isinstance(parsed, dict) and "rerank" in parsed: + try: + validated_rerank_request = QueryRequest.model_validate(parsed) + except ValueError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + if validated_rerank_request.rerank: + rerank_request = validated_rerank_request + has_remote_reranker = bool((config.nim_endpoints.rerank_invoke_url or "").strip()) + has_local_reranker = config.local_models.rerank.enabled + if not (has_remote_reranker or has_local_reranker): + raise HTTPException( + status_code=400, + detail=( + "Reranking is not configured. Set " + "nim_endpoints.rerank_invoke_url or enable " + "local_models.rerank.enabled on the main service." + ), + ) + candidate_k = rerank_request.rerank_top_k or max(rerank_request.top_k, 50) + forwarded = dict(parsed) + forwarded.pop("rerank", None) + forwarded.pop("rerank_top_k", None) + forwarded["top_k"] = candidate_k + query_body = json.dumps(forwarded).encode("utf-8") + except (UnicodeDecodeError, json.JSONDecodeError, AttributeError, TypeError): + agentic = False + + if agentic and not config.agentic.enabled: + raise HTTPException( + status_code=400, + detail="Agentic retrieval is not enabled in the service configuration. " + "Set agentic.enabled with llm_model and invoke_url.", + ) + vectordb_url = config.vectordb.vectordb_url.rstrip("/") + timeout = config.agentic.request_timeout_s if agentic else 60.0 try: - async with httpx.AsyncClient(timeout=60.0) as client: + from nemo_retriever.service.auth import authorized_scope, internal_auth_headers + + async with httpx.AsyncClient(timeout=timeout) as client: resp = await client.post( - target, - content=body, - headers={"Content-Type": "application/json"}, + f"{vectordb_url}/v1/query", + content=query_body, + headers={ + "Content-Type": "application/json", + "X-NRL-Scope": authorized_scope(request), + **internal_auth_headers(config.vectordb.internal_api_token), + }, ) - except Exception as exc: - logger.exception("Failed to proxy query to vectordb at %s", target) - raise HTTPException( - status_code=502, - detail=f"Failed to reach VectorDB service: {type(exc).__name__}: {exc}", + except httpx.HTTPError as exc: + logger.exception("Failed to proxy query to vectordb at %s", vectordb_url) + raise HTTPException(status_code=502, detail="VectorDB service is unavailable.") from exc + + if rerank_request is None or resp.status_code >= 400: + return Response( + content=resp.content, + status_code=resp.status_code, + media_type="application/json", ) + try: + response = QueryResponse.model_validate_json(resp.content) + queries = [rerank_request.query] if isinstance(rerank_request.query, str) else rerank_request.query + if len(response.results) != len(queries): + raise ValueError("VectorDB response count did not match the query count") + + from nemo_retriever.operators.rerank import rerank_hits + + local_rerank = config.local_models.rerank + use_remote_reranker = bool((config.nim_endpoints.rerank_invoke_url or "").strip()) + rerank_lock: asyncio.Lock | None = None + local_reranker = None + if not use_remote_reranker: + # The local model is service-owned and lazily loaded once. Serialize + # scoring too: HF/vLLM model instances are not safe to invoke from + # simultaneous request threads, and this avoids duplicate GPU work. + rerank_lock = getattr(request.app.state, "local_reranker_lock", None) + if rerank_lock is None: + rerank_lock = asyncio.Lock() + request.app.state.local_reranker_lock = rerank_lock + + async with rerank_lock: + local_reranker = getattr(request.app.state, "local_reranker", None) + if local_reranker is None: + from nemo_retriever.models import create_local_reranker + + local_reranker = await asyncio.to_thread( + create_local_reranker, + local_rerank.model_name, + backend=local_rerank.backend, + device=config.local_models.device, + hf_cache_dir=config.local_models.hf_cache_dir, + gpu_memory_utilization=local_rerank.gpu_memory_utilization, + ) + request.app.state.local_reranker = local_reranker + + async def _rerank_results() -> None: + rerank_kwargs: dict[str, Any] = {"top_n": rerank_request.top_k} + if use_remote_reranker: + rerank_kwargs.update( + rerank_invoke_url=config.nim_endpoints.rerank_invoke_url, + model_name=config.nim_endpoints.rerank_model_name or "nvidia/llama-nemotron-rerank-1b-v2", + api_key=config.nim_endpoints.api_key or "", + ) + else: + rerank_kwargs.update( + model=local_reranker, + model_name=local_rerank.model_name, + max_length=local_rerank.max_length, + batch_size=local_rerank.batch_size, + ) + for query_text, result in zip(queries, response.results): + result.hits = await asyncio.to_thread( + rerank_hits, + query_text, + result.hits, + **rerank_kwargs, + ) + + if rerank_lock is not None: + async with rerank_lock: + await _rerank_results() + else: + await _rerank_results() + except Exception as exc: + logger.exception("Failed to rerank VectorDB query results") + raise HTTPException(status_code=502, detail="Reranker service is unavailable.") from exc + return Response( - content=resp.content, + content=response.model_dump_json(), status_code=resp.status_code, media_type="application/json", ) @@ -1755,7 +2037,10 @@ async def job_callback(request: Request) -> JSONResponse: broker = None lease_record = None if _is_gateway(request): - from nemo_retriever.service.services.work_queue import StaleLease, get_work_broker + from nemo_retriever.service.services.work_queue import ( + StaleLease, + get_work_broker, + ) broker = get_work_broker() lease_id = body.get("lease_id") diff --git a/nemo_retriever/src/nemo_retriever/service/service_ingestor.py b/nemo_retriever/src/nemo_retriever/service/service_ingestor.py index c962f245f6..767c781d24 100644 --- a/nemo_retriever/src/nemo_retriever/service/service_ingestor.py +++ b/nemo_retriever/src/nemo_retriever/service/service_ingestor.py @@ -44,6 +44,7 @@ Fluent methods that *do* take effect by writing to the spec: * ``.extract(...)`` — per-request extraction knobs (DPI, OCR enable, …) +* ``.texts(...)`` — inline text payloads * ``.embed(...)`` — embedding model/dim overrides bounded by the operator's allow-list * ``.dedup(...)``, ``.split(...)``, ``.filter()`` — shape knobs @@ -76,14 +77,16 @@ import threading import time import warnings +from contextlib import nullcontext from io import BytesIO from pathlib import Path -from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple, Union +from typing import Any, AsyncIterator, Iterator, List, Optional, Self, Sequence, Tuple, Union import httpx from nemo_retriever.ingestor.results import ResultSchema, concat_ingest_results from nemo_retriever.ingestor import _merge_params, ingestor +from nemo_retriever.common.inline_text import inline_text_source_id, is_blank_inline_corpus, normalize_inline_texts from nemo_retriever.common.params import ( CaptionParams, IngestExecuteParams, @@ -92,6 +95,7 @@ VdbUploadParams, WebhookParams, ) +from nemo_retriever.service.client import InMemoryUpload, RetrieverServiceClient, UploadInput logger = logging.getLogger(__name__) @@ -102,6 +106,12 @@ "with bulky image/embedding values stripped during the deprecation window." ) +_RESULT_FETCH_TRANSIENT_ERRORS: tuple[type[Exception], ...] = ( + httpx.ConnectError, + httpx.ReadError, + httpx.RemoteProtocolError, +) + # ---------------------------------------------------------------------- # Result container @@ -200,6 +210,18 @@ def _normalize_files(files: Union[str, List[str], List[Path]]) -> list[Path]: return [Path(f) for f in files] +def _empty_service_result_dataframe(result_schema: ResultSchema) -> Any: + """Return an empty DataFrame matching the selected public result schema.""" + if result_schema == "legacy": + from nemo_retriever.common.modality.txt.split import empty_text_chunks_df + + return empty_text_chunks_df() + + import pandas as pd + + return pd.DataFrame(columns=["text", "source_id", "element_type", "page_number"]).astype({"page_number": "int64"}) + + # ---------------------------------------------------------------------- # Client-side mirror of service.models.pipeline_spec.PipelineSpec # ---------------------------------------------------------------------- @@ -216,6 +238,7 @@ def _normalize_files(files: Union[str, List[str], List[Path]]) -> list[Path]: "ocr_api_key", "table_structure_invoke_url", "nemotron_parse_invoke_url", + "nemotron_parse_model", "embed_invoke_url", "embedding_endpoint", "embed_model_provider_prefix", @@ -425,6 +448,7 @@ def __init__( self._max_concurrency = max_concurrency self._request_timeout_s = request_timeout_s self._api_token = (api_token or "").strip() or None + self._inline_texts: list[str] | None = None self._document_ids: list[str] = [] self._last_run_elapsed_s: float = 0.0 self._last_job_id: str | None = None @@ -446,20 +470,44 @@ def _record_stage(self, name: str) -> None: if name not in order: order.append(name) - def _fetch_document_result_data(self, document_id: str) -> list[dict[str, Any]]: + def _new_result_fetch_client(self) -> httpx.Client: + """Create a client for retained-result status requests.""" + return httpx.Client(timeout=self._request_timeout_s, headers=self._auth_headers) + + def _fetch_document_result_data( + self, + document_id: str, + *, + client: httpx.Client | None = None, + ) -> list[dict[str, Any]]: """Fetch ``result_data`` for *document_id* from the status endpoint. The status endpoint retains ``result_data`` through the job retention - window, so retrying this read is safe. + window, so retrying this read is safe. When the caller supplies a + client it is reused for the first attempt. A transient failure receives + one retry through a fresh client so a stale pooled connection cannot be + selected again. """ if not document_id: raise ValueError("_fetch_document_result_data(): empty document_id") + if client is None: + with self._new_result_fetch_client() as scoped_client: + return self._fetch_document_result_data(document_id, client=scoped_client) + url = f"{self._base_url}/v1/ingest/status/{document_id}" - with httpx.Client(timeout=self._request_timeout_s, headers=self._auth_headers) as client: + try: resp = client.get(url) - resp.raise_for_status() - body = resp.json() + except _RESULT_FETCH_TRANSIENT_ERRORS as exc: + logger.debug( + "Transient %s fetching retained result for %s; retrying on a fresh connection", + type(exc).__name__, + document_id, + ) + with self._new_result_fetch_client() as retry_client: + resp = retry_client.get(url) + resp.raise_for_status() + body = resp.json() return list(body.get("result_data") or []) def _write_result_data_to_disk(self, document_id: str, result_data: list[dict[str, Any]]) -> Path: @@ -483,7 +531,12 @@ def _write_result_data_to_disk(self, document_id: str, result_data: list[dict[st out_path.write_bytes(payload) return out_path - def _save_document_to_disk(self, document_id: str) -> Path: + def _save_document_to_disk( + self, + document_id: str, + *, + client: httpx.Client | None = None, + ) -> Path: """Fetch ``result_data`` for *document_id* and write a JSON artifact. Returns the path that was written. Raises if the document_id is @@ -491,7 +544,7 @@ def _save_document_to_disk(self, document_id: str) -> Path: """ if self._save_to_disk_dir is None: raise RuntimeError("_save_document_to_disk(): save_to_disk was never enabled") - result_data = self._fetch_document_result_data(document_id) + result_data = self._fetch_document_result_data(document_id, client=client) return self._write_result_data_to_disk(document_id, result_data) def _materialize_completed_document( @@ -499,11 +552,12 @@ def _materialize_completed_document( document_id: str, *, return_results: bool, + client: httpx.Client | None = None, ) -> list[dict[str, Any]] | None: """Fetch (once) and optionally persist rows for a completed document.""" if not return_results and self._save_to_disk_dir is None: return None - result_data = self._fetch_document_result_data(document_id) + result_data = self._fetch_document_result_data(document_id, client=client) if self._save_to_disk_dir is not None: self._write_result_data_to_disk(document_id, result_data) return result_data if return_results else None @@ -521,6 +575,8 @@ def _pipeline_payload( so the worker can short-circuit identically. """ spec = dict(self._pipeline_spec) + if self._has_mixed_inline_sources(): + spec["extraction_mode"] = "auto" spec["result_schema"] = result_schema spec["return_embeddings"] = bool(return_embeddings or spec.get("return_embeddings", False)) spec["return_images"] = bool(return_images or spec.get("return_images", False)) @@ -563,6 +619,11 @@ def files(self, documents: Union[str, List[str]]) -> "ServiceIngestor": self._documents.extend(documents) return self + def texts(self, texts: Union[str, Sequence[str]]) -> Self: + """Set raw inline text documents, optionally alongside file or buffer uploads.""" + self._inline_texts = normalize_inline_texts(texts) + return self + def buffers( self, buffers: Union[Tuple[str, BytesIO], List[Tuple[str, BytesIO]]], @@ -1048,6 +1109,25 @@ def webhook(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor": self._record_stage("webhook") return self + def _ingest_events_with_result_client( + self, + *, + retain_results: bool, + result_schema: ResultSchema, + return_embeddings: bool, + return_images: bool, + ) -> Iterator[tuple[dict[str, Any], httpx.Client | None]]: + """Yield ingest events while owning the optional shared result client.""" + client_context = self._new_result_fetch_client() if retain_results else nullcontext(None) + with client_context as result_client: + for evt in self.ingest_stream( + retain_results=retain_results, + result_schema=result_schema, + return_embeddings=return_embeddings, + return_images=return_images, + ): + yield evt, result_client + # ------------------------------------------------------------------ # Execution — sync materialized # ------------------------------------------------------------------ @@ -1110,6 +1190,19 @@ def ingest(self, params: Any = None, **kwargs: Any) -> Any: self._resolve_execute_flags(params, kwargs) ) del params, kwargs + if not self._documents and not self._buffers and is_blank_inline_corpus(self._inline_texts): + self._document_ids.clear() + self._last_run_elapsed_s = 0.0 + self._last_job_id = None + result = ServiceIngestResult() + if return_results: + result.dataframe = _empty_service_result_dataframe(result_schema) + if return_failures and return_traces: + return result, [], [] + if return_failures or return_traces: + return result, [] + return result + retain_results = return_results or self._save_to_disk_dir is not None if retain_results and result_schema == "legacy": warnings.warn(_LEGACY_RESULT_SCHEMA_DEPRECATION, DeprecationWarning, stacklevel=2) @@ -1122,7 +1215,7 @@ def ingest(self, params: Any = None, **kwargs: Any) -> Any: documents_failed = 0 total_uploaded = 0 - for evt in self.ingest_stream( + for evt, result_client in self._ingest_events_with_result_client( retain_results=retain_results, result_schema=result_schema, return_embeddings=return_embeddings, @@ -1186,6 +1279,7 @@ def ingest(self, params: Any = None, **kwargs: Any) -> Any: rows = self._materialize_completed_document( doc_id, return_results=return_results, + client=result_client, ) if rows is not None and return_results: rows_by_document[doc_id] = rows @@ -1398,15 +1492,13 @@ def _factory(): async def _aingest_stream_impl( self, - files: list[Path], + files: list[UploadInput], *, retain_results: bool = False, result_schema: ResultSchema = "legacy", return_embeddings: bool = False, return_images: bool = False, ) -> AsyncIterator[dict[str, Any]]: - from nemo_retriever.service.client import RetrieverServiceClient - client = RetrieverServiceClient( base_url=self._base_url, max_concurrency=self._max_concurrency, @@ -1512,9 +1604,15 @@ def cancel(self, job_id: str | None = None) -> dict[str, Any]: # Internals # ------------------------------------------------------------------ - def _collect_inputs(self) -> list[Path]: - """Gather both file paths and any in-memory buffers into Paths.""" - files = [Path(p) for p in self._documents] + def _has_mixed_inline_sources(self) -> bool: + return bool(self._inline_texts) and bool(self._documents or self._buffers) + + def _collect_inputs(self) -> list[UploadInput]: + """Gather filesystem and in-memory inputs for the service client.""" + if not self._documents and not self._buffers and is_blank_inline_corpus(self._inline_texts): + return [] + + files: list[UploadInput] = [Path(p) for p in self._documents] if self._buffers: import tempfile @@ -1525,4 +1623,15 @@ def _collect_inputs(self) -> list[Path]: target.write_bytes(buf.getvalue()) files.append(target) + for index, text in enumerate(self._inline_texts or []): + source_id = inline_text_source_id(index) + files.append( + InMemoryUpload( + filename=source_id, + content=text.encode("utf-8"), + content_type="text/plain; charset=utf-8", + classification_filename=f"inline-{index:08d}.txt", + ) + ) + return files diff --git a/nemo_retriever/src/nemo_retriever/service/services/job_tracker.py b/nemo_retriever/src/nemo_retriever/service/services/job_tracker.py index 2c3604a111..51661ba686 100644 --- a/nemo_retriever/src/nemo_retriever/service/services/job_tracker.py +++ b/nemo_retriever/src/nemo_retriever/service/services/job_tracker.py @@ -45,11 +45,12 @@ import time from datetime import datetime, timezone from enum import Enum -from typing import Any +from typing import Any, Callable from pydantic import Field from nemo_retriever.common.schemas.base import RichModel +from nemo_retriever.common.schemas.collections import IngestOperation logger = logging.getLogger(__name__) @@ -136,6 +137,7 @@ class DocumentRecord(RichModel): """ id: str + stable_document_id: str job_id: str status: DocumentStatus = DocumentStatus.PENDING submitted_at: str = "" @@ -145,8 +147,10 @@ class DocumentRecord(RichModel): result_rows: int | None = None result_data: list[dict[str, Any]] | None = None error: str | None = None - filename: str | None = None - """Original upload filename, surfaced in the dashboard UI.""" + filename: str | None = Field(default=None, description="Original upload filename surfaced in the dashboard") + collection_name: str | None = None + content_sha256: str | None = None + manifest_entry_id: str | None = None class JobAggregate(RichModel): @@ -174,13 +178,18 @@ class JobAggregate(RichModel): started_at: str | None = None finalized_at: str | None = None elapsed_s: float | None = None - label: str | None = None - """Optional client-supplied tag, e.g. ``"Q4-2026-corpus"``.""" + label: str | None = Field(default=None, description="Optional client-supplied job tag") metadata: dict[str, Any] = {} trace_id: str | None = None trace_context: dict[str, str] = Field(default_factory=dict) - retain_results: bool = False - """When false, :meth:`JobTracker.mark_completed` drops bulky ``result_data``.""" + retain_results: bool = Field(default=False, description="Keep completed document result payloads in memory") + collection_name: str | None = None + scope: str = "default" + operation: IngestOperation = IngestOperation.APPEND + target_document_id: str | None = None + idempotency_key: str | None = None + idempotency_fingerprint: str | None = None + document_manifest: list[dict[str, str]] = Field(default_factory=list) # ── eviction tunables (apply to terminal aggregates) ────────────────── @@ -249,6 +258,9 @@ def __init__( self._started_mono: dict[str, float] = {} # per-document elapsed timing self._job_started_mono: dict[str, float] = {} # per-job elapsed timing self._event_bus: Any = None + # Observers receive defensive snapshots after a genuine terminal + # document transition. This keeps all completion paths consistent. + self._terminal_observers: list[Callable[[DocumentRecord, JobAggregate | None], None]] = [] self._ttl_s = ttl_s self._stale_job_ttl_s = stale_job_ttl_s self._max_jobs = max_jobs @@ -257,6 +269,8 @@ def __init__( # counts) we last published, so we don't emit duplicate progress # events on every doc transition. self._progress_published: dict[str, int] = {} + self._idempotency: dict[tuple[str, str], tuple[str, str]] = {} + self._accepted_manifest_entries: dict[tuple[str, str], str] = {} self._progress_step: int = 10 # ── wiring ─────────────────────────────────────────────────────── @@ -265,6 +279,19 @@ def set_event_bus(self, bus: Any) -> None: """Attach an :class:`EventBus` so state transitions publish SSE events.""" self._event_bus = bus + def add_terminal_observer(self, observer: Callable[[DocumentRecord, JobAggregate | None], None]) -> None: + """Register a callback for real document terminal transitions. + + Args: + observer: Callback receiving defensive snapshots of the transitioned + document and its aggregate job, when available. + + Returns: + None. The callback runs outside the tracker lock after tracker + state is updated; duplicate and unknown callbacks do not invoke it. + """ + self._terminal_observers.append(observer) + def set_progress_step(self, step: int) -> None: """Override the progress-event cadence (default: every 10 docs).""" if step <= 0: @@ -283,6 +310,13 @@ def register_job( retain_results: bool = False, trace_id: str | None = None, trace_context: dict[str, str] | None = None, + collection_name: str | None = None, + scope: str = "default", + operation: IngestOperation = IngestOperation.APPEND, + target_document_id: str | None = None, + idempotency_key: str | None = None, + idempotency_fingerprint: str | None = None, + document_manifest: list[dict[str, str]] | None = None, ) -> JobAggregate: """Create a new :class:`JobAggregate` in ``pending`` state.""" if expected_documents <= 0: @@ -290,6 +324,15 @@ def register_job( with self._lock: now = datetime.now(timezone.utc) self._evict_locked(now=now) + if idempotency_key and idempotency_fingerprint: + entry = self._idempotency.get((scope, idempotency_key)) + if entry: + existing_job_id, previous_fingerprint = entry + if previous_fingerprint != idempotency_fingerprint: + raise JobFullError("Idempotency key was already used with a different request payload") + existing = self._get_live_job_locked(existing_job_id, now=now) + if existing is not None: + return existing.model_copy(deep=True) if job_id in self._jobs: raise JobTrackerError(f"Job {job_id!r} already exists") if len(self._jobs) >= self._max_jobs: @@ -309,9 +352,21 @@ def register_job( trace_id=trace_id, trace_context=dict(trace_context or {}), retain_results=retain_results, + collection_name=collection_name, + scope=scope, + operation=operation, + target_document_id=target_document_id, + idempotency_key=idempotency_key, + idempotency_fingerprint=idempotency_fingerprint, + document_manifest=list(document_manifest or []), ) agg.counts[DocumentStatus.PENDING.value] = 0 self._jobs[job_id] = agg + if idempotency_key and idempotency_fingerprint: + self._idempotency[(scope, idempotency_key)] = ( + job_id, + idempotency_fingerprint, + ) logger.info( "Job registered: %s (expected_documents=%d, label=%r)", job_id, @@ -371,6 +426,9 @@ def register_document( *, job_id: str, filename: str | None = None, + content_sha256: str | None = None, + stable_document_id: str | None = None, + manifest_entry_id: str | None = None, ) -> DocumentRecord: """Attach a new :class:`DocumentRecord` to *job_id*. @@ -379,11 +437,43 @@ def register_document( terminal state, or :class:`JobFullError` if the job is at capacity (``len(document_ids) == expected_documents``). """ + rec, _ = self.register_document_idempotent( + document_id, + job_id=job_id, + filename=filename, + content_sha256=content_sha256, + stable_document_id=stable_document_id, + manifest_entry_id=manifest_entry_id, + ) + return rec + + def register_document_idempotent( + self, + attempt_id: str, + *, + job_id: str, + filename: str | None = None, + content_sha256: str | None = None, + stable_document_id: str | None = None, + manifest_entry_id: str | None = None, + ) -> tuple[DocumentRecord, bool]: + """Register an upload or return its original accepted record. + + Manifest-entry replay is checked before terminal/capacity validation so a + client can safely replay every file after losing any upload response. + """ with self._lock: now = datetime.now(timezone.utc) agg = self._get_live_job_locked(job_id, now=now) if agg is None: raise JobNotFoundError(f"Job {job_id!r} not found") + if manifest_entry_id: + existing_attempt = self._accepted_manifest_entries.get((job_id, manifest_entry_id)) + if existing_attempt: + existing = self._documents[existing_attempt] + if existing.filename != filename or existing.content_sha256 != content_sha256: + raise JobFullError("Manifest entry was already accepted with different filename or content") + return existing.model_copy(deep=True), False if agg.status in _JOB_TERMINAL: raise JobFinalizedError( f"Job {job_id!r} has already finalized with status " @@ -395,22 +485,28 @@ def register_document( f"({agg.expected_documents} documents); rejected document " f"#{len(agg.document_ids) + 1}." ) - if document_id in self._documents: - raise JobTrackerError(f"Document {document_id!r} already registered.") + if attempt_id in self._documents: + raise JobTrackerError(f"Document attempt {attempt_id!r} already registered.") rec = DocumentRecord( - id=document_id, + id=attempt_id, + stable_document_id=stable_document_id or attempt_id, job_id=job_id, status=DocumentStatus.PENDING, submitted_at=_utcnow_iso(), filename=filename, + collection_name=agg.collection_name, + content_sha256=content_sha256, + manifest_entry_id=manifest_entry_id, ) - self._documents[document_id] = rec - agg.document_ids.append(document_id) + self._documents[attempt_id] = rec + agg.document_ids.append(attempt_id) + if manifest_entry_id: + self._accepted_manifest_entries[(job_id, manifest_entry_id)] = attempt_id agg.counts[DocumentStatus.PENDING.value] = agg.counts.get(DocumentStatus.PENDING.value, 0) + 1 self._reg_count += 1 if self._reg_count % _EVICTION_INTERVAL == 0: self._evict_locked(now=now) - return rec.model_copy(deep=True) + return rec.model_copy(deep=True), True def mark_processing(self, document_id: str) -> None: """Transition a document from ``pending`` → ``processing``. @@ -523,7 +619,7 @@ def _mark_terminal( "JobTracker.%s: no record of document %r — callback dropped (likely " "gateway-pod restart between upload acceptance and worker callback); " "client may hang waiting for an SSE event that will never arrive", - "mark_failed" if new_status == DocumentStatus.FAILED else "mark_completed", + ("mark_failed" if new_status == DocumentStatus.FAILED else "mark_completed"), document_id, ) return MarkOutcome.UNKNOWN_DOCUMENT @@ -549,6 +645,7 @@ def _mark_terminal( agg = self._jobs.get(rec.job_id) finalized_snapshot: JobAggregate | None = None progress_snapshot: JobAggregate | None = None + aggregate_snapshot: JobAggregate | None = None if agg is not None: terminal_count = agg.counts.get(DocumentStatus.COMPLETED.value, 0) + agg.counts.get( DocumentStatus.FAILED.value, 0 @@ -564,8 +661,16 @@ def _mark_terminal( if terminal_count - last_published >= self._progress_step: self._progress_published[rec.job_id] = terminal_count progress_snapshot = agg.model_copy(deep=True) + aggregate_snapshot = agg.model_copy(deep=True) + + # Phase 2: project and publish with the lock released. + for observer in tuple(self._terminal_observers): + try: + observer(doc_snapshot, aggregate_snapshot) + except Exception: + logger.exception("JobTracker terminal observer failed for document %s", document_id) - # Phase 2: publish events with the lock released. + # Events are published after observers see finalized state. self._publish_document_event(doc_snapshot) if progress_snapshot is not None: self._publish_job_event("job_progress", progress_snapshot) @@ -693,8 +798,18 @@ def _drop_job_locked(self, job_id: str) -> None: agg = self._jobs.pop(job_id, None) if agg is None: return + if agg.idempotency_key: + # Only release the key if this job still owns it; the map is written + # solely for jobs that carry a fingerprint, so a keyed job may point + # at an entry that now belongs to a later job. + entry_key = (agg.scope, agg.idempotency_key) + owner = self._idempotency.get(entry_key) + if owner is not None and owner[0] == job_id: + del self._idempotency[entry_key] for did in agg.document_ids: - self._documents.pop(did, None) + rec = self._documents.pop(did, None) + if rec and rec.manifest_entry_id: + self._accepted_manifest_entries.pop((job_id, rec.manifest_entry_id), None) self._started_mono.pop(did, None) self._job_started_mono.pop(job_id, None) self._progress_published.pop(job_id, None) @@ -707,7 +822,8 @@ def _publish_document_event(self, rec: DocumentRecord) -> None: event: dict[str, Any] = { "type": rec.status.value, "id": rec.id, - "document_id": rec.id, + "document_id": rec.stable_document_id, + "attempt_id": rec.id, "job_id": rec.job_id, "status": rec.status.value, "result_rows": rec.result_rows, diff --git a/nemo_retriever/src/nemo_retriever/service/services/metrics.py b/nemo_retriever/src/nemo_retriever/service/services/metrics.py index 6e7d76a67a..d1154dd4bc 100644 --- a/nemo_retriever/src/nemo_retriever/service/services/metrics.py +++ b/nemo_retriever/src/nemo_retriever/service/services/metrics.py @@ -33,10 +33,15 @@ import time from collections import deque from datetime import datetime, timezone +from typing import TYPE_CHECKING + from pydantic import ConfigDict, Field from nemo_retriever.common.schemas.base import RichModel +if TYPE_CHECKING: + from nemo_retriever.service.services.job_tracker import DocumentRecord, JobAggregate + logger = logging.getLogger(__name__) # ── Capacity knobs (tweak without touching the class) ──────────────── @@ -54,6 +59,7 @@ class PageMetric(RichModel): page_id: str document_id: str endpoint: str + job_id: str | None = None page_number: int | None = None file_size_bytes: int = 0 file_category: str = "" @@ -170,6 +176,9 @@ def __init__(self) -> None: self._jobs: dict[str, JobMetric] = {} self._documents: dict[str, DocumentMetric] = {} self._pages: deque[PageMetric] = deque(maxlen=MAX_RECENT_PAGES) + # Recent page details are bounded, but page terminal outcomes for + # active jobs must not be lost when those details are evicted. + self._pending_page_job_ids: dict[str, str] = {} # ── recording helpers ──────────────────────────────────────────── @@ -234,6 +243,7 @@ def record_page_accepted( *, page_id: str, document_id: str, + job_id: str | None = None, endpoint: str = "", page_number: int | None = None, file_size_bytes: int = 0, @@ -248,6 +258,7 @@ def record_page_accepted( PageMetric( page_id=page_id, document_id=document_id, + job_id=job_id, endpoint=endpoint, page_number=page_number, file_size_bytes=file_size_bytes, @@ -263,10 +274,13 @@ def record_page_accepted( "pages_submitted": doc.pages_submitted + 1, } ) - job_id = self._documents[document_id].job_id if document_id in self._documents else None - if job_id and job_id in self._jobs: - job = self._jobs[job_id] - self._jobs[job_id] = job.model_copy( + resolved_job_id = job_id or ( + self._documents[document_id].job_id if document_id in self._documents else None + ) + if resolved_job_id and resolved_job_id in self._jobs: + self._pending_page_job_ids[page_id] = resolved_job_id + job = self._jobs[resolved_job_id] + self._jobs[resolved_job_id] = job.model_copy( update={ "pages_total": job.pages_total + 1, } @@ -298,6 +312,66 @@ def record_page_completed(self, page_id: str) -> None: ) break + def record_terminal_transition(self, document: DocumentRecord, job: JobAggregate | None) -> None: + """Mirror an authoritative terminal document transition into metrics. + + Args: + document: Immutable snapshot of the document that reached a terminal + state. + job: Immutable snapshot of the document's aggregate job, if it is + still tracked. + + Returns: + None. The method updates in-memory metric records only. + """ + with self._lock: + if document.id in self._documents: + metric = self._documents[document.id] + self._documents[document.id] = metric.model_copy( + update={ + "status": document.status.value, + "completed_at": document.completed_at, + "processing_duration_s": document.elapsed_s, + "error": document.error, + } + ) + + for index, page in enumerate(self._pages): + if page.page_id == document.id: + self._pages[index] = page.model_copy( + update={ + "status": document.status.value, + "completed_at": document.completed_at, + "processing_duration_s": document.elapsed_s, + "error": document.error, + } + ) + break + + page_job_id = self._pending_page_job_ids.pop(document.id, None) + if job is not None and job.job_id in self._jobs: + metric_job = self._jobs[job.job_id] + # Observer calls may arrive out of order across concurrent + # worker completions; terminal counts are monotonic. + updates: dict[str, object] = { + "documents_completed": max(metric_job.documents_completed, job.counts.get("completed", 0)), + "documents_failed": max(metric_job.documents_failed, job.counts.get("failed", 0)), + } + if job.finalized_at is not None: + updates.update( + { + "status": job.status.value, + "completed_at": job.finalized_at, + "wall_duration_s": job.elapsed_s, + } + ) + if page_job_id == job.job_id: + if document.status.value == "completed": + updates["pages_completed"] = metric_job.pages_completed + 1 + elif document.status.value == "failed": + updates["pages_failed"] = metric_job.pages_failed + 1 + self._jobs[job.job_id] = metric_job.model_copy(update=updates) + # ── single-record lookups ──────────────────────────────────────── def get_job(self, job_id: str) -> JobMetric | None: diff --git a/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py b/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py index 3031de1d2a..98715ea8c3 100644 --- a/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py +++ b/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py @@ -21,6 +21,7 @@ from __future__ import annotations import asyncio +import copy import json import logging import multiprocessing as mp @@ -36,7 +37,7 @@ NimEndpointsConfig, ServiceConfig, ) - from nemo_retriever.service.services.pipeline_pool import WorkItem + from nemo_retriever.service.services.pipeline_pool import DocumentWriteContext, WorkItem logger = logging.getLogger(__name__) @@ -279,38 +280,72 @@ def shutdown_process_executors() -> None: logger.info("All pipeline process executors shut down") -def _post_rows_to_vectordb(rows: list[dict[str, Any]], vectordb_url: str, filename: str) -> None: - """Fire-and-forget POST of LanceDB rows to the vectordb service.""" +def _post_records_to_vectordb( + records: list[list[dict[str, Any]]], + vectordb_url: str, + filename: str, + *, + context: DocumentWriteContext, + job_id: str | None = None, + internal_api_token: str | None = None, +) -> None: + """Post canonical NRL record batches to VectorDB, preserving legacy best-effort behavior. + + Collection-managed writes are lifecycle-authoritative and therefore fail + the job when storage rejects the write. Legacy fixed-table writes retain + their historical best-effort behavior. + """ import json import urllib.request import urllib.error + from nemo_retriever.service.auth import internal_auth_headers - if not rows: + if not records or not any(records): + if context.collection_name: + raise ValueError(f"No vector rows were produced for collection document {filename}") return url = vectordb_url.rstrip("/") + "/internal/vectordb/write" - body = json.dumps({"rows": rows}).encode() + body = json.dumps( + { + "records": records, + "scope": context.scope, + "collection_name": context.collection_name, + "document_id": context.storage_document_id, + "job_id": job_id, + "filename": filename, + "content_sha256": context.content_sha256, + "document_version": context.resolved_version, + "operation": context.operation, + } + ).encode() req = urllib.request.Request( url, data=body, - headers={"Content-Type": "application/json"}, + headers={ + "Content-Type": "application/json", + **internal_auth_headers(internal_api_token), + }, method="POST", ) + record_count = sum(len(batch) for batch in records) try: with urllib.request.urlopen(req, timeout=30) as resp: logging.getLogger(__name__).info( - "Posted %d rows to vectordb for %s — HTTP %d", - len(rows), + "Posted %d records to vectordb for %s — HTTP %d", + record_count, filename, resp.status, ) except Exception as exc: logging.getLogger(__name__).warning( - "Failed to POST %d rows to vectordb for %s: %s", - len(rows), + "Failed to POST %d records to vectordb for %s: %s", + record_count, filename, exc, ) + if context.collection_name: + raise RuntimeError(f"Collection write failed for {filename}: {exc}") from exc _TRUST_OWNED_EXTRACT_KEYS: tuple[str, ...] = ( @@ -322,6 +357,7 @@ def _post_rows_to_vectordb(rows: list[dict[str, Any]], vectordb_url: str, filena "ocr_api_key", "table_structure_invoke_url", "nemotron_parse_invoke_url", + "nemotron_parse_model", ) _TRUST_OWNED_EMBED_KEYS: tuple[str, ...] = ( "embed_invoke_url", @@ -360,6 +396,24 @@ def _merge_server_owned( return merged +def _resolve_extract_params( + base_extract: dict[str, Any], + extract_override: dict[str, Any] | None, +) -> Any: + """Merge extraction settings while isolating Parse-only server fields.""" + from nemo_retriever.common.params import ExtractParams + + extract_kwargs = _merge_server_owned( + base_extract, + extract_override, + _TRUST_OWNED_EXTRACT_KEYS, + ) + if extract_kwargs.get("method", "pdfium") != "nemotron_parse": + extract_kwargs.pop("nemotron_parse_invoke_url", None) + extract_kwargs.pop("nemotron_parse_model", None) + return ExtractParams(**extract_kwargs) + + def _resolve_sidecar_in_spec(spec: dict[str, Any] | None) -> dict[str, Any] | None: """Resolve ``vdb_upload_params.meta_dataframe_id`` to in-band bytes. @@ -517,7 +571,6 @@ def _build_graph_ingestor_from_spec( ASRParams, CaptionParams, DedupParams, - ExtractParams, StoreParams, VdbUploadParams, WebhookParams, @@ -527,8 +580,7 @@ def _build_graph_ingestor_from_spec( extraction_mode = _resolve_service_extraction_mode(spec.get("extraction_mode", "auto"), filename) split_config = spec.get("split_config") - extract_kwargs = _merge_server_owned(base_extract, spec.get("extract_params"), _TRUST_OWNED_EXTRACT_KEYS) - extract_params = ExtractParams(**extract_kwargs) + extract_params = _resolve_extract_params(base_extract, spec.get("extract_params")) embed_override = spec.get("embed_params") embed_params = _resolve_embed_params(base_embed, embed_override) @@ -555,8 +607,6 @@ def _build_graph_ingestor_from_spec( if extraction_mode == "image": ingestor = ingestor.extract_image_files(extract_params, split_config=split_config) - elif extraction_mode == "text" and split_config is None: - ingestor = ingestor.extract_txt() elif extraction_mode == "html" and split_config is None: ingestor = ingestor.extract_html() else: @@ -655,6 +705,9 @@ def _run_pipeline_in_process( trace_context: dict[str, str] | None = None, pool_label: str | None = None, service_role: str | None = None, + write_context: DocumentWriteContext | None = None, + job_id: str | None = None, + internal_api_token: str | None = None, ) -> tuple[int, list[dict[str, Any]], float]: """Execute one pipeline run inside a child process. @@ -697,6 +750,10 @@ def _run_pipeline_in_process( ) result_df = ingestor.ingest() + _merge_document_metadata( + result_df, + write_context.document_metadata if write_context is not None else None, + ) finally: tracing.force_flush(timeout_millis=500) @@ -718,10 +775,18 @@ def _run_pipeline_in_process( # Skip the out-of-graph fan-out when the client already wired # IngestVdbOperator into the spec — that operator handles # persistence itself. - from nemo_retriever.common.vdb.lancedb_schema import build_lancedb_rows + from nemo_retriever.common.vdb.records import to_client_vdb_records + from nemo_retriever.service.services.pipeline_pool import DocumentWriteContext - lancedb_rows = build_lancedb_rows(result_df) - _post_rows_to_vectordb(lancedb_rows, vectordb_url, filename) + records = to_client_vdb_records(result_df) + _post_records_to_vectordb( + records, + vectordb_url, + filename, + context=write_context or DocumentWriteContext(), + job_id=job_id, + internal_api_token=internal_api_token, + ) result_options = pipeline_spec or {} result_schema = result_options.get("result_schema", "legacy") @@ -734,6 +799,40 @@ def _run_pipeline_in_process( return row_count, result_data, elapsed +def _merge_document_metadata(result: Any, document_metadata: dict[str, Any] | None) -> None: + """Merge request metadata into extracted rows without replacing parser fields.""" + if not document_metadata: + return + + # Validate and copy at the child-process boundary so storage receives no + # aliases or values that cannot be represented in JSON query hits. + canonical = json.loads(json.dumps(document_metadata, ensure_ascii=False)) + + def merge_row(row: Any) -> None: + if not isinstance(row, dict): + return + metadata = row.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + row["metadata"] = metadata + content_metadata = metadata.get("content_metadata") + if not isinstance(content_metadata, dict): + content_metadata = {} + metadata["content_metadata"] = content_metadata + for key, value in canonical.items(): + content_metadata.setdefault(key, copy.deepcopy(value)) + + if isinstance(result, list): + for row in result: + merge_row(row) + return + if hasattr(result, "iterrows"): + for index, row in result.iterrows(): + holder = {"metadata": row.get("metadata")} + merge_row(holder) + result.at[index, "metadata"] = holder["metadata"] + + def _local_model_runtime_kwargs(local: "LocalModelsConfig") -> dict[str, Any]: """Shared ``ModelRuntimeParams`` fields for in-pod HF stages.""" runtime: dict[str, Any] = {} @@ -786,6 +885,13 @@ def build_extract_params(nim: "NimEndpointsConfig", local: "LocalModelsConfig | kwargs["ocr_invoke_url"] = nim.ocr_invoke_url if nim.table_structure_invoke_url: kwargs["table_structure_invoke_url"] = nim.table_structure_invoke_url + if nim.nemotron_parse_invoke_url: + # ExtractParams validates that Parse-specific configuration and the + # extraction method are selected together. + kwargs["method"] = "nemotron_parse" + kwargs["nemotron_parse_invoke_url"] = nim.nemotron_parse_invoke_url + if nim.nemotron_parse_model: + kwargs["nemotron_parse_model"] = nim.nemotron_parse_model if nim.api_key: kwargs["api_key"] = nim.api_key @@ -977,6 +1083,7 @@ async def _work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: loop = asyncio.get_running_loop() resolved_spec = _resolve_sidecar_in_spec(item.pipeline_spec) + write_context = item.write.resolved(fallback_document_id=item.id) try: trace_context = _capture_trace_context_for_pipeline() @@ -994,6 +1101,9 @@ async def _work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: trace_context, label, config.mode, + write_context, + item.job_id, + config.vectordb.internal_api_token, ) except BrokenProcessPool: logger.error( diff --git a/nemo_retriever/src/nemo_retriever/service/services/pipeline_pool.py b/nemo_retriever/src/nemo_retriever/service/services/pipeline_pool.py index 967b9df72c..31d8b49bc3 100644 --- a/nemo_retriever/src/nemo_retriever/service/services/pipeline_pool.py +++ b/nemo_retriever/src/nemo_retriever/service/services/pipeline_pool.py @@ -35,6 +35,7 @@ from nemo_retriever.service.config import AuthConfig, PipelinePoolConfig, WorkQueueConfig from nemo_retriever.common.schemas.base import RichModel +from nemo_retriever.common.schemas.collections import IngestOperation from nemo_retriever.service.services.prometheus import ( POOL_CALLBACK_BACKPRESSURE_TOTAL, POOL_ACTIVE_SLOTS, @@ -90,6 +91,42 @@ def _safe_extract_trace_context(carrier: Mapping[str, str] | None, *, pool_name: return None +class DocumentWriteContext(RichModel): + """Durable write identity for one document, carried with its payload. + + Travels intact from upload admission through the gateway broker into the + pipeline, so no layer has to rebuild it from loose fields. ``None`` values + describe a legacy fixed-table write rather than a collection write. + """ + + scope: str = "default" + collection_name: str | None = None + operation: IngestOperation = IngestOperation.APPEND + content_sha256: str | None = None + document_version: str | None = None + storage_document_id: str | None = None + document_metadata: dict[str, Any] = Field(default_factory=dict) + + @property + def resolved_version(self) -> str: + """Durable version key: the explicit version, else the content digest.""" + return self.document_version or self.content_sha256 or "1" + + def resolved(self, *, fallback_document_id: str) -> "DocumentWriteContext": + """Return this context with a guaranteed storage document identity. + + Args: + fallback_document_id: Identity to adopt when the upload never + resolved one, as on the legacy per-page admission path. + + Returns: + This context, or a copy carrying the fallback identity. + """ + if self.storage_document_id: + return self + return self.model_copy(update={"storage_document_id": fallback_document_id}) + + class WorkItem(RichModel): """A unit of work submitted to a pool.""" @@ -104,6 +141,7 @@ class WorkItem(RichModel): # Owning job aggregate (J1+). Always set today since the only # admission path is /v1/ingest/job/{job_id}/document. job_id: str | None = None + write: DocumentWriteContext = Field(default_factory=DocumentWriteContext) retain_results: bool = False # Validated per-request pipeline overrides (PipelineSpec serialised # to a dict). ``None`` means: run the legacy startup-baked pipeline. @@ -731,6 +769,7 @@ def __init__( batch_work_fn: Callable[[WorkItem], Any] | None = None, work_queue_config: WorkQueueConfig | None = None, auth_config: AuthConfig | None = None, + internal_api_token: str | None = None, ) -> None: self._config = config self._mode = mode @@ -739,13 +778,16 @@ def __init__( pull_client = None if mode in ("realtime", "batch") and work_queue_config is not None: - from nemo_retriever.service.auth import auth_headers + from nemo_retriever.service.auth import auth_headers, internal_auth_headers from nemo_retriever.service.services.work_queue import GatewayWorkClient + headers = internal_auth_headers(internal_api_token) + if not headers: + headers = auth_headers(auth_config or AuthConfig()) pull_client = GatewayWorkClient( work_queue_config, pool=PoolType(mode), - headers=auth_headers(auth_config or AuthConfig()), + headers=headers, ) if mode in ("standalone", "realtime"): @@ -820,6 +862,7 @@ def init_pipeline_pool( batch_work_fn: Callable[[WorkItem], Any] | None = None, work_queue_config: WorkQueueConfig | None = None, auth_config: AuthConfig | None = None, + internal_api_token: str | None = None, ) -> PipelinePool: """Create and start the global pipeline pool (call once at startup). @@ -838,6 +881,7 @@ def init_pipeline_pool( batch_work_fn=batch_work_fn, work_queue_config=work_queue_config, auth_config=auth_config, + internal_api_token=internal_api_token, ) pool.start() _instance = pool diff --git a/nemo_retriever/src/nemo_retriever/service/services/work_queue.py b/nemo_retriever/src/nemo_retriever/service/services/work_queue.py index fd2db2e951..03fbcf8074 100644 --- a/nemo_retriever/src/nemo_retriever/service/services/work_queue.py +++ b/nemo_retriever/src/nemo_retriever/service/services/work_queue.py @@ -424,6 +424,7 @@ def claim_payload(self, record: WorkRecord, *, base_url: str) -> dict[str, Any]: "retain_results": record.retain_results, "pipeline_spec": record.pipeline_spec, "trace_context": record.trace_context, + "extra": record.extra, } @@ -487,6 +488,7 @@ async def claim(self) -> WorkItem | None: retain_results=bool(claim.get("retain_results")), pipeline_spec=claim.get("pipeline_spec"), trace_context=claim.get("trace_context") or {}, + **(claim.get("extra") or {}), lease_id=claim["lease_id"], lease_generation=claim["lease_generation"], delivery_attempt=claim["delivery_attempt"], diff --git a/nemo_retriever/src/nemo_retriever/service/utils/file_type.py b/nemo_retriever/src/nemo_retriever/service/utils/file_type.py index 367ce57b51..7aecefe6dd 100644 --- a/nemo_retriever/src/nemo_retriever/service/utils/file_type.py +++ b/nemo_retriever/src/nemo_retriever/service/utils/file_type.py @@ -11,6 +11,7 @@ from fastapi import HTTPException, UploadFile +from nemo_retriever.common.inline_text import is_inline_text_source from nemo_retriever.common.schemas.base import RichModel @@ -128,6 +129,9 @@ def infer_extraction_mode_from_filename(filename: str) -> str | None: routing text-like uploads through the PDF or audio-only graphs when the client leaves ``extraction_mode`` at the default ``"auto"``. """ + if is_inline_text_source(filename): + return "text" + dot = filename.rfind(".") suffix = filename[dot:].lower() if dot != -1 else "" entry = FileClassifier.SUFFIX_MAP.get(suffix) diff --git a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py index da2b3cf90c..abef9c917b 100644 --- a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py +++ b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py @@ -2,52 +2,59 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Standalone VectorDB microservice backed by LanceDB. +"""Standalone VectorDB service built on the backend-neutral VDB contract. -Provides three endpoints: - -- ``POST /internal/vectordb/write`` -- append embedding rows from ingest workers -- ``POST /v1/query`` -- embed query text and search the index -- ``GET /v1/health`` -- liveness probe - -Run with a remote NIM embed endpoint:: - - python -m nemo_retriever.service.vectordb_app \\ - --lancedb-uri /data/vectordb \\ - --embed-endpoint http://nemo-retriever-nim-embed-0...:8000/v1/embeddings \\ - --port 7671 - -Run with in-pod Hugging Face query embedding (requires ``[local]`` extras + GPU):: - - python -m nemo_retriever.service.vectordb_app \\ - --lancedb-uri /data/vectordb \\ - --local-embed \\ - --local-embed-backend hf \\ - --embed-model nvidia/llama-nemotron-embed-vl-1b-v2 \\ - --port 7671 +The HTTP layer owns transport, internal authentication, query embedding, and +error mapping. Concrete VDB implementations own persistence, collection +lifecycle behavior, and native retrieval semantics. """ from __future__ import annotations import argparse import asyncio +import hmac import logging +import os import threading +from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager +from pathlib import Path from typing import Any, AsyncIterator, Union -import lancedb import uvicorn -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI, Header, HTTPException, Query, Request, Response +from fastapi.responses import JSONResponse from pydantic import BaseModel from nemo_retriever.common.remote_auth import resolve_remote_api_key -from nemo_retriever.common.vdb.lancedb_capabilities import ( - LanceRetrievalMode, - LanceTableCapabilities, - inspect_lancedb_table_object, +from nemo_retriever.common.schemas.collections import ( + CollectionCreateRequest, + CollectionDeleteResult, + CollectionInfo, + CollectionPage, + CollectionUpdateRequest, + DocumentDeleteResult, + DocumentId, + DocumentInfo, + DocumentPage, + IngestOperation, ) +from nemo_retriever.common.vdb.adt_vdb import ( + CollectionWriteContext, + CollectionWriteResult, + UnsupportedVDBOperation, + VDB, + VDBInvalidRequest, + VDBResourceConflict, + VDBResourceNotFound, +) +from nemo_retriever.common.vdb.factory import get_vdb_op_cls +from nemo_retriever.common.vdb.records import RetrievalContractError +from nemo_retriever.operators.vdb import IngestVdbOperator, RetrieveVdbOperator from nemo_retriever.query.evidence import build_evidence_result +from nemo_retriever.service.agentic_query import run_agentic_query +from nemo_retriever.service.config import AgenticConfig from nemo_retriever.service.query_schema import ( EvidenceQueryResponse, EvidenceResult, @@ -58,19 +65,33 @@ logger = logging.getLogger(__name__) -# ── Request / response models ──────────────────────────────────────── +MAX_CONCURRENT_QUERIES = 4 +MAX_CONCURRENT_AGENTIC_QUERIES = 100 class WriteRequest(BaseModel): - rows: list[dict[str, Any]] + """Canonical internal payload emitted by an ingest worker.""" + + records: list[list[dict[str, Any]]] + scope: str | None = None + collection_name: str | None = None + document_id: DocumentId | None = None + job_id: str | None = None + filename: str | None = None + content_sha256: str | None = None + document_version: str | None = None + operation: IngestOperation = IngestOperation.APPEND class WriteResponse(BaseModel): + """Counts returned after an internal fixed-table or collection write.""" + written: int total_rows: int -# ── Embedding helpers ──────────────────────────────────────────────── +def _scope(value: str | None) -> str: + return (value or "default").strip() or "default" def _tensor_to_embedding_rows(tensor: Any) -> list[list[float]]: @@ -108,26 +129,16 @@ def _embed_queries_remote( ) -def _strategies_for_retrieval_mode(mode: LanceRetrievalMode | str) -> list[str]: - if mode == "hybrid": - return ["hybrid"] - return ["dense"] - - -# ── VectorDB state ─────────────────────────────────────────────────── - - class VectorDBState: - """Thread-safe wrapper around a LanceDB connection.""" + """Small service-owned wrapper for embedding and VDB operator dispatch.""" def __init__( self, - lancedb_uri: str, - table_name: str, + *, + vdb: VDB, embed_endpoint: str, embed_model: str, embed_api_key: str, - *, embed_model_provider_prefix: str | None = None, local_embed: bool = False, local_embed_backend: str = "hf", @@ -135,8 +146,9 @@ def __init__( device: str | None = None, gpu_memory_utilization: float = 0.45, ) -> None: - self.lancedb_uri = lancedb_uri - self.table_name = table_name + self.vdb = vdb + self.ingest_operator = IngestVdbOperator(vdb=vdb) + self.retrieve_operator = RetrieveVdbOperator(vdb=vdb) self.embed_endpoint = embed_endpoint self.embed_model = embed_model self.embed_model_provider_prefix = embed_model_provider_prefix @@ -146,17 +158,9 @@ def __init__( self.hf_cache_dir = hf_cache_dir self.device = device self.gpu_memory_utilization = gpu_memory_utilization - self._write_lock = threading.Lock() + self.query_semaphore = asyncio.Semaphore(MAX_CONCURRENT_QUERIES) self._embed_lock = threading.Lock() self._local_embedder: Any | None = None - self._db = lancedb.connect(uri=lancedb_uri) - self._table_exists = False - if self.table_name in self._db.list_tables().tables: - self._db.open_table(table_name) - self._table_exists = True - logger.info("Opened existing LanceDB table '%s' at %s", table_name, lancedb_uri) - else: - logger.info("LanceDB table '%s' does not exist yet at %s", table_name, lancedb_uri) @property def embed_mode(self) -> str: @@ -168,130 +172,8 @@ def embed_mode(self) -> str: @property def table_exists(self) -> bool: - return self._table_exists - - def _table_capabilities(self): - if not self._table_exists: - return None - table = self._db.open_table(self.table_name) - return inspect_lancedb_table_object(table) - - def resolve_effective_retrieval_mode(self, caps: LanceTableCapabilities | None = None) -> LanceRetrievalMode: - """Resolve retrieval mode from table capabilities (auto). - - Optional ``caps`` avoids re-opening the table when the caller already has it. - """ - if not self._table_exists: - return "dense" - - if caps is None: - caps = self._table_capabilities() - if caps is None: - raise ValueError( - f"Unable to inspect LanceDB table {self.table_name!r} at {self.lancedb_uri!r}: " - "capabilities could not be determined." - ) - - mode: LanceRetrievalMode = caps.retrieval_mode - if mode == "unknown": - raise ValueError( - f"LanceDB table {self.table_name!r} at {self.lancedb_uri!r} is not queryable: " - "no vector column or FTS index was detected." - ) - if mode == "sparse": - raise ValueError( - f"LanceDB table {self.table_name!r} at {self.lancedb_uri!r} has an FTS index but no vector " - "column; sparse-only retrieval is not supported by the VectorDB service." - ) - return mode - - def write_rows(self, rows: list[dict[str, Any]]) -> int: - """Append rows to the LanceDB table (creates table on first write).""" - if not rows: - return 0 - - from nemo_retriever.common.vdb.lancedb_schema import ( - create_or_append_lancedb_table, - infer_vector_dim, - lancedb_schema, - ) - - with self._write_lock: - # Decide create-vs-append on raw on-disk presence via ``list_tables`` - # rather than interpreting ``open_table`` errors, so a transient I/O - # failure cannot be misread as "table absent" and route execution - # into the destructive ``overwrite=True`` create path. - if not self._table_exists and self.table_name not in self._db.list_tables().tables: - dim = infer_vector_dim(rows) - if dim == 0: - logger.warning("Cannot infer vector dimension from rows; skipping write") - return 0 - schema = lancedb_schema(vector_dim=dim) - create_or_append_lancedb_table( - self._db, - self.table_name, - rows, - schema, - overwrite=True, - ) - self._table_exists = True - logger.info( - "Created LanceDB table '%s' with %d rows (dim=%d)", - self.table_name, - len(rows), - dim, - ) - else: - table = self._db.open_table(self.table_name) - table.add(rows) - self._table_exists = True - logger.info("Appended %d rows to table '%s'", len(rows), self.table_name) - - return len(rows) - - def total_rows(self) -> int: - if not self._table_exists: - return 0 - try: - table = self._db.open_table(self.table_name) - return table.count_rows() - except Exception: - logger.warning( - "Failed to count rows in LanceDB table '%s' at %s; reporting 0 to health", - self.table_name, - self.lancedb_uri, - exc_info=True, - ) - return 0 - - def search( - self, - vectors: list[list[float]], - query_texts: list[str], - top_k: int, - ) -> tuple[list[list[dict[str, Any]]], list[str]]: - """Search the LanceDB table with precomputed query vectors.""" - if not self._table_exists: - return [[] for _ in vectors], _strategies_for_retrieval_mode("dense") - - caps = self._table_capabilities() - mode = self.resolve_effective_retrieval_mode(caps) - strategies = _strategies_for_retrieval_mode(mode) - - from nemo_retriever.common.vdb.lancedb import LanceDB - from nemo_retriever.common.vdb.records import normalize_retrieval_results - - hybrid = mode == "hybrid" - vdb = LanceDB(uri=self.lancedb_uri, table_name=self.table_name, overwrite=False, hybrid=hybrid) - retrieval_kwargs: dict[str, Any] = {"top_k": top_k, "hybrid": hybrid} - if hybrid: - retrieval_kwargs["query_texts"] = query_texts - - if caps is not None and caps.vector_column and caps.vector_column != "vector": - retrieval_kwargs["vector_column_name"] = caps.vector_column - - raw_results = vdb.retrieval(vectors, **retrieval_kwargs) - return normalize_retrieval_results(raw_results), strategies + """Return whether the configured legacy table is available for queries.""" + return self.vdb.health().get("table_exists") is True def _get_local_embedder(self) -> Any: if self._local_embedder is None: @@ -312,7 +194,7 @@ def _get_local_embedder(self) -> Any: return self._local_embedder def embed_queries(self, texts: list[str]) -> list[list[float]]: - """Embed query texts via remote NIM or in-pod Hugging Face.""" + """Embed query texts via a remote endpoint or local model.""" if self.embed_endpoint: return _embed_queries_remote( texts, @@ -323,18 +205,47 @@ def embed_queries(self, texts: list[str]) -> list[list[float]]: ) if self.local_embed: with self._embed_lock: - embedder = self._get_local_embedder() - tensor = embedder.embed_queries(texts) + tensor = self._get_local_embedder().embed_queries(texts) return _tensor_to_embedding_rows(tensor) - raise RuntimeError("No embedding backend configured (remote endpoint or --local-embed).") + raise RuntimeError("No embedding backend configured") -# ── FastAPI app ────────────────────────────────────────────────────── +def _production_vdb( + *, + lancedb_uri: str, + table_name: str, + expiration_cleanup_enabled: bool, +) -> VDB: + """Construct the sole production VDB implementation for this service.""" + vdb_cls = get_vdb_op_cls("lancedb") + return vdb_cls( + uri=lancedb_uri, + table_name=table_name, + vector_dim=None, + overwrite=False, + build_index=False, + _service_table_schema=True, + expiration_cleanup_enabled=expiration_cleanup_enabled, + ) -_state: VectorDBState | None = None -_query_semaphore: asyncio.Semaphore | None = None -MAX_CONCURRENT_QUERIES = 4 +def _safe_backend_health(state: VectorDBState | None) -> dict[str, Any] | None: + if state is None: + return None + try: + health = state.vdb.health() + except Exception: + logger.exception("VectorDB backend health inspection failed") + return None + return dict(health) if isinstance(health, dict) else {} + + +def _legacy_strategies(health: dict[str, Any]) -> list[str]: + strategies = health.get("retrieval_strategies") + if isinstance(strategies, list) and all(isinstance(item, str) for item in strategies): + return list(strategies) + mode = health.get("effective_retrieval_mode") + return ["hybrid" if mode == "hybrid" else "dense"] def create_vectordb_app( @@ -350,15 +261,31 @@ def create_vectordb_app( hf_cache_dir: str | None = None, device: str | None = None, gpu_memory_utilization: float = 0.45, + internal_api_token: str | None = None, + reconciliation_interval_seconds: int = 60, + expiration_cleanup_enabled: bool = True, + vdb: VDB | None = None, + agentic_config: AgenticConfig | None = None, ) -> FastAPI: - """Build the VectorDB FastAPI application.""" + """Build the VectorDB FastAPI application around an injected VDB contract.""" + if reconciliation_interval_seconds < 0: + raise ValueError("reconciliation_interval_seconds must be non-negative") + + agentic_config = agentic_config or AgenticConfig() + state: VectorDBState | None = None + agentic_executor: ThreadPoolExecutor | None = None + agentic_slots: threading.BoundedSemaphore | None = None @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncIterator[None]: - global _state, _query_semaphore - _state = VectorDBState( + nonlocal state, agentic_executor, agentic_slots + backend = vdb or _production_vdb( lancedb_uri=lancedb_uri, table_name=table_name, + expiration_cleanup_enabled=expiration_cleanup_enabled, + ) + state = VectorDBState( + vdb=backend, embed_endpoint=embed_endpoint, embed_model=embed_model, embed_model_provider_prefix=embed_model_provider_prefix, @@ -369,83 +296,387 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: device=device, gpu_memory_utilization=gpu_memory_utilization, ) - _query_semaphore = asyncio.Semaphore(MAX_CONCURRENT_QUERIES) + app.state.vectordb_state = state + if agentic_config.enabled: + agentic_executor = ThreadPoolExecutor( + max_workers=MAX_CONCURRENT_AGENTIC_QUERIES, + thread_name_prefix="agentic-query", + ) + agentic_slots = threading.BoundedSemaphore(MAX_CONCURRENT_AGENTIC_QUERIES) + app.state.agentic_slots = agentic_slots logger.info( - "VectorDB service started: uri=%s table=%s embed_mode=%s max_concurrent_queries=%d", - lancedb_uri, - table_name, - _state.embed_mode, + "VectorDB service started: embed_mode=%s max_concurrent_queries=%d", + state.embed_mode, MAX_CONCURRENT_QUERIES, ) - if _state.embed_mode == "none": + if state.embed_mode == "none": logger.error( "VectorDB started without an embedding backend; /v1/query will " - "return HTTP 501 until --embed-endpoint or --local-embed is " - "configured." + "return HTTP 501 until --embed-endpoint or --local-embed is configured." ) - yield - _state = None - _query_semaphore = None - logger.info("VectorDB service stopped") + + async def reconciliation_loop() -> None: + while True: + try: + await asyncio.to_thread(backend.reconcile_collections) + except Exception: + logger.exception("VectorDB reconciliation iteration failed") + await asyncio.sleep(reconciliation_interval_seconds) + + reconciliation_task = ( + asyncio.create_task(reconciliation_loop()) if reconciliation_interval_seconds > 0 else None + ) + try: + yield + finally: + if reconciliation_task is not None: + reconciliation_task.cancel() + try: + await reconciliation_task + except asyncio.CancelledError: + pass + if agentic_executor is not None: + # Agentic LLM calls cannot be interrupted, so detach rather than + # blocking service shutdown on in-flight work. + agentic_executor.shutdown(wait=False, cancel_futures=True) + agentic_executor = None + agentic_slots = None + state = None + app.state.agentic_slots = None + app.state.vectordb_state = None + logger.info("VectorDB service stopped") app = FastAPI( title="NeMo Retriever VectorDB", - description="LanceDB-backed vector storage and retrieval", + description="Vector storage and retrieval through the VDB contract", version="1.0.0", lifespan=lifespan, ) + def require_state() -> VectorDBState: + if state is None: + raise HTTPException(503, "VectorDB not initialised") + return state + + @app.exception_handler(UnsupportedVDBOperation) + async def unsupported_operation(_request: Request, exc: UnsupportedVDBOperation) -> JSONResponse: + logger.info("Unsupported VDB operation: %s", exc) + return JSONResponse( + status_code=501, + content={"detail": "The configured VectorDB backend does not support this operation."}, + ) + + @app.exception_handler(VDBResourceNotFound) + async def resource_not_found(_request: Request, exc: VDBResourceNotFound) -> JSONResponse: + return JSONResponse(status_code=404, content={"detail": str(exc)}) + + @app.exception_handler(VDBResourceConflict) + async def resource_conflict(_request: Request, exc: VDBResourceConflict) -> JSONResponse: + return JSONResponse(status_code=409, content={"detail": str(exc)}) + + @app.exception_handler(VDBInvalidRequest) + async def invalid_request(_request: Request, exc: VDBInvalidRequest) -> JSONResponse: + return JSONResponse(status_code=422, content={"detail": str(exc)}) + + @app.exception_handler(RetrievalContractError) + async def retrieval_contract_failure(_request: Request, exc: RetrievalContractError) -> JSONResponse: + logger.exception("VectorDB retrieval contract violation", exc_info=exc) + return JSONResponse( + status_code=500, + content={"detail": "VectorDB retrieval contract violation."}, + ) + + @app.exception_handler(Exception) + async def unexpected_backend_failure(_request: Request, exc: Exception) -> JSONResponse: + logger.exception("Unexpected VectorDB service failure", exc_info=exc) + return JSONResponse( + status_code=500, + content={"detail": "VectorDB backend operation failed."}, + ) + + required_internal_token = (internal_api_token or "").strip() + + @app.middleware("http") + async def require_internal_credential(request: Request, call_next): + if request.url.path == "/v1/health" or not required_internal_token: + return await call_next(request) + supplied = request.headers.get("X-NRL-Internal-Token", "") + if not supplied or not hmac.compare_digest(supplied, required_internal_token): + return JSONResponse( + status_code=401, + content={"detail": "Missing or invalid internal credential."}, + ) + return await call_next(request) + @app.get("/v1/health", tags=["system"]) async def health() -> dict[str, Any]: - rows = _state.total_rows() if _state else 0 - effective_mode: str | None = None - if _state is not None and _state.table_exists: - try: - effective_mode = _state.resolve_effective_retrieval_mode() - except Exception: - # Health must never 500 (it backs k8s liveness/readiness probes), - # so report "unknown" for any failure — misconfiguration - # (ValueError) or transient I/O / LanceDB errors on open_table. - effective_mode = "unknown" - logger.warning( - "Failed to resolve effective retrieval mode for table '%s' at %s; reporting unknown to health", - table_name, - lancedb_uri, - exc_info=True, - ) + current = state + backend_health = _safe_backend_health(current) + if backend_health is None: + raise HTTPException(503, "VectorDB backend is unavailable") return { "status": "ok", - "table": table_name, - "total_rows": rows, - "table_exists": _state.table_exists if _state else False, - "embed_mode": _state.embed_mode if _state else "none", - "effective_retrieval_mode": effective_mode, + "total_rows": backend_health.pop("total_rows", 0), + "table_exists": backend_health.pop("table_exists", False), + "embed_mode": current.embed_mode if current else "none", + "effective_retrieval_mode": backend_health.pop("effective_retrieval_mode", None), + **backend_health, } + @app.get("/metrics", include_in_schema=False) + async def metrics() -> Response: + from prometheus_client import CollectorRegistry, Gauge, generate_latest + + registry = CollectorRegistry() + backend_health = _safe_backend_health(state) or {} + collections = backend_health.get("collections") or {} + cleanup = backend_health.get("cleanup") or {} + reconciliation = backend_health.get("reconciliation") or {} + collection_gauge = Gauge( + "nrl_vectordb_collections", + "Collection count by lifecycle status", + ["status"], + registry=registry, + ) + for status in ("active", "deleting", "expired"): + collection_gauge.labels(status=status).set(collections.get(status, 0)) + Gauge( + "nrl_vectordb_cleanup_pending", + "Pending lifecycle cleanup", + registry=registry, + ).set(cleanup.get("pending", 0)) + Gauge( + "nrl_vectordb_cleanup_oldest_age_seconds", + "Oldest pending cleanup age", + registry=registry, + ).set(cleanup.get("oldest_age_seconds", 0)) + Gauge( + "nrl_vectordb_reconciliation_successes_total", + "Successful reconciliations", + registry=registry, + ).set(reconciliation.get("successes", 0)) + Gauge( + "nrl_vectordb_reconciliation_failures_total", + "Failed reconciliations", + registry=registry, + ).set(reconciliation.get("failures", 0)) + Gauge( + "nrl_vectordb_open_table_cache", + "Open collection-table cache size", + registry=registry, + ).set(backend_health.get("open_table_cache_count", 0)) + return Response(generate_latest(registry), media_type="text/plain; version=0.0.4") + @app.post("/internal/vectordb/write", response_model=WriteResponse, tags=["internal"]) async def write(req: WriteRequest) -> WriteResponse: - if _state is None: - raise HTTPException(503, "VectorDB not initialised") - written = await asyncio.to_thread(_state.write_rows, req.rows) - return WriteResponse(written=written, total_rows=_state.total_rows()) + current = require_state() + context: CollectionWriteContext | None = None + if req.collection_name is not None: + missing = [ + name + for name, value in ( + ("document_id", req.document_id), + ("job_id", req.job_id), + ("filename", req.filename), + ("content_sha256", req.content_sha256), + ("document_version", req.document_version), + ) + if not value + ] + if missing: + raise VDBInvalidRequest("Collection writes require: " + ", ".join(missing)) + if not req.records or not any(req.records): + raise VDBInvalidRequest("Collection writes require at least one record") + context = CollectionWriteContext( + scope=_scope(req.scope), + collection_name=req.collection_name, + document_id=str(req.document_id), + document_version=str(req.document_version), + content_sha256=str(req.content_sha256), + filename=str(req.filename), + job_id=req.job_id, + operation=req.operation, + ) - @app.post("/v1/query", response_model=Union[QueryResponse, EvidenceQueryResponse], tags=["query"]) - async def query(req: QueryRequest) -> QueryResponse | EvidenceQueryResponse: - if _state is None: - raise HTTPException(503, "VectorDB not initialised") + result = await asyncio.to_thread( + current.ingest_operator.run, + req.records, + collection_context=context, + ) + if isinstance(result, CollectionWriteResult): + return WriteResponse(written=result.written, total_rows=result.total_rows) + backend_health = current.vdb.health() + return WriteResponse( + written=sum(len(batch) for batch in req.records), + total_rows=int(backend_health.get("total_rows", 0)), + ) - if _state.embed_mode == "none": + @app.post( + "/v1/collections", + response_model=CollectionInfo, + status_code=201, + tags=["collections"], + ) + async def create_collection( + req: CollectionCreateRequest, + x_nrl_scope: str | None = Header(None), + ) -> CollectionInfo: + backend = require_state().vdb + return await asyncio.to_thread( + backend.create_collection, + scope=_scope(x_nrl_scope), + request=req, + ) + + @app.get("/v1/collections", response_model=CollectionPage, tags=["collections"]) + async def list_collections( + limit: int = Query(100, ge=1, le=200), + continuation_token: str | None = None, + x_nrl_scope: str | None = Header(None), + ) -> CollectionPage: + backend = require_state().vdb + return await asyncio.to_thread( + backend.list_collections, + scope=_scope(x_nrl_scope), + limit=limit, + continuation_token=continuation_token, + ) + + @app.get("/v1/collections/{name}", response_model=CollectionInfo, tags=["collections"]) + async def get_collection( + name: str, + x_nrl_scope: str | None = Header(None), + ) -> CollectionInfo: + backend = require_state().vdb + return await asyncio.to_thread( + backend.get_collection, + scope=_scope(x_nrl_scope), + collection_name=name, + ) + + @app.patch("/v1/collections/{name}", response_model=CollectionInfo, tags=["collections"]) + async def update_collection( + name: str, + req: CollectionUpdateRequest, + x_nrl_scope: str | None = Header(None), + ) -> CollectionInfo: + backend = require_state().vdb + return await asyncio.to_thread( + backend.update_collection, + scope=_scope(x_nrl_scope), + collection_name=name, + request=req, + ) + + @app.delete( + "/v1/collections/{name}", + response_model=CollectionDeleteResult, + tags=["collections"], + ) + async def delete_collection( + response: Response, + name: str, + if_exists: bool = False, + x_nrl_scope: str | None = Header(None), + ) -> CollectionDeleteResult: + backend = require_state().vdb + result = await asyncio.to_thread( + backend.delete_collection, + scope=_scope(x_nrl_scope), + collection_name=name, + if_exists=if_exists, + ) + response.status_code = 202 if result.cleanup_pending else 200 + return result + + @app.get( + "/v1/collections/{name}/documents", + response_model=DocumentPage, + tags=["collections"], + ) + async def list_documents( + name: str, + limit: int = Query(100, ge=1, le=200), + continuation_token: str | None = None, + x_nrl_scope: str | None = Header(None), + ) -> DocumentPage: + backend = require_state().vdb + return await asyncio.to_thread( + backend.list_documents, + scope=_scope(x_nrl_scope), + collection_name=name, + limit=limit, + continuation_token=continuation_token, + ) + + @app.get( + "/v1/collections/{name}/documents/{document_id}", + response_model=DocumentInfo, + tags=["collections"], + ) + async def get_document( + name: str, + document_id: DocumentId, + x_nrl_scope: str | None = Header(None), + ) -> DocumentInfo: + backend = require_state().vdb + return await asyncio.to_thread( + backend.get_document, + scope=_scope(x_nrl_scope), + collection_name=name, + document_id=document_id, + ) + + @app.delete( + "/v1/collections/{name}/documents/{document_id}", + response_model=DocumentDeleteResult, + tags=["collections"], + ) + async def delete_document( + response: Response, + name: str, + document_id: DocumentId, + if_exists: bool = False, + x_nrl_scope: str | None = Header(None), + ) -> DocumentDeleteResult: + backend = require_state().vdb + result = await asyncio.to_thread( + backend.delete_document, + scope=_scope(x_nrl_scope), + collection_name=name, + document_id=document_id, + if_exists=if_exists, + ) + response.status_code = 202 if result.cleanup_pending else 200 + return result + + @app.post( + "/v1/query", + response_model=Union[QueryResponse, EvidenceQueryResponse], + tags=["query"], + ) + async def query( + req: QueryRequest, + x_nrl_scope: str | None = Header(None), + ) -> QueryResponse | EvidenceQueryResponse: + current = require_state() + if req.agentic: + if req.collection_name is not None: + raise UnsupportedVDBOperation("agentic collection retrieval") + return await _run_agentic_query(req) + + if current.embed_mode == "none": raise HTTPException( 501, "No embedding backend configured. Set --embed-endpoint for a remote " "NIM or --local-embed for in-pod Hugging Face query embedding.", ) - if not _state.table_exists: - raise HTTPException( - status_code=422, - detail="No data has been ingested yet. Ingest documents first, then query.", - ) + backend_health: dict[str, Any] = {} + if req.collection_name is None: + backend_health = current.vdb.health() + if backend_health.get("table_exists") is False: + raise VDBInvalidRequest("No data has been ingested yet. Ingest documents first, then query.") queries = req.query if isinstance(req.query, list) else [req.query] if not queries: @@ -453,47 +684,150 @@ async def query(req: QueryRequest) -> QueryResponse | EvidenceQueryResponse: return EvidenceQueryResponse(results=[]) return QueryResponse(results=[]) - try: - async with _query_semaphore: - vectors = await asyncio.to_thread(_state.embed_queries, queries) - hits_per_query, strategies = await asyncio.to_thread( - _state.search, + async with current.query_semaphore: + vectors = await asyncio.to_thread(current.embed_queries, queries) + if req.collection_name is not None: + result = await asyncio.to_thread( + current.retrieve_operator.run, vectors, - queries, - req.top_k, + scope=_scope(x_nrl_scope), + collection_name=req.collection_name, + query_texts=queries, + top_k=req.top_k, ) - except ValueError as exc: - raise HTTPException(status_code=422, detail=str(exc)) from exc + if not isinstance(result, tuple): + raise RetrievalContractError("Collection retrieval did not return strategies") + hits_per_query, strategies = result + else: + hits_per_query = await asyncio.to_thread( + current.retrieve_operator.run, + vectors, + query_texts=queries, + top_k=req.top_k, + hybrid=backend_health.get("effective_retrieval_mode") == "hybrid", + ) + if not isinstance(hits_per_query, list): + raise RetrievalContractError("Legacy retrieval returned an invalid shape") + strategies = _legacy_strategies(backend_health) if req.format == "evidence": return EvidenceQueryResponse( results=[EvidenceResult(**build_evidence_result(hits, strategies)) for hits in hits_per_query] ) - return QueryResponse(results=[QueryResult(hits=hits) for hits in hits_per_query]) - return app + async def _run_agentic_query(req: QueryRequest) -> QueryResponse: + """Run the blocking agentic workflow without consuming plain-query workers.""" + current = require_state() + if not agentic_config.enabled: + raise HTTPException( + status_code=400, + detail=( + "Agentic retrieval is not enabled in the VectorDB configuration. " + "Start with --agentic and a remote LLM invoke URL/model, or set " + "agentic.enabled in the service config." + ), + ) + if current.embed_mode == "none": + raise HTTPException( + 501, + "No embedding backend configured. Set --embed-endpoint for a remote " + "NIM or --local-embed for in-pod Hugging Face query embedding.", + ) + if current.embed_mode != "remote": + raise HTTPException(501, "Agentic service queries require a remote embedding endpoint.") + if not current.table_exists: + raise VDBInvalidRequest("No data has been ingested yet. Ingest documents first, then query.") + if req.top_k > agentic_config.backend_top_k: + raise VDBInvalidRequest( + f"top_k ({req.top_k}) cannot exceed the configured agentic " + f"backend_top_k ({agentic_config.backend_top_k})." + ) + + executor, slots = agentic_executor, agentic_slots + if executor is None or slots is None: + raise HTTPException(503, "Agentic retrieval workers are not running") + if not slots.acquire(blocking=False): + logger.warning( + "Rejecting agentic query: all %d agentic workers are busy", + MAX_CONCURRENT_AGENTIC_QUERIES, + ) + raise HTTPException( + status_code=503, + detail=( + f"All {MAX_CONCURRENT_AGENTIC_QUERIES} agentic retrieval workers are busy. " + "Retry once an in-flight query finishes." + ), + headers={"Retry-After": "30"}, + ) + assert isinstance(req.query, str) + try: + future = executor.submit( + run_agentic_query, + query=req.query, + top_k=req.top_k, + config=agentic_config, + lancedb_uri=lancedb_uri, + table_name=table_name, + embed_endpoint=current.embed_endpoint, + embed_model=current.embed_model, + embed_model_provider_prefix=current.embed_model_provider_prefix, + embed_api_key=current.embed_api_key, + ) + except RuntimeError as exc: + slots.release() + raise HTTPException(503, "VectorDB is shutting down") from exc -# ── CLI entry point ────────────────────────────────────────────────── + future.add_done_callback(lambda _future: slots.release()) + try: + return await asyncio.wrap_future(future) + except ValueError as exc: + raise VDBInvalidRequest(str(exc)) from exc + + return app def main() -> None: + internal_token = os.environ.get("NRL_INTERNAL_VDB_TOKEN", "") + if not internal_token and (token_file := os.environ.get("NRL_INTERNAL_VDB_TOKEN_FILE")): + internal_token = Path(token_file).read_text(encoding="utf-8").strip() + parser = argparse.ArgumentParser(description="NeMo Retriever VectorDB service") parser.add_argument("--lancedb-uri", default="/data/vectordb", help="LanceDB directory") - parser.add_argument("--table-name", default="nemo_retriever", help="LanceDB table name") + parser.add_argument("--table-name", default="nemo_retriever", help="Vector table name") parser.add_argument("--embed-endpoint", default="", help="Remote NIM/OpenAI-compatible embed URL") parser.add_argument("--embed-model", default="nvidia/llama-nemotron-embed-vl-1b-v2") - parser.add_argument("--embed-model-provider-prefix", default="", help="Optional LiteLLM provider prefix") + parser.add_argument( + "--embed-model-provider-prefix", + default="", + help="Optional LiteLLM provider prefix", + ) parser.add_argument( "--embed-api-key", default="", help="Remote embedding API key (defaults to NVIDIA_API_KEY, then NGC_API_KEY).", ) + parser.add_argument( + "--internal-api-token", + default=internal_token, + help="Dedicated internal credential (prefer NRL_INTERNAL_VDB_TOKEN from a Secret).", + ) + parser.add_argument( + "--reconciliation-interval-seconds", + type=int, + default=int(os.environ.get("NRL_RECONCILIATION_INTERVAL_SECONDS", "60")), + help="Lifecycle reconciliation interval; zero disables the local loop.", + ) + parser.add_argument( + "--disable-expiration-cleanup", + action="store_true", + help="Disable automatic collection expiration cleanup.", + ) parser.add_argument( "--local-embed", action="store_true", - help="Load Hugging Face embedder in-pod for /v1/query (requires [local] extras + GPU).", + help="Load a local embedder for /v1/query (requires local extras and a GPU).", ) parser.add_argument( "--local-embed-backend", @@ -502,13 +836,34 @@ def main() -> None: help="Backend for --local-embed (default: hf).", ) parser.add_argument("--hf-cache-dir", default="", help="Hugging Face model cache directory") - parser.add_argument("--device", default="", help="Torch device for --local-embed (e.g. cuda:0)") + parser.add_argument( + "--device", + default="", + help="Torch device for --local-embed (for example cuda:0)", + ) parser.add_argument( "--gpu-memory-utilization", type=float, default=0.45, help="vLLM GPU memory fraction when --local-embed-backend=vllm.", ) + parser.add_argument( + "--agentic", + action="store_true", + help="Enable agentic=true on POST /v1/query using the agentic retrieval workflow.", + ) + parser.add_argument("--agentic-llm-model", default="", help="Agentic retrieval chat model.") + parser.add_argument( + "--agentic-invoke-url", + default="", + help="OpenAI-compatible chat completions endpoint for agentic retrieval.", + ) + parser.add_argument("--agentic-reasoning-effort", default="high") + parser.add_argument("--agentic-backend-top-k", type=int, default=20) + parser.add_argument("--agentic-react-max-steps", type=int, default=50) + parser.add_argument("--agentic-text-truncation", type=int, default=0) + parser.add_argument("--agentic-temperature", type=float, default=0.0) + parser.add_argument("--agentic-request-timeout", type=float, default=1800.0) parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=7671) parser.add_argument("--log-level", default="info") @@ -534,6 +889,20 @@ def main() -> None: hf_cache_dir=args.hf_cache_dir or None, device=args.device or None, gpu_memory_utilization=args.gpu_memory_utilization, + internal_api_token=args.internal_api_token or None, + reconciliation_interval_seconds=args.reconciliation_interval_seconds, + expiration_cleanup_enabled=not args.disable_expiration_cleanup, + agentic_config=AgenticConfig( + enabled=args.agentic, + llm_model=args.agentic_llm_model or None, + invoke_url=args.agentic_invoke_url or None, + reasoning_effort=args.agentic_reasoning_effort or None, + backend_top_k=args.agentic_backend_top_k, + react_max_steps=args.agentic_react_max_steps, + text_truncation=args.agentic_text_truncation, + temperature=args.agentic_temperature, + request_timeout_s=args.agentic_request_timeout, + ), ) uvicorn.run(app, host=args.host, port=args.port) diff --git a/nemo_retriever/tests/query/test_evidence.py b/nemo_retriever/tests/query/test_evidence.py index 2cfdcd6327..f21bb00b1b 100644 --- a/nemo_retriever/tests/query/test_evidence.py +++ b/nemo_retriever/tests/query/test_evidence.py @@ -6,6 +6,8 @@ import json +import pytest + from nemo_retriever.query.evidence import build_evidence_result @@ -79,3 +81,29 @@ def test_visual_only_match_is_reported_when_text_evidence_remains() -> None: "n_docs_seen": 2, "thin_spots": ["single source", "visual-only matches omitted"], } + + +@pytest.mark.parametrize( + ("score_fields", "expected"), + [ + ({"distance": 0.17}, 0.17), + ({"_score": 0.81}, 0.81), + ({"_distance": 0.24}, 0.24), + ({}, 0.0), + ], +) +def test_reachable_ranking_values_are_preserved_in_evidence(score_fields: dict[str, float], expected: float) -> None: + result = build_evidence_result( + [ + { + "text": "answer-ready text", + "source": "report.pdf", + "page_number": 2, + "metadata": {"type": "text"}, + **score_fields, + } + ], + ["dense"], + ) + + assert result["evidence"][0]["score"] == expected diff --git a/nemo_retriever/tests/test_actor_operators.py b/nemo_retriever/tests/test_actor_operators.py index 1b652cc532..00305e5ef9 100644 --- a/nemo_retriever/tests/test_actor_operators.py +++ b/nemo_retriever/tests/test_actor_operators.py @@ -5,9 +5,11 @@ """Unit tests verifying all pipeline actors inherit from AbstractOperator.""" import inspect +import json from unittest.mock import MagicMock, patch import pandas as pd +import pytest from nemo_retriever.operators.abstract_operator import AbstractOperator @@ -441,6 +443,7 @@ def invoke_chat_completions_images(self, **kwargs): assert client.kwargs["model"] == NEMOTRON_PARSE_REMOTE_DEFAULT_MODEL assert client.kwargs["task_prompt"] == NEMOTRON_PARSE_DEFAULT_TASK_PROMPT assert client.kwargs["extra_body"] == {"max_tokens": 8192} + assert client.kwargs["repetition_penalty"] == 1.1 def test_remote_chat_completions_supports_legacy_tool_call_protocol(self): from nemo_retriever.operators.extract.parse.nemotron_parse import nemotron_parse_pages @@ -472,6 +475,7 @@ def invoke_chat_completions_images(self, **kwargs): "max_tokens": 8192, "tools": [{"type": "function", "function": {"name": "markdown_bbox"}}], } + assert client.kwargs["repetition_penalty"] == 1.1 def test_remote_chat_completions_does_not_treat_v1_10_as_legacy(self): from nemo_retriever.operators.extract.parse.nemotron_parse import nemotron_parse_pages @@ -499,6 +503,144 @@ def invoke_chat_completions_images(self, **kwargs): assert client.kwargs["task_prompt"] is not None assert client.kwargs["extra_body"] == {"max_tokens": 8192} + def test_hosted_build_contract_is_image_only_and_routes_nested_tool_json(self): + from nemo_retriever.operators.extract.parse.nemotron_parse import nemotron_parse_pages + + class _FakeNIMClient: + def __init__(self): + self.kwargs = None + + def invoke_chat_completions_images(self, **kwargs): + self.kwargs = kwargs + elements = [ + {"type": "Text", "bbox": {"xmin": 0, "ymin": 0, "xmax": 1, "ymax": 1}, "text": "Hosted text"}, + {"type": "Table", "bbox": {"xmin": 1, "ymin": 1, "xmax": 2, "ymax": 2}, "text": "A | B"}, + {"type": "Chart", "bbox": {"xmin": 2, "ymin": 2, "xmax": 3, "ymax": 3}, "text": "Chart text"}, + {"type": "Picture", "bbox": {"xmin": 3, "ymin": 3, "xmax": 4, "ymax": 4}, "text": "Picture text"}, + ] + return [json.dumps({"result": [elements]})] + + client = _FakeNIMClient() + df = pd.DataFrame({"page_image": [{"image_b64": "aW1hZ2U="}]}) + result = nemotron_parse_pages( + df, + invoke_url="https://integrate.api.nvidia.com/v1/chat/completions", + extract_text=True, + extract_tables=True, + extract_charts=True, + extract_infographics=True, + nim_client=client, + ) + + assert client.kwargs["model"] == "nvidia/nemotron-parse" + assert client.kwargs["task_prompt"] is None + assert client.kwargs["repetition_penalty"] is None + assert client.kwargs["extra_body"] == {"max_tokens": 8192} + assert result["text"].tolist() == ["Hosted text"] + assert len(result.at[0, "table"]) == 1 + assert len(result.at[0, "chart"]) == 1 + assert len(result.at[0, "infographic"]) == 1 + + @pytest.mark.parametrize( + ("endpoint", "model", "expected_model", "expected_profile"), + [ + ("https://integrate.api.nvidia.com/v1/chat/completions", None, "nvidia/nemotron-parse", "hosted_tool_call"), + ("http://parse:8000/v1/chat/completions", None, "nvidia/nemotron-parse-v1.2", "v1_2_tagged"), + ( + "http://parse:8000/v1/chat/completions", + "nvidia/nemotron-parse", + "nvidia/nemotron-parse", + "hosted_tool_call", + ), + ( + "http://parse:8000/v1/chat/completions", + "nvidia/nemotron-parse-v1.0", + "nvidia/nemotron-parse-v1.0", + "legacy_tool_call", + ), + ( + "http://parse:8000/v1/chat/completions", + "nvidia/nemotron-parse-v1.1", + "nvidia/nemotron-parse-v1.1", + "legacy_tool_call", + ), + ( + "http://parse:8000/v1/chat/completions", + "nvidia/nemotron-parse-v1.2", + "nvidia/nemotron-parse-v1.2", + "v1_2_tagged", + ), + ( + "http://parse:8000/v1/chat/completions", + "nvidia/nemotron-parse-v1.10", + "nvidia/nemotron-parse-v1.10", + "v1_2_tagged", + ), + ("http://parse:8000/v1/chat/completions", "custom/parse", "custom/parse", "v1_2_tagged"), + ], + ) + def test_contract_resolution(self, endpoint, model, expected_model, expected_profile): + from nemo_retriever.operators.extract.parse.nemotron_parse import _resolve_nemotron_parse_contract + + contract = _resolve_nemotron_parse_contract(endpoint, model) + + assert contract.model == expected_model + assert contract.profile.value == expected_profile + + def test_contract_resolution_rejects_mixed_endpoints_without_model(self): + from nemo_retriever.operators.extract.parse.nemotron_parse import _resolve_nemotron_parse_contract + + endpoints = "https://integrate.api.nvidia.com/v1/chat/completions," "http://parse:8000/v1/chat/completions" + with pytest.raises(ValueError, match="cannot mix NVIDIA Build and self-hosted"): + _resolve_nemotron_parse_contract(endpoints, None) + + contract = _resolve_nemotron_parse_contract(endpoints, "nvidia/nemotron-parse-v1.2") + assert contract.profile.value == "v1_2_tagged" + + def test_forced_v1_2_build_text_rejection_reports_contract_mismatch(self): + from nemo_retriever.operators.extract.parse.nemotron_parse import nemotron_parse_pages + + class _RejectingNIMClient: + def invoke_chat_completions_images(self, **kwargs): + raise RuntimeError("Content cannot be a plain string; model does not support text input") + + df = pd.DataFrame({"page_image": [{"image_b64": "aW1hZ2U="}]}) + result = nemotron_parse_pages( + df, + invoke_url="https://integrate.api.nvidia.com/v1/chat/completions", + nemotron_parse_model="nvidia/nemotron-parse-v1.2", + nim_client=_RejectingNIMClient(), + ) + + error = result.at[0, "nemotron_parse_v1_2"]["error"] + assert error["type"] == "ValueError" + assert "model/contract mismatch" in error["message"] + assert "nvidia/nemotron-parse" in error["message"] + assert "RuntimeError: Content cannot be a plain string" in error["traceback"] + assert "ValueError: Nemotron Parse model/contract mismatch" in error["traceback"] + + def test_image_wrapper_can_omit_repetition_penalty(self): + from nemo_retriever.models.nim.nim import NIMClient + + client = NIMClient(max_pool_workers=1) + try: + with patch.object(client, "invoke_chat_completions", return_value=["ok"]) as invoke: + result = client.invoke_chat_completions_images( + invoke_url="https://integrate.api.nvidia.com/v1/chat/completions", + image_b64_list=["aW1hZ2U="], + model="nvidia/nemotron-parse", + repetition_penalty=None, + extra_body={"max_tokens": 8192}, + ) + assert result == ["ok"] + payload_args = invoke.call_args.kwargs + assert payload_args["messages_list"][0][0]["content"] == [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}} + ] + assert payload_args["extra_body"] == {"max_tokens": 8192} + finally: + client.shutdown() + @patch( "nemo_retriever.operators.extract.parse.nemotron_parse.nemotron_parse_pages", side_effect=RuntimeError("boom") ) @@ -609,7 +751,7 @@ def test_inherits(self): def test_preprocess_empty(self): actor = self._make() result = actor.preprocess(pd.DataFrame()) - assert list(result.columns) == ["text", "path", "page_number", "metadata"] + assert list(result.columns) == ["text", "content", "path", "page_number", "metadata"] def test_postprocess_passthrough(self): actor = self._make() @@ -634,6 +776,20 @@ def test_call_delegates(self, mock_fn): result = actor(pd.DataFrame({"bytes": [b"hello"], "path": ["/a.txt"]})) pd.testing.assert_frame_equal(result, expected) + @patch( + "nemo_retriever.operators.extract.txt.ray_data.text_to_chunks_df", + side_effect=RuntimeError("tokenizer failed"), + ) + def test_inline_failure_is_logged_with_source(self, mock_fn, caplog): + actor = self._make() + + with caplog.at_level("WARNING", logger="nemo_retriever.operators.extract.txt.ray_data"): + result = actor.process(pd.DataFrame({"text": ["hello"], "path": ["inline://00000000"]})) + + assert result.empty + record = next(record for record in caplog.records if "inline://00000000" in record.getMessage()) + assert record.exc_info is not None + # --------------------------------------------------------------------------- # 11. HtmlSplitActor diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index e714904fb6..b6dfb6a5aa 100644 --- a/nemo_retriever/tests/test_agentic_eval.py +++ b/nemo_retriever/tests/test_agentic_eval.py @@ -5,7 +5,7 @@ from __future__ import annotations import json -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pandas as pd import pytest @@ -60,17 +60,6 @@ def query(self, query: str, *, top_k: int | None = None): "_score": 0.1, }, ] - hits.extend( - { - "source": f"/tmp/extra_{idx}.pdf", - "source_id": f"/tmp/extra_{idx}.pdf", - "page_number": idx + 3, - "pdf_page": f"extra_{idx}", - "text": f"extra document {idx}", - "_score": 0.05, - } - for idx in range(8) - ) return hits[:top_k] def queries(self, queries, *, top_k: int | None = None): @@ -99,138 +88,70 @@ def test_build_beir_run_from_ranked_doc_ids_rejects_length_mismatch(): build_beir_run_from_ranked_doc_ids(["q1", "q2"], [["d1"]]) -def test_agentic_config_validates_max_tokens(): - from nemo_retriever.query.agentic import AgenticRetrievalConfig +def _dispatch_chat_fn(react_response, selection_response): + """Fake in-process completion callable shared by both agents. + + The ReAct and selection agents share one injected ``chat_completion_fn``, so + the fake returns the selection response whenever the selection tool is offered + and the ReAct response otherwise. + """ - cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", max_tokens="128") + def fn(**kwargs): + tool_names = {(tool.get("function") or {}).get("name") for tool in (kwargs.get("tools") or [])} + if "log_selected_documents" in tool_names: + return selection_response + return react_response - assert cfg.max_tokens == 128 - with pytest.raises(ValueError, match="max_tokens"): - AgenticRetrievalConfig(llm_model="nemotron-8b", max_tokens=0) + return fn -@patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") -@patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") @patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) -def test_agentic_retriever_runs_graph_with_wrapped_retriever(mock_react_step, mock_selection_step): +def test_agentic_retriever_runs_graph_with_wrapped_retriever(): from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever - final_ids = ["doc_1", "other_2"] + [f"extra_{i}" for i in range(8)] - mock_react_step.return_value = _make_tool_call_response( - "final_results", - {"doc_ids": final_ids, "message": "done", "search_successful": "true"}, - ) - mock_selection_step.return_value = _make_tool_call_response( - "log_selected_documents", - {"doc_ids": ["doc_1"], "message": "doc_1 is best"}, + final_ids = ["doc_1"] + [f"extra_{i}" for i in range(9)] + chat_fn = _dispatch_chat_fn( + _make_tool_call_response( + "final_results", {"doc_ids": final_ids, "message": "done", "search_successful": "true"} + ), + _make_tool_call_response("log_selected_documents", {"doc_ids": ["doc_1"], "message": "doc_1 is best"}), ) - cfg = AgenticRetrievalConfig( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - max_tokens=77, - ) - result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) + # In-process path -> callable client backend; inject the fake completion fn. + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b") + with patch("nemo_retriever.query.agentic._build_agent_chat_completion_fn", return_value=chat_fn): + retriever = AgenticRetriever(cfg, match_mode="pdf_page") + result = retriever.retrieve(["0"], ["find doc"]) - assert mock_react_step.call_args.kwargs["max_tokens"] == 77 + assert "local_ingest_embed_backend" not in retriever._retriever.kwargs["embed_kwargs"] assert list(result.columns) == ["query_id", "doc_id", "rank", "message", "result_source"] assert result["query_id"].tolist() == ["0"] * 10 assert result["doc_id"].tolist()[0] == "doc_1" assert result["rank"].tolist() == list(range(1, 11)) -@patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") -@patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") @patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) -def test_agentic_retriever_honors_top_k(mock_react_step, mock_selection_step): +def test_agentic_retriever_honors_top_k(): """cfg.top_k drives the pipeline output count, not the hardcoded default of 10.""" from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever - final_ids = ["doc_1", "other_2"] + [f"extra_{i}" for i in range(3)] # exactly 5 - mock_react_step.return_value = _make_tool_call_response( - "final_results", - {"doc_ids": final_ids, "message": "done", "search_successful": "true"}, - ) - mock_selection_step.return_value = _make_tool_call_response( - "log_selected_documents", - {"doc_ids": ["doc_1"], "message": "doc_1 is best"}, + final_ids = ["doc_1"] + [f"extra_{i}" for i in range(4)] # exactly 5 + chat_fn = _dispatch_chat_fn( + _make_tool_call_response( + "final_results", {"doc_ids": final_ids, "message": "done", "search_successful": "true"} + ), + _make_tool_call_response("log_selected_documents", {"doc_ids": ["doc_1"], "message": "doc_1 is best"}), ) - cfg = AgenticRetrievalConfig( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=5, - ) - result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", top_k=5) + with patch("nemo_retriever.query.agentic._build_agent_chat_completion_fn", return_value=chat_fn): + result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) assert result["rank"].tolist() == list(range(1, 6)) # 5 rows, honoring top_k=5 -@patch("nemo_retriever.models.create_local_agent_llm") -@patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) -def test_agentic_retriever_builds_in_process_llm_lazily(mock_create_local_agent_llm): - from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever - - local_chat = MagicMock( - return_value=_make_tool_call_response( - "final_results", - {"doc_ids": ["doc_1"], "message": "done", "search_successful": "true"}, - ) - ) - mock_create_local_agent_llm.return_value = local_chat - - cfg = AgenticRetrievalConfig(top_k=1) - retriever = AgenticRetriever(cfg, match_mode="pdf_page") - - mock_create_local_agent_llm.assert_not_called() - - result = retriever.retrieve(["0"], ["find doc"]) - - mock_create_local_agent_llm.assert_called_once_with( - "nemotron-8b", - backend="vllm", - hf_cache_dir=None, - gpu_memory_utilization=0.8, - tensor_parallel_size=1, - max_model_len=None, - max_num_seqs=None, - ) - assert local_chat.call_count == 1 - assert result["doc_id"].tolist() == ["doc_1"] - - -@patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") -@patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") -@patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) -def test_agentic_retriever_rejects_partial_react_final_results(mock_react_step, mock_selection_step): - from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever - - mock_react_step.return_value = _make_tool_call_response( - "final_results", - {"doc_ids": ["doc_1"], "message": "partial but valid", "search_successful": "true"}, - ) - mock_selection_step.return_value = _make_tool_call_response( - "log_selected_documents", - {"doc_ids": ["doc_1"], "message": "selection should not run"}, - ) - - cfg = AgenticRetrievalConfig( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=5, - react_max_steps=1, - ) - result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) - - assert len(result) == 5 - assert "final_results" not in set(result["result_source"]) - mock_selection_step.assert_not_called() - - -@patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") -@patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") @patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) -def test_run_agentic_audio_recall_evaluation_computes_metrics(mock_react_step, mock_selection_step, tmp_path): +def test_run_agentic_audio_recall_evaluation_computes_metrics(tmp_path): from nemo_retriever.query.agentic import AgenticRetrievalConfig, run_agentic_audio_recall_evaluation query_csv = tmp_path / "queries.csv" @@ -245,21 +166,20 @@ def test_run_agentic_audio_recall_evaluation_computes_metrics(mock_react_step, m audio_doc_id = "clip 1.000000 3.000000" final_ids = [audio_doc_id] + [f"extra_{i}" for i in range(9)] - mock_react_step.return_value = _make_tool_call_response( - "final_results", - {"doc_ids": final_ids, "message": "done", "search_successful": "true"}, - ) - mock_selection_step.return_value = _make_tool_call_response( - "log_selected_documents", - {"doc_ids": [audio_doc_id], "message": "clip is best"}, + chat_fn = _dispatch_chat_fn( + _make_tool_call_response( + "final_results", {"doc_ids": final_ids, "message": "done", "search_successful": "true"} + ), + _make_tool_call_response("log_selected_documents", {"doc_ids": [audio_doc_id], "message": "clip is best"}), ) - cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions") - df_query, result, gold, retrieved, metrics = run_agentic_audio_recall_evaluation( - query_csv=query_csv, - cfg=cfg, - ks=(1, 5, 10), - ) + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b") + with patch("nemo_retriever.query.agentic._build_agent_chat_completion_fn", return_value=chat_fn): + df_query, result, gold, retrieved, metrics = run_agentic_audio_recall_evaluation( + query_csv=query_csv, + cfg=cfg, + ks=(1, 5, 10), + ) assert df_query["golden_answer"].tolist() == ["clip 0.000000 4.000000"] assert result["doc_id"].tolist()[0] == audio_doc_id @@ -268,21 +188,17 @@ def test_run_agentic_audio_recall_evaluation_computes_metrics(mock_react_step, m assert metrics["recall@1"] == 1.0 -@patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") -@patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") @patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) -def test_run_agentic_beir_evaluation_loads_queries_and_qrels(mock_react_step, mock_selection_step): +def test_run_agentic_beir_evaluation_loads_queries_and_qrels(): from nemo_retriever.query.agentic import AgenticRetrievalConfig, run_agentic_beir_evaluation from nemo_retriever.tools.recall.beir import BeirDataset - final_ids = ["doc", "other"] + [f"extra_{i}" for i in range(8)] - mock_react_step.return_value = _make_tool_call_response( - "final_results", - {"doc_ids": final_ids, "message": "done", "search_successful": "true"}, - ) - mock_selection_step.return_value = _make_tool_call_response( - "log_selected_documents", - {"doc_ids": ["doc"], "message": "doc is best"}, + final_ids = ["doc"] + [f"extra_{i}" for i in range(9)] + chat_fn = _dispatch_chat_fn( + _make_tool_call_response( + "final_results", {"doc_ids": final_ids, "message": "done", "search_successful": "true"} + ), + _make_tool_call_response("log_selected_documents", {"doc_ids": ["doc"], "message": "doc is best"}), ) beir_dataset = BeirDataset( @@ -291,9 +207,11 @@ def test_run_agentic_beir_evaluation_loads_queries_and_qrels(mock_react_step, mo queries=["find doc"], qrels={"q1": {"doc": 1}}, ) - cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions") + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b") - with patch("nemo_retriever.query.agentic.load_beir_dataset", return_value=beir_dataset) as mock_loader: + with patch("nemo_retriever.query.agentic._build_agent_chat_completion_fn", return_value=chat_fn), patch( + "nemo_retriever.query.agentic.load_beir_dataset", return_value=beir_dataset + ) as mock_loader: df_query, result, qrels, run, metrics = run_agentic_beir_evaluation( loader="vidore_hf", dataset_name="vidore_v3_finance_en", @@ -310,47 +228,45 @@ def test_run_agentic_beir_evaluation_loads_queries_and_qrels(mock_react_step, mo assert metrics["recall@1"] == 1.0 -def test_agentic_config_defaults_empty_in_process_llm_model_to_nemotron_8b(): - from nemo_retriever.query.agentic import AgenticRetrievalConfig +_REMOTE_URL = "http://localhost/v1/chat/completions" - cfg = AgenticRetrievalConfig(llm_model="") - assert cfg.llm_backend == "in_process" - assert cfg.local_llm_backend == "vllm" - assert cfg.llm_model == "nemotron-8b" - cfg = AgenticRetrievalConfig(llm_model=None) - assert cfg.llm_model == "nemotron-8b" - - -def test_agentic_config_requires_llm_model_for_openai_compatible(): +def test_agentic_config_requires_llm_model_on_remote_path(): from nemo_retriever.query.agentic import AgenticRetrievalConfig + # A model is required only on the remote (invoke_url) path; in-process runs + # default to the local model instead of raising. with pytest.raises(ValueError, match="llm_model"): - AgenticRetrievalConfig(llm_model="", invoke_url="http://localhost/v1/chat/completions") + AgenticRetrievalConfig(llm_model="", invoke_url=_REMOTE_URL) # None must not slip through as the literal string "None". with pytest.raises(ValueError, match="llm_model"): - AgenticRetrievalConfig(llm_model=None, invoke_url="http://localhost/v1/chat/completions") + AgenticRetrievalConfig(llm_model=None, invoke_url=_REMOTE_URL) -def test_agentic_config_rejects_custom_in_process_llm_model(): +def test_agentic_config_defaults_in_process_model_and_client(): from nemo_retriever.query.agentic import AgenticRetrievalConfig - with pytest.raises(ValueError, match="Custom in-process agent LLMs are not supported yet"): - AgenticRetrievalConfig(llm_model="custom/local-model") + # No invoke_url and no model -> local in-process default with the callable + # LLM client. + cfg = AgenticRetrievalConfig(llm_model="") + + assert cfg.llm_backend == "in_process" + assert cfg.llm_model == "nemotron-8b" + assert cfg.llm_client == "callable" def test_agentic_config_rejects_nonpositive_top_k(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="top_k"): - AgenticRetrievalConfig(llm_model="nemotron-8b", top_k=0) + AgenticRetrievalConfig(llm_model="m", invoke_url=_REMOTE_URL, top_k=0) def test_agentic_config_rejects_noninteger_top_k(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="top_k must be an integer"): - AgenticRetrievalConfig(llm_model="nemotron-8b", top_k=1.5) + AgenticRetrievalConfig(llm_model="m", invoke_url=_REMOTE_URL, top_k=1.5) def test_agentic_config_normalizes_integer_like_values(): @@ -358,22 +274,21 @@ def test_agentic_config_normalizes_integer_like_values(): cfg = AgenticRetrievalConfig( llm_model="m", - invoke_url="http://localhost/v1/chat/completions", + invoke_url=_REMOTE_URL, top_k="5.0", - backend_top_k="6.0", temperature="0.25", ) assert cfg.top_k == 5 - assert cfg.backend_top_k == 6 assert cfg.temperature == 0.25 -def test_agentic_config_rejects_backend_top_k_below_target(): +def test_agentic_config_allows_none_temperature(): from nemo_retriever.query.agentic import AgenticRetrievalConfig - with pytest.raises(ValueError, match="backend_top_k"): - AgenticRetrievalConfig(llm_model="nemotron-8b", backend_top_k=4, top_k=5) + cfg = AgenticRetrievalConfig(llm_model="m", invoke_url=_REMOTE_URL) + + assert cfg.temperature is None def test_agentic_config_rejects_nvidia_temperature_above_max(): @@ -387,20 +302,77 @@ def test_agentic_config_rejects_nvidia_temperature_above_max(): ) +def test_agentic_config_accepts_in_process_temperature_above_nvidia_limit(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + # In-process uses the OpenAI-compatible bound (2.0), so a value above the + # hosted-NVIDIA 1.0 cap is accepted. + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", temperature=1.5) + + assert cfg.llm_backend == "in_process" + assert cfg.temperature == pytest.approx(1.5) + + def test_agentic_config_rejects_nonfinite_temperature(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="temperature must be finite"): - AgenticRetrievalConfig(llm_model="nemotron-8b", temperature=float("nan")) + AgenticRetrievalConfig(llm_model="m", invoke_url=_REMOTE_URL, temperature=float("nan")) -def test_agentic_config_accepts_in_process_temperature_above_nvidia_limit(): +def test_agentic_config_defaults_client_to_callable(): from nemo_retriever.query.agentic import AgenticRetrievalConfig - cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", temperature=1.5) + # Remote transport (invoke_url set), client unset -> callable default. The + # same client serves both transports; only the injected completion callable + # differs. + cfg = AgenticRetrievalConfig(llm_model="m", invoke_url=_REMOTE_URL) - assert cfg.llm_backend == "in_process" - assert cfg.temperature == pytest.approx(1.5) + assert cfg.llm_backend == "openai_compatible" + assert cfg.llm_client == "callable" + + +def test_agentic_config_accepts_and_normalizes_known_client(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + cfg = AgenticRetrievalConfig(llm_model="m", invoke_url=_REMOTE_URL, llm_client=" litellm ") + + assert cfg.llm_client == "litellm" + + +def test_agentic_config_defaults_callable_client_in_process(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + # In-process transport, client unset or explicitly callable -> callable. + assert AgenticRetrievalConfig(llm_model="nemotron-8b").llm_client == "callable" + assert AgenticRetrievalConfig(llm_model="nemotron-8b", llm_client="callable").llm_client == "callable" + + +def test_agentic_config_rejects_remote_client_without_invoke_url(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + # A non-callable client is a remote client and needs invoke_url; no silent + # override to callable. + with pytest.raises(ValueError, match="in-process agentic runs use the 'callable' LLM client"): + AgenticRetrievalConfig(llm_model="nemotron-8b", llm_client="litellm") + + +def test_agentic_config_accepts_callable_client_with_invoke_url(): + # `callable` spans both transports: it wraps the in-process engine locally and + # the shared HTTP client remotely, so pairing it with an invoke_url is valid. + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + cfg = AgenticRetrievalConfig(llm_model="m", invoke_url=_REMOTE_URL, llm_client="callable") + + assert cfg.llm_backend == "openai_compatible" + assert cfg.llm_client == "callable" + + +def test_agentic_config_rejects_unknown_client(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + with pytest.raises(ValueError, match="llm_client must be one of"): + AgenticRetrievalConfig(llm_model="m", invoke_url=_REMOTE_URL, llm_client="bogus") def test_agentic_config_rejects_invalid_local_llm_backend(): diff --git a/nemo_retriever/tests/test_agentic_local_llm.py b/nemo_retriever/tests/test_agentic_local_llm.py index 308f8b5abf..aa5e5c78ab 100644 --- a/nemo_retriever/tests/test_agentic_local_llm.py +++ b/nemo_retriever/tests/test_agentic_local_llm.py @@ -292,3 +292,51 @@ def test_agentic_retriever_unload_noop_when_no_local_llm() -> None: retriever.unload() assert retriever._chat_completion_fn is None + + +def _offline_llm(sampling_params_cls: Any) -> Any: + """A VLLMAgentChatLLM with the engine faked out, so no vLLM install is needed.""" + from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + + completion = MagicMock(text="hello", finish_reason="stop", token_ids=[1, 2]) + completion.tool_calls = None + completion.tool_call = None + request_output = MagicMock(outputs=[completion], prompt_token_ids=[1]) + + llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) + llm._llm = MagicMock(chat=MagicMock(return_value=[request_output])) + llm._lock = threading.Lock() + llm._sampling_params_cls = sampling_params_cls + llm._model_path = "nvidia/Llama-3.1-Nemotron-Nano-8B-v1" + llm._max_tokens = 512 + llm._request_extras = {} + return llm + + +def test_local_llm_maps_unset_temperature_to_greedy() -> None: + # The agent forwards temperature=None to mean "unset". There is no provider to + # defer to in-process, and adopting vLLM's own sampling default would silently + # make every local benchmark non-deterministic -- so None means greedy here. + # Deliberately asymmetric with invoke_chat_completion_step, which omits the + # field so the remote provider's default applies. + sampling_params_cls = MagicMock() + _offline_llm(sampling_params_cls)(messages=[{"role": "user", "content": "q"}], temperature=None) + + assert sampling_params_cls.call_args.kwargs["temperature"] == 0.0 + + +def test_local_llm_forwards_an_explicit_temperature() -> None: + sampling_params_cls = MagicMock() + _offline_llm(sampling_params_cls)(messages=[{"role": "user", "content": "q"}], temperature=0.7) + + assert sampling_params_cls.call_args.kwargs["temperature"] == 0.7 + + +def test_local_llm_returns_an_openai_shaped_response() -> None: + # The callable contract is an OpenAI chat.completion dict; CallableLLMBackend + # parses this exact shape. + response = _offline_llm(MagicMock())(messages=[{"role": "user", "content": "q"}], temperature=None) + + assert response["choices"][0]["message"] == {"role": "assistant", "content": "hello"} + assert response["choices"][0]["finish_reason"] == "stop" + assert response["usage"]["total_tokens"] == 3 diff --git a/nemo_retriever/tests/test_agentic_operators.py b/nemo_retriever/tests/test_agentic_operators.py index c30e7986cf..58ed6234e4 100644 --- a/nemo_retriever/tests/test_agentic_operators.py +++ b/nemo_retriever/tests/test_agentic_operators.py @@ -4,18 +4,51 @@ """Smoke tests for the agentic retrieval operators. +The ReAct and selection operators delegate all agent logic to the vendored +``nemo_retriever._agentic.nemo_agent`` library, so these tests exercise the operators' +adapter responsibilities — DataFrame translation and the selection gate — by +mocking the ``nemo_agent`` entry points (``Agent.run_sync`` / +``SelectionAgent.select_sync``) rather than the LLM transport. + Run with: cd nemo_retriever && uv run pytest tests/test_agentic_operators.py -v """ from __future__ import annotations -import json from unittest.mock import MagicMock, patch import pandas as pd import pytest + +# --------------------------------------------------------------------------- +# Shared helpers: canned nemo_agent results +# --------------------------------------------------------------------------- + + +def _agent_result(*, final_doc_ids=None, retrieval_log=None, error_category=None): + """Build a canned ``AgentRunResult`` like ``Agent.run``/``SelectionAgent.select``.""" + from nemo_retriever._agentic.nemo_agent.results import AgentError, AgentRunResult + + error = AgentError(category=error_category, message="stub") if error_category else None + return AgentRunResult( + final_doc_ids=list(final_doc_ids or []), + retrieval_log=list(retrieval_log or []), + error=error, + ) + + +def _step(docs, query_type="agent"): + """One retrieval_log entry; ``docs`` is a list of (id, score, text) triples.""" + return { + "input": {"query": "q", "top_k": len(docs)}, + "tool_name": "retrieve", + "query_type": query_type, + "output": [{"id": did, "score": score, "text": text} for did, score, text in docs], + } + + # --------------------------------------------------------------------------- # RRFAggregatorOperator — pure pandas, no mocking needed # --------------------------------------------------------------------------- @@ -86,618 +119,276 @@ def test_carries_react_final_rank(self): q1 = result[result["query_id"] == "q1"].set_index("doc_id") assert int(q1.loc["d1", "react_final_rank"]) == 1 - def test_missing_column_raises(self): + def test_final_result_step_excluded_from_score(self): + """The synthetic final step must not contribute to the RRF score.""" from nemo_retriever.operators.graph_ops.rrf_aggregator_operator import RRFAggregatorOperator - op = RRFAggregatorOperator(k=60) - bad_df = pd.DataFrame({"query_id": ["q1"], "query_text": ["x"]}) - with pytest.raises(ValueError, match="missing required column"): - op.run(bad_df) - - -# --------------------------------------------------------------------------- -# Prompt rendering — pure Python, no mocking needed -# --------------------------------------------------------------------------- - - -class TestPromptRendering: - def test_react_prompt_no_extended_relevance(self): - from nemo_retriever.operators.graph_ops.react_agent_operator import _render_react_agent_prompt - - prompt = _render_react_agent_prompt(10, with_init_docs=True, extended_relevance=False) - assert "" in prompt - assert "" in prompt - assert "" in prompt - assert "RELEVANCE_DEFINITION" not in prompt - assert "exactly the 10 most relevant" in prompt - assert "TIP" in prompt - - def test_react_prompt_with_extended_relevance(self): - from nemo_retriever.operators.graph_ops.react_agent_operator import _render_react_agent_prompt - - prompt = _render_react_agent_prompt(5, with_init_docs=False, extended_relevance=True) - assert "RELEVANCE_DEFINITION" in prompt - assert "exactly the 5" in prompt - assert "TIP" not in prompt - - def test_selection_prompt_no_extended_relevance(self): - from nemo_retriever.operators.graph_ops.selection_agent_operator import _render_selection_prompt - - prompt = _render_selection_prompt(5, extended_relevance=False) - assert "" in prompt - assert "" in prompt - assert "THINKING TIPS" in prompt - assert "RELEVANCE_DEFINITION" not in prompt - assert "5 most relevant" in prompt - - def test_selection_prompt_with_extended_relevance(self): - from nemo_retriever.operators.graph_ops.selection_agent_operator import _render_selection_prompt - - prompt = _render_selection_prompt(5, extended_relevance=True) - assert "RELEVANCE_DEFINITION" in prompt - assert "As explained above" in prompt - - -# --------------------------------------------------------------------------- -# SelectionAgentOperator — mock invoke_chat_completion_step -# --------------------------------------------------------------------------- - - -def _make_tool_call_response(fn_name: str, fn_args: dict, tc_id: str = "call_1") -> dict: - """Build a canned /v1/chat/completions response with one tool call.""" - return { - "choices": [ + # d1 appears once as a retrieve hit (step 0, rank 1) and once as the + # final-results selection (step 1, rank 1, is_final_result=True). Only + # the retrieve hit should score; react_final_rank must still be recorded. + df = pd.DataFrame( { - "message": { - "content": None, - "tool_calls": [ - { - "id": tc_id, - "type": "function", - "function": {"name": fn_name, "arguments": json.dumps(fn_args)}, - } - ], - }, - "finish_reason": "tool_calls", + "query_id": ["q1", "q1"], + "query_text": ["q"] * 2, + "step_idx": [0, 1], + "doc_id": ["d1", "d1"], + "text": ["t1", "t1"], + "rank": [1, 1], + "is_final_result": [False, True], } - ] - } + ) + result = RRFAggregatorOperator(k=60).run(df) + row = result[result["doc_id"] == "d1"].iloc[0] + assert abs(row["rrf_score"] - 1 / (1 + 60)) < 1e-10 # only the retrieve hit + assert int(row["react_final_rank"]) == 1 + def test_final_only_doc_is_still_emitted(self): + """A doc that appears only in the final step keeps its react_final_rank.""" + from nemo_retriever.operators.graph_ops.rrf_aggregator_operator import RRFAggregatorOperator -def _make_raw_arguments_tool_call_response(fn_name: str, arguments: str, tc_id: str = "call_1") -> dict: - """Build a canned chat-completions response with raw function arguments.""" - return { - "choices": [ + df = pd.DataFrame( { - "message": { - "content": None, - "tool_calls": [ - { - "id": tc_id, - "type": "function", - "function": {"name": fn_name, "arguments": arguments}, - } - ], - }, - "finish_reason": "tool_calls", + "query_id": ["q1", "q1"], + "query_text": ["q"] * 2, + "step_idx": [0, 1], + "doc_id": ["d1", "d2"], # d2 appears ONLY in the final step + "text": ["t1", "t2"], + "rank": [1, 1], + "is_final_result": [False, True], } - ] - } + ) + result = RRFAggregatorOperator(k=60).run(df) + d2 = result[result["doc_id"] == "d2"] + assert len(d2) == 1 + assert d2.iloc[0]["rrf_score"] == 0.0 + assert int(d2.iloc[0]["react_final_rank"]) == 1 + def test_missing_column_raises(self): + from nemo_retriever.operators.graph_ops.rrf_aggregator_operator import RRFAggregatorOperator -class TestSelectionAgentOperator: - def _make_input(self): - return pd.DataFrame( - { - "query_id": ["q1", "q1", "q1"], - "query_text": ["What causes inflation?"] * 3, - "doc_id": ["d1", "d2", "d3"], - "text": ["monetary policy doc", "supply chain doc", "unrelated doc"], - } - ) + op = RRFAggregatorOperator(k=60) + bad_df = pd.DataFrame({"query_id": ["q1"], "query_text": ["x"]}) + with pytest.raises(ValueError, match="missing required column"): + op.run(bad_df) - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_happy_path_selects_docs(self, mock_step): - from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator - # LLM immediately calls log_selected_documents - mock_step.return_value = _make_tool_call_response( - "log_selected_documents", - {"doc_ids": ["d1", "d2"], "message": "d1 most relevant"}, - ) +# --------------------------------------------------------------------------- +# ReActAgentOperator — mock nemo_agent.Agent.run_sync +# --------------------------------------------------------------------------- - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=2, - ) - result = op.run(self._make_input()) - assert set(result.columns) >= {"query_id", "doc_id", "rank", "message", "result_source"} - assert result["query_id"].tolist() == ["q1", "q1"] - assert result["doc_id"].tolist() == ["d1", "d2"] - assert result["rank"].tolist() == [1, 2] - assert result["result_source"].tolist() == ["selection_agent", "selection_agent"] +class TestReActAgentOperator: + def _op(self, **kwargs): + from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_retries_when_selection_returns_invalid_doc_ids(self, mock_step): - from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + defaults = dict(llm_model="test-model", retriever_fn=lambda q, k: [], target_top_k=2) + defaults.update(kwargs) + return ReActAgentOperator(**defaults) - mock_step.side_effect = [ - _make_tool_call_response( - "log_selected_documents", - {"doc_ids": ["d1", "missing"], "message": "mixed valid and invalid"}, - ), - _make_tool_call_response( - "log_selected_documents", - {"doc_ids": ["d1", "d2"], "message": "corrected"}, - ), - ] + def _input(self): + return pd.DataFrame({"query_id": ["q1"], "query_text": ["What causes inflation?"]}) - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=2, - max_steps=2, + def test_retrieve_adapter_renames_and_coerces(self): + op = self._op( + retriever_fn=lambda q, k: [ + {"doc_id": "d1", "text": "t", "score": 0.5}, + {"id": "d2", "text": "u", "score": "0.4"}, # already-id + str score + {"doc_id": "", "text": "skip"}, # empty id dropped + ] ) - result = op.run(self._make_input()) - - assert mock_step.call_count == 2 - assert result["doc_id"].tolist() == ["d1", "d2"] - assert result["message"].tolist() == ["corrected", "corrected"] + out = op._retrieve_adapter("q", 5) + assert out == [ + {"id": "d1", "score": 0.5, "text": "t"}, + {"id": "d2", "score": 0.4, "text": "u"}, + ] - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_non_object_tool_arguments_are_reported_and_fall_back(self, mock_step): - from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + def test_translates_retrieval_log_and_final(self): + from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - mock_step.return_value = _make_raw_arguments_tool_call_response( - "log_selected_documents", json.dumps("doc_ids=d1") + mock_agent = MagicMock() + mock_agent.run_sync.return_value = _agent_result( + retrieval_log=[ + _step([("d1", 0.9, "A"), ("d2", 0.8, "B")], query_type="main"), + _step([("d2", 0.7, "B"), ("d3", 0.6, "C")]), + ], + final_doc_ids=["d2", "d1"], ) + op = self._op() + with patch.object(ReActAgentOperator, "_ensure_agent", return_value=mock_agent): + result = op.run(self._input()) + + assert set(result.columns) == { + "query_id", + "query_text", + "step_idx", + "doc_id", + "text", + "rank", + "has_valid_final_results", + "is_final_result", + } + assert result["has_valid_final_results"].all() + # retrieve steps 0 and 1 present + assert sorted(result[~result.is_final_result]["step_idx"].unique().tolist()) == [0, 1] + # synthetic final step carries final_doc_ids in order + finals = result[result.is_final_result].sort_values("rank") + assert finals["doc_id"].tolist() == ["d2", "d1"] + # query_id is bound on the agent run + assert mock_agent.run_sync.call_args.kwargs["query_id"] == "q1" + assert mock_agent.run_sync.call_args.kwargs["raw_log_dir"] is None + + def test_empty_final_doc_ids_no_synthetic_step(self): + from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=2, - max_steps=1, + mock_agent = MagicMock() + mock_agent.run_sync.return_value = _agent_result( + retrieval_log=[_step([("d1", 0.9, "A")], query_type="main")], + final_doc_ids=[], + error_category="max_steps", ) - result = op.run(self._make_input()) - - assert result["doc_id"].tolist() == ["d1", "d2"] - assert result["result_source"].tolist() == ["candidate_ranking", "candidate_ranking"] + op = self._op() + with patch.object(ReActAgentOperator, "_ensure_agent", return_value=mock_agent): + result = op.run(self._input()) - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_think_then_select(self, mock_step): - from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + assert not result["has_valid_final_results"].any() + assert not result["is_final_result"].any() + assert result["doc_id"].tolist() == ["d1"] # retrieval log preserved on failure - # First call: think; second call: log_selected_documents - mock_step.side_effect = [ - _make_tool_call_response("think", {"thought": "let me reason..."}), - _make_tool_call_response("log_selected_documents", {"doc_ids": ["d3"], "message": "only d3"}), + def test_empty_input_returns_full_schema(self): + op = self._op() + result = op.run(pd.DataFrame({"query_id": [], "query_text": []})) + assert list(result.columns) == [ + "query_id", + "query_text", + "step_idx", + "doc_id", + "text", + "rank", + "has_valid_final_results", + "is_final_result", ] + assert result.empty - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=1, - ) - result = op.run(self._make_input()) - - assert result["doc_id"].tolist() == ["d3"] - assert mock_step.call_count == 2 + def test_multiple_queries_preserve_input_order(self): + from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_injected_chat_completion_fn_replaces_http_call(self, mock_step): - from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + mock_agent = MagicMock() - local_chat = MagicMock( - return_value=_make_tool_call_response( - "log_selected_documents", - {"doc_ids": ["d1"], "message": "d1 is best"}, + def run_sync(query, *, query_id=None, raw_log_dir=None): + return _agent_result( + retrieval_log=[_step([(f"{query_id}d", 1.0, "x")], query_type="main")], + final_doc_ids=[], ) - ) - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=1, - max_tokens=234, - chat_completion_fn=local_chat, - ) - result = op.run(self._make_input()) + mock_agent.run_sync.side_effect = run_sync + op = self._op(num_concurrent=4) + df = pd.DataFrame({"query_id": ["qA", "qB", "qC"], "query_text": ["a", "b", "c"]}) + with patch.object(ReActAgentOperator, "_ensure_agent", return_value=mock_agent): + result = op.run(df) - mock_step.assert_not_called() - assert local_chat.call_count == 1 - assert local_chat.call_args.kwargs["max_tokens"] == 234 - assert result["doc_id"].tolist() == ["d1"] - - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_extended_relevance_in_prompt(self, mock_step): - from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + # Deterministic input order regardless of thread completion order. + assert result["query_id"].tolist() == ["qA", "qB", "qC"] + assert result["doc_id"].tolist() == ["qAd", "qBd", "qCd"] - captured_prompts = [] - def capture_and_respond(**kwargs): - captured_prompts.append(kwargs["messages"][0]["content"]) - return _make_tool_call_response("log_selected_documents", {"doc_ids": ["d1"], "message": "ok"}) - - mock_step.side_effect = capture_and_respond - - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=1, - extended_relevance=True, - ) - op.run(self._make_input()) +# --------------------------------------------------------------------------- +# SelectionAgentOperator — mock nemo_agent.SelectionAgent.select_sync +# --------------------------------------------------------------------------- - assert "RELEVANCE_DEFINITION" in captured_prompts[0] - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_final_results_policy_skips_selection_agent(self, mock_step): +class TestSelectionAgentOperator: + def _op(self, **kwargs): from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=2, - ) - df = pd.DataFrame( + defaults = dict(llm_model="test-model", top_k=2) + defaults.update(kwargs) + return SelectionAgentOperator(**defaults) + + def _rrf_frame(self, react_final_rank): + return pd.DataFrame( { "query_id": ["q1", "q1", "q1"], "query_text": ["What causes inflation?"] * 3, "doc_id": ["d1", "d2", "d3"], "text": ["doc one", "doc two", "doc three"], - "rrf_score": [0.1, 0.9, 0.8], - "react_final_rank": [2, None, 1], + "rrf_score": [0.9, 0.5, 0.7], + "react_final_rank": react_final_rank, } ) - result = op.run(df) + + def test_final_results_passthrough(self): + """Tier 1: a ReAct final list passes through; the selection agent is not run.""" + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + + mock_sel = MagicMock() + op = self._op(top_k=2) + df = self._rrf_frame(react_final_rank=[2, None, 1]) # d3 rank1, d1 rank2 + with patch.object(SelectionAgentOperator, "_ensure_agent", return_value=mock_sel): + result = op.run(df) assert result["doc_id"].tolist() == ["d3", "d1"] assert result["result_source"].tolist() == ["final_results", "final_results"] - mock_step.assert_not_called() + assert result["rank"].tolist() == [1, 2] + mock_sel.select_sync.assert_not_called() - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_result_policy_uses_rrf_before_selection_when_no_final_results(self, mock_step): + def test_selection_runs_when_no_final_results(self): + """Tier 2: no ReAct final list -> run the selection agent over RRF candidates.""" from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=2, - ) - df = pd.DataFrame( - { - "query_id": ["q1", "q1", "q1"], - "query_text": ["What causes inflation?"] * 3, - "doc_id": ["d1", "d2", "d3"], - "text": ["doc one", "doc two", "doc three"], - "rrf_score": [0.1, 0.9, 0.8], - "react_final_rank": [None, None, None], - } - ) - result = op.run(df) + mock_sel = MagicMock() + mock_sel.select_sync.return_value = _agent_result(final_doc_ids=["d3", "d2"]) + op = self._op(top_k=2) + df = self._rrf_frame(react_final_rank=[None, None, None]) + with patch.object(SelectionAgentOperator, "_ensure_agent", return_value=mock_sel): + result = op.run(df) - assert result["doc_id"].tolist() == ["d2", "d3"] - assert result["result_source"].tolist() == ["rrf", "rrf"] - mock_step.assert_not_called() - - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_empty_final_results_is_not_valid_and_falls_back_to_rrf(self, mock_step): + assert result["doc_id"].tolist() == ["d3", "d2"] + assert result["result_source"].tolist() == ["selection_agent", "selection_agent"] + mock_sel.select_sync.assert_called_once() + # scores side-table covers every candidate; candidates are RRF-descending. + call = mock_sel.select_sync.call_args + assert call.kwargs["scores"] == {"d1": 0.9, "d2": 0.5, "d3": 0.7} + assert [d["id"] for d in call.args[1]] == ["d1", "d3", "d2"] + + def test_falls_back_to_rrf_when_selection_returns_empty(self): + """Tier 3: selection failed/empty -> top RRF-ranked candidates.""" from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=2, - ) - df = pd.DataFrame( - { - "query_id": ["q1", "q1"], - "query_text": ["What causes inflation?"] * 2, - "doc_id": ["d1", "d2"], - "text": ["doc one", "doc two"], - "rrf_score": [0.9, 0.8], - "has_valid_final_results": [False, False], - "react_final_rank": [None, None], - } - ) - result = op.run(df) + mock_sel = MagicMock() + mock_sel.select_sync.return_value = _agent_result(final_doc_ids=[], error_category="max_steps") + op = self._op(top_k=2) + df = self._rrf_frame(react_final_rank=[None, None, None]) + with patch.object(SelectionAgentOperator, "_ensure_agent", return_value=mock_sel): + result = op.run(df) - assert result["doc_id"].tolist() == ["d1", "d2"] + # RRF-descending top 2: d1 (0.9), d3 (0.7) + assert result["doc_id"].tolist() == ["d1", "d3"] assert result["result_source"].tolist() == ["rrf", "rrf"] - mock_step.assert_not_called() - - -# --------------------------------------------------------------------------- -# ReActAgentOperator — mock retriever_fn + invoke_chat_completion_step -# --------------------------------------------------------------------------- - - -class TestReActAgentOperator: - def _make_input(self): - return pd.DataFrame( - { - "query_id": ["q1"], - "query_text": ["What causes inflation?"], - } - ) - - def _make_retriever(self, docs=None): - """Return a mock retriever_fn that returns canned docs.""" - if docs is None: - docs = [{"doc_id": "d1", "text": "monetary policy"}, {"doc_id": "d2", "text": "supply chains"}] - - def retriever_fn(query_text: str, top_k: int): - return docs[:top_k] - - return retriever_fn - - @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_simple_mode_retrieve_then_final(self, mock_step): - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - - # 1) agent calls retrieve("subquery"), 2) agent calls final_results - mock_step.side_effect = [ - _make_tool_call_response("retrieve", {"query": "inflation monetary policy"}), - _make_tool_call_response( - "final_results", - {"doc_ids": ["d1", "d2"], "message": "found them", "search_successful": "true"}, - ), - ] - - op = ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=self._make_retriever(), - user_msg_type="simple", - target_top_k=2, - ) - result = op.run(self._make_input()) + mock_sel.select_sync.assert_called_once() - assert set(result.columns) >= {"query_id", "query_text", "step_idx", "doc_id", "text", "rank"} - assert result["query_id"].unique().tolist() == ["q1"] - # step 0 is the retrieve tool call result - assert 0 in result["step_idx"].values - assert "d1" in result["doc_id"].values - - @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_injected_chat_completion_fn_replaces_http_call(self, mock_step): - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - - local_chat = MagicMock( - return_value=_make_tool_call_response( - "final_results", - {"doc_ids": ["d1"], "message": "ok", "search_successful": "true"}, - ) - ) - retriever = MagicMock(return_value=[{"doc_id": "d1", "text": "monetary policy"}]) - - op = ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=retriever, - user_msg_type="with_results", - target_top_k=1, - max_tokens=123, - chat_completion_fn=local_chat, - ) - - result = op.run(self._make_input()) - - mock_step.assert_not_called() - assert local_chat.call_count == 1 - assert local_chat.call_args.kwargs["max_tokens"] == 123 - assert result[result["doc_id"] == "d1"]["is_final_result"].astype(bool).any() - - @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_non_object_tool_arguments_are_reported_without_crashing(self, mock_step): - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - - mock_step.return_value = _make_raw_arguments_tool_call_response("retrieve", json.dumps("query=inflation")) - retriever = MagicMock(return_value=[{"doc_id": "d1", "text": "monetary policy"}]) - - op = ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=retriever, - user_msg_type="with_results", - target_top_k=1, - max_steps=1, - ) - result = op.run(self._make_input()) - - assert result["doc_id"].tolist() == ["d1"] - assert not result["is_final_result"].astype(bool).any() - - @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_with_results_mode_initial_retrieval(self, mock_step): - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - - # with_results=True → retriever_fn called once before LLM, then LLM immediately calls final_results - mock_step.return_value = _make_tool_call_response( - "final_results", - {"doc_ids": ["d1"], "message": "ok", "search_successful": "true"}, - ) - retriever = MagicMock(return_value=[{"doc_id": "d1", "text": "monetary policy"}]) - - op = ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=retriever, - user_msg_type="with_results", - target_top_k=1, - ) - result = op.run(self._make_input()) - - # retriever was called upfront (step_idx=0) before any LLM step - assert retriever.call_count >= 1 - assert 0 in result["step_idx"].values - - def test_backend_top_k_caps_fetch_depth_and_replays_seen_docs(self): - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - - calls = [] - docs = [ - {"doc_id": "d1", "text": "already seen", "score": 0.9}, - {"doc_id": "d2", "text": "new two", "score": 0.8}, - {"doc_id": "d3", "text": "new three", "score": 0.7}, - {"doc_id": "d4", "text": "outside backend cap", "score": 0.6}, - ] - - def retriever_fn(query_text, top_k): - calls.append((query_text, top_k)) - return docs[:top_k] - - op = ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=retriever_fn, - retriever_top_k=2, - backend_top_k=3, - ) - - result = op._call_retriever("inflation", {"d1"}, api_key=None) - - assert calls == [("inflation", 3)] - assert [doc["doc_id"] for doc in result] == ["d1", "d2", "d3"] - assert "retrieved before" in result[0]["text"] - assert result[1]["text"] == "new two" - - @pytest.mark.parametrize( - ("fn_args", "target_top_k"), - [ - ({"doc_ids": [1], "message": "bad id type", "search_successful": "true"}, 1), - ({"doc_ids": [], "message": "empty", "search_successful": "false"}, 1), - ({"doc_ids": [""], "message": "empty-string id", "search_successful": "true"}, 1), - ({"doc_ids": [" "], "message": "whitespace id", "search_successful": "true"}, 1), - ({"doc_ids": ["d1"], "message": "wrong count", "search_successful": "true"}, 2), - ({"doc_ids": ["missing"], "message": "hallucinated id", "search_successful": "true"}, 1), - ({"doc_ids": ["d1"], "message": "bad status", "search_successful": "yes"}, 1), - ], - ) - @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_invalid_final_results_are_rejected(self, mock_step, fn_args, target_top_k): - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - - mock_step.return_value = _make_tool_call_response("final_results", fn_args) - retriever = MagicMock(return_value=[{"doc_id": "d1", "text": "monetary policy"}]) - - op = ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=retriever, - user_msg_type="with_results", - target_top_k=target_top_k, - ) - result = op.run(self._make_input()) - - assert not result["is_final_result"].astype(bool).any() - assert not result["has_valid_final_results"].astype(bool).any() - - @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_output_row_structure(self, mock_step): - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - - mock_step.side_effect = [ - _make_tool_call_response("retrieve", {"query": "q"}), - _make_tool_call_response( - "final_results", {"doc_ids": ["d1"], "message": "ok", "search_successful": "true"} - ), - ] - - op = ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=self._make_retriever(), - user_msg_type="simple", - target_top_k=1, - ) - result = op.run(self._make_input()) - - assert (result["rank"] >= 1).all() - assert result["step_idx"].dtype in (int, "int64") - assert result["doc_id"].notna().all() - - @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_no_final_results_falls_back_to_retrieval_log(self, mock_step): - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - - mock_step.side_effect = [ - _make_tool_call_response("retrieve", {"query": "inflation monetary policy"}), - _make_tool_call_response("think", {"thought": "still reasoning"}), - ] - retriever = MagicMock( - return_value=[ - {"doc_id": "d1", "text": "monetary policy"}, - {"doc_id": "d2", "text": "supply chains"}, - ] - ) - - op = ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=retriever, - user_msg_type="simple", - target_top_k=2, - max_steps=2, - ) - result = op.run(self._make_input()) - - assert result["doc_id"].tolist() == ["d1", "d2"] - assert result["rank"].tolist() == [1, 2] - assert retriever.call_count == 1 - - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_pipeline_end_to_end_with_mocks(self, mock_react_step, mock_selection_step): - """Wire ReAct → RRF → Selection with mocks; verify final output shape. - - Each operator imports invoke_chat_completion_step into its own module - namespace, so both must be patched independently. - """ - from nemo_retriever.graph.executor import InprocessExecutor - from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - from nemo_retriever.operators.graph_ops.rrf_aggregator_operator import RRFAggregatorOperator + def test_selection_exception_falls_back_to_rrf(self): + """An unexpected error inside selection degrades to the RRF ranking.""" from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator - # ReAct: retrieve once, then final_results - mock_react_step.side_effect = [ - _make_tool_call_response("retrieve", {"query": "inflation"}), - _make_tool_call_response( - "final_results", {"doc_ids": ["d1"], "message": "ok", "search_successful": "true"} - ), - ] - # Selection: immediately log_selected_documents - mock_selection_step.return_value = _make_tool_call_response( - "log_selected_documents", {"doc_ids": ["d1"], "message": "d1 best"} - ) - - def retriever_fn(query_text, top_k): - return [{"doc_id": "d1", "text": "monetary policy"}, {"doc_id": "d2", "text": "supply chains"}] + mock_sel = MagicMock() + mock_sel.select_sync.side_effect = RuntimeError("boom") + op = self._op(top_k=2) + df = self._rrf_frame(react_final_rank=[None, None, None]) + with patch.object(SelectionAgentOperator, "_ensure_agent", return_value=mock_sel): + result = op.run(df) - pipeline = ( - ReActAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - retriever_fn=retriever_fn, - user_msg_type="simple", - target_top_k=1, - ) - >> RRFAggregatorOperator(k=60) - >> SelectionAgentOperator( - invoke_url="http://localhost/v1/chat/completions", - llm_model="test-model", - top_k=1, - ) - ) + assert result["doc_id"].tolist() == ["d1", "d3"] + assert result["result_source"].tolist() == ["rrf", "rrf"] - query_df = pd.DataFrame({"query_id": ["q1"], "query_text": ["What causes inflation?"]}) - result = InprocessExecutor(pipeline).ingest(query_df) + def test_message_column_empty(self): + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator - assert set(result.columns) >= {"query_id", "doc_id", "rank", "message", "result_source"} - assert result["query_id"].tolist() == ["q1"] - assert result["rank"].tolist() == [1] + mock_sel = MagicMock() + op = self._op(top_k=2) + df = self._rrf_frame(react_final_rank=[2, None, 1]) + with patch.object(SelectionAgentOperator, "_ensure_agent", return_value=mock_sel): + result = op.run(df) + assert result["message"].tolist() == ["", ""] # --------------------------------------------------------------------------- @@ -882,37 +573,303 @@ def test_non_dataframe_raises(self): op.preprocess([{"query_id": "q1"}]) +class TestBuildLLMForwarding: + """``_build_llm`` forwards temperature / parallel_tool_calls, preserving falsy values. + + Builds a real backend on the default (``callable``) path, so no LLM SDK is + required. ``temperature=0.0`` / ``parallel_tool_calls=False`` are real settings + and must not be collapsed to ``None`` (guards against an ``x or None`` + regression). + """ + + def test_react_forwards_falsy_sampling_args(self): + from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator + + op = ReActAgentOperator( + llm_model="gpt-4o-mini", + retriever_fn=lambda q, k: [], + temperature=0.0, + parallel_tool_calls=False, + ) + config = op._build_llm().config + assert config.temperature == 0.0 + assert config.parallel_tool_calls is False + + def test_selection_forwards_falsy_sampling_args(self): + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + + op = SelectionAgentOperator( + llm_model="gpt-4o-mini", + temperature=0.0, + parallel_tool_calls=False, + ) + config = op._build_llm().config + assert config.temperature == 0.0 + assert config.parallel_tool_calls is False + + # --------------------------------------------------------------------------- -# SelectionAgentOperator — max_steps exhausted fallback +# Callable LLM backend — the in-process (local vLLM) adapter seam # --------------------------------------------------------------------------- -class TestSelectionAgentMaxSteps: - @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") - def test_rrf_candidates_skip_selection_agent(self, mock_step): - from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator +class TestCallableBackendWiring: + """The operators forward an injected chat_completion_fn to the callable backend.""" - mock_step.return_value = _make_tool_call_response("think", {"thought": "still thinking..."}) + def test_build_llm_returns_callable_backend_when_completion_fn_set(self): + from nemo_retriever._agentic.nemo_agent.llm import CallableLLMBackend + from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator - op = SelectionAgentOperator( - llm_model="test-model", - invoke_url="http://localhost/v1/chat/completions", - top_k=2, - max_steps=3, + def fake_fn(**kwargs): + return {"choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]} + + op = ReActAgentOperator( + llm_model="nemotron-8b", + retriever_fn=lambda q, k: [], + backend="callable", + chat_completion_fn=fake_fn, ) - df = pd.DataFrame( - { - "query_id": ["q1", "q1", "q1"], - "query_text": ["What causes inflation?"] * 3, - "doc_id": ["d1", "d2", "d3"], - "text": ["doc one", "doc two", "doc three"], - "rrf_score": [0.2, 0.9, 0.5], + # No litellm import required on this path. + assert isinstance(op._build_llm(), CallableLLMBackend) + + def test_build_llm_uses_default_backend_without_completion_fn(self): + from nemo_retriever._agentic.nemo_agent.llm import CallableLLMBackend + from nemo_retriever.models.nim.chat_completions import invoke_chat_completion_step + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + + # No completion_fn on the default backend: the operator must inject the + # shared chat-completions client so a remote run works out of the box. + # Needs no LLM SDK installed. + op = SelectionAgentOperator(llm_model="gpt-4o") + llm = op._build_llm() + assert isinstance(llm, CallableLLMBackend) + assert llm._completion_fn is invoke_chat_completion_step + + def test_injected_completion_fn_wins_over_the_default_http_client(self): + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + + def fake_fn(**kwargs): + return {"choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]} + + op = SelectionAgentOperator(llm_model="nemotron-8b", chat_completion_fn=fake_fn) + assert op._build_llm()._completion_fn is fake_fn + + def test_build_llm_honors_an_explicit_backend(self): + pytest.importorskip("litellm") + from nemo_retriever._agentic.nemo_agent.llm import LiteLLMBackend + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + + op = SelectionAgentOperator(llm_model="gpt-4o", backend="litellm") + assert isinstance(op._build_llm(), LiteLLMBackend) + + +class TestCallableLLMBackend: + """The adapter maps an OpenAI chat.completion dict to a CompletionResult.""" + + @staticmethod + def _backend(fn): + from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config + + return create_llm(create_llm_config("callable", model="nemotron-8b"), completion_fn=fn) + + def test_parses_content_and_stubs_usage(self): + seen: dict = {} + + def fn(**kwargs): + seen.update(kwargs) + return { + "choices": [{"message": {"role": "assistant", "content": "answer"}, "finish_reason": "stop"}], + "usage": {"total_tokens": 5}, } + + result = self._backend(fn).completion(messages=[{"role": "user", "content": "q"}]) + assert result.message == {"role": "assistant", "content": "answer"} + assert result.finish_reason == "stop" + assert result.usage == {"total_tokens": 5} # reported usage is recorded, not discarded + # Headers/status never cross the callable boundary, so this stays empty. + assert result.extra_response_info == {} + assert seen["tool_choice"] == "none" # no tools -> suppress tool calls + + def test_passes_tools_and_forwards_tool_calls(self): + def fn(**kwargs): + assert kwargs["tool_choice"] == "auto" + assert kwargs["tools"] + return { + "choices": [ + { + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "retrieve", "arguments": "{}"}} + ], + }, + "finish_reason": "tool_calls", + } + ] + } + + result = self._backend(fn).completion( + messages=[{"role": "user", "content": "q"}], + tools=[{"type": "function", "function": {"name": "retrieve"}}], ) - result = op.run(df) + assert result.message["tool_calls"][0]["function"]["name"] == "retrieve" + assert result.finish_reason == "tool_calls" - assert result["doc_id"].tolist() == ["d2", "d3"] - assert result["rank"].tolist() == [1, 2] - assert result["message"].tolist() == ["Using RRF ranking."] * 2 - assert result["result_source"].tolist() == ["rrf", "rrf"] - mock_step.assert_not_called() + def test_requires_completion_fn(self): + from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config + + with pytest.raises(ValueError, match="completion_fn"): + create_llm(create_llm_config("callable", model="nemotron-8b")) + + def test_callable_failure_is_translated_to_llm_call_error(self): + from nemo_retriever._agentic.nemo_agent.llm import LLMCallError + + def fn(**kwargs): + raise RuntimeError("engine died") + + with pytest.raises(LLMCallError) as excinfo: + self._backend(fn).completion(messages=[{"role": "user", "content": "q"}]) + # The original exception stays chained so it remains diagnosable. + assert isinstance(excinfo.value.__cause__, RuntimeError) + + @pytest.mark.parametrize( + "bad_response", + [ + "not a dict", + {}, # no choices + {"choices": []}, # empty choices + {"choices": [{"finish_reason": "stop"}]}, # choice missing message + ], + ) + def test_malformed_response_is_translated_to_llm_call_error(self, bad_response): + from nemo_retriever._agentic.nemo_agent.llm import LLMCallError + + with pytest.raises(LLMCallError): + self._backend(lambda **kwargs: bad_response).completion(messages=[{"role": "user", "content": "q"}]) + + @staticmethod + def _seen_kwargs(backend, **call_kwargs): + seen: dict = {} + + def fn(**kwargs): + seen.update(kwargs) + return {"choices": [{"message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}]} + + backend(fn).completion(messages=[{"role": "user", "content": "q"}], **call_kwargs) + return seen + + def test_known_knobs_routed_into_extra_body(self): + from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config + + def backend(fn): + return create_llm( + create_llm_config("callable", model="nemotron-8b", parallel_tool_calls=False, reasoning_effort="high"), + completion_fn=fn, + ) + + seen = self._seen_kwargs(backend) + assert seen["extra_body"] == {"parallel_tool_calls": False, "reasoning_effort": "high"} + + def test_max_completion_tokens_override_aliases_max_tokens(self): + seen = self._seen_kwargs(self._backend, max_completion_tokens=64) + assert seen["max_tokens"] == 64 + assert "max_completion_tokens" not in seen # aliased, not passed through raw + + def test_unknown_override_is_passed_through_not_dropped(self): + seen = self._seen_kwargs(self._backend, top_p=0.25) + assert seen["top_p"] == 0.25 # rides along as a kwarg rather than being silently ignored + + def test_override_wins_over_config_knob_in_extra_body(self): + from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config + + def backend(fn): + return create_llm( + create_llm_config("callable", model="nemotron-8b", reasoning_effort="high"), + completion_fn=fn, + ) + + seen = self._seen_kwargs(backend, reasoning_effort="low") + assert seen["extra_body"]["reasoning_effort"] == "low" # override beats the config value + + def test_temperature_forwarded_as_none_when_unset(self): + seen = self._seen_kwargs(self._backend) # default config -> temperature is None + # Forwarded, not omitted: omitting it would let the callable's own default + # (0.0, i.e. greedy) apply, which is not what "unset" means. Passing None + # lets each callable decide — omit the field remotely, pick a concrete + # value in-process. + assert "temperature" in seen + assert seen["temperature"] is None + + def test_temperature_forwarded_when_set_including_zero(self): + from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config + + def backend(fn): + return create_llm( + create_llm_config("callable", model="nemotron-8b", temperature=0.0), + completion_fn=fn, + ) + + seen = self._seen_kwargs(backend) + assert seen["temperature"] == 0.0 # explicit falsy value is still forwarded + + def test_capture_raw_io_off_by_default(self): + result = self._backend( + lambda **kwargs: {"choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]} + ).completion(messages=[{"role": "user", "content": "q"}]) + assert result.raw_request is None + assert result.raw_response is None + + def test_capture_raw_io_populates_request_and_response(self): + from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config + + response = { + "choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"total_tokens": 3}, + } + backend = create_llm( + create_llm_config("callable", model="nemotron-8b", capture_raw_io=True), + completion_fn=lambda **kwargs: response, + ) + result = backend.completion(messages=[{"role": "user", "content": "q"}], max_completion_tokens=64) + assert result.raw_response == response + assert result.raw_response is not response # an independent snapshot, not an alias + assert result.raw_request["model"] == "nemotron-8b" + assert result.raw_request["messages"] == [{"role": "user", "content": "q"}] + assert result.raw_request["max_tokens"] == 64 # not redacted despite containing "token" + + def test_capture_raw_io_redacts_api_key(self): + from nemo_retriever._agentic.nemo_agent.llm import create_llm, create_llm_config + from nemo_retriever._agentic.nemo_agent.llm.callable_backend import _REDACTED + + backend = create_llm( + create_llm_config("callable", model="nemotron-8b", capture_raw_io=True), + completion_fn=lambda **kwargs: { + "choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}] + }, + ) + result = backend.completion(messages=[{"role": "user", "content": "q"}], api_key="super-secret") + assert result.raw_request["api_key"] == _REDACTED + assert result.raw_request["messages"] == [{"role": "user", "content": "q"}] # content preserved + + +class TestAgentConfigMode: + """``mode`` is retained as the extension point but only ``select`` is implemented.""" + + def test_select_mode_is_accepted(self): + from nemo_retriever._agentic.nemo_agent import AgentConfig + + assert AgentConfig(mode="select").mode == "select" + + def test_mode_defaults_to_select(self): + from nemo_retriever._agentic.nemo_agent import AgentConfig + + assert AgentConfig().mode == "select" + + def test_answer_mode_is_rejected(self): + from pydantic import ValidationError + + from nemo_retriever._agentic.nemo_agent import AgentConfig + + with pytest.raises(ValidationError): + AgentConfig(mode="answer") diff --git a/nemo_retriever/tests/test_agentic_source_hygiene.py b/nemo_retriever/tests/test_agentic_source_hygiene.py new file mode 100644 index 0000000000..ba8db1778b --- /dev/null +++ b/nemo_retriever/tests/test_agentic_source_hygiene.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Source hygiene checks for the private agent implementation.""" + +from __future__ import annotations + +from pathlib import Path + + +def _source_files() -> list[Path]: + root = Path(__file__).parents[1] / "src" / "nemo_retriever" + paths = list((root / "_agentic").rglob("*.py")) + paths += [ + root / "operators" / "graph_ops" / "react_agent_operator.py", + root / "operators" / "graph_ops" / "selection_agent_operator.py", + ] + return paths + + +def test_private_agent_sources_do_not_advertise_public_api_or_vendor_boundary(): + text = "\n".join(path.read_text(encoding="utf-8") for path in _source_files()) + + assert "Public surface:" not in text + assert "public API" not in text + assert "vendored" not in text diff --git a/nemo_retriever/tests/test_ci_workflows.py b/nemo_retriever/tests/test_ci_workflows.py index 89e2281d79..7b01efc85d 100644 --- a/nemo_retriever/tests/test_ci_workflows.py +++ b/nemo_retriever/tests/test_ci_workflows.py @@ -264,6 +264,23 @@ def test_dev_compose_helpers_are_feature_scoped(): assert "nemo_retriever/dev/compose/neo4j.compose.yaml" in neo4j_setup +def test_default_service_mode_compose_wires_optional_collection_auth(): + compose_path = REPO_ROOT / "nemo_retriever" / "dev" / "compose" / "service-mode.compose.yaml" + compose_text = compose_path.read_text(encoding="utf-8") + compose_data = yaml.safe_load(compose_text) + + retriever = compose_data["services"]["retriever"] + vectordb = compose_data["services"]["vectordb"] + assert retriever["environment"]["NRL_API_TOKEN"] == "${NRL_API_TOKEN:-}" + assert retriever["environment"]["NRL_INTERNAL_VDB_TOKEN"] == "${NRL_INTERNAL_VDB_TOKEN:-}" + assert vectordb["environment"]["NRL_INTERNAL_VDB_TOKEN"] == "${NRL_INTERNAL_VDB_TOKEN:-}" + + service_config = compose_data["configs"]["retriever_service_config"]["content"] + assert 'api_token: "${NRL_API_TOKEN:-}"' in service_config + assert "allow_unscoped_dev: true" in service_config + assert "use_graphic_elements" not in service_config + + def test_legacy_tools_harness_is_removed(): assert not (REPO_ROOT / "tools" / "harness").exists() diff --git a/nemo_retriever/tests/test_collection_management_api.py b/nemo_retriever/tests/test_collection_management_api.py new file mode 100644 index 0000000000..4c26a6cc65 --- /dev/null +++ b/nemo_retriever/tests/test_collection_management_api.py @@ -0,0 +1,898 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import hashlib +import threading +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import AsyncMock, Mock, patch + +import httpx +import lancedb +import pyarrow as pa +import pytest +from fastapi.testclient import TestClient + +from nemo_retriever import RetrieverServiceClient +from nemo_retriever.common.schemas.collections import CollectionCreateRequest +from nemo_retriever.common.schemas.requests import JobCreateRequest +from nemo_retriever.common.vdb.adt_vdb import ( + CollectionWriteContext, + VDBInvalidRequest, + VDBResourceNotFound, +) +from nemo_retriever.common.vdb.lancedb import LanceDB +from nemo_retriever.service.auth import ScopeAuthorizer +from nemo_retriever.service.app import create_app +from nemo_retriever.service.config import AuthConfig, LLMConfig, ServiceConfig, VectorDbConfig +from nemo_retriever.service.query_schema import QueryRequest +from nemo_retriever.service.errors import RetrieverServiceError +import nemo_retriever.service.client as client_module +from nemo_retriever.service.services.job_tracker import JobFullError, JobTracker +from nemo_retriever.service.vectordb_app import VectorDBState, create_vectordb_app + + +def _record(_chunk_id: str, _document_id: str, text: str, version: str = "v1") -> dict: + return { + "document_type": "text", + "metadata": { + "embedding": [1.0, 0.0], + "content": text, + "content_metadata": {"page_number": 1, "type": "text"}, + "source_metadata": { + "source_id": "report.pdf", + "source_name": "report.pdf", + }, + }, + } + + +def _context( + document_id: str, + *, + version: str = "v1", + collection_name: str = "research", + operation: str = "append", +) -> CollectionWriteContext: + return CollectionWriteContext( + scope="scope", + collection_name=collection_name, + document_id=document_id, + document_version=version, + content_sha256=version, + filename="report.pdf", + job_id=f"job-{version}", + operation=operation, + ) + + +def _vdb(tmp_path, *, table_name: str = "legacy") -> LanceDB: + return LanceDB( + uri=str(tmp_path), + table_name=table_name, + overwrite=False, + build_index=False, + ) + + +def test_collection_crud_scope_pagination_and_injection_rejection(tmp_path) -> None: + app = create_vectordb_app(lancedb_uri=str(tmp_path), embed_endpoint="http://embed") + with TestClient(app) as client: + headers = {"X-NRL-Scope": "workspace-a"} + assert client.post("/v1/collections", json={"name": "one"}, headers=headers).status_code == 201 + assert client.post("/v1/collections", json={"name": "two"}, headers=headers).status_code == 201 + page = client.get("/v1/collections?limit=1", headers=headers).json() + assert len(page["items"]) == 1 and page["next_token"] + other_scope = client.get("/v1/collections/one", headers={"X-NRL-Scope": "workspace-b"}) + assert other_scope.status_code == 404 + injected = client.post( + "/v1/query", + json={"query": "x", "collection_name": "one", "table_name": "secret"}, + headers=headers, + ) + assert injected.status_code == 422 + health = client.get("/v1/health").json() + assert "table" not in health and "workspace-a" not in str(health) + metrics = client.get("/metrics").text + assert "workspace-a" not in metrics and "nrl_vectordb_cleanup_pending" in metrics + assert client.delete("/v1/collections/one", headers=headers).status_code == 200 + repeated = client.delete("/v1/collections/one?if_exists=true", headers=headers).json() + assert repeated == { + "name": "one", + "scope": "workspace-a", + "existed": False, + "deleted": False, + "status": "deleted", + "cleanup_pending": False, + } + + +def test_append_replace_and_document_delete_are_collection_scoped(tmp_path) -> None: + backend = _vdb(tmp_path) + backend.create_collection( + scope="scope", + request=CollectionCreateRequest(name="research"), + ) + appended = backend.write_collection( + [[_record("a", "doc", "old"), _record("b", "doc", "obsolete")]], + context=_context("doc"), + ) + assert appended.total_rows == 2 + replaced = backend.write_collection( + [[_record("c", "doc", "new", "v2")]], + context=_context("doc", version="v2", operation="replace"), + ) + assert replaced.total_rows == 1 + assert ( + backend.get_document( + scope="scope", + collection_name="research", + document_id="doc", + ).document_version + == "v2" + ) + assert backend.delete_document( + scope="scope", + collection_name="research", + document_id="doc", + if_exists=False, + ).deleted + assert not backend.delete_document( + scope="scope", + collection_name="research", + document_id="doc", + if_exists=True, + ).deleted + + +def test_public_sdk_and_citation_ready_query(tmp_path) -> None: + app = create_vectordb_app(lancedb_uri=str(tmp_path), embed_endpoint="http://embed") + with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0]]): + with TestClient(app) as service: + + class InProcessClient(RetrieverServiceClient): + async def _arequest(self, method: str, path: str, **kwargs): + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url="http://test", + headers=self._auth_headers, + ) as client: + response = await client.request(method, path, **kwargs) + self._raise_for_response(response, f"{method} {path}") + return response.json() if response.content else None + + sdk = InProcessClient(scope="workspace") + collection = sdk.create_collection("research") + assert collection.name == "research" + write = service.post( + "/internal/vectordb/write", + json={ + "records": [[_record("chunk", "doc", "finding")]], + "scope": "workspace", + "collection_name": "research", + "document_id": "doc", + "job_id": "job-v1", + "filename": "report.pdf", + "content_sha256": "v1", + "document_version": "v1", + }, + ) + assert write.status_code == 200, write.text + sync_hits = sdk.query("finding", top_k=10, collection_name="research") + async_hits = asyncio.run(sdk.aquery("finding", top_k=10, collection_name="research")) + assert sync_hits[0].model_dump() == async_hits[0].model_dump() + assert sync_hits[0].chunk_id == hashlib.sha256(f"doc\0v1\0{0}".encode()).hexdigest() + assert sync_hits[0].text == "finding" + assert sync_hits[0].distance >= 0.0 + assert sync_hits[0].page_number == 1 + assert sync_hits[0].filename == "report.pdf" + assert sdk.list_documents("research").items[0].document_id == "doc" + + +def test_idempotency_replay_and_conflict() -> None: + tracker = JobTracker() + original = tracker.register_job( + "job", + expected_documents=1, + scope="workspace", + idempotency_key="request", + idempotency_fingerprint="same", + ) + replay = tracker.register_job( + "replay", + expected_documents=1, + scope="workspace", + idempotency_key="request", + idempotency_fingerprint="same", + ) + assert replay.job_id == original.job_id + + with pytest.raises(JobFullError): + tracker.register_job( + "conflict", + expected_documents=1, + scope="workspace", + idempotency_key="request", + idempotency_fingerprint="different", + ) + + +def test_idempotent_job_registration_is_atomic() -> None: + tracker = JobTracker() + + def register(index: int) -> str: + return tracker.register_job( + f"job-{index}", + expected_documents=1, + scope="workspace", + idempotency_key="request", + idempotency_fingerprint="same", + ).job_id + + with ThreadPoolExecutor(max_workers=8) as executor: + job_ids = list(executor.map(register, range(8))) + + assert len(set(job_ids)) == 1 + assert len(tracker.all_jobs()) == 1 + + +def test_target_document_id_rejects_query_metacharacters() -> None: + with pytest.raises(ValueError, match="target_document_id"): + JobCreateRequest( + expected_documents=1, + operation="replace", + target_document_id="doc' OR 1=1 --", + ) + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ({"expected_documents": 1, "operation": "replace", "target_document_id": "doc"}, "collection_name"), + ( + { + "expected_documents": 2, + "collection_name": "research", + "operation": "replace", + "target_document_id": "doc", + }, + "exactly one document", + ), + ({"expected_documents": 1, "collection_name": "research", "operation": "replace"}, "target_document_id"), + ({"expected_documents": 1, "operation": "append", "target_document_id": "doc"}, "append"), + ], +) +def test_job_create_request_enforces_collection_operation_invariants(payload, message) -> None: + with pytest.raises(ValueError, match=message): + JobCreateRequest(**payload) + + +def test_job_create_request_accepts_valid_replacement() -> None: + request = JobCreateRequest( + expected_documents=1, + collection_name="research", + operation="replace", + target_document_id="doc", + ) + + assert request.operation == "replace" + assert request.target_document_id == "doc" + + +def test_manifest_entry_replay_is_capacity_neutral_and_identity_stable() -> None: + tracker = JobTracker() + tracker.register_job( + "job", + expected_documents=1, + document_manifest=[ + { + "manifest_entry_id": "a" * 64, + "filename": "report.pdf", + "content_sha256": "b" * 64, + } + ], + ) + accepted, created = tracker.register_document_idempotent( + "attempt-1", + job_id="job", + stable_document_id="document-1", + filename="report.pdf", + content_sha256="b" * 64, + manifest_entry_id="a" * 64, + ) + assert created is True + tracker.mark_completed("attempt-1") + replay, created = tracker.register_document_idempotent( + "attempt-2", + job_id="job", + stable_document_id="document-2", + filename="report.pdf", + content_sha256="b" * 64, + manifest_entry_id="a" * 64, + ) + assert created is False + assert replay.id == accepted.id == "attempt-1" + assert replay.stable_document_id == accepted.stable_document_id == "document-1" + assert len(tracker.job_documents("job")) == 1 + + try: + tracker.register_document_idempotent( + "attempt-3", + job_id="job", + filename="report.pdf", + content_sha256="c" * 64, + manifest_entry_id="a" * 64, + ) + except JobFullError: + pass + else: + raise AssertionError("conflicting manifest replay must fail") + + +def test_raw_storage_selection_is_rejected_for_query_requests() -> None: + for key in ("table_name", "lancedb_uri", "uri", "physical_table"): + try: + QueryRequest.model_validate({"query": "x", key: "untrusted"}) + except ValueError: + pass + else: + raise AssertionError(f"query storage key {key} must be rejected") + + +def test_scope_authorizer_secret_mapping_and_internal_vectordb_token(tmp_path) -> None: + secret = tmp_path / "scope-tokens.json" + secret.write_text('{"tokens":[{"token":"alpha-token","scopes":["alpha"]}]}', encoding="utf-8") + authorizer = ScopeAuthorizer(AuthConfig(enabled=True, scope_token_file=str(secret), allow_unscoped_dev=False)) + assert authorizer.authorize("alpha-token", "alpha") == ("alpha", None) + assert authorizer.authorize("alpha-token", "beta") == (None, 401) + assert authorizer.authorize("invalid", "alpha") == (None, 401) + + app = create_vectordb_app( + lancedb_uri=str(tmp_path / "db"), + embed_endpoint="http://embed", + internal_api_token="internal-secret", + ) + with TestClient(app) as client: + assert client.get("/v1/health").status_code == 200 + assert client.get("/v1/collections").status_code == 401 + assert client.get("/v1/collections", headers={"X-NRL-Internal-Token": "wrong"}).status_code == 401 + assert client.get("/v1/collections", headers={"X-NRL-Internal-Token": "internal-secret"}).status_code == 200 + + +def test_scope_authorizer_defaults_to_unprotected_and_fails_closed_when_enabled() -> None: + assert ScopeAuthorizer(AuthConfig()).authorize("", "workspace-a") == ("workspace-a", None) + assert ScopeAuthorizer(AuthConfig(enabled=True)).authorize("", "workspace-a") == (None, 401) + assert ScopeAuthorizer(AuthConfig(enabled=True, allow_unscoped_dev=True)).authorize("", "workspace-a") == ( + "workspace-a", + None, + ) + + +def test_public_routes_accept_requests_when_auth_is_disabled() -> None: + app = create_app(ServiceConfig(mode="gateway", auth=AuthConfig())) + with TestClient(app) as client: + response = client.get("/v1/collections") + + # The route is reached and then rejects the disabled VectorDB, rather than + # being rejected by bearer authentication. + assert response.status_code == 404 + + +def test_service_routes_use_authorized_scope_not_raw_header() -> None: + app = create_app( + ServiceConfig( + mode="gateway", + auth=AuthConfig( + enabled=True, + api_token="alpha-token", + default_scope="alpha", + allow_unscoped_dev=False, + ), + vectordb=VectorDbConfig(internal_api_token="internal-secret"), + ) + ) + with TestClient(app) as client: + invalid_token = client.post("/v1/ingest/job", json={"expected_documents": 1}) + invalid_scope = client.post( + "/v1/ingest/job", + json={"expected_documents": 1}, + headers={"Authorization": "Bearer alpha-token", "X-NRL-Scope": "beta"}, + ) + assert invalid_token.status_code == invalid_scope.status_code == 401 + assert invalid_token.json() == invalid_scope.json() == {"detail": "Missing or invalid bearer token."} + created = client.post( + "/v1/ingest/job", + json={"expected_documents": 1}, + headers={"Authorization": "Bearer alpha-token", "X-NRL-Scope": "alpha"}, + ) + assert created.status_code == 201 + job_id = created.json()["job_id"] + assert ( + client.get( + f"/v1/ingest/job/{job_id}", + headers={"Authorization": "Bearer alpha-token", "X-NRL-Scope": "alpha"}, + ).status_code + == 200 + ) + assert ( + client.get( + "/v1/internal/document-result/missing", + headers={"Authorization": "Bearer alpha-token"}, + ).status_code + == 401 + ) + assert ( + client.get( + "/v1/internal/document-result/missing", + headers={"X-NRL-Internal-Token": "internal-secret"}, + ).status_code + == 404 + ) + + +def test_vectordb_proxy_failures_do_not_expose_internal_details(monkeypatch) -> None: + class FailingAsyncClient: + def __init__(self, **_kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def aclose(self): + return None + + async def request(self, *_args, **_kwargs): + raise httpx.ConnectError("private-vectordb:7671?token=sensitive") + + async def post(self, *_args, **_kwargs): + raise httpx.ConnectError("private-vectordb:7671?token=sensitive") + + monkeypatch.setattr(httpx, "AsyncClient", FailingAsyncClient) + app = create_app( + ServiceConfig( + mode="gateway", + auth=AuthConfig(allow_unscoped_dev=True), + vectordb=VectorDbConfig(enabled=True, vectordb_url="http://private-vectordb:7671"), + llm=LLMConfig(enabled=True), + ) + ) + with TestClient(app) as client: + responses = ( + client.get("/v1/collections"), + client.post("/v1/query", json={"query": "test"}), + client.post("/v1/answer", json={"query": "test"}), + ) + + for response in responses: + assert response.status_code == 502 + assert response.json() == {"detail": "VectorDB service is unavailable."} + assert "private-vectordb" not in response.text + assert "sensitive" not in response.text + + +def test_openapi_operation_ids_are_unique() -> None: + app = create_app(ServiceConfig()) + operation_ids = [ + operation["operationId"] + for path in app.openapi()["paths"].values() + for operation in path.values() + if isinstance(operation, dict) and "operationId" in operation + ] + assert len(operation_ids) == len(set(operation_ids)) + + +def test_sdk_replays_every_manifest_entry_after_idempotent_job_replay(tmp_path, monkeypatch) -> None: + first = tmp_path / "first.txt" + second = tmp_path / "second.txt" + first.write_text("one", encoding="utf-8") + second.write_text("two", encoding="utf-8") + + class FakeResponse: + status_code = 200 + content = b"{}" + text = "" + + def json(self): + return { + "job_id": "job", + "expected_documents": 2, + "status": "completed", + "created_at": "now", + "counts": {"completed": 2}, + } + + class FakeAsyncClient: + def __init__(self, **_kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def get(self, _url): + return FakeResponse() + + monkeypatch.setattr(client_module.httpx, "AsyncClient", FakeAsyncClient) + sdk = RetrieverServiceClient() + sdk._create_job = AsyncMock(return_value=client_module._CreatedJob("job")) + sdk._upload_one = AsyncMock(return_value={"status": "accepted"}) + + result = asyncio.run(sdk.asubmit_documents("research", [first, second], idempotency_key="key")) + assert result.job_id == "job" + assert sdk._upload_one.await_count == 2 + entry_ids = [call.kwargs["manifest_entry_id"] for call in sdk._upload_one.await_args_list] + expected = [] + for position, path in enumerate((first, second)): + digest = hashlib.sha256(path.read_bytes()).hexdigest() + expected.append(hashlib.sha256(f"{position}\0{path.name}\0{digest}".encode("utf-8")).hexdigest()) + assert entry_ids == expected + + +def test_sdk_wraps_malformed_sync_and_async_lifecycle_responses() -> None: + class MalformedClient(RetrieverServiceClient): + async def _arequest(self, *_args, **_kwargs): + return {} + + sdk = MalformedClient() + with pytest.raises(RetrieverServiceError, match="invalid response"): + sdk.create_collection("research") + with pytest.raises(RetrieverServiceError, match="invalid response"): + asyncio.run(sdk.acreate_collection("research")) + + +def test_expiration_is_timezone_aware_and_normalized() -> None: + with pytest.raises(ValueError, match="timezone"): + CollectionCreateRequest(name="bad", expires_at="2030-01-01T00:00:00") + with pytest.raises(ValueError, match="ISO-8601"): + CollectionCreateRequest(name="bad-type", expires_at=123) + request = CollectionCreateRequest(name="good", expires_at="2030-01-01T01:00:00+01:00") + assert request.expires_at == "2030-01-01T00:00:00+00:00" + + +def test_keyset_cursors_are_stable_and_context_bound(tmp_path) -> None: + backend = _vdb(tmp_path) + for name in ("a", "c"): + backend.create_collection(scope="scope", request=CollectionCreateRequest(name=name)) + + first = backend.list_collections( + scope="scope", + limit=1, + continuation_token=None, + ) + assert [item.name for item in first.items] == ["a"] + backend.create_collection(scope="scope", request=CollectionCreateRequest(name="b")) + second = backend.list_collections( + scope="scope", + limit=2, + continuation_token=first.next_token, + ) + assert [item.name for item in second.items] == ["b", "c"] + with pytest.raises(VDBInvalidRequest, match="context"): + backend.list_collections( + scope="other", + limit=1, + continuation_token=first.next_token, + ) + + backend.write_collection( + [[_record("one", "doc-1", "one")]], + context=_context("doc-1", collection_name="a"), + ) + backend.write_collection( + [[_record("two", "doc-2", "two")]], + context=_context("doc-2", collection_name="a"), + ) + documents = backend.list_documents( + scope="scope", + collection_name="a", + limit=1, + continuation_token=None, + ) + assert documents.next_token + with pytest.raises(VDBInvalidRequest, match="context"): + backend.list_documents( + scope="scope", + collection_name="b", + limit=1, + continuation_token=documents.next_token, + ) + + +def test_search_and_collection_delete_share_lifecycle_lock(tmp_path, monkeypatch) -> None: + backend = _vdb(tmp_path) + backend.create_collection( + scope="scope", + request=CollectionCreateRequest(name="research"), + ) + backend.write_collection( + [[_record("chunk", "doc", "searchable text")]], + context=_context("doc"), + ) + + search_entered = threading.Event() + release_search = threading.Event() + delete_entered = threading.Event() + original_retrieval = LanceDB.retrieval + + def blocked_retrieval(vdb, *args, **kwargs): + search_entered.set() + assert release_search.wait(5) + return original_retrieval(vdb, *args, **kwargs) + + def observed_drop_table(*args, **kwargs): + delete_entered.set() + + monkeypatch.setattr(LanceDB, "retrieval", blocked_retrieval) + monkeypatch.setattr(backend._get_collection_store()._db, "drop_table", observed_drop_table) + + with ThreadPoolExecutor(max_workers=2) as pool: + search = pool.submit( + backend.retrieve_collection, + [[1.0, 0.0]], + scope="scope", + collection_name="research", + query_texts=["searchable text"], + top_k=1, + ) + assert search_entered.wait(5) + delete = pool.submit( + backend.delete_collection, + scope="scope", + collection_name="research", + if_exists=False, + ) + + assert not delete_entered.wait(0.25) + release_search.set() + + results, strategies = search.result(timeout=5) + deleted = delete.result(timeout=30) + + assert results[0][0]["document_id"] == "doc" + assert strategies == ["dense"] + assert delete_entered.is_set() + assert deleted.deleted + + +def test_collection_searches_run_concurrently(tmp_path, monkeypatch) -> None: + backend = _vdb(tmp_path) + backend.create_collection( + scope="scope", + request=CollectionCreateRequest(name="research"), + ) + backend.write_collection( + [[_record("chunk", "doc", "searchable text")]], + context=_context("doc"), + ) + searches_entered = threading.Barrier(2) + + def synchronized_retrieval(_vdb, vectors, **_kwargs): + searches_entered.wait(timeout=5) + return [[] for _ in vectors] + + monkeypatch.setattr(LanceDB, "retrieval", synchronized_retrieval) + with ThreadPoolExecutor(max_workers=2) as pool: + searches = [ + pool.submit( + backend.retrieve_collection, + [[1.0, 0.0]], + scope="scope", + collection_name="research", + query_texts=["searchable text"], + top_k=1, + ) + for _ in range(2) + ] + assert [future.result(timeout=5)[0] for future in searches] == [[[]], [[]]] + + +def test_reconcile_catalog_scan_does_not_hold_the_write_lock(tmp_path, monkeypatch) -> None: + backend = _vdb(tmp_path) + store = backend._get_collection_store() + scan_entered = threading.Event() + release_scan = threading.Event() + original_rows = store._rows + + def blocked_rows(table_name, where=None, columns=None): + if table_name == "_nrl_documents": + assert where == "recovery_state != ''" + scan_entered.set() + assert release_scan.wait(5) + return original_rows(table_name, where, columns) + + write_lock_acquired = threading.Event() + + def acquire_write_lock() -> None: + with store._write_lock: + write_lock_acquired.set() + + monkeypatch.setattr(store, "_rows", blocked_rows) + with ThreadPoolExecutor(max_workers=2) as pool: + reconciliation = pool.submit(backend.reconcile_collections) + assert scan_entered.wait(5) + write_lock_probe = pool.submit(acquire_write_lock) + try: + assert write_lock_acquired.wait(2) + finally: + release_scan.set() + write_lock_probe.result(timeout=5) + assert reconciliation.result(timeout=5) == {"successes": 0, "failures": 0} + + +def test_replacement_marker_recovers_after_catalog_finalize_failure(tmp_path, monkeypatch) -> None: + backend = _vdb(tmp_path) + backend.create_collection( + scope="scope", + request=CollectionCreateRequest(name="research"), + ) + backend.write_collection( + [[_record("old", "doc", "old", "v1")]], + context=_context("doc"), + ) + store = backend._get_collection_store() + original_persist = store._persist_document_row + + def fail_finalize(row): + if row.get("current_document_version") == "v2": + raise RuntimeError("injected finalize failure") + return original_persist(row) + + monkeypatch.setattr(store, "_persist_document_row", fail_finalize) + with pytest.raises(RuntimeError, match="injected"): + backend.write_collection( + [ + [ + _record("new-1", "doc", "new first", "v2"), + _record("new-2", "doc", "new second", "v2"), + ] + ], + context=_context("doc", version="v2", operation="replace"), + ) + monkeypatch.setattr(store, "_persist_document_row", original_persist) + document = backend.get_document(scope="scope", collection_name="research", document_id="doc") + assert document.status == "replacing" + result = backend.reconcile_collections() + assert result["successes"] == 1 + document = backend.get_document(scope="scope", collection_name="research", document_id="doc") + assert document.document_version == "v2" + assert document.content_sha256 == "v2" + assert document.filename == "report.pdf" + assert document.chunk_count == 2 + # No pending job ID is persisted in the catalog schema, so recovery keeps + # the last successfully finalized job rather than fabricating an ID. + assert document.job_id == "job-v1" + assert document.status == "completed" + table = store._open_table(store._resolved_table("scope", "research")) + rows = table.search().to_list() + versions = {row["document_version"] for row in rows} + assert versions == {"v2"} + assert len(rows) == 2 + + +def test_collection_deletion_does_not_delete_external_artifacts(tmp_path) -> None: + artifact = tmp_path / "artifacts" / "external" / "image.png" + artifact.parent.mkdir(parents=True) + artifact.write_bytes(b"image") + + backend = _vdb(tmp_path / "db") + backend.create_collection( + scope="scope", + request=CollectionCreateRequest(name="research"), + ) + backend.write_collection( + [[_record("chunk", "doc", "owned")]], + context=_context("doc"), + ) + deleted = backend.delete_collection(scope="scope", collection_name="research", if_exists=False) + + assert deleted.deleted + assert artifact.read_bytes() == b"image" + + +def test_expired_collection_uses_retryable_deletion_and_health_is_aggregate_only( + tmp_path, +) -> None: + backend = _vdb(tmp_path, table_name="secret-legacy-table") + backend.create_collection( + scope="tenant-secret", + request=CollectionCreateRequest( + name="expired", + expires_at="2000-01-01T00:00:00Z", + ), + ) + assert backend.reconcile_collections()["successes"] == 1 + with pytest.raises(VDBResourceNotFound): + backend.get_collection( + scope="tenant-secret", + collection_name="expired", + ) + health = backend.health() + assert health["catalog"]["schema_version"] == 2 + assert "tenant-secret" not in str(health) + assert "secret-legacy-table" not in str(health) + + +def test_catalog_startup_fails_fast_on_missing_required_columns(tmp_path) -> None: + db = lancedb.connect(str(tmp_path)) + db.create_table( + "_nrl_collections", + schema=pa.schema( + [ + pa.field("scope", pa.string()), + pa.field("name", pa.string()), + pa.field("physical_table", pa.string()), + pa.field("status", pa.string()), + pa.field("description", pa.string()), + pa.field("metadata_json", pa.string()), + pa.field("created_at", pa.string()), + pa.field("updated_at", pa.string()), + pa.field("expires_at", pa.string()), + ] + ), + ) + with pytest.raises(RuntimeError, match="missing required columns"): + _vdb(tmp_path).list_collections( + scope="scope", + limit=1, + continuation_token=None, + ) + backend = _vdb(tmp_path) + with pytest.raises(RuntimeError, match="missing required columns"): + backend.list_collections(scope="scope", limit=1, continuation_token=None) + with pytest.raises(RuntimeError, match="Collection catalog initialization failed"): + backend.health() + + db.drop_table("_nrl_collections") + assert backend.list_collections(scope="scope", limit=1, continuation_token=None).items == [] + assert backend.health()["catalog"]["initialized"] is True + + +def test_catalog_startup_fails_fast_on_incompatible_schema(tmp_path) -> None: + db = lancedb.connect(str(tmp_path)) + db.create_table( + "_nrl_collections", + schema=pa.schema( + [ + pa.field("scope", pa.int64()), + pa.field("name", pa.string()), + pa.field("physical_table", pa.string()), + pa.field("status", pa.string()), + pa.field("description", pa.string()), + pa.field("metadata_json", pa.string()), + pa.field("created_at", pa.string()), + pa.field("updated_at", pa.string()), + pa.field("expires_at", pa.string()), + ] + ), + ) + with pytest.raises(RuntimeError, match="Incompatible"): + _vdb(tmp_path).list_collections( + scope="scope", + limit=1, + continuation_token=None, + ) + + +def test_catalog_startup_does_not_recreate_an_unreadable_existing_table( + tmp_path, +) -> None: + backend = _vdb(tmp_path) + backend.list_collections(scope="scope", limit=1, continuation_token=None) + store = backend._get_collection_store() + db = Mock(wraps=store._db) + db.list_tables.return_value = store._db.list_tables() + db.open_table.side_effect = RuntimeError("catalog unreadable") + store._db = db + + with pytest.raises(RuntimeError, match="catalog unreadable"): + store._ensure_catalogs() + + db.create_table.assert_not_called() diff --git a/nemo_retriever/tests/test_common_datetools.py b/nemo_retriever/tests/test_common_datetools.py new file mode 100644 index 0000000000..8dc93ba940 --- /dev/null +++ b/nemo_retriever/tests/test_common_datetools.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from nemo_retriever.common.api.util.converters.datetools import normalize_timezone_aware_iso8601_to_utc + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("2030-01-01T00:00:00Z", "2030-01-01T00:00:00+00:00"), + ("2030-01-01T01:00:00+01:00", "2030-01-01T00:00:00+00:00"), + ("2030-01-01T00:00:00.123456+00:00", "2030-01-01T00:00:00.123456+00:00"), + ], +) +def test_normalize_timezone_aware_iso8601_to_utc(value: str, expected: str) -> None: + assert normalize_timezone_aware_iso8601_to_utc(value) == expected + + +@pytest.mark.parametrize( + ("value", "message"), + [ + ("not-a-timestamp", "ISO-8601"), + (123, "ISO-8601"), + ("2030-01-01T00:00:00", "timezone"), + ], +) +def test_normalize_timezone_aware_iso8601_to_utc_rejects_invalid_input(value: object, message: str) -> None: + with pytest.raises(ValueError, match=message): + normalize_timezone_aware_iso8601_to_utc(value) diff --git a/nemo_retriever/tests/test_compose_nim_2_contracts.py b/nemo_retriever/tests/test_compose_nim_2_contracts.py new file mode 100644 index 0000000000..58fe5405f8 --- /dev/null +++ b/nemo_retriever/tests/test_compose_nim_2_contracts.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compose contracts for the Object Detection and OCR 2.0 NIMs.""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + +ROOT = Path(__file__).resolve().parents[1] +COMPOSE = ROOT / "dev/compose/service-mode.compose.yaml" +CORE_PRESET = ROOT / "dev/compose/presets/nims-core.env" + + +def _compose() -> dict: + return yaml.safe_load(COMPOSE.read_text(encoding="utf-8")) + + +def test_extraction_nims_use_distinct_native_models_and_model_store_paths() -> None: + services = _compose()["services"] + expected = { + "nim-page-elements": ( + "nvidia/nemotron-page-elements-v3", + "page-elements", + "NIM_PAGE_ELEMENTS_CACHE_PATH", + ), + "nim-table-structure": ( + "nvidia/nemotron-table-structure-v1", + "table-structure", + "NIM_TABLE_STRUCTURE_CACHE_PATH", + ), + "nim-ocr": ("nvidia/nemotron-ocr-v2", "ocr", "NIM_OCR_CACHE_PATH"), + } + + for service_name, (model_name, model_dir, path_variable) in expected.items(): + service = services[service_name] + env = service["environment"] + assert env["NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER"] == "ngc" + assert env["NIM_ENGINE_MODEL_NAME"] == model_name + assert env["NIM_ENGINE_MODEL_PATH"] == f"${{{path_variable}:-/model-store}}/{model_dir}" + assert service["volumes"] == [ + f"nim_{service_name.removeprefix('nim-').replace('-', '_')}_cache:${{{path_variable}:-/model-store}}" + ] + assert service["healthcheck"]["test"] == [ + "CMD", + "curl", + "--fail", + "--silent", + "http://localhost:8000/v1/health/ready", + ] + assert not any(name.startswith("NIM_TRITON_") for name in env) + + assert services["nim-ocr"]["environment"]["NIM_ENGINE_MODEL_VARIANT"] == "multilingual" + + +def test_compose_and_core_preset_use_nim_2_routes() -> None: + compose = _compose() + config = compose["configs"]["retriever_service_config"]["content"] + preset = CORE_PRESET.read_text(encoding="utf-8") + + routes = ("/v1/page-elements", "/v1/table-structure", "/v1/ocr") + for route in routes: + assert route in config + assert route in preset + assert "/v1/infer" not in preset diff --git a/nemo_retriever/tests/test_create_local_embedder.py b/nemo_retriever/tests/test_create_local_embedder.py index 8a72bff595..3793d65adf 100644 --- a/nemo_retriever/tests/test_create_local_embedder.py +++ b/nemo_retriever/tests/test_create_local_embedder.py @@ -4,14 +4,21 @@ """Unit tests for nemo_retriever.models.create_local_embedder factory.""" +import json import sys import warnings +from pathlib import Path from types import ModuleType from unittest.mock import MagicMock import pytest -from nemo_retriever.models import create_local_embedder, create_local_query_embedder +from nemo_retriever.models import ( + EmbedModelSpec, + create_local_embedder, + create_local_query_embedder, + is_vl_embed_model, +) @pytest.fixture(autouse=True) @@ -44,9 +51,57 @@ def _patch_embedders(monkeypatch): monkeypatch.setitem(sys.modules, "nemo_retriever.models.local.llama_nemotron_embed_1b_v2_hf_embedder", text_hf_mod) monkeypatch.setitem(sys.modules, "nemo_retriever.models.local.llama_nemotron_embed_vl_1b_v2_embedder", vl_mod) + def resolve_spec(model_id, *, revision=None, hf_cache_dir=None): + config_path = Path(model_id) / "config.json" + if config_path.is_file(): + config = json.loads(config_path.read_text(encoding="utf-8")) + family = "vl" if config["model_type"] == "llama_nemotron_vl" else "text" + quantization = config.get("quantization_config") or {} + requires_vllm = quantization.get("quant_method") == "modelopt" + dimension_config = (config.get("llm_config") or {}) if family == "vl" else config + output_dimension = dimension_config.get("hidden_size", 2048) + resolved_revision = None + else: + family = "vl" if "embed-vl" in model_id or "vlm-embed" in model_id else "text" + quantization = {} + requires_vllm = False + output_dimension = 4096 if "8b" in model_id.lower() else 2048 + resolved_revision = revision or "a" * 40 + instruction_model = model_id == "nvidia/llama-embed-nemotron-8b" + return EmbedModelSpec( + model_id=model_id, + revision=resolved_revision, + family=family, + output_dimension=output_dimension, + query_prefix=( + "Instruct: Given a question, retrieve passages that answer the question\nQuery: " + if instruction_model + else "query: " + ), + document_prefix="" if instruction_model else "passage: ", + quantization=quantization.get("quant_algo"), + requires_vllm=requires_vllm, + ) + + monkeypatch.setattr("nemo_retriever.models.resolve_embed_model_spec", resolve_spec) + yield fake_text_vllm, fake_text_hf, fake_vl_hf, fake_vl_vllm +@pytest.fixture +def local_checkpoint(tmp_path): + """Create a local checkpoint whose architecture is declared by config.json.""" + + def create(model_type, *, modelopt=False): + config = {"model_type": model_type} + if modelopt: + config["quantization_config"] = {"quant_method": "modelopt", "quant_algo": "FP8"} + (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8") + return tmp_path + + return create + + # --------------------------------------------------------------------------- # create_local_embedder — default model (VL, since _DEFAULT_EMBED_MODEL is VL) # --------------------------------------------------------------------------- @@ -74,6 +129,18 @@ def test_alias_resolved_to_text_embedder(_patch_embedders): assert result is fake_text_vllm.return_value +@pytest.mark.parametrize( + ("model_name", "expected"), + [ + pytest.param(None, True, id="default"), + pytest.param("llama-3.2-nemoretriever-1b-vlm-embed-v1", True, id="legacy-vl-alias"), + pytest.param("nvidia/llama-nemotron-embed-1b-v2", False, id="text"), + ], +) +def test_is_vl_embed_model_preserves_legacy_contract(model_name, expected): + assert is_vl_embed_model(model_name) is expected + + def test_default_model_explicit_vllm_backend(_patch_embedders): _, _, _, fake_vl_vllm = _patch_embedders result = create_local_embedder(backend="vllm") @@ -121,6 +188,16 @@ def test_unknown_model_passes_through(_patch_embedders): assert kw["model_id"] == "custom-org/my-embed-model" +def test_8b_prompt_metadata_is_forwarded_to_text_loader(_patch_embedders): + fake_text_vllm, _, _, _ = _patch_embedders + + create_local_embedder("nvidia/llama-embed-nemotron-8b") + + kw = fake_text_vllm.call_args.kwargs + assert kw["query_prefix"].startswith("Instruct:") + assert kw["document_prefix"] == "" + + # --------------------------------------------------------------------------- # create_local_embedder — VL model # --------------------------------------------------------------------------- @@ -221,6 +298,67 @@ def test_query_embedder_vl_vllm_uses_vllm_vl(_patch_embedders): assert result is fake_vl_vllm.return_value +# --------------------------------------------------------------------------- +# Local checkpoint directories — routed by config.json +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("model_type", "backend", "embedder_index"), + [ + pytest.param("llama_nemotron_vl", "vllm", 3, id="vl-vllm"), + pytest.param("llama_nemotron_vl", "hf", 2, id="vl-hf"), + pytest.param("llama_bidirec", "vllm", 0, id="text-vllm"), + pytest.param("llama_bidirec", "hf", 1, id="text-hf"), + ], +) +def test_local_checkpoint_routes_from_config( + local_checkpoint, + _patch_embedders, + model_type, + backend, + embedder_index, +): + checkpoint = local_checkpoint(model_type) + + result = create_local_embedder(str(checkpoint), backend=backend) + + selected = _patch_embedders[embedder_index] + selected.assert_called_once() + assert selected.call_args.kwargs["model_id"] == str(checkpoint) + assert selected.call_args.kwargs["revision"] is None + assert result is selected.return_value + + +def test_modelopt_checkpoint_rejects_hf_before_loader(local_checkpoint, _patch_embedders): + checkpoint = local_checkpoint("llama_nemotron_vl", modelopt=True) + with pytest.raises(ValueError, match="requires backend='vllm'"): + create_local_embedder(str(checkpoint), backend="hf") + + +def test_relative_local_checkpoint_routes_locally(tmp_path, monkeypatch, _patch_embedders): + _, fake_text_hf, _, _ = _patch_embedders + checkpoint = tmp_path / "my-checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text('{"model_type":"llama_bidirec"}', encoding="utf-8") + monkeypatch.chdir(tmp_path) + + create_local_embedder("./my-checkpoint", backend="hf") + + assert fake_text_hf.call_args.kwargs["model_id"] == "./my-checkpoint" + + +def test_query_embedder_forwards_recorded_revision(_patch_embedders): + _, _, fake_vl_hf, _ = _patch_embedders + revision = "b" * 40 + create_local_query_embedder( + "nvidia/llama-nemotron-embed-vl-1b-v2", + backend="hf", + revision=revision, + ) + assert fake_vl_hf.call_args.kwargs["revision"] == revision + + # --------------------------------------------------------------------------- # Real-class smoke test (requires torch; skipped if not installed) # --------------------------------------------------------------------------- diff --git a/nemo_retriever/tests/test_embed_model_spec.py b/nemo_retriever/tests/test_embed_model_spec.py new file mode 100644 index 0000000000..c43010cbc2 --- /dev/null +++ b/nemo_retriever/tests/test_embed_model_spec.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import huggingface_hub +import pytest + +from nemo_retriever.models.embed_model_spec import ( + resolve_embed_model_revision, + resolve_embed_model_spec, +) +from nemo_retriever.models.hf_model_registry import HF_MODEL_REVISIONS + + +def _write_config(tmp_path, config): + path = tmp_path / "config.json" + path.write_text(json.dumps(config), encoding="utf-8") + return path + + +def _write_prompt_config(tmp_path, prompts): + path = tmp_path / "config_sentence_transformers.json" + path.write_text(json.dumps({"prompts": prompts}), encoding="utf-8") + return path + + +def _text_config(**overrides): + config = { + "model_type": "llama_bidirec", + "architectures": ["LlamaBidirectionalModel"], + "hidden_size": 2048, + "pooling": "avg", + } + config.update(overrides) + return config + + +def _vl_config(**overrides): + config = { + "model_type": "llama_nemotron_vl", + "architectures": ["LlamaNemotronVLModel"], + "llm_config": {"hidden_size": 2048}, + "pooling": "avg", + } + config.update(overrides) + return config + + +def test_local_text_checkpoint_is_resolved_from_model_type(tmp_path): + _write_config(tmp_path, _text_config()) + + spec = resolve_embed_model_spec(str(tmp_path)) + + assert spec.family == "text" + assert spec.output_dimension == 2048 + assert spec.revision is None + assert spec.requires_vllm is False + + +@pytest.mark.parametrize( + ("config", "family", "quantization"), + [ + pytest.param( + _vl_config(quantization_config={"quant_method": "modelopt", "quant_algo": "FP8"}), + "vl", + "FP8", + id="vl-fp8", + ), + pytest.param( + _text_config(quantization_config={"quant_method": "modelopt", "quant_algo": "NVFP4"}), + "text", + "NVFP4", + id="text-nvfp4", + ), + ], +) +def test_local_modelopt_checkpoint_requires_vllm(tmp_path, config, family, quantization): + _write_config(tmp_path, config) + + spec = resolve_embed_model_spec(str(tmp_path)) + + assert spec.family == family + assert spec.quantization == quantization + assert spec.requires_vllm is True + + +def test_local_checkpoint_rejects_unsupported_architecture(tmp_path): + _write_config(tmp_path, {"model_type": "bert"}) + + with pytest.raises(ValueError, match="supported Nemotron embed model types"): + resolve_embed_model_spec(str(tmp_path)) + + +def test_text_reranker_architecture_is_rejected(tmp_path): + config = _text_config(architectures=["LlamaBidirectionalForSequenceClassification"]) + _write_config(tmp_path, config) + + with pytest.raises(ValueError, match="unsupported architectures"): + resolve_embed_model_spec(str(tmp_path)) + + +@pytest.mark.parametrize( + ("config", "family"), + [ + pytest.param(_text_config(hidden_size=4096), "text", id="text"), + pytest.param(_vl_config(llm_config={"hidden_size": 4096}), "vl", id="vl"), + ], +) +def test_4096_wide_embedder_is_supported(tmp_path, config, family): + _write_config(tmp_path, config) + + spec = resolve_embed_model_spec(str(tmp_path)) + + assert spec.family == family + assert spec.output_dimension == 4096 + + +def test_invalid_embedding_dimension_is_rejected(tmp_path): + _write_config(tmp_path, _text_config(hidden_size=0)) + + with pytest.raises(ValueError, match="invalid embedding dimension 0"): + resolve_embed_model_spec(str(tmp_path)) + + +def test_ministral_checkpoint_is_rejected(tmp_path): + _write_config( + tmp_path, + { + "model_type": "ministral3", + "architectures": ["Ministral3Model"], + "hidden_size": 2048, + "is_causal": False, + "pooling": "avg", + }, + ) + + with pytest.raises(ValueError, match="unsupported model_type 'ministral3'"): + resolve_embed_model_spec(str(tmp_path)) + + +def test_checkpoint_prompt_metadata_is_resolved(tmp_path): + _write_config(tmp_path, _text_config(hidden_size=4096)) + _write_prompt_config( + tmp_path, + { + "query": "Instruct: Retrieve relevant passages\nQuery: ", + "document": "", + }, + ) + + spec = resolve_embed_model_spec(str(tmp_path)) + + assert spec.query_prefix == "Instruct: Retrieve relevant passages\nQuery: " + assert spec.document_prefix == "" + + +def test_non_average_pooling_is_rejected(tmp_path): + _write_config(tmp_path, _text_config(pooling="last")) + + with pytest.raises(ValueError, match="unsupported pooling 'last'"): + resolve_embed_model_spec(str(tmp_path)) + + +def test_local_directory_requires_config_json(tmp_path): + with pytest.raises(ValueError, match="does not contain config.json"): + resolve_embed_model_spec(str(tmp_path)) + + +@pytest.mark.parametrize( + ("model_id", "config", "family", "dimension", "requires_vllm"), + [ + ("nvidia/llama-3.2-nv-embedqa-1b-v2", _text_config(), "text", 2048, False), + ("nvidia/llama-nemotron-embed-1b-v2", _text_config(), "text", 2048, False), + ("nvidia/llama-nemotron-embed-vl-1b-v2", _vl_config(), "vl", 2048, False), + ( + "nvidia/llama-nemotron-embed-vl-1b-v2-fp8", + _vl_config(quantization_config={"quant_method": "modelopt", "quant_algo": "FP8"}), + "vl", + 2048, + True, + ), + ("nvidia/llama-embed-nemotron-8b", _text_config(hidden_size=4096), "text", 4096, False), + ("nvidia/llama-nv-embed-reasoning-3b", _text_config(hidden_size=3072), "text", 3072, False), + ], +) +def test_supported_hub_models_resolve_immutable_revisions( + monkeypatch, + tmp_path, + model_id, + config, + family, + dimension, + requires_vllm, +): + _write_config(tmp_path, config) + resolved_sha = "c" * 40 + model_info_calls = [] + + def model_info(_self, requested_model_id, revision=None): + model_info_calls.append((requested_model_id, revision)) + return SimpleNamespace(sha=resolved_sha) + + monkeypatch.setattr(huggingface_hub.HfApi, "model_info", model_info) + monkeypatch.setattr(huggingface_hub, "hf_hub_download", lambda **kwargs: str(tmp_path / "config.json")) + + spec = resolve_embed_model_spec(model_id) + + assert spec.family == family + assert spec.output_dimension == dimension + assert spec.requires_vllm is requires_vllm + if model_id in HF_MODEL_REVISIONS: + assert spec.revision == HF_MODEL_REVISIONS[model_id] + assert model_info_calls == [] + else: + assert spec.revision == resolved_sha + assert model_info_calls == [(model_id, None)] + + +def test_compatible_custom_hub_model_is_pinned_before_config_load(monkeypatch, tmp_path): + _write_config(tmp_path, _vl_config()) + resolved_sha = "c" * 40 + calls = {"downloads": []} + + def model_info(_self, model_id, revision=None): + calls["model_info"] = (model_id, revision) + return SimpleNamespace(sha=resolved_sha) + + def download(**kwargs): + calls["downloads"].append(kwargs) + return str(tmp_path / "config.json") + + monkeypatch.setattr(huggingface_hub.HfApi, "model_info", model_info) + monkeypatch.setattr(huggingface_hub, "hf_hub_download", download) + + cache_dir = str(tmp_path / "hub-cache") + spec = resolve_embed_model_spec("acme/fine-tuned-nemotron", hf_cache_dir=cache_dir) + + assert spec.family == "vl" + assert spec.revision == resolved_sha + assert calls["model_info"] == ("acme/fine-tuned-nemotron", None) + assert {call["filename"] for call in calls["downloads"]} == { + "config.json", + "config_sentence_transformers.json", + } + assert all(call["revision"] == resolved_sha for call in calls["downloads"]) + assert all(call["cache_dir"] == cache_dir for call in calls["downloads"]) + + +def test_revision_only_resolution_avoids_config_download_for_registered_model(monkeypatch): + model_id = "nvidia/llama-nemotron-embed-vl-1b-v2" + + def fail_download(*args, **kwargs): + raise AssertionError("revision-only resolution must not download config.json") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", fail_download) + + assert resolve_embed_model_revision(model_id, None) == HF_MODEL_REVISIONS[model_id] diff --git a/nemo_retriever/tests/test_embed_params.py b/nemo_retriever/tests/test_embed_params.py index db2b334c99..119417946a 100644 --- a/nemo_retriever/tests/test_embed_params.py +++ b/nemo_retriever/tests/test_embed_params.py @@ -55,16 +55,16 @@ def test_image_modalities_constant(): assert isinstance(IMAGE_MODALITIES, frozenset) -def test_build_embed_option_kwargs_applies_remote_model_provider_prefix(): +def test_build_embed_option_kwargs_defers_remote_model_provider_prefix(): kwargs = build_embed_option_kwargs( "https://litellm.example.com/v1/embeddings", "nvidia/llama-nemotron-embed-vl-1b-v2", embed_model_provider_prefix="nvidia", ) - assert kwargs["model_name"] == "nvidia/nvidia/llama-nemotron-embed-vl-1b-v2" - assert kwargs["embed_model_name"] == "nvidia/nvidia/llama-nemotron-embed-vl-1b-v2" - assert "embed_model_provider_prefix" not in kwargs + assert kwargs["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert kwargs["embed_model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert kwargs["embed_model_provider_prefix"] == "nvidia" def test_build_embed_option_kwargs_leaves_model_unchanged_without_prefix(): @@ -77,38 +77,30 @@ def test_build_embed_option_kwargs_leaves_model_unchanged_without_prefix(): assert kwargs["embed_model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" -def test_build_embed_option_kwargs_prefix_supports_other_vendor_namespaces(): - kwargs = build_embed_option_kwargs( - "https://litellm.example.com/v1/embeddings", - "mistral/embed-small", - embed_model_provider_prefix="acme", - ) - - assert kwargs["model_name"] == "acme/mistral/embed-small" - assert kwargs["embed_model_name"] == "acme/mistral/embed-small" - - -def test_build_embed_option_kwargs_prefix_supports_bare_model_name(): +def test_build_embed_option_kwargs_keeps_provider_prefix_separate_without_endpoint(): kwargs = build_embed_option_kwargs( - "https://litellm.example.com/v1/embeddings", - "nv-embedqa-e5-v5", + None, + "nvidia/llama-nemotron-embed-vl-1b-v2", embed_model_provider_prefix="nvidia", ) - assert kwargs["model_name"] == "nvidia/nv-embedqa-e5-v5" - assert kwargs["embed_model_name"] == "nvidia/nv-embedqa-e5-v5" + assert kwargs["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert kwargs["embed_model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert kwargs["embed_model_provider_prefix"] == "nvidia" -def test_build_embed_option_kwargs_prefix_is_remote_only(): +def test_build_embed_option_kwargs_retains_prefix_when_model_is_omitted(): kwargs = build_embed_option_kwargs( + "https://inference-api.nvidia.com/v1", None, - "nvidia/llama-nemotron-embed-vl-1b-v2", embed_model_provider_prefix="nvidia", ) - assert kwargs["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" - assert kwargs["embed_model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" - assert "embed_model_provider_prefix" not in kwargs + assert kwargs == { + "embed_invoke_url": "https://inference-api.nvidia.com/v1", + "embedding_endpoint": "https://inference-api.nvidia.com/v1", + "embed_model_provider_prefix": "nvidia", + } # =================================================================== diff --git a/nemo_retriever/tests/test_harness_agentic_eval.py b/nemo_retriever/tests/test_harness_agentic_eval.py index b47dd26e61..caab76c09a 100644 --- a/nemo_retriever/tests/test_harness_agentic_eval.py +++ b/nemo_retriever/tests/test_harness_agentic_eval.py @@ -47,11 +47,11 @@ def test_query_override_paths_include_agentic_fields() -> None: "query.agentic_local_max_model_len", "query.agentic_local_max_num_seqs", "query.agentic_reasoning_effort", - "query.agentic_backend_top_k", "query.agentic_react_max_steps", "query.agentic_text_truncation", "query.agentic_num_concurrent", "query.agentic_temperature", + "query.agentic_llm_client", ): assert key in QUERY_OVERRIDE_PATHS @@ -69,11 +69,11 @@ def test_build_query_request_populates_agentic() -> None: "agentic_local_max_model_len": 8192, "agentic_local_max_num_seqs": 4, "agentic_reasoning_effort": "high", - "agentic_backend_top_k": 25, "agentic_react_max_steps": 12, "agentic_text_truncation": 4000, "agentic_num_concurrent": 4, "agentic_temperature": 0.5, + "agentic_llm_client": "litellm", } ), "", @@ -89,17 +89,18 @@ def test_build_query_request_populates_agentic() -> None: assert agentic.local_max_model_len == 8192 assert agentic.local_max_num_seqs == 4 assert agentic.reasoning_effort == "high" - assert agentic.backend_top_k == 25 assert agentic.react_max_steps == 12 assert agentic.text_truncation == 4000 assert agentic.num_concurrent == 4 assert agentic.temperature == pytest.approx(0.5) + assert agentic.llm_client == "litellm" def test_build_query_request_agentic_defaults_when_absent() -> None: request = build_query_request(_resolved({"top_k": 10}), "") assert request.agentic == QueryAgenticOptions() assert request.agentic.enabled is False + assert request.agentic.temperature is None def test_build_agentic_config_maps_request_and_top_k_override() -> None: @@ -109,7 +110,6 @@ def test_build_agentic_config_maps_request_and_top_k_override() -> None: enabled=True, llm_model="test-model", invoke_url="http://localhost/v1/chat/completions", - backend_top_k=20, num_concurrent=4, temperature=0.0, ), @@ -118,7 +118,6 @@ def test_build_agentic_config_maps_request_and_top_k_override() -> None: assert cfg.llm_model == "test-model" assert cfg.llm_backend == "openai_compatible" assert cfg.top_k == 10 # harness sets this to the deepest BEIR k - assert cfg.backend_top_k == 20 assert cfg.num_concurrent == 4 @@ -197,7 +196,7 @@ def test_run_beir_queries_routes_to_agentic(tmp_path) -> None: def test_run_beir_queries_invalid_agentic_config_is_structured_failure(tmp_path) -> None: - # An invalid agentic config (backend_top_k below the target top_k = max(ks)) + # An invalid agentic config (temperature above the hosted-NVIDIA max of 1.0) # must surface as a structured HarnessRunError, not a raw ValueError. from nemo_retriever.harness.contracts import EXIT_INVALID, HarnessRunError @@ -207,7 +206,15 @@ def test_run_beir_queries_invalid_agentic_config_is_structured_failure(tmp_path) "query": {}, } request = build_query_request( - _resolved({"top_k": 10, "agentic": True, "agentic_llm_model": "nemotron-8b", "agentic_backend_top_k": 5}), + _resolved( + { + "top_k": 10, + "agentic": True, + "agentic_llm_model": "nemotron-8b", + "agentic_invoke_url": "https://integrate.api.nvidia.com/v1/chat/completions", + "agentic_temperature": 1.5, + } + ), "", ) dataset = BeirDataset(dataset_name="demo", query_ids=["q1"], queries=["t1"], qrels={"q1": {"d1": 1}}) diff --git a/nemo_retriever/tests/test_harness_helm_nightly.py b/nemo_retriever/tests/test_harness_helm_nightly.py index 05fdc37659..af7db000c2 100644 --- a/nemo_retriever/tests/test_harness_helm_nightly.py +++ b/nemo_retriever/tests/test_harness_helm_nightly.py @@ -134,10 +134,12 @@ def test_service_dry_run_writes_plans_without_network(service_execution, monkeyp def test_service_ingest_and_beir_use_shared_results(service_execution, monkeypatch, tmp_path: Path) -> None: + ingest_calls = [] monkeypatch.setattr( service_execution, "execute_service_ingest_request", - lambda _request: SimpleNamespace(to_summary_dict=lambda: {"n_rows": 12}), + lambda _request, **kwargs: ingest_calls.append(kwargs) + or SimpleNamespace(to_summary_dict=lambda: {"n_rows": 12}), ) monkeypatch.setattr( service_execution, @@ -151,6 +153,7 @@ def test_service_ingest_and_beir_use_shared_results(service_execution, monkeypat ) assert outcome.exit_code == 0 + assert ingest_calls == [{"return_results": False}] assert outcome.results["summary_metrics"]["rows_processed"] == 12 assert outcome.results["summary_metrics"]["query_count"] == 1 assert outcome.results["summary_metrics"]["recall_5"] == 0.9 @@ -199,7 +202,7 @@ def test_service_metric_gate_failure_uses_standard_exit_code(service_execution, monkeypatch.setattr( service_execution, "execute_service_ingest_request", - lambda _request: SimpleNamespace(to_summary_dict=lambda: {"n_rows": 1}), + lambda _request, **_kwargs: SimpleNamespace(to_summary_dict=lambda: {"n_rows": 1}), ) outcome = run_prepared_benchmark( _prepared(tmp_path, requirements=("files==2",)), @@ -211,7 +214,7 @@ def test_service_metric_gate_failure_uses_standard_exit_code(service_execution, def test_service_ingest_failure_is_concise(service_execution, monkeypatch, tmp_path: Path) -> None: - def fail(_request): + def fail(_request, **_kwargs): raise RuntimeError("service unavailable\ntraceback details") monkeypatch.setattr(service_execution, "execute_service_ingest_request", fail) diff --git a/nemo_retriever/tests/test_helm_auth.py b/nemo_retriever/tests/test_helm_auth.py new file mode 100644 index 0000000000..bf99117f2a --- /dev/null +++ b/nemo_retriever/tests/test_helm_auth.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Helm coverage for public and internal Secret-backed authentication.""" + +from __future__ import annotations + +import subprocess + +import pytest + +from tests.test_helm_shared_results import _render + + +def _deployments(documents: list[dict]) -> dict[str, dict]: + return { + document["metadata"]["labels"]["app.kubernetes.io/component"]: document + for document in documents + if document.get("kind") == "Deployment" + } + + +def _container_env(deployment: dict) -> dict[str, dict]: + container = deployment["spec"]["template"]["spec"]["containers"][0] + return {item["name"]: item for item in container.get("env", [])} + + +def _auth_args(*, split: bool = False) -> tuple[str, ...]: + args = ( + "--set", + "serviceConfig.vectordb.enabled=true", + "--set", + "serviceConfig.nimEndpoints.embedInvokeUrl=http://embed.invalid/v1/embeddings", + "--set", + "serviceConfig.vectordb.internalAuth.enabled=true", + "--set", + "serviceConfig.vectordb.internalAuth.existingSecret.name=nrl-internal-auth", + "--set", + "serviceConfig.auth.enabled=true", + "--set", + "serviceConfig.auth.scopeTokenSecret.name=nrl-public-auth", + ) + if split: + return ( + *args, + "--set", + "topology.mode=split", + "--set", + "serviceMonitor.autoEnableInSplitMode=false", + ) + return args + + +def test_auth_is_disabled_by_default_for_unprotected_standalone_deployments() -> None: + documents = _render() + deployment = next(iter(_deployments(documents).values())) + assert "NRL_INTERNAL_VDB_TOKEN" not in _container_env(deployment) + assert "NRL_SCOPE_TOKEN_FILE" not in _container_env(deployment) + config = next( + item["data"]["retriever-service.yaml"] + for item in documents + if item.get("kind") == "ConfigMap" and "retriever-service.yaml" in item.get("data", {}) + ) + assert "api_token: null" in config + assert 'default_scope: "default"' in config + assert "enabled: false" in config + assert "allow_unscoped_dev: false" in config + + +def test_standalone_wires_separate_public_and_internal_secrets() -> None: + documents = _render(*_auth_args()) + deployments = _deployments(documents) + assert set(deployments) >= {"service", "vectordb"} + + for component in ("service", "vectordb"): + internal = _container_env(deployments[component])["NRL_INTERNAL_VDB_TOKEN"]["valueFrom"]["secretKeyRef"] + assert internal == { + "name": "nrl-internal-auth", + "key": "token", + "optional": False, + } + + service = deployments["service"] + scope_env = _container_env(service)["NRL_SCOPE_TOKEN_FILE"] + assert scope_env["value"] == "/var/run/secrets/nemo-retriever/auth/scope-tokens.json" + pod_spec = service["spec"]["template"]["spec"] + scope_volume = next(item for item in pod_spec["volumes"] if item["name"] == "scope-token") + assert scope_volume["secret"] == { + "secretName": "nrl-public-auth", + "defaultMode": 288, + "items": [{"key": "scope-tokens.json", "path": "scope-tokens.json"}], + } + config = next( + item["data"]["retriever-service.yaml"] + for item in documents + if item.get("kind") == "ConfigMap" and "retriever-service.yaml" in item.get("data", {}) + ) + assert "api_token: null" in config + assert "nrl-public-auth" not in config + assert "nrl-internal-auth" not in config + + +def test_split_mounts_public_secret_only_on_gateway_and_internal_secret_everywhere() -> None: + deployments = _deployments(_render(*_auth_args(split=True))) + assert set(deployments) >= {"gateway", "realtime", "batch", "vectordb"} + for component in ("gateway", "realtime", "batch", "vectordb"): + assert ( + _container_env(deployments[component])["NRL_INTERNAL_VDB_TOKEN"]["valueFrom"]["secretKeyRef"]["name"] + == "nrl-internal-auth" + ) + + assert "NRL_SCOPE_TOKEN_FILE" in _container_env(deployments["gateway"]) + for component in ("realtime", "batch", "vectordb"): + assert "NRL_SCOPE_TOKEN_FILE" not in _container_env(deployments[component]) + pod_spec = deployments[component]["spec"]["template"]["spec"] + assert all(item["name"] != "scope-token" for item in pod_spec.get("volumes", [])) + + +def test_internal_auth_requires_existing_secret_name() -> None: + with pytest.raises(subprocess.CalledProcessError) as error: + _render("--set", "serviceConfig.vectordb.internalAuth.enabled=true") + assert "internalAuth.existingSecret.name" in error.value.stderr + + +def test_internal_auth_requires_existing_secret_key() -> None: + with pytest.raises(subprocess.CalledProcessError) as error: + _render( + "--set", + "serviceConfig.vectordb.internalAuth.enabled=true", + "--set", + "serviceConfig.vectordb.internalAuth.existingSecret.name=nrl-internal-auth", + "--set-string", + "serviceConfig.vectordb.internalAuth.existingSecret.key=", + ) + assert "internalAuth.existingSecret.key" in error.value.stderr + + +def test_public_auth_requires_existing_secret_key() -> None: + with pytest.raises(subprocess.CalledProcessError) as error: + _render( + "--set", + "serviceConfig.auth.enabled=true", + "--set", + "serviceConfig.auth.scopeTokenSecret.name=nrl-public-auth", + "--set-string", + "serviceConfig.auth.scopeTokenSecret.key=", + ) + assert "scopeTokenSecret.key" in error.value.stderr + + +def test_inline_public_token_requires_explicit_insecure_gate() -> None: + with pytest.raises(subprocess.CalledProcessError) as error: + _render( + "--set", + "serviceConfig.auth.enabled=true", + "--set-string", + "serviceConfig.auth.apiToken=sentinel-public-token", + ) + assert "allowInsecureInlineApiToken=true" in error.value.stderr + + documents = _render( + "--set", + "serviceConfig.auth.enabled=true", + "--set-string", + "serviceConfig.auth.apiToken=sentinel-public-token", + "--set", + "serviceConfig.auth.allowInsecureInlineApiToken=true", + ) + config = next( + item["data"]["retriever-service.yaml"] + for item in documents + if item.get("kind") == "ConfigMap" and "retriever-service.yaml" in item.get("data", {}) + ) + assert "api_token: sentinel-public-token" in config + + +def test_inline_and_secret_public_auth_are_mutually_exclusive() -> None: + with pytest.raises(subprocess.CalledProcessError) as error: + _render( + "--set", + "serviceConfig.auth.enabled=true", + "--set-string", + "serviceConfig.auth.apiToken=sentinel-public-token", + "--set", + "serviceConfig.auth.allowInsecureInlineApiToken=true", + "--set", + "serviceConfig.auth.scopeTokenSecret.name=nrl-public-auth", + ) + assert "mutually exclusive" in error.value.stderr + + +def test_enabled_public_auth_requires_a_credential_source() -> None: + with pytest.raises(subprocess.CalledProcessError) as error: + _render("--set", "serviceConfig.auth.enabled=true") + assert "auth.enabled=true requires" in error.value.stderr + + +def test_public_credentials_require_auth_to_be_enabled() -> None: + with pytest.raises(subprocess.CalledProcessError) as error: + _render("--set", "serviceConfig.auth.scopeTokenSecret.name=nrl-public-auth") + assert "require serviceConfig.auth.enabled=true" in error.value.stderr diff --git a/nemo_retriever/tests/test_helm_embed_model_provider_prefix.py b/nemo_retriever/tests/test_helm_embed_model_provider_prefix.py new file mode 100644 index 0000000000..adb72db50d --- /dev/null +++ b/nemo_retriever/tests/test_helm_embed_model_provider_prefix.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Helm rendering for the remote embedding provider prefix.""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path +from unittest import SkipTest + +import yaml + +CHART = Path(__file__).resolve().parents[1] / "helm" + + +def _render(*extra_args: str) -> dict: + helm = shutil.which("helm") + if helm is None: + raise SkipTest("`helm` binary not available in this environment.") + + command = [ + helm, + "template", + "embed-provider-prefix-test", + str(CHART), + "--set", + "nims.enabled=false", + "--set", + "serviceConfig.nimEndpoints.embedInvokeUrl=http://embed:8000/v1/embeddings", + *extra_args, + ] + completed = subprocess.run(command, check=True, capture_output=True, text=True) + documents = [document for document in yaml.safe_load_all(completed.stdout) if document] + configmap = next( + document + for document in documents + if document.get("kind") == "ConfigMap" and "retriever-service.yaml" in document.get("data", {}) + ) + return yaml.safe_load(configmap["data"]["retriever-service.yaml"]) + + +def test_default_provider_prefix_preserves_null_schema_keys() -> None: + config = _render() + + assert "embed_model_provider_prefix" in config["nim_endpoints"] + assert config["nim_endpoints"]["embed_model_provider_prefix"] is None + assert "embed_model_provider_prefix" in config["vectordb"] + assert config["vectordb"]["embed_model_provider_prefix"] is None + + +def test_provider_prefix_override_is_rendered_in_both_sections() -> None: + config = _render("--set", "serviceConfig.vectordb.embedModelProviderPrefix=nvidia") + + assert config["nim_endpoints"]["embed_model_provider_prefix"] == "nvidia" + assert config["vectordb"]["embed_model_provider_prefix"] == "nvidia" diff --git a/nemo_retriever/tests/test_helm_nemotron_parse_endpoint.py b/nemo_retriever/tests/test_helm_nemotron_parse_endpoint.py new file mode 100644 index 0000000000..d0c357181e --- /dev/null +++ b/nemo_retriever/tests/test_helm_nemotron_parse_endpoint.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression tests for Nemotron Parse service endpoint Helm wiring.""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path +from typing import Sequence +from unittest import SkipTest + +import yaml + + +REPO_ROOT = Path(__file__).resolve().parents[2] +CHART_DIR = REPO_ROOT / "nemo_retriever" / "helm" +VALUES_PATH = CHART_DIR / "values.yaml" +CONFIGMAP_PATH = CHART_DIR / "templates" / "configmap.yaml" +PARSE_SERVICE = "nemotron-parse" +PARSE_PATH = "/v1/chat/completions" + + +def _helm_template( + extra_args: Sequence[str] = (), + *, + api_versions: Sequence[str] = (), +) -> str: + helm = shutil.which("helm") + if helm is None: + raise SkipTest("`helm` binary not available") + command = [ + helm, + "template", + "parse-endpoint", + str(CHART_DIR), + "--set", + "ngcImagePullSecret.create=false", + "--set", + "ngcApiSecret.create=false", + ] + for version in api_versions: + command += ["--api-versions", version] + command += list(extra_args) + completed = subprocess.run(command, check=False, capture_output=True, text=True) + assert completed.returncode == 0, completed.stderr + return completed.stdout + + +def test_values_expose_parse_endpoint_and_model_overrides() -> None: + values = yaml.safe_load(VALUES_PATH.read_text()) + endpoints = values["serviceConfig"]["nimEndpoints"] + assert endpoints["nemotronParseInvokeUrl"] == "" + assert endpoints["nemotronParseModel"] == "" + + +def test_configmap_resolves_and_renders_parse_fields() -> None: + template = CONFIGMAP_PATH.read_text() + assert '"key" "nemotron_parse"' in template + assert f'"serviceName" "{PARSE_SERVICE}"' in template + assert '"configKey" "nemotronParseInvokeUrl"' in template + assert f'"invokePath" "{PARSE_PATH}"' in template + assert "nemotronParseModel" in template + assert "nemotron_parse_invoke_url:" in template + assert "nemotron_parse_model:" in template + + +def test_helm_template_parse_null_when_disabled() -> None: + rendered = _helm_template( + ("--set", "nimOperator.nemotron_parse.enabled=false"), + api_versions=("apps.nvidia.com/v1alpha1",), + ) + assert "nemotron_parse_invoke_url: null" in rendered + assert "nemotron_parse_model: null" in rendered + + +def test_helm_template_autowires_operator_parse_endpoint() -> None: + rendered = _helm_template( + ("--set", "nimOperator.nemotron_parse.enabled=true"), + api_versions=("apps.nvidia.com/v1alpha1",), + ) + assert f'nemotron_parse_invoke_url: "http://{PARSE_SERVICE}:8000{PARSE_PATH}"' in rendered + assert "nemotron_parse_model: null" in rendered + + +def test_helm_template_explicit_hosted_endpoint_and_model_win() -> None: + hosted_url = "https://integrate.api.nvidia.com/v1/chat/completions" + hosted_model = "nvidia/nemotron-parse" + rendered = _helm_template( + ( + "--set", + "nimOperator.nemotron_parse.enabled=true", + "--set", + f"serviceConfig.nimEndpoints.nemotronParseInvokeUrl={hosted_url}", + "--set", + f"serviceConfig.nimEndpoints.nemotronParseModel={hosted_model}", + ), + api_versions=("apps.nvidia.com/v1alpha1",), + ) + assert f'nemotron_parse_invoke_url: "{hosted_url}"' in rendered + assert f'nemotron_parse_model: "{hosted_model}"' in rendered + + +def test_helm_template_parse_endpoint_renders_in_split_mode() -> None: + rendered = _helm_template( + ( + "--set", + "nimOperator.nemotron_parse.enabled=true", + "--set", + "topology.mode=split", + ), + api_versions=("apps.nvidia.com/v1alpha1",), + ) + expected = f"http://{PARSE_SERVICE}:8000{PARSE_PATH}" + assert rendered.count(expected) >= 3 diff --git a/nemo_retriever/tests/test_helm_nim_2_contracts.py b/nemo_retriever/tests/test_helm_nim_2_contracts.py new file mode 100644 index 0000000000..56e9fdb999 --- /dev/null +++ b/nemo_retriever/tests/test_helm_nim_2_contracts.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Permanent Helm runtime and API contracts for extraction NIM 2.0.""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path +from unittest import SkipTest + +import yaml + +CHART = Path(__file__).resolve().parents[1] / "helm" +EXPECTED_MODELS = { + "nemotron-page-elements-v3": ( + "nvidia/nemotron-page-elements-v3", + "/model-store/page-elements", + ), + "nemotron-table-structure-v1": ( + "nvidia/nemotron-table-structure-v1", + "/model-store/table-structure", + ), + "nemotron-ocr-v2": ("nvidia/nemotron-ocr-v2", "/model-store/ocr"), +} + + +def _render() -> list[dict]: + helm = shutil.which("helm") + if helm is None: + raise SkipTest("`helm` binary not available") + command = [ + helm, + "template", + "nrl-nim-2-contracts", + str(CHART), + "--set", + "ngcImagePullSecret.create=false", + "--set", + "ngcApiSecret.create=false", + "--api-versions", + "apps.nvidia.com/v1alpha1", + ] + completed = subprocess.run(command, check=True, capture_output=True, text=True) + return [document for document in yaml.safe_load_all(completed.stdout) if isinstance(document, dict)] + + +def _find(documents: list[dict], kind: str, name: str) -> dict: + return next( + document + for document in documents + if document.get("kind") == kind and document.get("metadata", {}).get("name") == name + ) + + +def test_extraction_nims_select_distinct_native_models() -> None: + documents = _render() + + for name, (model_name, model_path) in EXPECTED_MODELS.items(): + service = _find(documents, "NIMService", name) + env = {item["name"]: item.get("value") for item in service["spec"]["env"]} + assert env["NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER"] == "ngc" + assert env["NIM_ENGINE_MODEL_NAME"] == model_name + assert env["NIM_ENGINE_MODEL_PATH"] == model_path + assert env["NIM_PERFORMANCE_MODE"] == "0" + assert env["NIM_PIPELINE_MAX_BATCH_SIZE"] == "1" + assert "NIM_TRITON_MAX_BATCH_SIZE" not in env + + ocr_env = { + item["name"]: item.get("value") for item in _find(documents, "NIMService", "nemotron-ocr-v2")["spec"]["env"] + } + assert ocr_env["NIM_ENGINE_MODEL_VARIANT"] == "multilingual" + + +def test_operator_managed_urls_use_nim_2_contracts() -> None: + documents = _render() + rendered_config = "\n".join( + document.get("data", {}).get("retriever-service.yaml", "") + for document in documents + if document.get("kind") == "ConfigMap" + ) + + assert 'page_elements_invoke_url: "http://nemotron-page-elements-v3:8000/v1/page-elements"' in rendered_config + assert 'table_structure_invoke_url: "http://nemotron-table-structure-v1:8000/v1/table-structure"' in rendered_config + assert 'ocr_invoke_url: "http://nemotron-ocr-v2:8000/v1/ocr"' in rendered_config diff --git a/nemo_retriever/tests/test_helm_nimcache_model_profile.py b/nemo_retriever/tests/test_helm_nimcache_model_profile.py index 5c6e77530c..70f7c69a08 100644 --- a/nemo_retriever/tests/test_helm_nimcache_model_profile.py +++ b/nemo_retriever/tests/test_helm_nimcache_model_profile.py @@ -265,7 +265,7 @@ def test_default_render_uses_ocr_v2_nim(self) -> None: ) self.assertEqual( ocr_cache["spec"]["source"]["ngc"]["modelPuller"], - "nvcr.io/nim/nvidia/nemotron-ocr-v2:1.4.0", + "nvcr.io/nim/nvidia/nemotron-ocr-v2:2.0.1", ) ocr_service = next( @@ -277,12 +277,12 @@ def test_default_render_uses_ocr_v2_nim(self) -> None: ocr_service["spec"]["image"]["repository"], "nvcr.io/nim/nvidia/nemotron-ocr-v2", ) - self.assertEqual(ocr_service["spec"]["image"]["tag"], "1.4.0") + self.assertEqual(ocr_service["spec"]["image"]["tag"], "2.0.1") configmaps = [doc for doc in docs if doc.get("kind") == "ConfigMap"] self.assertTrue( any( - 'ocr_invoke_url: "http://nemotron-ocr-v2:8000/v1/infer"' + 'ocr_invoke_url: "http://nemotron-ocr-v2:8000/v1/ocr"' in doc.get("data", {}).get("retriever-service.yaml", "") for doc in configmaps ), diff --git a/nemo_retriever/tests/test_helm_nimservice_resources.py b/nemo_retriever/tests/test_helm_nimservice_resources.py index f6ce7ac958..3f9fd7d01d 100644 --- a/nemo_retriever/tests/test_helm_nimservice_resources.py +++ b/nemo_retriever/tests/test_helm_nimservice_resources.py @@ -5,7 +5,7 @@ """Regression tests for NIMService GPU resource rendering. The NIM Operator does not reliably populate ``spec.resources.limits.nvidia.com/gpu`` -from the model profile on all tested versions (for example v3.1.1 on A100/H100). +from the model profile on all tested versions (for example v3.1.2 on A100/H100). The chart therefore defaults to rendering ``nvidia.com/gpu: 1`` via ``nimOperator.nimServiceGpuLimit``. @@ -27,10 +27,10 @@ ("llama-nemotron-embed-vl-1b-v2.yaml", "vlm_embed"), ("llama-nemotron-rerank-vl-1b-v2.yaml", "rerankqa"), ("nemotron-3-nano-omni-30b-a3b-reasoning.yaml", "nemotron_3_nano_omni_30b_a3b_reasoning"), - ("nemotron-ocr-v2.yaml", "ocr"), ("nemotron-page-elements-v3.yaml", "page_elements"), - ("nemotron-parse.yaml", "nemotron_parse"), ("nemotron-table-structure-v1.yaml", "table_structure"), + ("nemotron-ocr-v2.yaml", "ocr"), + ("nemotron-parse.yaml", "nemotron_parse"), ) diff --git a/nemo_retriever/tests/test_helm_optional_nims_disabled_by_default.py b/nemo_retriever/tests/test_helm_optional_nims_disabled_by_default.py index 6ca457690c..e010542fe0 100644 --- a/nemo_retriever/tests/test_helm_optional_nims_disabled_by_default.py +++ b/nemo_retriever/tests/test_helm_optional_nims_disabled_by_default.py @@ -176,7 +176,7 @@ def test_values_omni_enabled_defaults_to_false(self) -> None: Omni 30B is the heaviest NIM in the chart (~62 GiB BF16 weights, ~80 GB on-disk NIM cache, requires its own ≥ 80 GiB GPU). It must not deploy on a "default" install — that contradicts the docs and - the README's [Recommended minimal install (26.05)] guidance. + the README's [Recommended minimal install (26.08)] guidance. """ values = _read_required_file(_VALUES_YAML) value = _enabled_value_for_block(values, _OMNI_BLOCK) @@ -303,15 +303,15 @@ def test_readme_image_table_pins_vl_rerank_sku(self) -> None: """ readme = _read_required_file(_README_MD) self.assertIn( - f"{_RERANK_VL_REPOSITORY}:1.10.0", + f"{_RERANK_VL_REPOSITORY}:2.3.0", readme, - "README mirror-image table must list the VL reranker " f"`{_RERANK_VL_REPOSITORY}:1.10.0`.", + "README mirror-image table must list the VL reranker " f"`{_RERANK_VL_REPOSITORY}:2.3.0`.", ) self.assertNotIn( - f"{_RERANK_TEXT_REPOSITORY}:1.10.0", + f"{_RERANK_TEXT_REPOSITORY}:2.3.0", readme, "README mirror-image table must not list the text-only " - f"rerank SKU `{_RERANK_TEXT_REPOSITORY}:1.10.0` — that " + f"rerank SKU `{_RERANK_TEXT_REPOSITORY}:2.3.0` — that " "would silently degrade multimodal reranking for air-gapped " "mirror setups.", ) @@ -357,12 +357,12 @@ def test_readme_minimal_install_no_longer_disables_parse_or_omni(self) -> None: # Find the heredoc-style minimal install command. The recipe # ends with `audio.enabled=false`; the block above that is what # we inspect. - marker = "Recommended minimal install (26.05)" + marker = "Recommended minimal install (26.08)" idx = readme.find(marker) self.assertNotEqual( idx, -1, - "README must keep a `Recommended minimal install (26.05)` " + "README must keep a `Recommended minimal install (26.08)` " "section even after the defaults flip — it documents the " "two flags that are still needed (`rerankqa` + `audio`).", ) diff --git a/nemo_retriever/tests/test_helm_tracing_zipkin.py b/nemo_retriever/tests/test_helm_tracing_zipkin.py index b0325ee063..7b35077575 100644 --- a/nemo_retriever/tests/test_helm_tracing_zipkin.py +++ b/nemo_retriever/tests/test_helm_tracing_zipkin.py @@ -28,10 +28,10 @@ "llama-nemotron-embed-vl-1b-v2", "llama-nemotron-rerank-vl-1b-v2", "nemotron-3-nano-omni-30b-a3b-reasoning", - "nemotron-ocr-v2", "nemotron-page-elements-v3", - "nemotron-parse", "nemotron-table-structure-v1", + "nemotron-ocr-v2", + "nemotron-parse", } @@ -234,9 +234,9 @@ def test_null_otel_env_maps_render_as_empty_maps() -> None: service_env = _env_values(_deployment_env(_find(docs, "Deployment", FULLNAME))) assert service_env["OTEL_EXPORTER_OTLP_ENDPOINT"] == f"http://{OTEL_NAME}:4317" - page_env = _env_values(_nim_env(_find(docs, "NIMService", "nemotron-page-elements-v3"))) - assert page_env["NIM_ENABLE_OTEL"] == "true" - assert page_env["NIM_OTEL_EXPORTER_OTLP_ENDPOINT"] == f"http://{OTEL_NAME}:4318" + page_elements_env = _env_values(_nim_env(_find(docs, "NIMService", "nemotron-page-elements-v3"))) + assert page_elements_env["NIM_ENABLE_OTEL"] == "true" + assert page_elements_env["NIM_OTEL_EXPORTER_OTLP_ENDPOINT"] == f"http://{OTEL_NAME}:4318" rerank_env = _env_values(_nim_env(_find(docs, "NIMService", "llama-nemotron-rerank-vl-1b-v2"))) assert rerank_env["NIM_ENABLE_OTEL"] == "true" @@ -709,8 +709,8 @@ def test_chart_wide_nim_otel_disable_omits_managed_env() -> None: _assert_unique_env_names(env) assert chart_managed_names.isdisjoint(values) - table_values = _env_values(_nim_env(_find(docs, "NIMService", "nemotron-table-structure-v1"))) - assert table_values["NIM_TRITON_CUDA_MEMORY_POOL_MB"] == "2048" + page_elements_values = _env_values(_nim_env(_find(docs, "NIMService", "nemotron-page-elements-v3"))) + assert page_elements_values["NIM_PIPELINE_MAX_BATCH_SIZE"] == "1" def test_per_nim_otel_endpoint_overrides_chart_endpoint() -> None: @@ -723,9 +723,9 @@ def test_per_nim_otel_endpoint_overrides_chart_endpoint() -> None: ) page_elements = _find(docs, "NIMService", "nemotron-page-elements-v3") - page_values = _env_values(_nim_env(page_elements)) - assert page_values["NIM_OTEL_EXPORTER_OTLP_ENDPOINT"] == "http://chart-otel:4318" - assert page_values["TRITON_OTEL_URL"] == "http://chart-otel:4318/v1/traces" + page_elements_values = _env_values(_nim_env(page_elements)) + assert page_elements_values["NIM_OTEL_EXPORTER_OTLP_ENDPOINT"] == "http://chart-otel:4318" + assert page_elements_values["TRITON_OTEL_URL"] == "http://chart-otel:4318/v1/traces" rerank = _find(docs, "NIMService", "llama-nemotron-rerank-vl-1b-v2") rerank_values = _env_values(_nim_env(rerank)) @@ -739,10 +739,10 @@ def test_chart_nim_otel_env_endpoint_drives_triton_url() -> None: ) page_elements = _find(docs, "NIMService", "nemotron-page-elements-v3") - page_values = _env_values(_nim_env(page_elements)) + page_elements_values = _env_values(_nim_env(page_elements)) - assert page_values["NIM_OTEL_EXPORTER_OTLP_ENDPOINT"] == "http://env-otel:4318" - assert page_values["TRITON_OTEL_URL"] == "http://env-otel:4318/v1/traces" + assert page_elements_values["NIM_OTEL_EXPORTER_OTLP_ENDPOINT"] == "http://env-otel:4318" + assert page_elements_values["TRITON_OTEL_URL"] == "http://env-otel:4318/v1/traces" def test_per_nim_otel_env_endpoint_drives_triton_url() -> None: @@ -767,10 +767,10 @@ def test_nim_otel_env_triton_url_override_is_preserved() -> None: ) page_elements = _find(docs, "NIMService", "nemotron-page-elements-v3") - page_values = _env_values(_nim_env(page_elements)) + page_elements_values = _env_values(_nim_env(page_elements)) - assert page_values["NIM_OTEL_EXPORTER_OTLP_ENDPOINT"] == "http://env-otel:4318" - assert page_values["TRITON_OTEL_URL"] == "http://explicit-triton/v1/traces" + assert page_elements_values["NIM_OTEL_EXPORTER_OTLP_ENDPOINT"] == "http://env-otel:4318" + assert page_elements_values["TRITON_OTEL_URL"] == "http://explicit-triton/v1/traces" def test_existing_nim_env_endpoint_drives_triton_url_without_duplicate_endpoint() -> None: diff --git a/nemo_retriever/tests/test_ingest_interface.py b/nemo_retriever/tests/test_ingest_interface.py index 2d098ef6c4..c99b6d6a77 100644 --- a/nemo_retriever/tests/test_ingest_interface.py +++ b/nemo_retriever/tests/test_ingest_interface.py @@ -24,6 +24,14 @@ ) +class _InlineTextTokenizer: + def encode(self, text: str, add_special_tokens: bool = False) -> list[str]: + return text.split() + + def decode(self, ids: list[str], skip_special_tokens: bool = True) -> str: + return " ".join(ids) + + def _graph_node_names(graph) -> list[str]: names: list[str] = [] @@ -112,6 +120,175 @@ def test_create_ingestor_rejects_unknown_run_modes() -> None: create_ingestor(run_mode="parallel") # type: ignore[arg-type] +def test_texts_accepts_scalar(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda *args, **kwargs: _InlineTextTokenizer() + ) + + result = create_ingestor(run_mode="inprocess").texts("first").ingest() + + assert result["text"].tolist() == ["first"] + assert result["path"].tolist() == ["inline://00000000"] + + +def test_texts_replaces_prior_inline_corpus(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda *args, **kwargs: _InlineTextTokenizer() + ) + + result = create_ingestor(run_mode="inprocess").texts("first").texts(["second", "third"]).ingest() + + assert result["text"].tolist() == ["second", "third"] + assert result["path"].tolist() == ["inline://00000000", "inline://00000001"] + + +@pytest.mark.parametrize("values", [["valid", None], ["valid", 3], [object()]]) +def test_texts_rejects_non_string_values_with_index(values) -> None: + bad_index = next(index for index, value in enumerate(values) if not isinstance(value, str)) + + with pytest.raises(TypeError, match=rf"texts\[{bad_index}\] must be a string"): + create_ingestor(run_mode="inprocess").texts(values) + + +def test_texts_rejects_non_sequence_input() -> None: + with pytest.raises(TypeError, match="string or sequence of strings"): + create_ingestor(run_mode="inprocess").texts(iter(["one", "two"])) + + +@pytest.mark.parametrize("files_first", [True, False]) +def test_texts_can_mix_with_text_files( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + files_first: bool, +) -> None: + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda *args, **kwargs: _InlineTextTokenizer() + ) + document = tmp_path / "document.txt" + document.write_text("document", encoding="utf-8") + + ingestor = create_ingestor(run_mode="inprocess") + if files_first: + ingestor.files([str(document)]).texts(["inline"]) + else: + ingestor.texts(["inline"]).files([str(document)]) + + result = ingestor.extract(split_config={"text": {"max_tokens": 10}}).ingest() + + assert result["text"].tolist() == ["document", "inline"] + assert result["path"].tolist() == [str(document.resolve()), "inline://00000000"] + + +def test_texts_can_mix_with_text_buffers(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda *args, **kwargs: _InlineTextTokenizer() + ) + + result = ( + create_ingestor(run_mode="inprocess") + .texts(["inline"]) + .buffers(("document.txt", BytesIO(b"document"))) + .extract(split_config={"text": {"max_tokens": 10}}) + .ingest() + ) + + assert result["text"].tolist() == ["document", "inline"] + assert result["path"].tolist() == [ + str(Path("document.txt").resolve()), + "inline://00000000", + ] + + +def test_texts_allow_explicit_text_extraction(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda *args, **kwargs: _InlineTextTokenizer() + ) + + result = ( + create_ingestor(run_mode="inprocess") + .texts(["one two"]) + .extract(extraction_mode="text", text_params=TextChunkParams(max_tokens=1)) + .ingest() + ) + + assert result["text"].tolist() == ["one", "two"] + + +def test_texts_plan_alongside_other_modalities_regardless_of_call_order(tmp_path: Path) -> None: + image = tmp_path / "scan.bmp" + image.write_bytes(b"bmp") + + images_first = GraphIngestor(run_mode="inprocess").files([str(image)]).extract_image_files().texts(["inline"]) + texts_first = GraphIngestor(run_mode="inprocess").texts(["inline"]).files([str(image)]).extract_image_files() + + for ingestor in (images_first, texts_first): + branches = ingestor._plan_default_extraction_branches() + assert [(branch.family, branch.input_paths) for branch in branches] == [ + ("image", (str(image),)), + ("txt", ("inline://00000000",)), + ] + + empty_inline = GraphIngestor(run_mode="inprocess").files([str(image)]).texts([]) + assert [(branch.family, branch.input_paths) for branch in empty_inline._plan_default_extraction_branches()] == [ + ("image", (str(image),)), + ] + + explicit_image = GraphIngestor(run_mode="inprocess").files([str(image)]).texts([]).extract_image_files() + assert explicit_image._plan_default_extraction_branches() is None + + +def test_empty_and_blank_inline_corpus_short_circuits_graph(monkeypatch: pytest.MonkeyPatch) -> None: + ingestor = create_ingestor(run_mode="inprocess").texts(["", " \n"]) + monkeypatch.setattr( + "nemo_retriever.ingestor.graph_ingestor.build_graph", + lambda *args, **kwargs: pytest.fail("empty inline corpus should not execute the graph"), + ) + + result = ingestor.embed().vdb_upload().ingest() + + assert result.empty + assert list(result.columns) == ["text", "content", "path", "page_number", "metadata"] + + +@pytest.mark.parametrize("inline_texts", [[], ["", " \n"]]) +def test_empty_inline_text_does_not_short_circuit_file_ingestion( + inline_texts: list[str], + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", + lambda *args, **kwargs: _InlineTextTokenizer(), + ) + document = tmp_path / "document.txt" + document.write_text("document", encoding="utf-8") + + result = create_ingestor(run_mode="inprocess").files([str(document)]).texts(inline_texts).ingest() + + assert result["text"].tolist() == ["document"] + assert result["path"].tolist() == [str(document.resolve())] + + +def test_empty_batch_inline_corpus_returns_dataframe_without_starting_ray(monkeypatch: pytest.MonkeyPatch) -> None: + ingestor = create_ingestor(run_mode="batch").texts(["", " \n"]) + monkeypatch.setattr( + ingestor, + "_ensure_batch_runtime", + lambda: pytest.fail("empty inline corpus should not start Ray"), + ) + monkeypatch.setattr( + "nemo_retriever.ingestor.graph_ingestor.build_graph", + lambda *args, **kwargs: pytest.fail("empty inline corpus should not execute the graph"), + ) + + result = ingestor.ingest() + + assert isinstance(result, pd.DataFrame) + assert ingestor.get_dataset() is result + assert result.empty + assert list(result.columns) == ["text", "content", "path", "page_number", "metadata"] + + def test_graph_ingestor_action_methods_materialize_default_params() -> None: ingestor = GraphIngestor(run_mode="inprocess") @@ -121,8 +298,8 @@ def test_graph_ingestor_action_methods_materialize_default_params() -> None: ingestor.extract_image_files() assert isinstance(ingestor._extract_params, ExtractParams) - ingestor.extract_txt() - assert isinstance(ingestor._text_params, TextChunkParams) + ingestor.extract(split_config={"text": {"max_tokens": 512}}) + assert isinstance(ingestor._split_config["text"], TextChunkParams) ingestor.extract_html() assert isinstance(ingestor._html_params, HtmlChunkParams) @@ -271,56 +448,24 @@ def test_extract_default_rejects_unknown_input_type(tmp_path) -> None: ingestor.ingest() -def test_extract_default_treats_markdown_as_plain_text(tmp_path) -> None: - document = tmp_path / "README.md" - document.write_text("# Heading\n\nBody text\n", encoding="utf-8") - - result = GraphIngestor(run_mode="inprocess", show_progress=False).files([str(document)]).extract().ingest() - - assert result["text"].tolist() == ["# Heading\n\nBody text\n"] - assert result["path"].tolist() == [str(document.resolve())] - - -def test_extract_txt_accepts_json_as_plain_text(tmp_path) -> None: - document = tmp_path / "payload.json" - document.write_text('{"message": "hello"}\n', encoding="utf-8") - - result = GraphIngestor(run_mode="inprocess", show_progress=False).files([str(document)]).extract_txt().ingest() - - assert result["text"].tolist() == ['{"message": "hello"}\n'] - assert result["path"].tolist() == [str(document.resolve())] +def test_extract_audio_does_not_enable_post_extraction_chunking_by_default() -> None: + audio_ingestor = GraphIngestor(run_mode="inprocess").extract_audio() + assert audio_ingestor._split_config["audio"] is None -def test_extract_default_accepts_shell_script_buffer_as_plain_text() -> None: - content = b"#!/bin/sh\necho hello\n" +def test_extract_split_configures_automatically_routed_text(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda *args, **kwargs: _InlineTextTokenizer() + ) result = ( GraphIngestor(run_mode="inprocess", show_progress=False) - .buffers(("setup.sh", BytesIO(content))) - .extract() + .texts(["one two three"]) + .extract(split_config={"text": {"max_tokens": 2}}) .ingest() ) - assert result["text"].tolist() == [content.decode()] - assert result["path"].tolist() == [str(Path("setup.sh").resolve())] - - -def test_typed_shortcuts_preserve_legacy_no_default_chunking() -> None: - """Typed shortcuts (extract_audio, extract_txt, ...) must NOT enable default - split_config chunking. Default-ON is reserved for the unified .extract() - path. extract_txt(custom_params) must propagate custom_params via the - text_params fallback. - """ - # extract_audio without split_config: no audio chunking. - audio_ingestor = GraphIngestor(run_mode="inprocess").extract_audio() - assert audio_ingestor._split_config["audio"] is None - - # extract_txt(custom): _split_config["text"] stays None so the operator - # falls back to self.text_params (= custom) in _effective_chunk_params. - custom = TextChunkParams(max_tokens=512) - txt_ingestor = GraphIngestor(run_mode="inprocess").extract_txt(custom) - assert txt_ingestor._split_config["text"] is None - assert txt_ingestor._text_params is custom + assert result["text"].tolist() == ["one two", "three"] def test_graph_ingestor_return_failures_returns_service_tuples_from_path(monkeypatch) -> None: diff --git a/nemo_retriever/tests/test_ingest_service.py b/nemo_retriever/tests/test_ingest_service.py index 54272f0632..4ee5bd6c84 100644 --- a/nemo_retriever/tests/test_ingest_service.py +++ b/nemo_retriever/tests/test_ingest_service.py @@ -12,7 +12,7 @@ from nemo_retriever.common.schemas.pipeline_spec import PipelineSpec from nemo_retriever.ingest.service import ServiceIngestRequest, build_service_ingestor, execute_service_ingest_request from nemo_retriever.service.config import PipelineOverridesConfig -from nemo_retriever.service.service_ingestor import ServiceIngestor +from nemo_retriever.service.service_ingestor import ServiceIngestor, ServiceIngestResult def test_build_service_ingestor_wires_extract_embed_and_chunking(tmp_path: Path) -> None: @@ -93,8 +93,63 @@ def test_execute_service_ingest_request_raises_for_document_failures(monkeypatch failed_result = SimpleNamespace(failures=[("doc.pdf", "HTTP 400: invalid request")]) monkeypatch.setattr( "nemo_retriever.ingest.service.build_service_ingestor", - lambda _request: SimpleNamespace(ingest=lambda: failed_result), + lambda _request: SimpleNamespace(ingest=lambda **_kwargs: failed_result), ) with pytest.raises(RuntimeError, match=r"failed for 1 document\(s\).+doc.pdf.+HTTP 400"): execute_service_ingest_request(request) + + +def test_execute_service_ingest_request_materializes_results_by_default( + monkeypatch, + tmp_path: Path, +) -> None: + request = ServiceIngestRequest(documents=[str(tmp_path / "doc.pdf")], input_type="pdf") + result = SimpleNamespace(dataframe=[{"row": 1}, {"row": 2}], failures=[]) + captured: dict[str, object] = {} + + def ingest(**kwargs): + captured.update(kwargs) + return result + + monkeypatch.setattr( + "nemo_retriever.ingest.service.build_service_ingestor", + lambda _request: SimpleNamespace(ingest=ingest), + ) + + execution = execute_service_ingest_request(request) + + assert captured == {"return_results": True} + assert execution.n_rows == 2 + assert execution.result_n_rows == 2 + + +def test_execute_service_ingest_request_can_disable_materialization_and_count_event_rows( + monkeypatch, + tmp_path: Path, +) -> None: + request = ServiceIngestRequest(documents=[str(tmp_path / "doc.pdf")], input_type="pdf") + result = ServiceIngestResult( + [ + {"event": "document_complete", "document_id": "a", "status": "completed", "result_rows": 2}, + {"event": "document_complete", "document_id": "b", "status": "failed", "result_rows": 99}, + {"event": "document_complete", "document_id": "c", "status": "completed", "result_rows": 3}, + ] + ) + captured: dict[str, object] = {} + + def ingest(**kwargs): + captured.update(kwargs) + return result + + monkeypatch.setattr( + "nemo_retriever.ingest.service.build_service_ingestor", + lambda _request: SimpleNamespace(ingest=ingest), + ) + + execution = execute_service_ingest_request(request, return_results=False) + + assert captured == {"return_results": False} + assert execution.n_rows == 5 + assert execution.result_n_rows == 5 + assert execution.to_summary_dict()["result_n_rows"] == 5 diff --git a/nemo_retriever/tests/test_inline_text_ingest.py b/nemo_retriever/tests/test_inline_text_ingest.py new file mode 100644 index 0000000000..a0618c757c --- /dev/null +++ b/nemo_retriever/tests/test_inline_text_ingest.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +import pandas as pd +import pytest + +from nemo_retriever.common.params import TextChunkParams +from nemo_retriever.graph import Graph +from nemo_retriever.ingestor.graph_ingestor import GraphIngestor +from nemo_retriever.operators.abstract_operator import AbstractOperator +from nemo_retriever.operators.extract.txt.ray_data import TxtSplitActor + + +class _MockTokenizer: + def encode(self, text: str, add_special_tokens: bool = False) -> list[str]: + return text.split() + + def decode(self, ids: list[str], skip_special_tokens: bool = True) -> str: + return " ".join(ids) + + +class _FakeEmbedOperator(AbstractOperator): + def preprocess(self, data: Any, **kwargs: Any) -> Any: + return data + + def process(self, data: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: + result = data.copy() + result["fake_embedding"] = [[1.0]] * len(result) + return result + + def postprocess(self, data: Any, **kwargs: Any) -> Any: + return data + + +class _FakeVdbOperator(AbstractOperator): + def preprocess(self, data: Any, **kwargs: Any) -> Any: + return data + + def process(self, data: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: + result = data.copy() + result["stored"] = True + return result + + def postprocess(self, data: Any, **kwargs: Any) -> Any: + return data + + +def test_inline_text_runs_split_embed_and_vdb_graph(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda *args, **kwargs: _MockTokenizer() + ) + captured: dict[str, Any] = {} + + def fake_build_graph(**kwargs: Any) -> Graph: + captured.update(kwargs) + return ( + Graph() + >> TxtSplitActor(params=kwargs["split_config"]["text"]) + >> _FakeEmbedOperator() + >> _FakeVdbOperator() + ) + + monkeypatch.setattr("nemo_retriever.ingestor.graph_ingestor.build_graph", fake_build_graph) + + result = ( + GraphIngestor(run_mode="inprocess", show_progress=False) + .texts(["one two three", "one two three"]) + .extract(split_config={"text": {"max_tokens": 2}}) + .embed() + .vdb_upload() + .ingest() + ) + + assert result["text"].tolist() == ["one two", "three", "one two", "three"] + assert result["path"].tolist() == [ + "inline://00000000", + "inline://00000000", + "inline://00000001", + "inline://00000001", + ] + assert result["stored"].tolist() == [True, True, True, True] + assert captured["embed_params"] is not None + assert captured["vdb_upload_params"] is not None + + +@pytest.mark.integration +def test_batch_inline_text_matches_text_file(tmp_path) -> None: + ray = pytest.importorskip("ray") + pytest.importorskip("transformers") + text = "one two three four five" + document = tmp_path / "document.txt" + document.write_text(text, encoding="utf-8") + params = TextChunkParams(max_tokens=2) + + try: + file_result = ( + GraphIngestor(run_mode="batch", show_progress=False) + .files([str(document)]) + .extract(split_config={"text": params}) + .ingest() + ) + inline_result = ( + GraphIngestor(run_mode="batch", show_progress=False) + .texts([text]) + .extract(split_config={"text": params}) + .ingest() + ) + finally: + ray.shutdown() + + assert file_result["text"].tolist() == inline_result["text"].tolist() + assert file_result["page_number"].tolist() == inline_result["page_number"].tolist() + + +@pytest.mark.integration +def test_batch_inline_text_ingests_alongside_text_file(tmp_path) -> None: + ray = pytest.importorskip("ray") + pytest.importorskip("transformers") + document = tmp_path / "document.txt" + document.write_text("from file", encoding="utf-8") + + try: + result = ( + GraphIngestor(run_mode="batch", show_progress=False) + .files([str(document)]) + .texts(["from inline"]) + .extract(split_config={"text": {"max_tokens": 10}}) + .ingest() + ) + finally: + ray.shutdown() + + assert set(zip(result["path"], result["text"])) == { + (str(document.resolve()), "from file"), + ("inline://00000000", "from inline"), + } diff --git a/nemo_retriever/tests/test_internal_auth.py b/nemo_retriever/tests/test_internal_auth.py new file mode 100644 index 0000000000..5cc5539061 --- /dev/null +++ b/nemo_retriever/tests/test_internal_auth.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dedicated internal-token normalization regressions.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import pytest + +from nemo_retriever.service.auth import internal_auth_headers +from nemo_retriever.service.config import load_config +from nemo_retriever.service.services.pipeline_executor import _post_records_to_vectordb +from nemo_retriever.service.services.pipeline_pool import DocumentWriteContext + + +def test_internal_auth_headers_strip_secret_whitespace() -> None: + assert internal_auth_headers(" internal-secret\n") == {"X-NRL-Internal-Token": "internal-secret"} + assert internal_auth_headers(" \n\t") == {} + + +def test_load_config_strips_internal_token_environment( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = tmp_path / "service.yaml" + config_path.write_text("vectordb:\n enabled: true\n", encoding="utf-8") + monkeypatch.setenv("NRL_INTERNAL_VDB_TOKEN", " internal-secret\n") + + assert load_config(str(config_path)).vectordb.internal_api_token == "internal-secret" + + +def test_pipeline_vectordb_request_uses_normalized_internal_header() -> None: + captured: dict[str, Any] = {} + + class Response: + status = 200 + + def __enter__(self) -> "Response": + return self + + def __exit__(self, *_args: object) -> None: + return None + + def urlopen(request: Any, timeout: int) -> Response: + captured["request"] = request + captured["timeout"] = timeout + return Response() + + with patch("urllib.request.urlopen", urlopen): + _post_records_to_vectordb( + [[{"text": "row"}]], + "http://vectordb:7671", + "document.txt", + internal_api_token=" internal-secret\n", + context=DocumentWriteContext(), + ) + + assert captured["timeout"] == 30 + assert captured["request"].get_header("X-nrl-internal-token") == "internal-secret" diff --git a/nemo_retriever/tests/test_lancedb_capabilities.py b/nemo_retriever/tests/test_lancedb_capabilities.py index 0507130930..778b582720 100644 --- a/nemo_retriever/tests/test_lancedb_capabilities.py +++ b/nemo_retriever/tests/test_lancedb_capabilities.py @@ -14,10 +14,17 @@ import nemo_retriever.graph.retriever as retriever_module # noqa: E402 from nemo_retriever.common.vdb.lancedb_capabilities import LanceTableCapabilities, inspect_lancedb_table # noqa: E402 +from nemo_retriever.common.vdb.lancedb import LanceDB # noqa: E402 from nemo_retriever.graph.retriever import Retriever # noqa: E402 +from nemo_retriever.operators.vdb import RetrieveVdbOperator # noqa: E402 -def _create_vector_table(uri: str, table_name: str, *, fts: bool = False) -> None: +def _create_vector_table( + uri: str, + table_name: str, + *, + fts: bool = False, +) -> None: schema = pa.schema( [ pa.field("vector", pa.list_(pa.float32(), 2)), @@ -94,6 +101,53 @@ def test_detector_returns_dense_for_vector_only_table(tmp_path) -> None: assert caps.retrieval_mode == "dense" +def test_lancedb_metadata_round_trip_drives_query_model(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + uri = str(tmp_path / "db") + records = [ + [ + { + "document_type": "text", + "metadata": { + "content": "alpha safety manual", + "embedding": [1.0, 0.0], + "content_metadata": {"id": "alpha"}, + "source_metadata": {"source_id": "alpha.pdf"}, + }, + } + ] + ] + LanceDB( + uri=uri, + table_name="docs", + vector_dim=2, + build_index=False, + embedding_model_name="nvidia/llama-nemotron-embed-vl-1b-v2", + ).run(records) + + operator = RetrieveVdbOperator(vdb_op="lancedb", vdb_kwargs={"uri": uri, "table_name": "docs"}) + + assert operator.get_index_metadata("embedding_model_name") == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert operator.get_index_metadata("retrieval_mode") == "dense" + + captured_embed_kwargs: dict[str, Any] = {} + + def capture_query_model( + _self: Retriever, + _query_texts: list[str], + *, + embed_extra: dict[str, Any] | None, + **_kwargs: Any, + ) -> list[list[dict[str, Any]]]: + captured_embed_kwargs.update(embed_extra or {}) + return [[{"text": "alpha safety manual", "source": "alpha.pdf"}]] + + monkeypatch.setattr(Retriever, "_execute_queries_graph", capture_query_model) + + Retriever(vdb_kwargs={"uri": uri, "table_name": "docs"}).query("alpha", top_k=1) + + assert captured_embed_kwargs["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + + def test_detector_returns_hybrid_for_vector_plus_fts_table(tmp_path) -> None: uri = str(tmp_path / "db") _create_vector_table(uri, "hybrid", fts=True) diff --git a/nemo_retriever/tests/test_lancedb_collections.py b/nemo_retriever/tests/test_lancedb_collections.py new file mode 100644 index 0000000000..3606cbaf32 --- /dev/null +++ b/nemo_retriever/tests/test_lancedb_collections.py @@ -0,0 +1,1129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Focused tests for LanceDB's optional collection capabilities.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor + +import hashlib +import json +import math +import threading +from dataclasses import replace + +import lancedb +import pytest + +import nemo_retriever.common.vdb.lancedb_collections as collections_module +from nemo_retriever.common.schemas.collections import ( + CollectionCreateRequest, + CollectionUpdateRequest, +) +from nemo_retriever.common.vdb.adt_vdb import ( + CollectionWriteContext, + UnsupportedVDBOperation, + VDBInvalidRequest, + VDBResourceConflict, + VDBResourceNotFound, +) +from nemo_retriever.common.vdb.lancedb_capabilities import LanceTableCapabilities +from nemo_retriever.common.vdb.lancedb import ( + LanceDB, + _create_lancedb_results, + _to_service_lancedb_rows, +) +from nemo_retriever.common.vdb.lancedb_collections import ( + _collection_rows, + _encode_cursor, + _normalize_collection_results, + _public_collection_hit, +) +from nemo_retriever.common.vdb.records import RetrievalContractError + + +def _context( + *, + version: str = "v1", + operation: str = "append", + document_id: str = "document-a", +) -> CollectionWriteContext: + return CollectionWriteContext( + scope="workspace-a", + collection_name="collection-a", + document_id=document_id, + document_version=version, + content_sha256=f"sha-{version}", + filename="source.pdf", + job_id="job-a", + operation=operation, + ) + + +def _records( + *, + text: str = "first chunk", + vector: list[float] | None = None, + page_number: int = 2, +) -> list[list[dict]]: + return [ + [ + { + "document_type": "text", + "metadata": { + "embedding": vector or [1.0, 0.0], + "content": text, + "content_metadata": { + "type": "text", + "page_number": page_number, + "stored_image_uri": "file:///artifacts/page-2.png", + "bbox_xyxy_norm": [0.1, 0.2, 0.8, 0.9], + }, + "source_metadata": { + "source_id": "/inputs/source.pdf", + "source_name": "source.pdf", + }, + }, + } + ] + ] + + +def _service_records( + content_type: str, + *, + text: str, + vector: list[float], +) -> list[list[dict]]: + records = _records(text=text, vector=vector, page_number=4) + metadata = records[0][0]["metadata"] + metadata["content_metadata"].update( + { + "type": content_type, + "stored_image_uri": "s3://bucket/figure.png", + "bbox_xyxy_norm": [0.1, 0.2, 0.8, 0.9], + } + ) + metadata["source_metadata"] = { + "source_id": "/inputs/report.pdf", + "source_name": "report.pdf", + "custom_source_field": "preserved", + } + return records + + +def _backend_with_collection(tmp_path) -> LanceDB: + backend = LanceDB( + uri=str(tmp_path / "lancedb"), + table_name="legacy", + vector_dim=2, + build_index=False, + ) + backend.create_collection( + scope="workspace-a", + request=CollectionCreateRequest(name="collection-a"), + ) + return backend + + +def _fail_document_finalize(monkeypatch, store): + original_persist = store._persist_document_row + + def fail_completed(row): + if row.get("status") == "completed": + raise RuntimeError("injected catalog finalize failure") + return original_persist(row) + + monkeypatch.setattr(store, "_persist_document_row", fail_completed) + return original_persist + + +def test_collection_row_conversion_preserves_identity_and_provenance(): + records = _records() + records[0][0]["metadata"]["content_metadata"]["type"] = "table_caption" + rows = _collection_rows(records, context=_context()) + + assert len(rows) == 1 + row = rows[0] + assert row["text"] == "first chunk" + assert row["filename"] == "source.pdf" + assert row["page_number"] == 2 + assert row["pdf_page"] == "source_2" + assert row["source_id"] == "/inputs/source.pdf" + assert json.loads(row["source"])["source_name"] == "source.pdf" + assert row["content_type"] == "table" + assert json.loads(row["metadata"])["type"] == "table" + assert json.loads(row["bbox_xyxy_norm"]) == [0.1, 0.2, 0.8, 0.9] + assert row["stored_image_uri"] == "file:///artifacts/page-2.png" + assert row["document_id"] == "document-a" + assert row["document_version"] == "v1" + assert row["content_sha256"] == "sha-v1" + assert row["chunk_id"] == hashlib.sha256(b"document-a\x00v1\x000").hexdigest() + + +@pytest.mark.parametrize( + ("raw_type", "expected_type"), + [("table_caption", "table"), ("chart_caption", "chart")], +) +def test_service_row_adapter_preserves_multimodal_provenance(raw_type, expected_type): + narrow_rows, counts = _create_lancedb_results( + _service_records(raw_type, text="caption", vector=[1.0, 0.0, 0.0]), + expected_dim=None, + ) + + rows = _to_service_lancedb_rows(narrow_rows) + + assert counts["accepted"] == 1 + assert len(rows) == 1 + row = rows[0] + assert row["content_type"] == expected_type + assert row["filename"] == "report.pdf" + assert row["page_number"] == 4 + assert row["pdf_page"] == "report_4" + assert row["source_id"] == "/inputs/report.pdf" + assert json.loads(row["source"])["custom_source_field"] == "preserved" + assert json.loads(row["metadata"])["type"] == expected_type + assert row["stored_image_uri"] == "s3://bucket/figure.png" + assert json.loads(row["bbox_xyxy_norm"]) == [0.1, 0.2, 0.8, 0.9] + + +def test_collection_hit_preserves_native_dense_distance(): + hit = { + "text": "chunk", + "_score": 42.0, + "_distance": 0.125, + } + + public = _public_collection_hit(hit) + + assert public["distance"] == 0.125 + assert public["text"] == "chunk" + assert not {"_score", "_distance"} & public.keys() + + +@pytest.mark.parametrize( + ("content_type", "page_number"), + [("audio", 3), ("video", 3), ("video_frame", 3), ("text", -1)], +) +def test_collection_hit_does_not_expose_non_document_pages(content_type: str, page_number: int): + public = _public_collection_hit( + { + "text": "chunk", + "content_type": content_type, + "page_number": page_number, + "pdf_page": "document_3", + "_distance": 0.125, + } + ) + + assert public["page_number"] is None + assert public["pdf_page"] == "" + + +@pytest.mark.parametrize("bad_value", [None, True, math.nan, math.inf, -math.inf, "not-a-number"]) +def test_collection_hit_rejects_missing_or_invalid_native_distance(bad_value): + with pytest.raises(RetrievalContractError): + _public_collection_hit({"text": "chunk", "_distance": bad_value}) + + +@pytest.mark.parametrize( + ("mode", "expected_error"), + [ + ("hybrid", UnsupportedVDBOperation), + ("sparse", UnsupportedVDBOperation), + ("unknown", RetrievalContractError), + ], +) +def test_collection_retrieval_mode_error_classification(mode, expected_error): + store = object.__new__(collections_module.LanceDBCollectionStore) + capabilities = LanceTableCapabilities( + has_vector=mode in {"dense", "hybrid"}, + has_fts=mode in {"hybrid", "sparse"}, + retrieval_mode=mode, + vector_column="vector" if mode in {"dense", "hybrid"} else None, + text_column="text" if mode in {"hybrid", "sparse"} else None, + ) + + with pytest.raises(expected_error): + store._resolve_effective_retrieval_mode("collection-table", capabilities) + + +def test_collection_write_enforces_append_and_replace_invariants(tmp_path): + backend = _backend_with_collection(tmp_path) + + with pytest.raises(VDBResourceNotFound, match="Document not found"): + backend.write_collection(_records(), context=_context(operation="replace")) + + backend.write_collection(_records(), context=_context()) + with pytest.raises(VDBResourceConflict, match="use replace"): + backend.write_collection(_records(text="new version"), context=_context(version="v2")) + with pytest.raises(VDBResourceConflict, match="content does not match"): + backend.write_collection( + _records(text="different content"), + context=replace(_context(), content_sha256="different-sha"), + ) + + document = backend.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + ) + assert document.document_version == "v1" + store = backend._get_collection_store() + table_name = store._resolved_table("workspace-a", "collection-a") + assert store._open_table(table_name).count_rows() == 1 + + +def test_collection_writes_refresh_sliding_expiration(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + store = backend._get_collection_store() + collection = store._collection_row("workspace-a", "collection-a") + assert collection is not None + collection.update( + { + "updated_at": "2030-01-01T00:00:00+00:00", + "expires_at": "2030-01-02T00:00:00+00:00", + } + ) + store._persist_collection_row(collection) + + monkeypatch.setattr(collections_module, "_now", lambda: "2030-01-01T12:00:00+00:00") + backend.write_collection(_records(), context=_context()) + appended = backend.get_collection(scope="workspace-a", collection_name="collection-a") + assert appended.updated_at == "2030-01-01T12:00:00+00:00" + assert appended.expires_at == "2030-01-02T12:00:00+00:00" + + monkeypatch.setattr(collections_module, "_now", lambda: "2030-01-01T18:00:00+00:00") + backend.write_collection( + _records(text="replacement"), + context=_context(version="v2", operation="replace"), + ) + replaced = backend.get_collection(scope="workspace-a", collection_name="collection-a") + assert replaced.updated_at == "2030-01-01T18:00:00+00:00" + assert replaced.expires_at == "2030-01-02T18:00:00+00:00" + + +def test_collection_activity_without_expiration_only_updates_timestamp(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + monkeypatch.setattr(collections_module, "_now", lambda: "2030-01-01T12:00:00+00:00") + + backend.write_collection(_records(), context=_context()) + + collection = backend.get_collection(scope="workspace-a", collection_name="collection-a") + assert collection.updated_at == "2030-01-01T12:00:00+00:00" + assert collection.expires_at is None + + +def test_failed_or_empty_collection_write_does_not_refresh_expiration(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + store = backend._get_collection_store() + collection = store._collection_row("workspace-a", "collection-a") + assert collection is not None + collection.update( + { + "updated_at": "2030-01-01T00:00:00+00:00", + "expires_at": "2030-01-02T00:00:00+00:00", + } + ) + store._persist_collection_row(collection) + monkeypatch.setattr(collections_module, "_now", lambda: "2030-01-01T12:00:00+00:00") + + with pytest.raises(VDBInvalidRequest, match="no writable vector rows"): + backend.write_collection( + [[{"document_type": "text", "metadata": {"content": "no vector"}}]], + context=_context(), + ) + empty = backend.write_collection([], context=_context()) + assert empty.written == 0 + + unchanged = backend.get_collection(scope="workspace-a", collection_name="collection-a") + assert unchanged.updated_at == "2030-01-01T00:00:00+00:00" + assert unchanged.expires_at == "2030-01-02T00:00:00+00:00" + + +def test_collection_updates_preserve_replace_or_clear_expiration_window(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + store = backend._get_collection_store() + collection = store._collection_row("workspace-a", "collection-a") + assert collection is not None + collection.update( + { + "updated_at": "2030-01-01T00:00:00+00:00", + "expires_at": "2030-01-02T00:00:00+00:00", + } + ) + store._persist_collection_row(collection) + + monkeypatch.setattr(collections_module, "_now", lambda: "2030-01-01T12:00:00+00:00") + preserved = backend.update_collection( + scope="workspace-a", + collection_name="collection-a", + request=CollectionUpdateRequest(description="active"), + ) + assert preserved.updated_at == "2030-01-01T12:00:00+00:00" + assert preserved.expires_at == "2030-01-02T12:00:00+00:00" + + monkeypatch.setattr(collections_module, "_now", lambda: "2030-01-01T13:00:00+00:00") + replaced = backend.update_collection( + scope="workspace-a", + collection_name="collection-a", + request=CollectionUpdateRequest(expires_at="2030-01-04T13:00:00+00:00"), + ) + assert replaced.updated_at == "2030-01-01T13:00:00+00:00" + assert replaced.expires_at == "2030-01-04T13:00:00+00:00" + + monkeypatch.setattr(collections_module, "_now", lambda: "2030-01-01T14:00:00+00:00") + cleared = backend.update_collection( + scope="workspace-a", + collection_name="collection-a", + request=CollectionUpdateRequest(expires_at=None), + ) + assert cleared.updated_at == "2030-01-01T14:00:00+00:00" + assert cleared.expires_at is None + + +def test_collection_lifecycle_is_lazy_and_restart_safe(tmp_path): + uri = str(tmp_path / "lancedb") + backend = LanceDB( + uri=uri, + table_name="legacy", + vector_dim=2, + build_index=False, + ) + + assert backend._collection_store is None + assert lancedb.connect(uri).list_tables().tables == [] + assert backend.health()["catalog"]["initialized"] is False + assert backend._collection_store is None + assert lancedb.connect(uri).list_tables().tables == [] + + created = backend.create_collection( + scope="workspace-a", + request=CollectionCreateRequest(name="collection-a"), + ) + assert created.name == "collection-a" + for invalid_last in ([], ["a", "b"]): + with pytest.raises(VDBInvalidRequest, match="continuation token"): + backend.list_collections( + scope="workspace-a", + limit=10, + continuation_token=_encode_cursor("collections", "workspace-a", None, invalid_last), + ) + expected_table = "nrl_" + hashlib.sha256(b"workspace-a\x00collection-a").hexdigest()[:40] + assert {"_nrl_collections", "_nrl_documents"} <= set(lancedb.connect(uri).list_tables().tables) + + backend.create_collection( + scope="workspace-a", + request=CollectionCreateRequest(name="state-test"), + ) + store = backend._get_collection_store() + state_row = store._collection_row("workspace-a", "state-test") + assert state_row is not None + state_row["status"] = "deleting" + store._persist_collection_row(state_row) + with pytest.raises(VDBInvalidRequest, match="deleting"): + backend.update_collection( + scope="workspace-a", + collection_name="state-test", + request=CollectionUpdateRequest(description="blocked"), + ) + state_row["status"] = "active" + state_row["expires_at"] = "2000-01-01T00:00:00+00:00" + store._persist_collection_row(state_row) + with pytest.raises(VDBInvalidRequest, match="expired"): + backend.update_collection( + scope="workspace-a", + collection_name="state-test", + request=CollectionUpdateRequest(description="blocked"), + ) + state_row["expires_at"] = "" + store._persist_collection_row(state_row) + backend.delete_collection( + scope="workspace-a", + collection_name="state-test", + if_exists=False, + ) + + with pytest.raises(VDBInvalidRequest): + backend.write_collection( + [[{"document_type": "text", "metadata": {"content": "no vector"}}]], + context=_context(), + ) + + result = backend.write_collection(_records(), context=_context()) + assert result.written == 1 + assert result.total_rows == 1 + + document_row = store._document_rows("workspace-a", "collection-a", "document-a")[0] + document_row.update( + { + "content_sha256": "stale-hash", + "pending_document_version": "v1", + "recovery_state": "replacing", + } + ) + store._persist_document_row(document_row) + with store._write_lock: + assert store._reconcile_document_row_locked(document_row, expected_table) + assert ( + backend.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + ).content_sha256 + == "sha-v1" + ) + assert expected_table in lancedb.connect(uri).list_tables().tables + + hits, strategies = backend.retrieve_collection( + [[1.0, 0.0]], + scope="workspace-a", + collection_name="collection-a", + query_texts=["first"], + top_k=1, + ) + assert strategies == ["dense"] + assert hits[0][0]["document_id"] == "document-a" + assert hits[0][0]["distance"] >= 0.0 + assert "_distance" not in hits[0][0] + + restarted = LanceDB(uri=uri, table_name="legacy", vector_dim=2, build_index=False) + assert restarted._collection_store is None + assert ( + restarted.get_collection( + scope="workspace-a", + collection_name="collection-a", + ).name + == "collection-a" + ) + assert ( + restarted.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + ).chunk_count + == 1 + ) + + replacement = restarted.write_collection( + _records(text="replacement", vector=[0.0, 1.0], page_number=3), + context=_context(version="v2", operation="replace"), + ) + assert replacement.written == 1 + assert replacement.total_rows == 1 + assert ( + restarted.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + ).document_version + == "v2" + ) + + deleted_document = restarted.delete_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + if_exists=False, + ) + assert deleted_document.deleted is True + + deleted_collection = restarted.delete_collection( + scope="workspace-a", + collection_name="collection-a", + if_exists=False, + ) + assert deleted_collection.deleted is True + assert expected_table not in lancedb.connect(uri).list_tables().tables + + +def test_initial_append_reconciles_after_catalog_finalize_failure(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + store = backend._get_collection_store() + collection = store._collection_row("workspace-a", "collection-a") + assert collection is not None + collection.update( + { + "updated_at": "2030-01-01T00:00:00+00:00", + "expires_at": "2030-01-02T00:00:00+00:00", + } + ) + store._persist_collection_row(collection) + monkeypatch.setattr(collections_module, "_now", lambda: "2030-01-01T12:00:00+00:00") + _fail_document_finalize(monkeypatch, store) + + with pytest.raises(RuntimeError, match="injected catalog finalize failure"): + backend.write_collection(_records(), context=_context()) + + marker = store._document_rows("workspace-a", "collection-a", "document-a")[0] + assert marker["recovery_state"] == "appending" + assert marker["pending_document_version"] == "v1" + with pytest.raises(VDBResourceNotFound): + backend.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + ) + assert ( + backend.list_documents( + scope="workspace-a", + collection_name="collection-a", + limit=10, + continuation_token=None, + ).items + == [] + ) + + pending_hits, strategies = backend.retrieve_collection( + [[1.0, 0.0]], + scope="workspace-a", + collection_name="collection-a", + query_texts=["first"], + top_k=1, + ) + assert strategies == ["dense"] + assert pending_hits == [[]] + table_name = store._resolved_table("workspace-a", "collection-a") + assert store._open_table(table_name).count_rows() == 1 + + restarted = LanceDB( + uri=str(tmp_path / "lancedb"), + table_name="legacy", + vector_dim=2, + build_index=False, + ) + assert restarted.reconcile_collections() == {"successes": 1, "failures": 0} + document = restarted.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + ) + assert document.status == "completed" + assert document.document_version == "v1" + assert document.chunk_count == 1 + collection = restarted.get_collection(scope="workspace-a", collection_name="collection-a") + assert collection.updated_at == "2030-01-01T12:00:00+00:00" + assert collection.expires_at == "2030-01-02T12:00:00+00:00" + assert restarted._get_collection_store()._open_table(table_name).count_rows() == 1 + visible_hits, strategies = restarted.retrieve_collection( + [[1.0, 0.0]], + scope="workspace-a", + collection_name="collection-a", + query_texts=["first"], + top_k=1, + ) + assert strategies == ["dense"] + assert visible_hits[0][0]["document_id"] == "document-a" + + +def test_recovery_retries_activity_refresh_without_reextending_expiration(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + backend.write_collection(_records(), context=_context()) + store = backend._get_collection_store() + collection = store._collection_row("workspace-a", "collection-a") + assert collection is not None + collection.update( + { + "updated_at": "2030-01-01T00:00:00+00:00", + "expires_at": "2030-01-02T00:00:00+00:00", + } + ) + store._persist_collection_row(collection) + document = store._document_rows("workspace-a", "collection-a", "document-a")[0] + document.update( + { + "pending_document_version": "v1", + "recovery_state": "appending", + "status": "appending", + } + ) + store._persist_document_row(document) + activity_at = "2030-01-01T12:00:00+00:00" + monkeypatch.setattr(collections_module, "_now", lambda: activity_at) + + refreshes = 0 + original_persist_collection = store._persist_collection_row + + def count_activity_refresh(row): + nonlocal refreshes + refreshes += 1 + return original_persist_collection(row) + + monkeypatch.setattr(store, "_persist_collection_row", count_activity_refresh) + original_persist_document = store._persist_document_row + marker_clear_failed = False + + def fail_first_marker_clear(row): + nonlocal marker_clear_failed + if not marker_clear_failed and row.get("recovery_state") == "" and row.get("updated_at") == activity_at: + marker_clear_failed = True + raise RuntimeError("injected marker clear failure") + return original_persist_document(row) + + monkeypatch.setattr(store, "_persist_document_row", fail_first_marker_clear) + + assert backend.reconcile_collections() == {"successes": 0, "failures": 1} + marker = store._document_rows("workspace-a", "collection-a", "document-a")[0] + assert marker["recovery_state"] == "refreshing_collection_activity" + refreshed = backend.get_collection(scope="workspace-a", collection_name="collection-a") + assert refreshed.updated_at == activity_at + assert refreshed.expires_at == "2030-01-02T12:00:00+00:00" + + assert backend.reconcile_collections() == {"successes": 1, "failures": 0} + assert refreshes == 1 + assert store._document_rows("workspace-a", "collection-a", "document-a")[0]["recovery_state"] == "" + + +def test_reconciliation_filters_recoverable_documents_before_scan_limit(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + backend.write_collection( + _records(text="completed", vector=[1.0, 0.0]), + context=_context(document_id="document-completed"), + ) + store = backend._get_collection_store() + original_persist = _fail_document_finalize(monkeypatch, store) + with pytest.raises(RuntimeError, match="injected catalog finalize failure"): + backend.write_collection( + _records(text="pending", vector=[0.0, 1.0]), + context=_context(document_id="document-pending"), + ) + monkeypatch.setattr(store, "_persist_document_row", original_persist) + monkeypatch.setattr(collections_module, "_CATALOG_SCAN_LIMIT", 1) + + assert backend.reconcile_collections() == {"successes": 1, "failures": 0} + assert ( + backend.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-pending", + ).status + == "completed" + ) + + +def test_reconciliation_filters_expired_collections_before_scan_limit(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + backend.create_collection( + scope="workspace-a", + request=CollectionCreateRequest( + name="collection-expired", + expires_at="2000-01-01T00:00:00Z", + ), + ) + monkeypatch.setattr(collections_module, "_CATALOG_SCAN_LIMIT", 1) + + assert backend.reconcile_collections() == {"successes": 1, "failures": 0} + with pytest.raises(VDBResourceNotFound): + backend.get_collection(scope="workspace-a", collection_name="collection-expired") + + +def test_pending_initial_append_does_not_hide_completed_documents(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + backend.write_collection( + _records(text="completed", vector=[1.0, 0.0]), + context=_context(document_id="document-completed"), + ) + store = backend._get_collection_store() + _fail_document_finalize(monkeypatch, store) + + with pytest.raises(RuntimeError, match="injected catalog finalize failure"): + backend.write_collection( + _records(text="pending", vector=[0.0, 1.0]), + context=_context(document_id="document-pending"), + ) + + hits, strategies = backend.retrieve_collection( + [[1.0, 0.0]], + scope="workspace-a", + collection_name="collection-a", + query_texts=["completed"], + top_k=10, + ) + + assert strategies == ["dense"] + assert [hit["document_id"] for hit in hits[0]] == ["document-completed"] + + +def test_initial_append_retry_does_not_duplicate_committed_chunks(tmp_path, monkeypatch): + backend = _backend_with_collection(tmp_path) + store = backend._get_collection_store() + original_persist = _fail_document_finalize(monkeypatch, store) + + with pytest.raises(RuntimeError, match="injected catalog finalize failure"): + backend.write_collection(_records(), context=_context()) + + table_name = store._resolved_table("workspace-a", "collection-a") + assert store._open_table(table_name).count_rows() == 1 + monkeypatch.setattr(store, "_persist_document_row", original_persist) + + result = backend.write_collection(_records(), context=_context()) + assert result.written == 1 + assert result.total_rows == 1 + assert store._open_table(table_name).count_rows() == 1 + assert ( + backend.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + ).chunk_count + == 1 + ) + + +def test_initial_append_marker_without_chunks_is_removed_by_reconciliation( + tmp_path, + monkeypatch, +): + backend = _backend_with_collection(tmp_path) + store = backend._get_collection_store() + original_write = collections_module.create_or_append_lancedb_table + + def fail_before_chunk_write(*args, **kwargs): + raise RuntimeError("injected chunk write failure") + + monkeypatch.setattr( + collections_module, + "create_or_append_lancedb_table", + fail_before_chunk_write, + ) + with pytest.raises(RuntimeError, match="injected chunk write failure"): + backend.write_collection(_records(), context=_context()) + + marker = store._document_rows("workspace-a", "collection-a", "document-a")[0] + assert marker["recovery_state"] == "appending" + with pytest.raises(VDBResourceNotFound): + backend.get_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + ) + monkeypatch.setattr( + collections_module, + "create_or_append_lancedb_table", + original_write, + ) + + assert backend.reconcile_collections() == {"successes": 1, "failures": 0} + assert store._document_rows("workspace-a", "collection-a", "document-a") == [] + + +@pytest.mark.parametrize( + "raw_results", + [None, {}, [[], []], [{}], [[object()]]], +) +def test_collection_results_reject_invalid_cardinality_or_hit_shapes(raw_results): + with pytest.raises(RetrievalContractError): + _normalize_collection_results(raw_results, expected_queries=1) + + +def test_legacy_table_can_explicitly_infer_vector_dimension(tmp_path): + uri = str(tmp_path / "inferred-lancedb") + backend = LanceDB( + uri=uri, + table_name="legacy", + vector_dim=None, + overwrite=False, + build_index=False, + ) + + backend.run(_records(vector=[1.0, 0.0, 0.0])) + + table = lancedb.connect(uri).open_table("legacy") + assert table.schema.field("vector").type.list_size == 3 + assert table.schema.names == ["vector", "text", "metadata", "source", "id"] + assert table.count_rows() == 1 + assert LanceDB(uri=str(tmp_path / "default")).vector_dim == 2048 + + +def test_service_table_schema_survives_append_restart_and_query(tmp_path): + uri = str(tmp_path / "service-lancedb") + common = { + "uri": uri, + "table_name": "legacy", + "vector_dim": None, + "overwrite": False, + "build_index": False, + "_service_table_schema": True, + } + backend = LanceDB(**common) + backend.run(_service_records("table_caption", text="table caption", vector=[1.0, 0.0, 0.0])) + + restarted = LanceDB(**common) + restarted.run(_service_records("chart_caption", text="chart caption", vector=[0.0, 1.0, 0.0])) + + table = lancedb.connect(uri).open_table("legacy") + assert table.schema.field("vector").type.list_size == 3 + assert {"content_type", "stored_image_uri", "bbox_xyxy_norm"} <= set(table.schema.names) + assert table.count_rows() == 2 + + results = restarted.retrieval([[1.0, 0.0, 0.0]], top_k=2) + hits = {hit["text"]: hit for hit in results[0]} + for text, content_type in (("table caption", "table"), ("chart caption", "chart")): + hit = hits[text] + assert hit["content_type"] == content_type + assert hit["filename"] == "report.pdf" + assert hit["page_number"] == 4 + assert hit["pdf_page"] == "report_4" + assert hit["source_id"] == "/inputs/report.pdf" + assert json.loads(hit["source"])["custom_source_field"] == "preserved" + assert json.loads(hit["metadata"])["type"] == content_type + assert hit["stored_image_uri"] == "s3://bucket/figure.png" + assert json.loads(hit["bbox_xyxy_norm"]) == [0.1, 0.2, 0.8, 0.9] + + +def test_collection_query_is_not_blocked_by_unrelated_write(tmp_path, monkeypatch): + backend = LanceDB( + uri=str(tmp_path / "query-during-write"), + table_name="legacy", + vector_dim=2, + overwrite=False, + build_index=False, + ) + for name in ("collection-a", "collection-b"): + backend.create_collection(scope="workspace-a", request=CollectionCreateRequest(name=name)) + backend.write_collection( + _records(text="collection b"), + context=replace(_context(), collection_name="collection-b", document_id="document-b"), + ) + + write_entered = threading.Event() + allow_write_to_finish = threading.Event() + query_entered = threading.Event() + original_create = collections_module.create_or_append_lancedb_table + + def blocking_create(*args, **kwargs): + write_entered.set() + if not allow_write_to_finish.wait(timeout=5): + raise TimeoutError("test did not release blocked collection write") + return original_create(*args, **kwargs) + + def immediate_retrieval(vectors, **kwargs): + query_entered.set() + return [[] for _ in vectors] + + monkeypatch.setattr(collections_module, "create_or_append_lancedb_table", blocking_create) + monkeypatch.setattr(backend, "retrieval", immediate_retrieval) + + with ThreadPoolExecutor(max_workers=2) as pool: + write_future = pool.submit(backend.write_collection, _records(), context=_context()) + assert write_entered.wait(timeout=5) + query_future = pool.submit( + backend.retrieve_collection, + [[1.0, 0.0]], + scope="workspace-a", + collection_name="collection-b", + query_texts=["query"], + top_k=1, + ) + try: + assert query_entered.wait(timeout=1) + query_future.result(timeout=1) + assert not write_future.done() + finally: + allow_write_to_finish.set() + write_future.result(timeout=5) + + +def test_collection_delete_waits_for_active_write(tmp_path, monkeypatch): + backend = LanceDB( + uri=str(tmp_path / "delete-during-write"), + table_name="legacy", + vector_dim=2, + overwrite=False, + build_index=False, + ) + backend.create_collection(scope="workspace-a", request=CollectionCreateRequest(name="collection-a")) + + write_entered = threading.Event() + allow_write_to_finish = threading.Event() + original_create = collections_module.create_or_append_lancedb_table + + def blocking_create(*args, **kwargs): + write_entered.set() + if not allow_write_to_finish.wait(timeout=5): + raise TimeoutError("test did not release blocked collection write") + return original_create(*args, **kwargs) + + monkeypatch.setattr(collections_module, "create_or_append_lancedb_table", blocking_create) + + with ThreadPoolExecutor(max_workers=2) as pool: + write_future = pool.submit(backend.write_collection, _records(), context=_context()) + assert write_entered.wait(timeout=5) + delete_future = pool.submit( + backend.delete_collection, + scope="workspace-a", + collection_name="collection-a", + if_exists=False, + ) + try: + with pytest.raises(TimeoutError): + delete_future.result(timeout=0.2) + finally: + allow_write_to_finish.set() + write_future.result(timeout=5) + result = delete_future.result(timeout=5) + + assert result.deleted is True + + +def test_collection_reconciliation_waits_for_active_write(tmp_path, monkeypatch): + backend = LanceDB( + uri=str(tmp_path / "reconcile-during-write"), + table_name="legacy", + vector_dim=2, + overwrite=False, + build_index=False, + ) + backend.create_collection(scope="workspace-a", request=CollectionCreateRequest(name="collection-a")) + + write_entered = threading.Event() + allow_write_to_finish = threading.Event() + original_create = collections_module.create_or_append_lancedb_table + + def blocking_create(*args, **kwargs): + write_entered.set() + if not allow_write_to_finish.wait(timeout=5): + raise TimeoutError("test did not release blocked collection write") + return original_create(*args, **kwargs) + + monkeypatch.setattr(collections_module, "create_or_append_lancedb_table", blocking_create) + + with ThreadPoolExecutor(max_workers=2) as pool: + write_future = pool.submit(backend.write_collection, _records(), context=_context()) + assert write_entered.wait(timeout=5) + reconcile_future = pool.submit(backend.reconcile_collections) + try: + with pytest.raises(TimeoutError): + reconcile_future.result(timeout=0.2) + finally: + allow_write_to_finish.set() + write_future.result(timeout=5) + result = reconcile_future.result(timeout=5) + + assert result == {"successes": 0, "failures": 0} + + +def test_collection_writes_remain_serialized(tmp_path, monkeypatch): + backend = LanceDB( + uri=str(tmp_path / "serialized-writes"), + table_name="legacy", + vector_dim=2, + overwrite=False, + build_index=False, + ) + for name in ("collection-a", "collection-b"): + backend.create_collection(scope="workspace-a", request=CollectionCreateRequest(name=name)) + + first_write_entered = threading.Event() + second_write_entered = threading.Event() + allow_first_write_to_finish = threading.Event() + call_lock = threading.Lock() + call_count = 0 + original_create = collections_module.create_or_append_lancedb_table + + def blocking_first_create(*args, **kwargs): + nonlocal call_count + with call_lock: + call_count += 1 + current_call = call_count + if current_call == 1: + first_write_entered.set() + if not allow_first_write_to_finish.wait(timeout=5): + raise TimeoutError("test did not release first collection write") + else: + second_write_entered.set() + return original_create(*args, **kwargs) + + monkeypatch.setattr(collections_module, "create_or_append_lancedb_table", blocking_first_create) + + second_context = replace(_context(), collection_name="collection-b", document_id="document-b") + with ThreadPoolExecutor(max_workers=2) as pool: + first_future = pool.submit(backend.write_collection, _records(), context=_context()) + assert first_write_entered.wait(timeout=5) + second_future = pool.submit(backend.write_collection, _records(), context=second_context) + try: + assert not second_write_entered.wait(timeout=0.2) + assert not second_future.done() + finally: + allow_first_write_to_finish.set() + first_future.result(timeout=5) + second_future.result(timeout=5) + + assert second_write_entered.is_set() + + +def test_document_delete_waits_for_active_collection_query(tmp_path, monkeypatch): + uri = str(tmp_path / "concurrent-lancedb") + backend = LanceDB( + uri=uri, + table_name="legacy", + vector_dim=2, + overwrite=False, + build_index=False, + ) + backend.create_collection( + scope="workspace-a", + request=CollectionCreateRequest(name="collection-a"), + ) + backend.write_collection(_records(), context=_context()) + + query_entered = threading.Event() + allow_query_to_finish = threading.Event() + query_finished = threading.Event() + delete_finished = threading.Event() + query_errors: list[BaseException] = [] + delete_errors: list[BaseException] = [] + + def blocking_retrieval(vectors, **kwargs): + query_entered.set() + if not allow_query_to_finish.wait(timeout=5): + raise TimeoutError("test did not release blocked collection query") + return [[] for _ in vectors] + + monkeypatch.setattr(backend, "retrieval", blocking_retrieval) + + def query_target(): + try: + backend.retrieve_collection( + [[1.0, 0.0]], + scope="workspace-a", + collection_name="collection-a", + query_texts=["query"], + top_k=1, + ) + except BaseException as exc: # pragma: no cover - asserted below + query_errors.append(exc) + finally: + query_finished.set() + + def delete_target(): + try: + backend.delete_document( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + if_exists=False, + ) + except BaseException as exc: # pragma: no cover - asserted below + delete_errors.append(exc) + finally: + delete_finished.set() + + query_thread = threading.Thread(target=query_target) + delete_thread = threading.Thread(target=delete_target) + query_thread.start() + assert query_entered.wait(timeout=5) + delete_thread.start() + try: + assert not delete_finished.wait(timeout=0.2) + finally: + allow_query_to_finish.set() + query_thread.join(timeout=5) + delete_thread.join(timeout=5) + + assert query_finished.is_set() + assert delete_finished.is_set() + assert query_errors == [] + assert delete_errors == [] diff --git a/nemo_retriever/tests/test_lancedb_write_policy.py b/nemo_retriever/tests/test_lancedb_write_policy.py index df820592bc..d4c3d80f62 100644 --- a/nemo_retriever/tests/test_lancedb_write_policy.py +++ b/nemo_retriever/tests/test_lancedb_write_policy.py @@ -84,6 +84,138 @@ def test_append_same_records_twice_doubles_row_count(tmp_path: Path, caplog: pyt assert "Append mode does not deduplicate" in caplog.text +def test_append_with_matching_embedding_model_succeeds(tmp_path: Path) -> None: + model_name = "nvidia/embedding-model-a" + LanceDB( + uri=str(tmp_path), + table_name="t", + vector_dim=2, + embedding_model_name=model_name, + create_index=False, + ).run(_records()) + + LanceDB( + uri=str(tmp_path), + table_name="t", + vector_dim=2, + embedding_model_name=model_name, + overwrite=False, + create_index=False, + ).run(_records()) + + assert _count_rows(tmp_path) == 2 + + +def test_embedding_model_revision_is_recorded_and_readable(tmp_path: Path) -> None: + op = LanceDB( + uri=str(tmp_path), + table_name="t", + vector_dim=2, + embedding_model_name="nvidia/embedding-model-a", + embedding_model_revision="a" * 40, + create_index=False, + ) + + op.run(_records()) + + assert op.get_index_metadata("embedding_model_name") == "nvidia/embedding-model-a" + assert op.get_index_metadata("embedding_model_revision") == "a" * 40 + + +def test_vector_dimension_can_be_inferred_from_model_output(tmp_path: Path) -> None: + op = LanceDB( + uri=str(tmp_path), + table_name="t", + vector_dim=None, + embedding_model_name="nvidia/llama-embed-nemotron-8b", + create_index=False, + ) + + op.run(_records(vector=[0.0] * 4096)) + + table = lancedb.connect(str(tmp_path)).open_table("t") + schema = table.schema() if callable(table.schema) else table.schema + assert schema.field("vector").type.list_size == 4096 + + +def test_append_with_inferred_dimension_uses_existing_table_schema(tmp_path: Path) -> None: + kwargs = { + "uri": str(tmp_path), + "table_name": "t", + "vector_dim": None, + "embedding_model_name": "nvidia/llama-embed-nemotron-8b", + "create_index": False, + } + LanceDB(**kwargs).run(_records(vector=[0.0] * 4096)) + + LanceDB(**kwargs, overwrite=False).run(_records(vector=[0.0] * 4096)) + + assert _count_rows(tmp_path) == 2 + + +def test_append_with_mismatched_embedding_model_fails_before_write(tmp_path: Path) -> None: + LanceDB( + uri=str(tmp_path), + table_name="t", + vector_dim=2, + embedding_model_name="nvidia/embedding-model-a", + create_index=False, + ).run(_records()) + + op = LanceDB( + uri=str(tmp_path), + table_name="t", + vector_dim=2, + embedding_model_name="nvidia/embedding-model-b", + overwrite=False, + create_index=False, + ) + + with pytest.raises(ValueError, match="cannot append vectors"): + op.run(_records()) + + assert _count_rows(tmp_path) == 1 + + +@pytest.mark.parametrize( + ("stored_revision", "incoming_revision", "error_pattern", "expected_rows"), + [ + pytest.param("a" * 40, "b" * 40, "cannot append vectors from revision", 1, id="mismatch"), + pytest.param("a" * 40, None, "without a known revision", 1, id="missing"), + pytest.param("a" * 40, "a" * 40, None, 2, id="matching"), + pytest.param(None, "a" * 40, None, 2, id="legacy-table"), + ], +) +def test_append_revision_compatibility( + tmp_path: Path, + stored_revision: str | None, + incoming_revision: str | None, + error_pattern: str | None, + expected_rows: int, +) -> None: + common = { + "uri": str(tmp_path), + "table_name": "t", + "vector_dim": 2, + "embedding_model_name": "nvidia/embedding-model-a", + "create_index": False, + } + LanceDB(**common, embedding_model_revision=stored_revision).run(_records()) + incoming = LanceDB( + **common, + embedding_model_revision=incoming_revision, + overwrite=False, + ) + + if error_pattern is not None: + with pytest.raises(ValueError, match=error_pattern): + incoming.run(_records()) + else: + incoming.run(_records()) + + assert _count_rows(tmp_path) == expected_rows + + def test_append_incompatible_schema_raises_clear_error(tmp_path: Path) -> None: LanceDB(uri=str(tmp_path), table_name="t", vector_dim=3, create_index=False).run(_records(vector=[1.0, 0.0, 0.0])) diff --git a/nemo_retriever/tests/test_nemo_agent_callable_contract.py b/nemo_retriever/tests/test_nemo_agent_callable_contract.py new file mode 100644 index 0000000000..3c0b8d66c7 --- /dev/null +++ b/nemo_retriever/tests/test_nemo_agent_callable_contract.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""The contract between CallableLLMBackend and the completion callables it drives. + +Two callables ship in this repo — the hosted ``invoke_chat_completion_step`` and +the in-process ``VLLMAgentChatLLM`` — and the backend must speak to both through +one keyword set. These tests pin that seam plus the request shaping that used to +live in the retired HTTP backend. +""" + +from __future__ import annotations + +import inspect + +import pytest +from pydantic import ValidationError + +from nemo_retriever._agentic.nemo_agent.llm import ( + CallableLLMConfig, + ContextLimitError, + LLMCallError, + create_llm, + create_llm_config, +) + +_OK = { + "choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, +} + + +def _backend(fn, **config_kwargs): + config_kwargs.setdefault("model", "m") + return create_llm(create_llm_config("callable", **config_kwargs), completion_fn=fn) + + +def _seen(**config_kwargs): + """Run one completion against a recording callable and return its kwargs.""" + seen: dict = {} + + def fn(**kwargs): + seen.update(kwargs) + return _OK + + _backend(fn, **config_kwargs).completion(messages=[{"role": "user", "content": "q"}]) + return seen + + +class TestKeywordContract: + #: Every keyword the backend promises to send. Both shipped callables accept + #: all of them; a callable that does not is a contract violation, and a loud + #: TypeError is the intended failure. + CONTRACT = frozenset( + { + "invoke_url", + "messages", + "model", + "api_key", + "tools", + "tool_choice", + "timeout_s", + "temperature", + "max_tokens", + "extra_body", + "max_retries", + "max_429_retries", + } + ) + + def test_every_contract_keyword_is_sent(self): + assert set(_seen()) == self.CONTRACT + + @pytest.mark.parametrize( + "target", + [ + "nemo_retriever.models.nim.chat_completions.invoke_chat_completion_step", + "nemo_retriever.models.local.agent_llm.VLLMAgentChatLLM.__call__", + ], + ) + def test_both_shipped_callables_accept_the_contract(self, target): + module_path, _, attr = target.rpartition(".") + if attr == "__call__": + module_path, _, cls_name = module_path.rpartition(".") + module = pytest.importorskip(module_path) + fn = getattr(module, cls_name).__call__ + else: + module = pytest.importorskip(module_path) + fn = getattr(module, attr) + + accepted = set(inspect.signature(fn).parameters) + missing = self.CONTRACT - accepted + assert not missing, f"{target} does not accept {sorted(missing)}" + + +class TestRetryKnobs: + def test_defaults_are_bounded_for_agent_use(self): + config = CallableLLMConfig(model="m") + assert config.max_retries == 3 + assert config.max_429_retries == 6 + assert config.timeout_s == 120.0 + # The callable owns retrying, so the base template must make one attempt. + assert config.rate_limit_max_retries == 0 + + def test_zero_max_retries_is_rejected(self): + # A callable that loops `while attempt < max_retries` issues ZERO requests + # at 0 and then reports retries-exhausted, which reads as an endpoint + # failure. The bound makes that unreachable. + with pytest.raises(ValidationError): + CallableLLMConfig(model="m", max_retries=0) + + @pytest.mark.parametrize("field", ["max_429_retries"]) + def test_other_retry_bounds(self, field): + with pytest.raises(ValidationError): + CallableLLMConfig(model="m", **{field: 0}) + + def test_non_positive_timeout_is_rejected(self): + with pytest.raises(ValidationError): + CallableLLMConfig(model="m", timeout_s=0) + + def test_knobs_are_forwarded_to_the_callable(self): + seen = _seen(timeout_s=30.0, max_retries=2, max_429_retries=4) + assert seen["timeout_s"] == 30.0 + assert seen["max_retries"] == 2 + assert seen["max_429_retries"] == 4 + + +class TestApiKeyResolution: + def test_env_indirection_resolved_once_at_construction(self, monkeypatch): + monkeypatch.setenv("AGENT_TEST_KEY", " sk-live ") + assert _seen(api_key="os.environ/AGENT_TEST_KEY")["api_key"] == "sk-live" + + def test_literal_key_passes_through(self): + assert _seen(api_key="sk-literal")["api_key"] == "sk-literal" + + def test_absent_key_is_none_not_empty_string(self): + # An empty Bearer header is worse than none at all. + assert _seen()["api_key"] is None + + def test_missing_env_var_fails_at_build_time(self, monkeypatch): + # Must fail while the pipeline is being constructed, not several hundred + # agent steps into a run. + monkeypatch.delenv("AGENT_TEST_MISSING", raising=False) + with pytest.raises(ValueError, match="AGENT_TEST_MISSING"): + _backend(lambda **kw: _OK, api_key="os.environ/AGENT_TEST_MISSING") + + +class TestRequestShaping: + def test_block_list_content_is_normalized_to_a_string(self): + # The agent builds every message as a content-block list, including the + # system message. Endpoints that only accept string content for some roles + # would otherwise see a shape this repo has never exercised. + seen: dict = {} + + def fn(**kwargs): + seen.update(kwargs) + return _OK + + _backend(fn).completion( + messages=[ + {"role": "system", "content": [{"type": "text", "text": "S"}]}, + {"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]}, + ] + ) + assert seen["messages"][0]["content"] == "S" + assert seen["messages"][1]["content"] == "a\nb" + + def test_private_keys_are_stripped_before_normalization(self): + seen: dict = {} + + def fn(**kwargs): + seen.update(kwargs) + return _OK + + _backend(fn).completion(messages=[{"role": "user", "content": "q", "__reasoning__": "secret"}]) + assert "__reasoning__" not in seen["messages"][0] + + def test_caller_messages_are_not_mutated(self): + messages = [{"role": "user", "content": [{"type": "text", "text": "q"}], "__reasoning__": "r"}] + _backend(lambda **kw: _OK).completion(messages=messages) + assert messages[0]["content"] == [{"type": "text", "text": "q"}] + assert messages[0]["__reasoning__"] == "r" + + def test_base_url_is_forwarded_verbatim_as_invoke_url(self): + url = "https://integrate.api.nvidia.com/v1/chat/completions" + assert _seen(base_url=url)["invoke_url"] == url + + def test_tool_choice_suppressed_without_tools(self): + assert _seen()["tool_choice"] == "none" + + +class TestResponseHandling: + def test_usage_is_recorded_not_discarded(self): + result = _backend(lambda **kw: _OK).completion(messages=[{"role": "user", "content": "q"}]) + assert result.usage == {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} + + def test_extra_response_info_is_always_empty(self): + # Structural: the callable returns a decoded body, so headers/status never + # cross the boundary. Pinned so nobody "fixes" it by plumbing them through. + result = _backend(lambda **kw: _OK).completion(messages=[{"role": "user", "content": "q"}]) + assert result.extra_response_info == {} + + def test_padded_finish_reason_is_stripped(self): + # The agent loop treats any finish reason outside ("stop", "tool_calls") as + # terminal, so a padded value would end an otherwise healthy run. + body = {"choices": [{"message": {"role": "assistant", "content": "x"}, "finish_reason": " tool_calls "}]} + result = _backend(lambda **kw: body).completion(messages=[{"role": "user", "content": "q"}]) + assert result.finish_reason == "tool_calls" + + def test_missing_finish_reason_defaults_to_stop(self): + body = {"choices": [{"message": {"role": "assistant", "content": "x"}}]} + result = _backend(lambda **kw: body).completion(messages=[{"role": "user", "content": "q"}]) + assert result.finish_reason == "stop" + + def test_block_list_response_content_is_coerced(self): + blocks = [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}] + body = { + "choices": [ + { + "message": {"role": "assistant", "content": blocks}, + "finish_reason": "stop", + } + ] + } + result = _backend(lambda **kw: body).completion(messages=[{"role": "user", "content": "q"}]) + assert result.message["content"] == "a\nb" + + @pytest.mark.parametrize( + "body", + [None, {}, {"choices": []}, {"choices": [None]}, {"choices": [{}]}, "not-a-dict"], + ) + def test_malformed_responses_raise_llm_call_error_not_key_error(self, body): + with pytest.raises(LLMCallError): + _backend(lambda **kw: body).completion(messages=[{"role": "user", "content": "q"}]) + + +class TestErrorTranslation: + def test_callable_failures_are_classified(self): + def boom(**kwargs): + raise ValueError("prompt is longer than the maximum model length of 8192") + + with pytest.raises(ContextLimitError): + _backend(boom).completion(messages=[{"role": "user", "content": "q"}]) + + def test_library_errors_keep_their_subclass(self): + # Re-wrapping would erase the class the selection agent's shrink loop + # branches on. + def boom(**kwargs): + raise ContextLimitError("already ours") + + with pytest.raises(ContextLimitError, match="already ours"): + _backend(boom).completion(messages=[{"role": "user", "content": "q"}]) + + def test_original_exception_stays_chained(self): + original = RuntimeError("wire down") + + def boom(**kwargs): + raise original + + with pytest.raises(LLMCallError) as excinfo: + _backend(boom).completion(messages=[{"role": "user", "content": "q"}]) + assert excinfo.value.__cause__ is original + + def test_cancellation_is_not_swallowed(self): + # BaseException must propagate untouched or task cancellation breaks. + def boom(**kwargs): + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + _backend(boom).completion(messages=[{"role": "user", "content": "q"}]) + + +class TestCaptureRawIO: + def test_credentials_and_url_are_redacted(self): + backend = _backend( + lambda **kw: _OK, + api_key="sk-secret", + base_url="https://user:pw@endpoint.invalid/v1/chat/completions?token=leaked", + capture_raw_io=True, + ) + result = backend.completion(messages=[{"role": "user", "content": "q"}]) + dumped = str(result.raw_request) + assert "sk-secret" not in dumped + assert "pw" not in dumped + assert "token=leaked" not in dumped + assert "endpoint.invalid" in dumped + + def test_disabled_by_default(self): + result = _backend(lambda **kw: _OK).completion(messages=[{"role": "user", "content": "q"}]) + assert result.raw_request is None + assert result.raw_response is None diff --git a/nemo_retriever/tests/test_nemo_agent_error_classification.py b/nemo_retriever/tests/test_nemo_agent_error_classification.py new file mode 100644 index 0000000000..9cbd99849b --- /dev/null +++ b/nemo_retriever/tests/test_nemo_agent_error_classification.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Exception -> error-hierarchy mapping for the callable LLM backend. + +The agent branches on these classes: ``loop.py`` records ContextLimitError and +ContentPolicyError as expected outcomes, and the selection agent retries a +shrunken candidate list only on a context limit. A misclassification therefore +silently changes agent behavior, so the table below is the contract. +""" + +from __future__ import annotations + +import json + +import pytest +import requests + +from nemo_retriever._agentic.nemo_agent.llm import ( + ContentPolicyError, + ContextLimitError, + LLMCallError, + RateLimitError, +) +from nemo_retriever._agentic.nemo_agent.llm.error_classification import ( + _parse_retry_after, + classify_call_exception, + classify_http_error, +) + +_URL = "https://user:secret@endpoint.invalid/v1/chat/completions?api_key=leaked" + + +def _response(status: int, text: str = "", headers: dict | None = None) -> requests.Response: + response = requests.Response() + response.status_code = status + response._content = text.encode() + response.url = _URL + response.headers.update(headers or {}) + return response + + +def _http_error(status: int, text: str = "", headers: dict | None = None) -> requests.HTTPError: + return requests.HTTPError(f"HTTP {status}", response=_response(status, text, headers)) + + +class TestResponseCarryingErrors: + def test_429_becomes_a_rate_limit_error_with_retry_after(self): + error = classify_call_exception(_http_error(429, "slow down", {"Retry-After": "12"})) + assert isinstance(error, RateLimitError) + assert error.retry_after == 12.0 + + def test_error_response_objects_are_falsy_and_still_classified(self): + # requests.Response.__bool__ is `status_code < 400`, so EVERY response this + # module exists to classify is falsy. A truthiness check instead of + # `is None` would skip classification for all of them. + response = _response(429, "slow down") + assert not response + assert isinstance(classify_call_exception(requests.HTTPError("x", response=response)), RateLimitError) + + @pytest.mark.parametrize("code", ["context_length_exceeded", "string_above_max_length"]) + def test_structured_context_limit_codes(self, code): + body = '{"error": {"code": "%s", "message": "too long"}}' % code + assert isinstance(classify_call_exception(_http_error(400, body)), ContextLimitError) + + @pytest.mark.parametrize( + "body", + [ + "This model's maximum context length is 8192 tokens", + "Please reduce the length of the messages", + "max_tokens must be at least 1, got -42", + ], + ) + def test_prose_context_limit_markers(self, body): + assert isinstance(classify_call_exception(_http_error(400, body)), ContextLimitError) + + def test_content_policy(self): + assert isinstance(classify_call_exception(_http_error(400, "blocked by content policy")), ContentPolicyError) + + def test_unclassified_4xx_stays_a_plain_llm_call_error(self): + error = classify_call_exception(_http_error(401, "invalid api key")) + assert type(error) is LLMCallError + + def test_message_redacts_url_credentials_and_query(self): + message = str(classify_call_exception(_http_error(500, "boom"))) + assert "secret" not in message + assert "api_key=leaked" not in message + assert "endpoint.invalid" in message + + def test_message_excerpts_a_huge_body(self): + message = str(classify_call_exception(_http_error(500, "x" * 50_000))) + assert len(message) < 3_000 + assert message.endswith("...") + + +class TestErrorsWithoutAResponse: + """No response object: timeouts, transport errors, and in-process backends.""" + + @pytest.mark.parametrize( + "exc", + [ + TimeoutError("Request timed out after 3 attempts."), + RuntimeError("Failed to get a successful response after 3 retries."), + ValueError("invoke_url is required"), + ], + ) + def test_degrade_to_plain_llm_call_error(self, exc): + assert type(classify_call_exception(exc)) is LLMCallError + + def test_prose_still_recovers_a_context_limit(self): + # An in-process engine reports overflow as a bare ValueError with no + # response to inspect. Recovering ContextLimitError from the text is what + # keeps the selection agent's shrink-and-retry loop alive locally. + exc = ValueError("The decoder prompt (length 9000) is longer than the maximum model length of 8192.") + assert isinstance(classify_call_exception(exc), ContextLimitError) + + +class TestTotality: + """The classifier must never raise — a bug here would mask the real failure.""" + + def test_a_response_whose_text_raises_is_survived(self): + class Hostile: + status_code = 500 + + @property + def text(self): + raise RuntimeError("content already consumed") + + exc = requests.HTTPError("x") + exc.response = Hostile() + assert isinstance(classify_call_exception(exc), LLMCallError) + + def test_a_non_integer_status_code_falls_back_to_prose(self): + class Weird: + status_code = "not-a-number" + + exc = RuntimeError("maximum context length exceeded") + exc.response = Weird() + assert isinstance(classify_call_exception(exc), ContextLimitError) + + def test_an_exception_whose_str_raises_is_survived(self): + class Explosive(Exception): + def __str__(self): + raise RuntimeError("nope") + + assert isinstance(classify_call_exception(Explosive()), LLMCallError) + + def test_library_errors_pass_through_with_their_subclass_intact(self): + # Re-wrapping would erase the class the agent branches on. + original = ContextLimitError("prompt too long") + assert classify_call_exception(original) is original + + +# ---------------------------------------------------------------------- +# Pure helpers, called directly. These were previously exercised through the +# retired HTTP backend; the logic moved here, so the coverage moves with it. +# ---------------------------------------------------------------------- + + +class TestClassifyHttpError: + """``classify_http_error`` on its own, without an exception wrapper.""" + + def _classify(self, status, body, headers=None): + text = body if isinstance(body, str) else json.dumps(body) + parsed = None if isinstance(body, str) else body + return classify_http_error(status, text, parsed, headers or {}, _URL) + + def test_rate_limit(self): + error = self._classify(429, {"error": {"message": "slow down"}}, {"Retry-After": "7"}) + assert isinstance(error, RateLimitError) + assert error.retry_after == 7.0 + + def test_context_limit_structured_code(self): + assert isinstance(self._classify(400, {"error": {"code": "context_length_exceeded"}}), ContextLimitError) + + def test_context_limit_structured_type(self): + # `type` is checked as well as `code`: providers populate one or the other. + assert isinstance(self._classify(400, {"error": {"type": "string_above_max_length"}}), ContextLimitError) + + @pytest.mark.parametrize( + "prose", + [ + "This model's maximum context length is 8192 tokens", + "The input is longer than the maximum model length", + "Please reduce the length of the messages", + "max_tokens must be at least 1, got -37", + ], + ) + def test_context_limit_prose(self, prose): + assert isinstance(self._classify(400, {"error": {"message": prose}}), ContextLimitError) + + def test_content_policy_structured(self): + assert isinstance(self._classify(400, {"error": {"code": "content_policy_violation"}}), ContentPolicyError) + + @pytest.mark.parametrize( + "prose", + ["blocked by the content filter", "guardrail intervened", "violates our content policy"], + ) + def test_content_policy_prose(self, prose): + assert isinstance(self._classify(400, {"error": {"message": prose}}), ContentPolicyError) + + def test_non_json_body_still_classified(self): + # body_json is None when the error body is not JSON; prose matching runs + # against the raw text, which is a superset of it. + error = classify_http_error(400, "maximum context length exceeded", None, {}, _URL) + assert isinstance(error, ContextLimitError) + + def test_unclassified_is_plain_llm_call_error(self): + assert type(self._classify(401, {"error": {"message": "invalid api key"}})) is LLMCallError + + def test_structured_code_wins_over_absent_prose(self): + # The structured check runs first precisely because provider wording drifts. + assert isinstance(self._classify(400, {"error": {"code": "context_length_exceeded"}}), ContextLimitError) + + +class TestParseRetryAfter: + """Header parsing for the 429 path. + + The base class accepts any finite ``>= 0`` and caps at its own ceiling, so a + bogus value must be rejected *here* or it would either burn the whole + rate-limit budget instantly or stall every retry. + """ + + def test_delta_seconds(self): + assert _parse_retry_after({"Retry-After": "12"}) == 12.0 + + def test_case_insensitive(self): + assert _parse_retry_after({"retry-after": "3"}) == 3.0 + + def test_http_date_far_future_is_discarded(self): + assert _parse_retry_after({"Retry-After": "Wed, 21 Oct 2099 07:28:00 GMT"}) is None + + def test_zero_is_discarded(self): + # Would otherwise burn the base class's whole rate-limit budget instantly. + assert _parse_retry_after({"Retry-After": "0"}) is None + + def test_negative_is_discarded(self): + assert _parse_retry_after({"Retry-After": "-5"}) is None + + def test_absurdly_large_is_discarded(self): + assert _parse_retry_after({"Retry-After": "99999"}) is None + + def test_garbage_is_discarded(self): + assert _parse_retry_after({"Retry-After": "soon"}) is None + + def test_absent(self): + assert _parse_retry_after({}) is None + + def test_no_headers_at_all(self): + assert _parse_retry_after(None) is None + + +class TestCredentialScrubbing: + """Messages this module builds must not carry a URL's credentials. + + These land in ``AgentError.message`` and from there in the run trajectory, so + a query-string API key would be persisted alongside ordinary results. Bounds + disclosure in what we surface; the original is still chained on ``__cause__``. + """ + + def test_transport_error_query_string_is_scrubbed(self): + # The regression that motivated this: a connection error carries no + # response, so it takes the prose path — and urllib3's message keeps the + # path and query even though it drops the scheme and host. + exc = requests.ConnectionError( + "HTTPSConnectionPool(host='h', port=443): Max retries exceeded with " + "url: /v1/chat/completions?api_key=SECRET123 (Caused by NewConnectionError(...))" + ) + message = str(classify_call_exception(exc)) + assert "SECRET123" not in message + # Still diagnosable: the failure mode survives scrubbing. + assert "Max retries exceeded" in message + + def test_full_url_in_a_prose_message_is_scrubbed(self): + exc = RuntimeError("failed calling https://user:pw@h.invalid/v1/chat?api_key=SECRET123") + message = str(classify_call_exception(exc)) + assert "SECRET123" not in message + assert "user:pw" not in message + assert "h.invalid" in message + + def test_scrubbing_does_not_cost_a_classification(self): + # The message is scrubbed but the marker matching runs against the raw + # text, so a URL sitting next to a marker cannot suppress the verdict. + exc = ValueError("prompt is longer than the maximum model length at https://h/v1?api_key=SECRET123") + error = classify_call_exception(exc) + assert isinstance(error, ContextLimitError) + assert "SECRET123" not in str(error) + + def test_response_body_echoing_a_url_is_scrubbed(self): + # Defense in depth: some providers echo the request back in the error body. + error = classify_call_exception(_http_error(400, "bad request to https://h/v1?api_key=SECRET123")) + assert "SECRET123" not in str(error) + + def test_original_exception_is_still_chained_unscrubbed(self): + # Scrubbing is about what we surface, not about destroying evidence. The + # backend attaches __cause__, so a debugger still sees the real message. + original = requests.ConnectionError("url: /v1?api_key=SECRET123") + assert "SECRET123" in str(original) diff --git a/nemo_retriever/tests/test_nemo_agent_llm_config.py b/nemo_retriever/tests/test_nemo_agent_llm_config.py new file mode 100644 index 0000000000..21638f2bf7 --- /dev/null +++ b/nemo_retriever/tests/test_nemo_agent_llm_config.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for ``nemo_agent.llm.create_llm_config`` — the kwargs-filtering config factory. + +These tests intentionally avoid instantiating a real ``LiteLLMBackend`` with a +valid config (which would import the optional ``litellm`` dependency). They only +exercise config *building* and the wrong-config-type failure path, both of which +run without ``litellm`` installed (the failure path raises before the lazy +``import litellm`` in ``LiteLLMBackend.__init__``). +""" + +from __future__ import annotations + +import logging + +import pytest +from pydantic import ValidationError + +from nemo_retriever._agentic.nemo_agent.llm import ( + BaseLLMConfig, + CallableLLMBackend, + CallableLLMConfig, + LiteLLMBackend, + LiteLLMConfig, + create_llm, + create_llm_config, + get_available_backends, +) + +#: Per-backend kwargs needed on top of ``model`` to build a valid config. Every +#: registered backend is currently satisfied by ``model`` alone — ``callable`` +#: leaves ``base_url`` optional because an in-process callable has no endpoint — +#: but the indirection stays so adding a backend with a required field is a +#: one-line change here rather than an edit at every call site. +_MINIMAL_KWARGS: dict[str, dict[str, str]] = {} + + +def _factory_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + """Rendered warning messages emitted by the factory itself.""" + return [r.getMessage() for r in caplog.records if r.getMessage().startswith("create_llm_config:")] + + +class TestBackendSelection: + def test_litellm_selected(self): + config = create_llm_config("litellm", model="m") + assert isinstance(config, LiteLLMConfig) + assert config.backend == "litellm" + assert config.model == "m" + + def test_unknown_backend_raises_valueerror(self): + with pytest.raises(ValueError) as exc: + create_llm_config("bogus", model="m") + message = str(exc.value) + assert "bogus" in message + # Message lists the available backends. + assert "litellm" in message + + def test_get_available_backends_reflects_registry(self): + # Public accessor over the registry: the wider library reads this instead + # of hard-coding a backend list, so new registrations propagate for free. + backends = get_available_backends() + assert isinstance(backends, tuple) + assert set(backends) == {"callable", "litellm"} + # Sorted, so callers can build stable, deterministic messages/choices. + assert list(backends) == sorted(backends) + # Every advertised name is actually selectable by the config factory. + for name in backends: + config = create_llm_config(name, model="m", **_MINIMAL_KWARGS.get(name, {})) + assert config.backend == name + + def test_create_llm_rejects_bare_base_config(self): + # BaseLLMConfig.backend defaults to "callable", but every backend requires + # its own config subclass, so a bare base config must still fail fast — at + # the config-type check rather than the registry lookup. + with pytest.raises(TypeError, match="CallableLLMConfig"): + create_llm(BaseLLMConfig(model="x"), completion_fn=lambda **kw: {}) + + def test_base_config_default_backend_is_registered(self): + # The declared library default must actually resolve — a default naming an + # unregistered backend would make the declaration a guaranteed ValueError. + assert BaseLLMConfig.model_fields["backend"].default in get_available_backends() + + +class TestDropAndWarn: + def test_unsupported_field_dropped_with_warning(self, caplog): + # drop_params is a LiteLLM-only field; callable (base-only) does not have it. + with caplog.at_level(logging.WARNING): + config = create_llm_config("callable", model="m", drop_params=True) + assert not hasattr(config, "drop_params") + assert any("drop_params" in w for w in _factory_warnings(caplog)) + + def test_warns_even_when_dropped_value_is_none(self, caplog): + # None is a valid value a caller may intend to set; a drop is still a drop. + with caplog.at_level(logging.WARNING): + create_llm_config("callable", model="m", cache_control=None) + assert any("cache_control" in w for w in _factory_warnings(caplog)) + + def test_supported_field_kept_no_warning(self, caplog): + # Pass the NON-default value: drop_params defaults to True, so asserting + # True here would pass even if the kwarg were silently ignored. + with caplog.at_level(logging.WARNING): + config = create_llm_config("litellm", model="m", drop_params=False) + assert config.drop_params is False + assert _factory_warnings(caplog) == [] + + def test_litellm_drop_params_defaults_true(self): + # The operators rely on this default rather than passing it explicitly — + # passing it would warn on every non-litellm backend, which has no such field. + assert create_llm_config("litellm", model="m").drop_params is True + + def test_warning_lists_keys_not_values(self, caplog): + # Values (potential secrets) must never be logged — only field names. + with caplog.at_level(logging.WARNING): + create_llm_config("callable", model="m", prompt_cache_key="SENSITIVE-VALUE") + joined = " ".join(_factory_warnings(caplog)) + assert "prompt_cache_key" in joined + assert "SENSITIVE-VALUE" not in joined + + +class TestValidationPreserved: + def test_missing_required_field_raises(self): + with pytest.raises(ValidationError): + create_llm_config("litellm") # no model + + def test_bad_type_for_known_field_raises(self): + with pytest.raises(ValidationError): + create_llm_config("litellm", model="m", max_completion_tokens="not-an-int") + + +class TestReasoningEffortHoisted: + @pytest.mark.parametrize("backend", ["litellm", "callable"]) + def test_reasoning_effort_supported_on_every_backend(self, backend, caplog): + with caplog.at_level(logging.WARNING): + config = create_llm_config(backend, model="m", reasoning_effort="high", **_MINIMAL_KWARGS.get(backend, {})) + assert config.reasoning_effort == "high" + # It is a base field now, so it is never dropped and never warned about. + assert not any("reasoning_effort" in w for w in _factory_warnings(caplog)) + + def test_reasoning_effort_declared_on_base_config(self): + assert "reasoning_effort" in BaseLLMConfig.model_fields + + +class TestConfigClsPairing: + def test_litellm_pairing(self): + assert LiteLLMBackend.config_cls is LiteLLMConfig + + def test_callable_pairing(self): + assert CallableLLMBackend.config_cls is CallableLLMConfig + + def test_wrong_config_type_rejected_before_litellm_import(self): + # Passing another backend's config to LiteLLMBackend must raise TypeError + # from the centralized base check, which runs before the lazy + # `import litellm`. + with pytest.raises(TypeError): + LiteLLMBackend(CallableLLMConfig(model="m")) + + +class TestLiteLLMCompletionKwargs: + """``temperature`` / ``parallel_tool_calls`` reach ``completion_kwargs`` only when set. + + Constructs a real ``LiteLLMBackend``, which imports litellm; skipped when + litellm is not installed. + """ + + def test_forwarded_when_set(self): + pytest.importorskip("litellm") + backend = LiteLLMBackend(LiteLLMConfig(model="gpt-4o-mini", temperature=0.3, parallel_tool_calls=True)) + assert backend.completion_kwargs["temperature"] == 0.3 + assert backend.completion_kwargs["parallel_tool_calls"] is True + + def test_omitted_when_none(self): + pytest.importorskip("litellm") + backend = LiteLLMBackend(LiteLLMConfig(model="gpt-4o-mini")) + assert "temperature" not in backend.completion_kwargs + assert "parallel_tool_calls" not in backend.completion_kwargs + + def test_parallel_tool_calls_dropped_when_no_tools(self): + pytest.importorskip("litellm") + backend = LiteLLMBackend(LiteLLMConfig(model="gpt-4o-mini", parallel_tool_calls=True)) + messages = [{"role": "user", "content": "hi"}] + no_tools = backend._prepare_request(messages, tools=None) + assert "parallel_tool_calls" not in no_tools + assert "tool_choice" not in no_tools + with_tools = backend._prepare_request( + messages, + tools=[{"type": "function", "function": {"name": "f", "parameters": {}}}], + ) + assert with_tools["parallel_tool_calls"] is True + + +class TestFalsyValuesPreserved: + """0.0 / False are real settings, not "unset" — the factory must keep them.""" + + def test_zero_temperature_and_false_parallel_tool_calls_kept(self): + config = create_llm_config("litellm", model="m", temperature=0.0, parallel_tool_calls=False) + assert config.temperature == 0.0 + assert config.parallel_tool_calls is False + + +class TestLiteLLMBaseUrlNormalization: + """LiteLLMConfig strips a trailing /chat/completions (litellm appends it itself).""" + + def test_strips_chat_completions_suffix(self): + c = LiteLLMConfig(model="m", base_url="https://integrate.api.nvidia.com/v1/chat/completions") + assert c.base_url == "https://integrate.api.nvidia.com/v1" + + def test_leaves_clean_base_untouched(self): + c = LiteLLMConfig(model="m", base_url="https://integrate.api.nvidia.com/v1") + assert c.base_url == "https://integrate.api.nvidia.com/v1" + + def test_none_is_preserved(self): + assert LiteLLMConfig(model="m", base_url=None).base_url is None + + def test_tolerates_trailing_slash(self): + assert LiteLLMConfig(model="m", base_url="https://x/v1/chat/completions/").base_url == "https://x/v1" + + def test_idempotent(self): + once = LiteLLMConfig(model="m", base_url="https://x/v1/chat/completions").base_url + twice = LiteLLMConfig(model="m", base_url=once).base_url + assert once == twice == "https://x/v1" + + def test_applied_via_create_llm_config(self): + c = create_llm_config("litellm", model="m", base_url="https://x/v1/chat/completions") + assert c.base_url == "https://x/v1" + + def test_callable_base_url_not_normalized(self): + # Normalization is litellm-specific; other backends keep the URL verbatim. + c = create_llm_config("callable", model="m", base_url="https://x/v1/chat/completions") + assert c.base_url == "https://x/v1/chat/completions" + + +class TestCallableBaseUrl: + """base_url is the callable's endpoint: optional, and used exactly as given.""" + + def test_optional_because_an_in_process_callable_has_no_endpoint(self): + # Neither the config nor the backend can tell a remote callable from an + # in-process one, so this cannot be required. A remote callable handed + # None fails on its first call with its own error. + assert CallableLLMConfig(model="m").base_url is None + + def test_used_verbatim(self): + # The opposite of LiteLLMConfig, which strips the suffix because litellm + # re-appends it. Here the value is forwarded as-is. + url = "https://integrate.api.nvidia.com/v1/chat/completions" + assert CallableLLMConfig(model="m", base_url=url).base_url == url + + def test_no_trailing_slash_cleanup(self): + url = "https://x/v1/chat/completions/" + assert CallableLLMConfig(model="m", base_url=url).base_url == url + + def test_non_positive_retry_budget_rejected(self): + # 0 would make a `while attempt < max_retries` client issue zero requests + # and then report retries-exhausted, which reads as an endpoint failure. + with pytest.raises(ValidationError): + CallableLLMConfig(model="m", max_retries=0) diff --git a/nemo_retriever/tests/test_nemo_agent_llm_helpers.py b/nemo_retriever/tests/test_nemo_agent_llm_helpers.py new file mode 100644 index 0000000000..ba5eb0a641 --- /dev/null +++ b/nemo_retriever/tests/test_nemo_agent_llm_helpers.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared helpers used by the LLM backends. + +Redaction and key resolution are security-adjacent: a regression here leaks a +credential into a log or an agent trajectory rather than failing a test loudly. +""" + +from __future__ import annotations + +import pytest + +from nemo_retriever._agentic.nemo_agent.llm.helpers import ( + BODY_EXCERPT_CHARS, + excerpt, + redact_url, + redact_urls_in_text, + resolve_api_key, +) + + +class TestRedactUrl: + def test_drops_userinfo_query_and_fragment(self): + out = redact_url("https://user:pw@host:8443/v1/chat/completions?token=abc#frag") + assert out == "https://host:8443/v1/chat/completions" + + def test_keeps_host_and_path(self): + url = "https://integrate.api.nvidia.com/v1/chat/completions" + assert redact_url(url) == url + + def test_each_url_in_a_comma_separated_list_is_redacted(self): + # The endpoint option accepts a comma-separated list. A bare urlsplit would + # leave every URL after the first sitting in `.path` with its userinfo + # intact, so each segment has to be redacted independently. + out = redact_url("https://a:pw@h1/v1/chat/completions, https://b:pw2@h2/v1/chat/completions") + assert out == "https://h1/v1/chat/completions,https://h2/v1/chat/completions" + assert "pw" not in out + + @pytest.mark.parametrize("value", ["", None]) + def test_empty_input_does_not_raise(self, value): + assert redact_url(value) == "" + + def test_unparseable_input_degrades_instead_of_raising(self): + # Redaction runs on error paths; raising here would mask the real failure. + assert isinstance(redact_url("http://[unclosed"), str) + + +class TestRedactUrlsInText: + """Scrubbing URLs out of free text we did not format ourselves.""" + + def test_full_url_in_prose_loses_userinfo_and_query(self): + out = redact_urls_in_text("failed calling https://user:pw@h.invalid/v1/chat?api_key=SECRET after 3 tries") + assert "SECRET" not in out + assert "user:pw" not in out + assert "https://h.invalid/v1/chat" in out + assert out.endswith("after 3 tries") + + def test_bare_path_query_is_scrubbed(self): + # urllib3 connection errors carry no scheme or host, so the URL pattern + # never matches them — yet the query string, where a credential actually + # lives, is right there. This is the case that motivated the helper. + text = "HTTPSConnectionPool(host='h', port=443): Max retries exceeded with url: /v1/chat?api_key=SECRET" + out = redact_urls_in_text(text) + assert "SECRET" not in out + assert "?" in out + + def test_trailing_sentence_punctuation_is_not_swallowed(self): + out = redact_urls_in_text("could not reach https://h.invalid/v1?k=SECRET. Retrying.") + assert "SECRET" not in out + assert out.endswith(". Retrying.") + + def test_prose_question_mark_is_left_alone(self): + # The `=` requirement is what keeps ordinary prose intact. + text = "Did the endpoint respond? No route to host." + assert redact_urls_in_text(text) == text + + def test_multiple_urls_all_scrubbed(self): + out = redact_urls_in_text("tried https://a.invalid/v1?k=S1 then https://b.invalid/v1?k=S2") + assert "S1" not in out and "S2" not in out + + def test_text_without_urls_is_unchanged(self): + text = "maximum context length is 8192 tokens" + assert redact_urls_in_text(text) == text + + @pytest.mark.parametrize("value", ["", None]) + def test_empty_input(self, value): + assert redact_urls_in_text(value) == "" + + +class TestExcerpt: + def test_short_text_passes_through_unchanged(self): + assert excerpt("boom") == "boom" + + def test_long_text_is_truncated_with_an_ellipsis(self): + out = excerpt("x" * (BODY_EXCERPT_CHARS + 500)) + assert len(out) == BODY_EXCERPT_CHARS + 3 + assert out.endswith("...") + + def test_exactly_at_the_limit_is_not_truncated(self): + out = excerpt("x" * BODY_EXCERPT_CHARS) + assert out == "x" * BODY_EXCERPT_CHARS + + @pytest.mark.parametrize("value", [None, ""]) + def test_empty_becomes_an_empty_string(self, value): + assert excerpt(value) == "" + + def test_non_string_is_coerced(self): + assert excerpt(404) == "404" + + +class TestResolveApiKey: + def test_literal_key_passes_through(self): + assert resolve_api_key("sk-literal") == "sk-literal" + + def test_env_indirection_is_followed_and_stripped(self, monkeypatch): + monkeypatch.setenv("HELPER_TEST_KEY", " sk-from-env\n") + assert resolve_api_key("os.environ/HELPER_TEST_KEY") == "sk-from-env" + + def test_whitespace_around_the_indirection_is_tolerated(self, monkeypatch): + monkeypatch.setenv("HELPER_TEST_KEY", "sk-from-env") + assert resolve_api_key(" os.environ/HELPER_TEST_KEY ") == "sk-from-env" + + @pytest.mark.parametrize("value", [None, "", " "]) + def test_absent_key_becomes_an_empty_string(self, value): + assert resolve_api_key(value) == "" + + def test_missing_variable_names_the_variable_and_how_to_set_it(self, monkeypatch): + monkeypatch.delenv("HELPER_TEST_MISSING", raising=False) + with pytest.raises(ValueError) as excinfo: + resolve_api_key("os.environ/HELPER_TEST_MISSING") + message = str(excinfo.value) + assert "HELPER_TEST_MISSING" in message + assert "export HELPER_TEST_MISSING=" in message + + def test_missing_variable_does_not_chain_a_confusing_key_error(self, monkeypatch): + # `from None`: a KeyError on __cause__ reads like an internal bug rather + # than a configuration problem the caller can fix. + monkeypatch.delenv("HELPER_TEST_MISSING", raising=False) + with pytest.raises(ValueError) as excinfo: + resolve_api_key("os.environ/HELPER_TEST_MISSING") + assert excinfo.value.__cause__ is None diff --git a/nemo_retriever/tests/test_nim_chat_completion_step.py b/nemo_retriever/tests/test_nim_chat_completion_step.py new file mode 100644 index 0000000000..b5b91d7374 --- /dev/null +++ b/nemo_retriever/tests/test_nim_chat_completion_step.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Payload shaping for ``invoke_chat_completion_step``. + +This is the completion callable the agent operators inject for remote runs, so +its request body is now an agent-facing contract. The ``temperature`` handling in +particular is load-bearing: the agent forwards ``None`` to mean *unset*, and +sending ``0.0`` instead would silently force greedy decoding on every remote run. +""" + +from __future__ import annotations + +import pytest + +from nemo_retriever.models.nim import chat_completions + +_URL = "https://endpoint.invalid/v1/chat/completions" +_MESSAGES = [{"role": "user", "content": "q"}] + + +@pytest.fixture +def sent(monkeypatch): + """Capture the kwargs handed to the underlying poster.""" + captured: dict = {} + + def fake_post(**kwargs): + captured.update(kwargs) + return {"choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]} + + monkeypatch.setattr(chat_completions, "_post_with_retries", fake_post) + return captured + + +class TestTemperature: + def test_none_omits_the_field_entirely(self, sent): + # "Unset" must mean the endpoint/model default applies. Sending 0.0 would + # force greedy decoding the caller never asked for. + chat_completions.invoke_chat_completion_step(invoke_url=_URL, messages=_MESSAGES, temperature=None) + assert "temperature" not in sent["payload"] + + def test_explicit_zero_is_still_sent(self, sent): + # 0.0 is a real, meaningful value (greedy) and must not be confused with unset. + chat_completions.invoke_chat_completion_step(invoke_url=_URL, messages=_MESSAGES, temperature=0.0) + assert sent["payload"]["temperature"] == 0.0 + + def test_default_is_unchanged_for_existing_callers(self, sent): + chat_completions.invoke_chat_completion_step(invoke_url=_URL, messages=_MESSAGES) + assert sent["payload"]["temperature"] == 0.0 + + def test_extra_body_can_still_override(self, sent): + chat_completions.invoke_chat_completion_step( + invoke_url=_URL, messages=_MESSAGES, temperature=None, extra_body={"temperature": 0.7} + ) + assert sent["payload"]["temperature"] == 0.7 + + +class TestPayloadShape: + def test_tools_and_tool_choice_travel_together(self, sent): + tools = [{"type": "function", "function": {"name": "f", "parameters": {}}}] + chat_completions.invoke_chat_completion_step( + invoke_url=_URL, messages=_MESSAGES, tools=tools, tool_choice="auto" + ) + assert sent["payload"]["tools"] == tools + assert sent["payload"]["tool_choice"] == "auto" + + def test_tool_choice_is_dropped_without_tools(self, sent): + # Suppressing tools is expressed by omitting them, so a stray tool_choice + # on a tool-less request would be meaningless at best. + chat_completions.invoke_chat_completion_step(invoke_url=_URL, messages=_MESSAGES, tool_choice="none") + assert "tool_choice" not in sent["payload"] + assert "tools" not in sent["payload"] + + def test_max_tokens_omitted_when_none(self, sent): + chat_completions.invoke_chat_completion_step(invoke_url=_URL, messages=_MESSAGES, max_tokens=None) + assert "max_tokens" not in sent["payload"] + + def test_extra_body_is_merged_top_level(self, sent): + chat_completions.invoke_chat_completion_step( + invoke_url=_URL, messages=_MESSAGES, extra_body={"reasoning_effort": "high"} + ) + assert sent["payload"]["reasoning_effort"] == "high" + + def test_api_key_becomes_a_bearer_header(self, sent): + chat_completions.invoke_chat_completion_step(invoke_url=_URL, messages=_MESSAGES, api_key=" sk-live ") + assert sent["headers"]["Authorization"] == "Bearer sk-live" + + def test_no_authorization_header_without_a_key(self, sent): + chat_completions.invoke_chat_completion_step(invoke_url=_URL, messages=_MESSAGES, api_key=None) + assert "Authorization" not in sent["headers"] + + def test_retry_budget_is_forwarded(self, sent): + chat_completions.invoke_chat_completion_step( + invoke_url=_URL, messages=_MESSAGES, timeout_s=30.0, max_retries=2, max_429_retries=4 + ) + assert (sent["timeout_s"], sent["max_retries"], sent["max_429_retries"]) == (30.0, 2, 4) + + def test_first_url_of_a_comma_separated_list_is_used(self, sent): + chat_completions.invoke_chat_completion_step( + invoke_url=f"{_URL},https://second.invalid/v1/chat/completions", messages=_MESSAGES + ) + assert sent["invoke_url"] == _URL + + def test_empty_invoke_url_is_rejected(self): + with pytest.raises(ValueError, match="invoke_url is required"): + chat_completions.invoke_chat_completion_step(invoke_url="", messages=_MESSAGES) diff --git a/nemo_retriever/tests/test_nv_ingest_vdb_operator.py b/nemo_retriever/tests/test_nv_ingest_vdb_operator.py index b7fd63080f..d39fddd418 100644 --- a/nemo_retriever/tests/test_nv_ingest_vdb_operator.py +++ b/nemo_retriever/tests/test_nv_ingest_vdb_operator.py @@ -10,18 +10,43 @@ import pandas as pd import pytest -from nemo_retriever.common.vdb.adt_vdb import VDB +from nemo_retriever.common.vdb.adt_vdb import ( + CollectionWriteContext, + CollectionWriteResult, + VDB, +) +from nemo_retriever.common.vdb.records import RetrievalContractError from nemo_retriever.operators.vdb import IngestVdbOperator, RetrieveVdbOperator from nemo_retriever.operators import vdb as vdb_operator_module from nemo_retriever.operators.vdb import PutVdbOperator -class FakeVDB(VDB): +class _CollectionContractStub(VDB): + """Implement required collection methods that these operator tests do not exercise.""" + + def _unexpected_collection_operation(self, *args: Any, **kwargs: Any) -> Any: + raise AssertionError("Unexpected collection operation") + + create_collection = _unexpected_collection_operation + get_collection = _unexpected_collection_operation + list_collections = _unexpected_collection_operation + update_collection = _unexpected_collection_operation + delete_collection = _unexpected_collection_operation + get_document = _unexpected_collection_operation + list_documents = _unexpected_collection_operation + delete_document = _unexpected_collection_operation + write_collection = _unexpected_collection_operation + retrieve_collection = _unexpected_collection_operation + + +class FakeVDB(_CollectionContractStub): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self.run_calls: list[Any] = [] self.retrieval_calls: list[tuple[Any, dict[str, Any]]] = [] self.put_calls: list[tuple[Any, dict[str, Any]]] = [] + self.write_collection_calls: list[tuple[Any, CollectionWriteContext]] = [] + self.retrieve_collection_calls: list[tuple[Any, dict[str, Any]]] = [] def create_index(self, **kwargs: Any) -> None: return None @@ -55,6 +80,58 @@ def put(self, records: list, **kwargs: Any) -> dict[str, Any]: self.put_calls.append((records, dict(kwargs))) return {"put": sum(len(b) for b in records)} + def write_collection(self, records: list, *, context: CollectionWriteContext) -> CollectionWriteResult: + self.write_collection_calls.append((records, context)) + return CollectionWriteResult(written=sum(len(batch) for batch in records), total_rows=7) + + def retrieve_collection( + self, + vectors: list, + *, + scope: str, + collection_name: str, + query_texts: list[str], + top_k: int, + **kwargs: Any, + ) -> tuple[list[list[dict[str, Any]]], list[str]]: + call_kwargs = { + "scope": scope, + "collection_name": collection_name, + "query_texts": query_texts, + "top_k": top_k, + **kwargs, + } + self.retrieve_collection_calls.append((vectors, call_kwargs)) + return [ + [ + { + "chunk_id": "chunk-1", + "document_id": "document-1", + "text": "retrieved chunk", + "distance": 0.12, + "filename": "doc-a.pdf", + "page_number": 1, + "content_type": "table", + "source": "doc-a.pdf", + "source_id": "doc-a.pdf", + "stored_image_uri": "file:///tmp/page.png", + "bbox": [0, 0, 1, 1], + "metadata": {}, + "physical_table": "private-table", + "lancedb_uri": "/private/vector-store", + } + ] + ], ["dense"] + + +class InvalidCollectionVDB(FakeVDB): + def __init__(self, result: Any) -> None: + super().__init__() + self.result = result + + def retrieve_collection(self, vectors: list, **kwargs: Any) -> Any: + return self.result + def _graph_rows() -> list[dict[str, Any]]: return [ @@ -146,7 +223,9 @@ def test_ingest_operator_converts_graph_rows_to_client_vdb_records() -> None: "text", [pytest.param("", id="empty"), pytest.param(" \n\t ", id="whitespace")], ) -def test_ingest_operator_retains_embedded_blank_image_row_without_text_fidelity(text: str) -> None: +def test_ingest_operator_retains_embedded_blank_image_row_without_text_fidelity( + text: str, +) -> None: vdb = FakeVDB() operator = IngestVdbOperator(vdb=vdb) data = [ @@ -184,7 +263,9 @@ def test_ingest_operator_retains_embedded_blank_image_row_without_text_fidelity( pytest.param(np.ones((2, 2), dtype=np.uint8), id="numpy"), ], ) -def test_ingest_operator_noncanonical_image_payload_fails_closed_without_truthiness(image_payload: Any) -> None: +def test_ingest_operator_noncanonical_image_payload_fails_closed_without_truthiness( + image_payload: Any, +) -> None: vdb = FakeVDB() operator = IngestVdbOperator(vdb=vdb) data = [ @@ -202,7 +283,9 @@ def test_ingest_operator_noncanonical_image_payload_fails_closed_without_truthin @pytest.mark.parametrize("uri_field", ["_stored_image_uri", "stored_image_uri"]) -def test_ingest_operator_retains_image_only_row_with_stored_image_uri(uri_field: str) -> None: +def test_ingest_operator_retains_image_only_row_with_stored_image_uri( + uri_field: str, +) -> None: vdb = FakeVDB() operator = IngestVdbOperator(vdb=vdb) data = [ @@ -282,14 +365,36 @@ def test_retrieve_operator_delegates_vectors_to_retrieval() -> None: } ] ] - assert vdb.retrieval_calls == [([[0.1, 0.2]], {"collection_name": "docs", "model_name": "embedder", "top_k": 3})] + assert vdb.retrieval_calls == [ + ( + [[0.1, 0.2]], + {"collection_name": "docs", "model_name": "embedder", "top_k": 3}, + ) + ] + + +def test_retrieve_operator_reads_index_metadata_from_any_vdb() -> None: + class MetadataVDB(FakeVDB): + def get_index_metadata(self, key: str, **kwargs: Any) -> str | None: + assert kwargs == {"collection_name": "docs"} + return {"embedding_model_name": "acme/embed", "retrieval_mode": "dense"}.get(key) + + operator = RetrieveVdbOperator(vdb=MetadataVDB(), vdb_kwargs={"collection_name": "docs"}) + + assert operator.get_index_metadata("embedding_model_name") == "acme/embed" + assert operator.get_index_metadata("retrieval_mode") == "dense" def test_retrieve_operator_forwards_runtime_query_texts() -> None: vdb = FakeVDB() operator = RetrieveVdbOperator( vdb=vdb, - vdb_kwargs={"collection_name": "docs", "model_name": "embedder", "hybrid": True, "query_texts": ["stale"]}, + vdb_kwargs={ + "collection_name": "docs", + "model_name": "embedder", + "hybrid": True, + "query_texts": ["stale"], + }, ) operator.process([[0.1, 0.2]], top_k=3, query_texts=["current"]) @@ -332,7 +437,12 @@ def test_retrieve_operator_does_not_forward_query_texts_for_dense_retrieval() -> operator.process([[0.1, 0.2]], top_k=3, query_texts=["current"]) - assert vdb.retrieval_calls == [([[0.1, 0.2]], {"collection_name": "docs", "model_name": "embedder", "top_k": 3})] + assert vdb.retrieval_calls == [ + ( + [[0.1, 0.2]], + {"collection_name": "docs", "model_name": "embedder", "top_k": 3}, + ) + ] def test_constructor_requires_exactly_one_vdb_source() -> None: @@ -348,7 +458,7 @@ def test_constructor_requires_exactly_one_vdb_source() -> None: # ────────────────────────────────────────────────────────────────────────────── -class _StubPutVDB(VDB): +class _StubPutVDB(_CollectionContractStub): """VDB subclass that intentionally does NOT override ``put``. Used to exercise the construction-time guard in @@ -469,3 +579,112 @@ def test_put_operator_merges_sidecar_metadata_into_records_before_put() -> None: # Sidecar column merged in alongside the per-row ``page_number``. assert merged_content_meta["category"] == "legal" assert merged_content_meta["page_number"] == 7 + + +def test_ingest_operator_preserves_canonical_batches_for_collection_write() -> None: + vdb = FakeVDB() + operator = IngestVdbOperator(vdb=vdb) + records = [ + [ + { + "document_type": "text", + "metadata": { + "embedding": [0.1, 0.2], + "content": "canonical chunk", + "content_metadata": {"page_number": 9}, + "source_metadata": {"source_id": "/tmp/canonical.pdf"}, + }, + } + ] + ] + context = CollectionWriteContext( + scope="tenant-a", + collection_name="papers", + document_id="document-1", + document_version="version-1", + content_sha256="a" * 64, + filename="canonical.pdf", + job_id="job-1", + operation="replace", + ) + + result = operator.process(records, collection_context=context) + + assert result == CollectionWriteResult(written=1, total_rows=7) + assert vdb.write_collection_calls == [(records, context)] + assert vdb.write_collection_calls[0][0] is records + assert vdb.run_calls == [] + + +def test_ingest_operator_rejects_empty_collection_write() -> None: + operator = IngestVdbOperator(vdb=FakeVDB()) + context = CollectionWriteContext( + scope="tenant-a", + collection_name="papers", + document_id="document-1", + document_version="version-1", + content_sha256="a" * 64, + filename="empty.pdf", + ) + + with pytest.raises(ValueError, match="at least one canonical VDB record"): + operator.process([], collection_context=context) + + +def test_retrieve_operator_dispatches_explicit_collection_context() -> None: + vdb = FakeVDB() + operator = RetrieveVdbOperator(vdb=vdb, vdb_kwargs={"model_name": "embedder"}) + + hits, strategies = operator.process( + [[0.1, 0.2]], + scope="tenant-a", + collection_name="papers", + top_k=3, + query_texts=["current"], + ) + + assert strategies == ["dense"] + assert hits[0][0]["text"] == "retrieved chunk" + assert "physical_table" not in hits[0][0] + assert "lancedb_uri" not in hits[0][0] + assert vdb.retrieve_collection_calls == [ + ( + [[0.1, 0.2]], + { + "scope": "tenant-a", + "collection_name": "papers", + "query_texts": ["current"], + "top_k": 3, + "model_name": "embedder", + }, + ) + ] + + +@pytest.mark.parametrize( + "result", + [ + None, + ([], ["dense"]), + ([[42]], ["dense"]), + ([[{"text": "missing collection identity"}]], ["dense"]), + ([[{"chunk_id": "chunk-1"}]], []), + ], +) +def test_retrieve_operator_rejects_malformed_collection_results(result: Any) -> None: + operator = RetrieveVdbOperator(vdb=InvalidCollectionVDB(result)) + + with pytest.raises(RetrievalContractError): + operator.process( + [[0.1, 0.2]], + scope="tenant-a", + collection_name="papers", + query_texts=["query"], + ) + + +def test_retrieve_operator_rejects_partial_collection_context() -> None: + operator = RetrieveVdbOperator(vdb=FakeVDB()) + + with pytest.raises(ValueError, match="both scope and collection_name"): + operator.process([[0.1, 0.2]], scope="tenant-a") diff --git a/nemo_retriever/tests/test_ocr_cross_page_batching.py b/nemo_retriever/tests/test_ocr_cross_page_batching.py new file mode 100644 index 0000000000..0c48bbd555 --- /dev/null +++ b/nemo_retriever/tests/test_ocr_cross_page_batching.py @@ -0,0 +1,418 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Correctness tests for local OCR crop batching across page rows.""" + +from __future__ import annotations + +import base64 +import io +from typing import Any + +import pandas as pd +from PIL import Image + +from nemo_retriever.common.params import RemoteRetryParams +from nemo_retriever.operators.extract.ocr.gpu_ocr import OCRActor + + +def _page_png_b64(crop_id: int) -> str: + image = Image.new("RGB", (16, 16), color=(crop_id, crop_id, crop_id)) + buffer = io.BytesIO() + image.save(buffer, format="PNG") + return base64.b64encode(buffer.getvalue()).decode("ascii") + + +def _detection(detection_id: str, label_name: str, bbox: list[float]) -> dict[str, Any]: + return { + "detection_id": detection_id, + "label_name": label_name, + "bbox_xyxy_norm": bbox, + } + + +def _page(page_id: str, crop_id: int, detections: list[dict[str, Any]]) -> dict[str, Any]: + return { + "page_id": page_id, + "metadata": {}, + "page_image": {"image_b64": _page_png_b64(crop_id)}, + "page_elements_v3": {"detections": detections}, + } + + +def _ocr_prediction(crop_id: int) -> list[dict[str, Any]]: + return [ + { + "left": 0.0, + "right": 1.0, + "upper": 0.0, + "lower": 1.0, + "text": f"crop-{crop_id}", + } + ] + + +class _RecordingListModel: + """Record the list-input contract used by the persistent local model.""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def invoke(self, crops: Any, *, merge_level: str) -> list[list[dict[str, Any]]]: + assert isinstance(crops, list) + crop_ids = [int(crop[0, 0, 0]) for crop in crops] + self.calls.append( + { + "crop_count": len(crops), + "crop_ids": crop_ids, + "merge_level": merge_level, + } + ) + return [_ocr_prediction(crop_id) for crop_id in crop_ids] + + +class _FallbackIsolationModel: + """Force batched failure, then fail one crop during per-item isolation.""" + + def __init__(self, failing_crop_id: int) -> None: + self.failing_crop_id = failing_crop_id + self.calls: list[dict[str, Any]] = [] + + def invoke(self, crops: Any, *, merge_level: str) -> Any: + is_batch = isinstance(crops, list) + crop_arrays = crops if is_batch else [crops] + crop_ids = [int(crop[0, 0, 0]) for crop in crop_arrays] + self.calls.append( + { + "input_kind": "batch" if is_batch else "single", + "crop_ids": crop_ids, + "merge_level": merge_level, + } + ) + if is_batch: + raise RuntimeError("force per-crop fallback") + if crop_ids[0] == self.failing_crop_id: + raise ValueError(f"bad crop {self.failing_crop_id}") + return _ocr_prediction(crop_ids[0]) + + +class _WrongCountFallbackModel: + """Return too few batched results, then succeed during per-item fallback.""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def invoke(self, crops: Any, *, merge_level: str) -> Any: + is_batch = isinstance(crops, list) + crop_arrays = crops if is_batch else [crops] + crop_ids = [int(crop[0, 0, 0]) for crop in crop_arrays] + self.calls.append( + { + "input_kind": "batch" if is_batch else "single", + "crop_ids": crop_ids, + "merge_level": merge_level, + } + ) + if is_batch: + return [_ocr_prediction(crop_ids[0])] + return _ocr_prediction(crop_ids[0]) + + +def _local_actor( + model: Any, + *, + inference_batch_size: int, + extract_text: bool = False, + extract_tables: bool = False, + extract_charts: bool = False, + extract_infographics: bool = False, +) -> OCRActor: + actor = object.__new__(OCRActor) + actor._graph_init_kwargs = {} + actor.ocr_kwargs = { + "extract_text": extract_text, + "extract_tables": extract_tables, + "extract_charts": extract_charts, + "extract_infographics": extract_infographics, + "use_table_structure": False, + "request_timeout_s": 120.0, + "inference_batch_size": inference_batch_size, + } + actor._remote_retry = RemoteRetryParams() + actor._model = model + actor._nim_client = None + return actor + + +def test_local_actor_batches_compatible_crops_across_page_rows() -> None: + model = _RecordingListModel() + actor = _local_actor(model, inference_batch_size=2, extract_charts=True) + chart_bboxes = [[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 1.0, 1.0]] + batch = pd.DataFrame( + [ + _page("page-A", 11, [_detection("chart-A", "chart", chart_bboxes[0])]), + _page("page-B", 22, [_detection("chart-B", "chart", chart_bboxes[1])]), + ] + ) + + result = actor(batch) + + assert { + "model_calls": model.calls, + "rows": [{"page_id": row.page_id, "chart": row.chart} for row in result.itertuples(index=False)], + } == { + "model_calls": [ + { + "crop_count": 2, + "crop_ids": [11, 22], + "merge_level": "paragraph", + } + ], + "rows": [ + { + "page_id": "page-A", + "chart": [{"bbox_xyxy_norm": chart_bboxes[0], "text": "crop-11"}], + }, + { + "page_id": "page-B", + "chart": [{"bbox_xyxy_norm": chart_bboxes[1], "text": "crop-22"}], + }, + ], + } + + +def test_local_actor_separates_merge_levels_chunks_and_preserves_detection_order() -> None: + model = _RecordingListModel() + actor = _local_actor( + model, + inference_batch_size=2, + extract_text=True, + extract_tables=True, + extract_charts=True, + extract_infographics=True, + ) + table_a = [0.0, 0.0, 0.4, 0.4] + chart_a = [0.5, 0.0, 1.0, 0.4] + title_a = [0.0, 0.5, 1.0, 0.7] + table_b1 = [0.0, 0.0, 0.3, 0.3] + table_b2 = [0.35, 0.0, 0.65, 0.3] + infographic_b = [0.0, 0.5, 1.0, 1.0] + page_a = _page( + "page-A", + 11, + [ + _detection("table-A", "table", table_a), + _detection("chart-A", "chart", chart_a), + _detection("title-A", "title", title_a), + ], + ) + page_a["metadata"] = {"needs_ocr_for_text": True} + page_a["text"] = "native-A" + page_b = _page( + "page-B", + 22, + [ + _detection("table-B1", "table", table_b1), + _detection("table-B2", "table", table_b2), + _detection("infographic-B", "infographic", infographic_b), + ], + ) + page_b["text"] = "native-B" + empty_page = _page("page-empty", 33, []) + empty_page["text"] = "native-empty" + batch = pd.DataFrame( + [ + page_a, + page_b, + empty_page, + { + "page_id": "page-malformed", + "metadata": {"error": {"stage": "page_render", "message": "missing raster"}}, + "page_image": None, + "page_elements_v3": {"detections": [_detection("chart-malformed", "chart", [0.0, 0.0, 1.0, 1.0])]}, + "text": "native-malformed", + }, + ] + ) + + result = actor(batch) + + assert { + "model_calls": model.calls, + "page_A": { + "table_bboxes": [entry["bbox_xyxy_norm"] for entry in result.at[0, "table"]], + "chart": result.at[0, "chart"], + "text": result.at[0, "text"], + "metadata": { + "num_detections": result.at[0, "ocr"]["num_detections"], + "counts_by_label": result.at[0, "ocr"]["counts_by_label"], + }, + }, + "page_B": { + "table_bboxes": [entry["bbox_xyxy_norm"] for entry in result.at[1, "table"]], + "infographic": result.at[1, "infographic"], + "text": result.at[1, "text"], + "metadata": { + "num_detections": result.at[1, "ocr"]["num_detections"], + "counts_by_label": result.at[1, "ocr"]["counts_by_label"], + }, + }, + "empty_page": { + "table": result.at[2, "table"], + "chart": result.at[2, "chart"], + "infographic": result.at[2, "infographic"], + "text": result.at[2, "text"], + "num_detections": result.at[2, "ocr"]["num_detections"], + }, + "malformed_page": { + "table": result.at[3, "table"], + "chart": result.at[3, "chart"], + "infographic": result.at[3, "infographic"], + "text": result.at[3, "text"], + "error": result.at[3, "ocr"]["error"], + "num_detections": result.at[3, "ocr"]["num_detections"], + }, + } == { + "model_calls": [ + {"crop_count": 2, "crop_ids": [11, 22], "merge_level": "word"}, + {"crop_count": 1, "crop_ids": [22], "merge_level": "word"}, + {"crop_count": 2, "crop_ids": [11, 11], "merge_level": "paragraph"}, + {"crop_count": 1, "crop_ids": [22], "merge_level": "paragraph"}, + ], + "page_A": { + "table_bboxes": [table_a], + "chart": [{"bbox_xyxy_norm": chart_a, "text": "crop-11"}], + "text": "crop-11", + "metadata": { + "num_detections": 3, + "counts_by_label": {"table": 1, "chart": 1, "text": 1}, + }, + }, + "page_B": { + "table_bboxes": [table_b1, table_b2], + "infographic": [{"bbox_xyxy_norm": infographic_b, "text": "crop-22"}], + "text": "native-B", + "metadata": { + "num_detections": 3, + "counts_by_label": {"table": 2, "infographic": 1}, + }, + }, + "empty_page": { + "table": [], + "chart": [], + "infographic": [], + "text": "native-empty", + "num_detections": 0, + }, + "malformed_page": { + "table": [], + "chart": [], + "infographic": [], + "text": "native-malformed", + "error": {"stage": "page_render", "message": "missing raster"}, + "num_detections": 0, + }, + } + + +def test_local_actor_isolates_one_failed_crop_to_its_source_row() -> None: + model = _FallbackIsolationModel(failing_crop_id=22) + actor = _local_actor(model, inference_batch_size=3, extract_charts=True) + bbox = [0.0, 0.0, 1.0, 1.0] + batch = pd.DataFrame( + [ + _page("page-A", 11, [_detection("chart-A", "chart", bbox)]), + _page("page-B", 22, [_detection("chart-B", "chart", bbox)]), + _page("page-C", 33, [_detection("chart-C", "chart", bbox)]), + ] + ) + + result = actor(batch) + + assert { + "model_calls": model.calls, + "charts": result["chart"].tolist(), + "errors": [meta["error"] for meta in result["ocr"]], + } == { + "model_calls": [ + { + "input_kind": "batch", + "crop_ids": [11, 22, 33], + "merge_level": "paragraph", + }, + { + "input_kind": "single", + "crop_ids": [11], + "merge_level": "paragraph", + }, + { + "input_kind": "single", + "crop_ids": [22], + "merge_level": "paragraph", + }, + { + "input_kind": "single", + "crop_ids": [33], + "merge_level": "paragraph", + }, + ], + "charts": [ + [{"bbox_xyxy_norm": bbox, "text": "crop-11"}], + [], + [{"bbox_xyxy_norm": bbox, "text": "crop-33"}], + ], + "errors": [ + None, + { + "stage": "ocr_page_elements", + "type": "ValueError", + "message": "bad crop 22", + "traceback": result.at[1, "ocr"]["error"]["traceback"], + }, + None, + ], + } + + +def test_local_actor_falls_back_when_batch_result_count_is_wrong() -> None: + model = _WrongCountFallbackModel() + actor = _local_actor(model, inference_batch_size=2, extract_charts=True) + bbox = [0.0, 0.0, 1.0, 1.0] + batch = pd.DataFrame( + [ + _page("page-A", 11, [_detection("chart-A", "chart", bbox)]), + _page("page-B", 22, [_detection("chart-B", "chart", bbox)]), + ] + ) + + result = actor(batch) + + assert { + "model_calls": model.calls, + "charts": result["chart"].tolist(), + "errors": [meta["error"] for meta in result["ocr"]], + } == { + "model_calls": [ + { + "input_kind": "batch", + "crop_ids": [11, 22], + "merge_level": "paragraph", + }, + { + "input_kind": "single", + "crop_ids": [11], + "merge_level": "paragraph", + }, + { + "input_kind": "single", + "crop_ids": [22], + "merge_level": "paragraph", + }, + ], + "charts": [ + [{"bbox_xyxy_norm": bbox, "text": "crop-11"}], + [{"bbox_xyxy_norm": bbox, "text": "crop-22"}], + ], + "errors": [None, None], + } diff --git a/nemo_retriever/tests/test_operator_flags_and_cpu_actors.py b/nemo_retriever/tests/test_operator_flags_and_cpu_actors.py index 92d7195ca9..26f147a91f 100644 --- a/nemo_retriever/tests/test_operator_flags_and_cpu_actors.py +++ b/nemo_retriever/tests/test_operator_flags_and_cpu_actors.py @@ -322,7 +322,13 @@ def test_uses_default_invoke_url(self, mock_probe): from nemo_retriever.operators.embed.cpu_operator import _BatchEmbedCPUActor from nemo_retriever.common.params import EmbedParams - actor = _BatchEmbedCPUActor(params=EmbedParams(model_name="test-model", api_key="test-key")) + actor = _BatchEmbedCPUActor( + params=EmbedParams( + model_name="test-model", + embed_model_provider_prefix="nvidia", + api_key="test-key", + ) + ) assert actor._model is None assert "integrate.api.nvidia.com" in actor._kwargs["embedding_endpoint"] mock_probe.assert_called_once_with( @@ -331,7 +337,7 @@ def test_uses_default_invoke_url(self, mock_probe): prefix="_BatchEmbedCPUActor", api_key="test-key", post_url=actor.DEFAULT_EMBED_INVOKE_URL, - post_body={"input": [], "model": "test-model"}, + post_body={"input": [], "model": "nvidia/test-model"}, ) def test_default_hosted_endpoint_requires_api_key(self): diff --git a/nemo_retriever/tests/test_page_elements_postprocess.py b/nemo_retriever/tests/test_page_elements_postprocess.py new file mode 100644 index 0000000000..0a28550a57 --- /dev/null +++ b/nemo_retriever/tests/test_page_elements_postprocess.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from nemo_retriever.common.modality.page_elements.shared import ( + _apply_final_score_filter, + _apply_page_elements_v3_postprocess, +) + + +def test_rejected_structured_box_cannot_suppress_surviving_title() -> None: + raw_detections = [ + { + "bbox_xyxy_norm": [0.081876, 0.077799, 0.738124, 0.111654], + "label": 2, + "label_name": "title", + "score": 0.912088, + }, + { + "bbox_xyxy_norm": [0.056368, 0.004097, 0.965700, 0.972466], + "label": 0, + "label_name": "table", + "score": 0.052635, + }, + ] + + postprocessed = _apply_page_elements_v3_postprocess(raw_detections) + final_detections = _apply_final_score_filter(postprocessed) + + assert [detection["label_name"] for detection in final_detections] == ["title"] + + +def test_surviving_structured_box_can_still_absorb_title() -> None: + raw_detections = [ + { + "bbox_xyxy_norm": [0.081876, 0.077799, 0.738124, 0.111654], + "label": 2, + "label_name": "title", + "score": 0.912088, + }, + { + "bbox_xyxy_norm": [0.056368, 0.004097, 0.965700, 0.972466], + "label": 0, + "label_name": "table", + "score": 0.9, + }, + ] + + postprocessed = _apply_page_elements_v3_postprocess(raw_detections) + final_detections = _apply_final_score_filter(postprocessed) + + assert [detection["label_name"] for detection in final_detections] == ["table"] diff --git a/nemo_retriever/tests/test_params_models.py b/nemo_retriever/tests/test_params_models.py index 408655a24f..7332faf9db 100644 --- a/nemo_retriever/tests/test_params_models.py +++ b/nemo_retriever/tests/test_params_models.py @@ -18,6 +18,24 @@ def test_fps_zero_rejected(self) -> None: class TestExtractParams: + def test_parse_specific_configuration_requires_parse_method(self) -> None: + for field, value in ( + ("nemotron_parse_invoke_url", "http://parse:8000/v1/chat/completions"), + ("nemotron_parse_model", "nvidia/nemotron-parse"), + ): + with pytest.raises(ValidationError, match="method='nemotron_parse'"): + ExtractParams(**{field: value}) + + def test_normal_and_selected_parse_configurations_are_valid(self) -> None: + assert ExtractParams().method == "pdfium" + assert ExtractParams(invoke_url="http://generic").method == "pdfium" + params = ExtractParams( + method="nemotron_parse", + nemotron_parse_invoke_url="https://integrate.api.nvidia.com/v1/chat/completions", + nemotron_parse_model="nvidia/nemotron-parse", + ) + assert params.method == "nemotron_parse" + def test_graphic_elements_controls_are_removed(self) -> None: assert "use_graphic_elements" not in ExtractParams.model_fields assert "graphic_elements_invoke_url" not in ExtractParams.model_fields diff --git a/nemo_retriever/tests/test_pipeline_graph.py b/nemo_retriever/tests/test_pipeline_graph.py index f54f1cbf01..3e1d87ec6f 100644 --- a/nemo_retriever/tests/test_pipeline_graph.py +++ b/nemo_retriever/tests/test_pipeline_graph.py @@ -72,6 +72,40 @@ def test_text_build_graph_does_not_use_modal_content_reshape() -> None: assert "ExplodeContentToRows" not in _graph_node_names(graph) +def test_batch_graph_forwards_resolvable_hosted_parse_contract() -> None: + from nemo_retriever.operators.extract.parse.nemotron_parse import _resolve_nemotron_parse_contract + + endpoint = "https://integrate.api.nvidia.com/v1/chat/completions" + model = "nvidia/nemotron-parse" + graph = build_graph( + extraction_mode="pdf", + extract_params=ExtractParams( + method="nemotron_parse", + nemotron_parse_invoke_url=endpoint, + nemotron_parse_model=model, + ), + ) + + nodes: list[Node] = [] + + def collect(node: Node) -> None: + nodes.append(node) + for child in node.children: + collect(child) + + for root in graph.roots: + collect(root) + parse_node = next(node for node in nodes if node.operator.__class__.__name__ == "NemotronParseActor") + + assert parse_node.operator_kwargs["nemotron_parse_invoke_url"] == endpoint + assert parse_node.operator_kwargs["nemotron_parse_model"] == model + contract = _resolve_nemotron_parse_contract( + parse_node.operator_kwargs["nemotron_parse_invoke_url"], + parse_node.operator_kwargs["nemotron_parse_model"], + ) + assert (contract.model, contract.profile.value) == (model, "hosted_tool_call") + + def test_auto_extract_extension_sets_share_manifest_registry() -> None: assert PDF_EXTENSIONS == INPUT_TYPE_EXTENSIONS["pdf"] | INPUT_TYPE_EXTENSIONS["doc"] assert TEXT_EXTENSIONS == INPUT_TYPE_EXTENSIONS["txt"] @@ -902,7 +936,7 @@ def run(self, data): ) def _fake_resolve(operator_class, resources, operator_kwargs=None): - calls.append((operator_class.__name__, resources)) + calls.append((operator_class.__name__, resources, operator_kwargs)) return _IdentityStage monkeypatch.setattr( @@ -913,16 +947,31 @@ def _fake_resolve(operator_class, resources, operator_kwargs=None): lambda: Resources(cpu_count=8, gpu_count=1), ) + endpoint = "https://integrate.api.nvidia.com/v1/chat/completions" + model = "nvidia/nemotron-parse" op = MultiTypeExtractCPUActor( extraction_mode="pdf", - extract_params=ExtractParams(method="nemotron_parse"), + extract_params=ExtractParams( + method="nemotron_parse", + nemotron_parse_invoke_url=endpoint, + nemotron_parse_model=model, + ), ) batch_df = pd.DataFrame({"path": ["/tmp/test.pdf"]}) result = op._run_pdf_pipeline(batch_df) pd.testing.assert_frame_equal(result, batch_df) - assert [name for name, _resources in calls] == ["NemotronParseActor"] + assert [name for name, _resources, _kwargs in calls] == ["NemotronParseActor"] + parse_kwargs = calls[0][2] + assert parse_kwargs["nemotron_parse_invoke_url"] == endpoint + assert parse_kwargs["nemotron_parse_model"] == model + from nemo_retriever.operators.extract.parse.nemotron_parse import _resolve_nemotron_parse_contract + + contract = _resolve_nemotron_parse_contract( + parse_kwargs["nemotron_parse_invoke_url"], parse_kwargs["nemotron_parse_model"] + ) + assert (contract.model, contract.profile.value) == (model, "hosted_tool_call") class TestFileListLoaderOperator: diff --git a/nemo_retriever/tests/test_private_agent_packaging.py b/nemo_retriever/tests/test_private_agent_packaging.py new file mode 100644 index 0000000000..aaa9101100 --- /dev/null +++ b/nemo_retriever/tests/test_private_agent_packaging.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Packaging contracts for the agent implementation.""" + +from __future__ import annotations + +import tomllib +from pathlib import Path + + +def test_private_agent_prompt_resources_are_declared_for_distribution(): + project = tomllib.loads((Path(__file__).parents[1] / "pyproject.toml").read_text(encoding="utf-8")) + + dependencies = project["project"]["dependencies"] + package_data = project["tool"]["setuptools"]["package-data"] + + assert any(dependency.lower().startswith("jinja2") for dependency in dependencies) + assert package_data["nemo_retriever._agentic.nemo_agent.prompts"] == ["templates/**/*.j2"] diff --git a/nemo_retriever/tests/test_retriever_queries.py b/nemo_retriever/tests/test_retriever_queries.py index 2dd2384520..af6a5d19fb 100644 --- a/nemo_retriever/tests/test_retriever_queries.py +++ b/nemo_retriever/tests/test_retriever_queries.py @@ -37,6 +37,9 @@ def _make_retriever(**overrides: Any) -> Retriever: "embed_kwargs": {"model_name": "embedder", "embed_model_name": "embedder"}, } defaults.update(overrides) + embed_kwargs = dict(defaults["embed_kwargs"]) + embed_kwargs.setdefault("local_ingest_embed_backend", "hf") + defaults["embed_kwargs"] = embed_kwargs return Retriever(**defaults) @@ -84,15 +87,126 @@ def test_merge_embed_params_per_call_overrides(self) -> None: p = r._merge_embed_params({"model_name": "call"}) assert p.model_name == "call" - def test_local_query_embedding_defaults_to_hf(self) -> None: - p = _make_retriever()._merge_embed_params() - assert p.local_ingest_embed_backend == "hf" + def test_index_model_keeps_constructor_provider_prefix(self) -> None: + retriever = _make_retriever( + embed_kwargs={ + "embed_invoke_url": "https://embed.example.com/v1/embeddings", + "embed_model_provider_prefix": "nvidia", + } + ) + + resolved = retriever._resolve_embed_kwargs("nvidia/llama-nemotron-embed-vl-1b-v2", None) + params = retriever._merge_embed_params(resolved) + + assert params.model_name == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert params.embed_model_provider_prefix == "nvidia" + + def test_index_model_revision_is_forwarded_to_query_embedder(self) -> None: + retriever = _make_retriever(embed_kwargs={}) + + resolved = retriever._resolve_embed_kwargs( + "acme/fine-tuned-nemotron", + None, + "a" * 40, + ) + + assert resolved["embed_model_name"] == "acme/fine-tuned-nemotron" + assert resolved["embed_model_revision"] == "a" * 40 + + def test_explicit_model_override_does_not_reuse_index_revision(self) -> None: + retriever = _make_retriever(embed_kwargs={"embed_model_name": "acme/override"}) + + resolved = retriever._resolve_embed_kwargs( + "acme/index-model", + None, + "a" * 40, + ) + + assert resolved["embed_model_name"] == "acme/override" + assert "embed_model_revision" not in resolved + + def test_runtime_model_change_clears_configured_revision(self) -> None: + retriever = _make_retriever(embed_kwargs={"embed_model_name": "acme/model-a", "embed_model_revision": "a" * 40}) + + resolved = retriever._resolve_embed_kwargs(None, {"embed_model_name": "acme/model-b"}) + params = retriever._merge_embed_params(resolved) + + assert params.embed_model_name == "acme/model-b" + assert params.embed_model_revision is None + + def test_runtime_same_model_keeps_configured_revision(self) -> None: + retriever = _make_retriever(embed_kwargs={"embed_model_name": "acme/model-a", "embed_model_revision": "a" * 40}) + + resolved = retriever._resolve_embed_kwargs(None, {"embed_model_name": "acme/model-a"}) + params = retriever._merge_embed_params(resolved) + + assert params.embed_model_revision == "a" * 40 + + def test_explicit_runtime_revision_overrides_index_revision(self) -> None: + retriever = _make_retriever(embed_kwargs={}) + + resolved = retriever._resolve_embed_kwargs( + "acme/index-model", + {"embed_model_revision": "b" * 40}, + "a" * 40, + ) + params = retriever._merge_embed_params(resolved) + + assert params.embed_model_name == "acme/index-model" + assert params.embed_model_revision == "b" * 40 + + @pytest.mark.parametrize(("requires_vllm", "expected_backend"), [(False, "hf"), (True, "vllm")]) + @patch("nemo_retriever.graph.retriever.resolve_embed_model_spec") + def test_local_query_embedding_selects_compatible_backend( + self, resolve_spec: MagicMock, requires_vllm: bool, expected_backend: str + ) -> None: + resolve_spec.return_value = MagicMock(requires_vllm=requires_vllm, revision="a" * 40) + retriever = _make_retriever( + embed_kwargs={ + "model_name": "embedder", + "embed_model_name": "embedder", + "local_ingest_embed_backend": None, + } + ) + + p = retriever._merge_embed_params() + + assert p.local_ingest_embed_backend == expected_backend + assert p.embed_model_revision == "a" * 40 + resolve_spec.assert_called_once_with("embedder", revision=None) + + @patch("nemo_retriever.graph.retriever.resolve_embed_model_spec") + def test_index_model_revision_drives_automatic_backend(self, resolve_spec: MagicMock) -> None: + revision = "b" * 40 + resolve_spec.return_value = MagicMock(requires_vllm=True, revision=revision) + retriever = _make_retriever(embed_kwargs={"local_ingest_embed_backend": None}) + resolved = retriever._resolve_embed_kwargs("acme/index-model", None, revision) + + p = retriever._merge_embed_params(resolved) + + assert p.embed_model_name == "acme/index-model" + assert p.embed_model_revision == revision + assert p.local_ingest_embed_backend == "vllm" + resolve_spec.assert_called_once_with("acme/index-model", revision=revision) def test_local_query_embedding_backend_can_be_overridden(self) -> None: r = _make_retriever(embed_kwargs={"local_ingest_embed_backend": "vllm"}) p = r._merge_embed_params() assert p.local_ingest_embed_backend == "vllm" + @patch("nemo_retriever.graph.retriever.resolve_embed_model_spec") + def test_remote_embedding_skips_local_model_resolution(self, resolve_spec: MagicMock) -> None: + retriever = _make_retriever( + embed_kwargs={ + "embed_invoke_url": "https://embed.example.com/v1/embeddings", + "local_ingest_embed_backend": None, + } + ) + + retriever._merge_embed_params() + + resolve_spec.assert_not_called() + def test_rerank_inflates_retrieval_top_k(self, monkeypatch: pytest.MonkeyPatch) -> None: graph = _install_mock_graph(monkeypatch, [[{"text": "x"}]]) retriever = _make_retriever(top_k=3, rerank=True, rerank_kwargs={"refine_factor": 4}) diff --git a/nemo_retriever/tests/test_root_cli_workflow.py b/nemo_retriever/tests/test_root_cli_workflow.py index 42a13a6cd2..fe7bc7b50a 100644 --- a/nemo_retriever/tests/test_root_cli_workflow.py +++ b/nemo_retriever/tests/test_root_cli_workflow.py @@ -161,6 +161,8 @@ def fake_create_ingestor(**kwargs: Any) -> Any: "uri": "lancedb", "table_name": "nemo-retriever", "overwrite": True, + "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", + "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } assert "Ingested 1 file(s) → 7 row(s) in LanceDB lancedb/nemo-retriever." in result.output @@ -201,6 +203,8 @@ def test_root_ingest_without_mode_accepts_local_options_before_documents(monkeyp "uri": "/tmp/default-lancedb", "table_name": "nemo-retriever", "overwrite": False, + "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", + "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } @@ -240,6 +244,7 @@ def embed(self, params=None, **_kwargs): return self def ingest(self, *args: Any, **kwargs: Any): + captured["ingest_kwargs"] = kwargs return self monkeypatch.setattr(service_ingestor_module, "ServiceIngestor", _FakeServiceIngestor) @@ -287,6 +292,7 @@ def ingest(self, *args: Any, **kwargs: Any): assert captured["dedup_params"].iou_threshold == 0.6 assert captured["caption_params"].context_text_max_chars == 12 assert captured["embed_params"].embed_granularity == "page" + assert captured["ingest_kwargs"] == {"return_results": True} assert "through retriever service http://retriever-service:7670" in result.output @@ -394,6 +400,8 @@ def fake_create_ingestor(**kwargs: Any) -> Any: "uri": "/tmp/lancedb", "table_name": "docs", "overwrite": True, + "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", + "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } assert "Ingested 2 file(s) → 12 row(s) in LanceDB /tmp/lancedb/docs." in result.output @@ -412,6 +420,8 @@ def test_root_ingest_append_forwards_overwrite_false(monkeypatch, tmp_path) -> N "uri": "lancedb", "table_name": "nemo-retriever", "overwrite": False, + "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", + "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } @@ -509,8 +519,13 @@ def fake_create_ingestor(**_kwargs: Any) -> Any: assert isinstance(embed_params, EmbedParams) assert embed_params.embed_invoke_url == "http://embed:8000/v1/embeddings" assert embed_params.embedding_endpoint == "http://embed:8000/v1/embeddings" - assert embed_params.model_name == "nvidia/nvidia/llama-nemotron-embed-1b-v2" - assert embed_params.embed_model_name == "nvidia/nvidia/llama-nemotron-embed-1b-v2" + assert embed_params.model_name == "nvidia/llama-nemotron-embed-1b-v2" + assert embed_params.embed_model_name == "nvidia/llama-nemotron-embed-1b-v2" + assert embed_params.embed_model_provider_prefix == "nvidia" + vdb_kwargs = fake_ingestor.vdb_upload.call_args.args[0].vdb_kwargs + assert vdb_kwargs["embedding_model_name"] == "nvidia/llama-nemotron-embed-1b-v2" + assert vdb_kwargs["vector_dim"] is None + assert "embedding_model_revision" not in vdb_kwargs def test_root_ingest_passes_embedding_overrides_without_stage_flags(monkeypatch, tmp_path) -> None: @@ -1694,6 +1709,8 @@ def test_root_ingest_index_mode_hybrid_passes_hybrid_into_vdb_kwargs(monkeypatch "table_name": "docs", "overwrite": True, "hybrid": True, + "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", + "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } @@ -1756,5 +1773,7 @@ def test_root_ingest_index_mode_sparse_skips_embedding_and_writes_fts_table(monk table = lancedb.connect(str(tmp_path / "db")).open_table("sparse_docs") assert "vector" not in table.schema.names assert table.schema.metadata[b"retrieval_mode"] == b"sparse" + assert table.schema.metadata[b"nemo_retriever.retrieval_mode"] == b"sparse" + assert b"nemo_retriever.embedding_model_name" not in table.schema.metadata index_names = {index.name.lower() for index in table.list_indices()} assert any("text" in name or "fts" in name for name in index_names) diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index f03c9e824b..6b116155fd 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -152,8 +152,9 @@ def query(self, query: str, **_kwargs: Any) -> list[dict[str, Any]]: "embed_kwargs": { "embed_invoke_url": "http://embed:8000/v1/embeddings", "embedding_endpoint": "http://embed:8000/v1/embeddings", - "model_name": "nvidia/nvidia/llama-nemotron-embed-1b-v2", - "embed_model_name": "nvidia/nvidia/llama-nemotron-embed-1b-v2", + "model_name": "nvidia/llama-nemotron-embed-1b-v2", + "embed_model_name": "nvidia/llama-nemotron-embed-1b-v2", + "embed_model_provider_prefix": "nvidia", }, } ] @@ -399,6 +400,10 @@ def unload(self) -> None: assert "llm_backend" not in cfg assert cfg["local_llm_backend"] == "vllm" assert cfg["llm_model"] == "nemotron-8b" + # --agentic-llm-client is optional and unset here; the callable default is + # resolved in AgenticRetrievalConfig.__post_init__ (faked here), so the + # pre-resolution kwarg is still None. + assert cfg["llm_client"] is None assert cfg["temperature"] == 1.25 # --top-k is honored end-to-end: plumbed into the agentic config (drives the # ReAct target / RRF / selection cut), not just applied as a post-filter. @@ -474,6 +479,116 @@ def unload(self) -> None: assert config_calls[-1]["invoke_url"] == "http://localhost:8000/v1/chat/completions" +def test_root_query_agentic_llm_client_override_plumbed_into_config(monkeypatch) -> None: + """`--agentic-llm-client` selects the LLM client wired into AgenticRetrievalConfig.""" + import pandas as pd + + import nemo_retriever.query.agentic as agentic_retrieval + + config_calls: list[dict[str, Any]] = [] + + class FakeConfig: + def __init__(self, **kwargs: Any) -> None: + config_calls.append(kwargs) + + class FakeAgenticRetriever: + def __init__(self, cfg: Any) -> None: + self.cfg = cfg + + def retrieve(self, query_ids: Any, query_texts: Any) -> Any: + return pd.DataFrame([{"query_id": "0", "doc_id": "a.pdf", "rank": 1, "result_source": "rrf"}]) + + def unload(self) -> None: + return None + + monkeypatch.setattr(agentic_retrieval, "AgenticRetrievalConfig", FakeConfig) + monkeypatch.setattr(agentic_retrieval, "AgenticRetriever", FakeAgenticRetriever) + + result = RUNNER.invoke( + cli_main.app, + [ + "query", + "q", + "--agentic", + "--agentic-llm-model", + "m", + "--agentic-invoke-url", + "http://localhost:8000/v1/chat/completions", + "--agentic-llm-client", + "litellm", + ], + ) + + assert result.exit_code == 0 + assert config_calls[0]["llm_client"] == "litellm" + + +def test_root_query_agentic_accepts_callable_client_with_invoke_url(monkeypatch) -> None: + """`callable` is valid remotely: it wraps the shared HTTP client, not just the local engine.""" + import pandas as pd + + import nemo_retriever.query.agentic as agentic_retrieval + + config_calls: list[dict[str, Any]] = [] + + class FakeConfig: + def __init__(self, **kwargs: Any) -> None: + config_calls.append(kwargs) + + class FakeAgenticRetriever: + def __init__(self, cfg: Any) -> None: + self.cfg = cfg + + def retrieve(self, query_ids: Any, query_texts: Any) -> Any: + return pd.DataFrame([{"query_id": "0", "doc_id": "a.pdf", "rank": 1, "result_source": "rrf"}]) + + def unload(self) -> None: + return None + + monkeypatch.setattr(agentic_retrieval, "AgenticRetrievalConfig", FakeConfig) + monkeypatch.setattr(agentic_retrieval, "AgenticRetriever", FakeAgenticRetriever) + + result = RUNNER.invoke( + cli_main.app, + [ + "query", + "q", + "--agentic", + "--agentic-llm-model", + "m", + "--agentic-invoke-url", + "http://localhost:8000/v1/chat/completions", + "--agentic-llm-client", + "callable", + ], + ) + + assert result.exit_code == 0 + assert config_calls[0]["llm_client"] == "callable" + + +def test_root_query_agentic_rejects_remote_client_without_invoke_url() -> None: + """A client that can only talk to an endpoint still needs one.""" + result = RUNNER.invoke( + cli_main.app, + ["query", "q", "--agentic", "--agentic-llm-model", "m", "--agentic-llm-client", "litellm"], + ) + + assert result.exit_code == 1 + assert "requires --agentic-invoke-url" in result.output + + +def test_root_query_agentic_rejects_unknown_client() -> None: + """An unsupported --agentic-llm-client fails fast before any retrieval work.""" + result = RUNNER.invoke( + cli_main.app, + ["query", "q", "--agentic", "--agentic-llm-model", "m", "--agentic-llm-client", "bogus"], + ) + + assert result.exit_code == 1 + assert "agentic_llm_client must be one of" in result.output + + def test_root_query_agentic_plumbs_rerank_into_config(monkeypatch) -> None: """`--rerank` with `--agentic` wires the reranker config into AgenticRetrievalConfig (reranker model + endpoint + backend), so the agent's retrieval backend reranks.""" @@ -658,11 +773,13 @@ def test_root_query_local_help_shows_retrieval_mode_not_hybrid() -> None: assert "--hybrid" not in result.output -def test_root_query_local_help_names_default_models() -> None: +def test_root_query_local_help_describes_model_resolution() -> None: result = RUNNER.invoke(cli_main.app, ["query", "q", "--help"]) assert result.exit_code == 0 - assert "Default embedding model" in result.output + assert "Embedding model: read from the selected table" in result.output + assert "legacy tables" in result.output + assert "fall back" in result.output assert VL_EMBED_MODEL in result.output assert "Default local reranker model" in result.output assert VL_RERANK_MODEL in result.output diff --git a/nemo_retriever/tests/test_service_agentic_query.py b/nemo_retriever/tests/test_service_agentic_query.py new file mode 100644 index 0000000000..fb41de7129 --- /dev/null +++ b/nemo_retriever/tests/test_service_agentic_query.py @@ -0,0 +1,399 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from unittest.mock import PropertyMock, patch + +import pytest +from fastapi.testclient import TestClient +from pydantic import ValidationError + +import nemo_retriever.service.vectordb_app as vectordb_module +from nemo_retriever.service.app import create_app +from nemo_retriever.service.agentic_query import ( + agentic_ranked_to_hits, + build_agentic_query_request, +) +from nemo_retriever.service.config import ( + AgenticConfig, + AuthConfig, + LoggingConfig, + PipelinePoolConfig, + ServiceConfig, + VectorDbConfig, +) +from nemo_retriever.service.query_schema import ( + MAX_AGENTIC_QUERY_CHARS, + QueryRequest, + QueryResponse, + QueryResult, +) +from nemo_retriever.service.vectordb_app import VectorDBState, create_vectordb_app + + +def test_agentic_service_config_requires_remote_model_and_endpoint() -> None: + with pytest.raises(ValidationError, match="agentic.invoke_url"): + AgenticConfig(enabled=True, llm_model="model") + with pytest.raises(ValidationError, match="agentic.llm_model"): + AgenticConfig( + enabled=True, + invoke_url="https://llm.example/v1/chat/completions", + ) + + +def test_query_request_agentic_requires_hits_format() -> None: + with pytest.raises(ValidationError, match="single query string"): + QueryRequest(query=["a", "b"], agentic=True) + with pytest.raises(ValidationError, match="non-empty"): + QueryRequest(query=" ", agentic=True) + with pytest.raises(ValidationError, match="format='hits'"): + QueryRequest(query="q", agentic=True, format="evidence") + + +def test_build_agentic_query_request_maps_server_owned_configuration() -> None: + request = build_agentic_query_request( + query="revenue trend", + top_k=3, + config=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + backend_top_k=25, + react_max_steps=7, + ), + lancedb_uri="/indexes/finance", + table_name="finance", + embed_endpoint="https://embed.example/v1/embeddings", + embed_model="embed-model", + embed_model_provider_prefix="openai", + embed_api_key="embed-key", + ) + + assert request.query == "revenue trend" + assert request.retrieval.top_k == 3 + assert request.storage.lancedb_uri == "/indexes/finance" + assert request.storage.table_name == "finance" + assert request.embed.embed_invoke_url == "https://embed.example/v1/embeddings" + assert request.embed.embed_model_name == "embed-model" + assert request.embed.embed_model_provider_prefix == "openai" + assert request.embed.embed_api_key == "embed-key" + assert request.agentic.enabled is True + assert request.agentic.llm_model == "model" + assert request.agentic.invoke_url == "https://llm.example/v1/chat/completions" + assert request.agentic.backend_top_k == 25 + assert request.agentic.react_max_steps == 7 + + +def test_agentic_ranked_to_hits_maps_doc_id_onto_query_hit_envelope() -> None: + hits = agentic_ranked_to_hits([{"rank": 1, "doc_id": "report.pdf", "result_source": "selection_agent"}]) + assert hits == [ + { + "text": None, + "metadata": {"result_source": "selection_agent", "rank": 1}, + "source": "report.pdf", + "source_id": None, + "path": None, + "page_number": None, + "pdf_basename": None, + "pdf_page": None, + } + ] + + +def test_agentic_ranked_to_hits_rejects_blank_doc_id() -> None: + with pytest.raises(ValueError, match="missing a non-empty doc_id"): + agentic_ranked_to_hits([{"rank": 1, "doc_id": "", "result_source": "selection_agent"}]) + + +def test_agentic_query_flag_rejected_when_disabled(tmp_path) -> None: + app = create_vectordb_app( + lancedb_uri=str(tmp_path), + embed_endpoint="https://embed.example/v1/embeddings", + ) + + with TestClient(app) as client: + response = client.post("/v1/query", json={"query": "q", "agentic": True}) + + assert response.status_code == 400 + assert "not enabled" in response.json()["detail"] + + +def test_agentic_query_rejects_collection_target(tmp_path) -> None: + app = create_vectordb_app( + lancedb_uri=str(tmp_path), + embed_endpoint="https://embed.example/v1/embeddings", + agentic_config=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + ), + ) + + with TestClient(app) as client: + response = client.post( + "/v1/query", + json={"query": "q", "agentic": True, "collection_name": "workspace"}, + ) + + assert response.status_code == 501 + + +def test_agentic_true_runs_react_workflow_on_v1_query(tmp_path) -> None: + app = create_vectordb_app( + lancedb_uri=str(tmp_path), + table_name="finance", + embed_endpoint="https://embed.example/v1/embeddings", + embed_model="embed-model", + agentic_config=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + ), + ) + expected = QueryResponse( + results=[ + QueryResult( + hits=agentic_ranked_to_hits( + [ + { + "rank": 1, + "doc_id": "report.pdf", + "result_source": "selection_agent", + } + ] + ) + ) + ], + query_mode="agentic", + ) + + with ( + patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), + patch.object(vectordb_module, "run_agentic_query", return_value=expected) as run_query, + TestClient(app) as client, + ): + response = client.post( + "/v1/query", + json={"query": "revenue trend", "top_k": 3, "agentic": True}, + ) + + assert response.status_code == 200 + assert response.json() == { + "results": [ + { + "hits": [ + { + "text": None, + "metadata": {"result_source": "selection_agent", "rank": 1}, + "source": "report.pdf", + "source_id": None, + "path": None, + "page_number": None, + "pdf_basename": None, + "pdf_page": None, + } + ] + } + ], + "query_mode": "agentic", + } + assert run_query.call_args.kwargs["query"] == "revenue trend" + assert run_query.call_args.kwargs["top_k"] == 3 + assert run_query.call_args.kwargs["lancedb_uri"] == str(tmp_path) + assert run_query.call_args.kwargs["table_name"] == "finance" + assert run_query.call_args.kwargs["embed_api_key"] == "" + + +def test_agentic_query_rejects_top_k_above_backend_depth(tmp_path) -> None: + app = create_vectordb_app( + lancedb_uri=str(tmp_path), + embed_endpoint="https://embed.example/v1/embeddings", + agentic_config=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + backend_top_k=5, + ), + ) + + with ( + patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), + TestClient(app) as client, + ): + response = client.post( + "/v1/query", + json={"query": "revenue trend", "top_k": 6, "agentic": True}, + ) + + assert response.status_code == 422 + assert "cannot exceed" in response.json()["detail"] + + +def test_agentic_query_rejects_query_above_length_limit(tmp_path) -> None: + app = create_vectordb_app( + lancedb_uri=str(tmp_path), + embed_endpoint="https://embed.example/v1/embeddings", + agentic_config=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + ), + ) + + with ( + patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), + patch.object(vectordb_module, "run_agentic_query") as run_query, + TestClient(app) as client, + ): + response = client.post( + "/v1/query", + json={"query": "x" * (MAX_AGENTIC_QUERY_CHARS + 1), "agentic": True}, + ) + + assert response.status_code == 422 + run_query.assert_not_called() + + +def test_agentic_query_slots_are_bounded_and_released_by_the_worker(tmp_path) -> None: + """Capacity follows the worker thread, not the caller: a saturated pool sheds + load with 503 instead of queueing behind non-cancellable ReAct work, and a + completed query returns its slot.""" + app = create_vectordb_app( + lancedb_uri=str(tmp_path), + embed_endpoint="https://embed.example/v1/embeddings", + agentic_config=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + ), + ) + expected = QueryResponse(results=[QueryResult(hits=[])], query_mode="agentic") + + with ( + patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), + patch.object(vectordb_module, "run_agentic_query", return_value=expected), + TestClient(app) as client, + ): + slots = app.state.agentic_slots + assert slots is not None + + for _ in range(vectordb_module.MAX_CONCURRENT_AGENTIC_QUERIES): + assert slots.acquire(blocking=False) is True + + busy = client.post("/v1/query", json={"query": "revenue trend", "agentic": True}) + + assert busy.status_code == 503 + assert busy.headers["Retry-After"] == "30" + query_semaphore = app.state.vectordb_state.query_semaphore + assert query_semaphore.locked() is False + + slots.release() + accepted = client.post("/v1/query", json={"query": "revenue trend", "agentic": True}) + + assert accepted.status_code == 200 + assert slots.acquire(blocking=False) is True + + +def test_gateway_proxies_agentic_flag_to_vectordb( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + async def _stub_work(_item): + return 0, [] + + monkeypatch.setattr( + "nemo_retriever.service.services.pipeline_executor.create_realtime_work_fn", + lambda _config: _stub_work, + ) + monkeypatch.setattr( + "nemo_retriever.service.services.pipeline_executor.create_batch_work_fn", + lambda _config: _stub_work, + ) + config = ServiceConfig( + mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), + logging=LoggingConfig(file=str(tmp_path / "service.log")), + pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), + vectordb=VectorDbConfig( + enabled=True, + vectordb_url="http://vectordb:7671", + ), + agentic=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + request_timeout_s=321.0, + ), + ) + seen: dict[str, object] = {} + + class _FakeResponse: + status_code = 200 + content = json.dumps({"results": [{"hits": []}]}).encode() + + class _FakeAsyncClient: + def __init__(self, *args, **kwargs) -> None: + seen["timeout"] = kwargs["timeout"] + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + async def post(self, url: str, **kwargs) -> _FakeResponse: + seen["url"] = url + seen["body"] = json.loads(kwargs["content"]) + return _FakeResponse() + + monkeypatch.setattr("httpx.AsyncClient", _FakeAsyncClient) + + with TestClient(create_app(config)) as client: + response = client.post( + "/v1/query", + json={"query": "revenue trend", "top_k": 3, "agentic": True}, + ) + + assert response.status_code == 200 + assert response.json() == {"results": [{"hits": []}]} + assert seen == { + "timeout": 321.0, + "url": "http://vectordb:7671/v1/query", + "body": {"query": "revenue trend", "top_k": 3, "agentic": True}, + } + + +def test_service_rejects_agentic_flag_when_not_configured( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + async def _stub_work(_item): + return 0, [] + + monkeypatch.setattr( + "nemo_retriever.service.services.pipeline_executor.create_realtime_work_fn", + lambda _config: _stub_work, + ) + monkeypatch.setattr( + "nemo_retriever.service.services.pipeline_executor.create_batch_work_fn", + lambda _config: _stub_work, + ) + config = ServiceConfig( + mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), + logging=LoggingConfig(file=str(tmp_path / "service.log")), + pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), + vectordb=VectorDbConfig(enabled=True, vectordb_url="http://vectordb:7671"), + ) + + with TestClient(create_app(config)) as client: + response = client.post( + "/v1/query", + json={"query": "revenue trend", "agentic": True}, + ) + + assert response.status_code == 400 + assert "not enabled" in response.json()["detail"] diff --git a/nemo_retriever/tests/test_service_answer_generation.py b/nemo_retriever/tests/test_service_answer_generation.py index 76e0d5504e..babb055c30 100644 --- a/nemo_retriever/tests/test_service_answer_generation.py +++ b/nemo_retriever/tests/test_service_answer_generation.py @@ -16,7 +16,14 @@ from nemo_retriever.models.llm.types import GenerationResult, JudgeResult from nemo_retriever.service.app import create_app -from nemo_retriever.service.config import LLMConfig, LoggingConfig, PipelinePoolConfig, ServiceConfig, VectorDbConfig +from nemo_retriever.service.config import ( + AuthConfig, + LLMConfig, + LoggingConfig, + PipelinePoolConfig, + ServiceConfig, + VectorDbConfig, +) def test_llm_config_defaults_to_reasoning_enabled_for_external_provider_safety() -> None: @@ -50,6 +57,7 @@ async def _stub_work(_item): cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), logging=LoggingConfig(file=str(tmp_path / "service.log")), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), vectordb=VectorDbConfig(enabled=True, vectordb_url="http://vectordb:7671"), @@ -137,6 +145,7 @@ async def post(self, url: str, **kwargs) -> _FakeResponse: { "url": "http://vectordb:7671/v1/query", "json": {"query": "What generates answers?", "top_k": 2}, + "headers": {"X-NRL-Scope": "default"}, } ] from_kwargs.assert_called_once_with( @@ -311,6 +320,7 @@ async def _stub_work(_item): app = create_app( ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), logging=LoggingConfig(file=str(tmp_path / "service.log")), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), vectordb=VectorDbConfig(enabled=True, vectordb_url="http://vectordb:7671"), @@ -341,6 +351,7 @@ async def _stub_work(_item): app = create_app( ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), logging=LoggingConfig(file=str(tmp_path / "service.log")), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), vectordb=VectorDbConfig(enabled=False), diff --git a/nemo_retriever/tests/test_service_caption.py b/nemo_retriever/tests/test_service_caption.py index e463474971..a82f8f7e69 100644 --- a/nemo_retriever/tests/test_service_caption.py +++ b/nemo_retriever/tests/test_service_caption.py @@ -28,6 +28,7 @@ from nemo_retriever.common.params import CaptionParams from nemo_retriever.service.app import create_app from nemo_retriever.service.config import ( + AuthConfig, NimEndpointsConfig, PipelineOverridesConfig, PipelinePoolConfig, @@ -318,6 +319,7 @@ async def _stub_work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), nim_endpoints=NimEndpointsConfig( caption_invoke_url="http://caption.svc/v1", diff --git a/nemo_retriever/tests/test_service_client_compat.py b/nemo_retriever/tests/test_service_client_compat.py index b0394653b0..def7f057c0 100644 --- a/nemo_retriever/tests/test_service_client_compat.py +++ b/nemo_retriever/tests/test_service_client_compat.py @@ -47,13 +47,14 @@ import pytest from nemo_retriever.service.client import ( + DocumentTracker, + InMemoryUpload, RetrieverServiceClient, RetrieverServiceCompatibilityError, _compat_error_message, _is_api_mismatch_status, ) - # ---------------------------------------------------------------------- # Helpers: drive _create_job and _upload_one directly against MockTransport # ---------------------------------------------------------------------- @@ -174,13 +175,13 @@ async def _call() -> None: assert request_paths == ["/v1/ingest/job"], request_paths -def test_create_job_500_still_raises_generic_http_status_error() -> None: +def test_create_job_500_raises_typed_service_error() -> None: """A real server error must NOT be misreported as a version mismatch. If the deployed service is the right version but transiently broken (500/503/etc.) the SDK should surface the existing - :class:`httpx.HTTPStatusError` so retry/alerting logic in the - caller still triggers as before. Mis-coding this as a + typed service error so retry/alerting logic can use one SDK error + hierarchy. Mis-coding this as a ``RetrieverServiceCompatibilityError`` would hide a real outage. """ @@ -193,7 +194,9 @@ async def _call() -> None: async with httpx.AsyncClient(transport=_make_transport(_handler)) as client: await rc._create_job(client, expected_documents=1) - with pytest.raises(httpx.HTTPStatusError) as ei: + from nemo_retriever.service.errors import RetrieverServiceError + + with pytest.raises(RetrieverServiceError) as ei: _run_async(_call()) assert "HTTP 500" in str(ei.value) @@ -242,11 +245,127 @@ async def _first_event() -> dict[str, object]: } +def test_sse_keepalive_after_uploads_falls_back_to_polling( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A missed terminal SSE event cannot block a completed ingestion forever.""" + + rc = RetrieverServiceClient(base_url="http://nrl:7670") + pending = {"doc-1"} + uploads_done = asyncio.Event() + uploads_done.set() + tracker = DocumentTracker() + fallback_calls: list[set[str]] = [] + + class FakeResponse: + status_code = 200 + + async def aiter_lines(self): + yield ": keepalive" + + class FakeStream: + async def __aenter__(self): + return FakeResponse() + + async def __aexit__(self, *_args): + return None + + class FakeClient: + def stream(self, *_args, **_kwargs): + return FakeStream() + + async def fallback(_client, current, _tracker, _on_event): + fallback_calls.append(set(current)) + current.clear() + + monkeypatch.setattr(rc, "_bulk_poll_fallback", fallback) + _run_async( + rc._consume_sse( + FakeClient(), + pending, + uploads_done, + tracker, + job_id="JOB-1", + ) + ) + assert fallback_calls == [{"doc-1"}] + assert pending == set() + + +def test_materialized_ingest_tracks_attempt_id_but_reports_stable_document_id( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Status tracking uses the attempt id returned by the job upload route.""" + + rc = RetrieverServiceClient(base_url="http://nrl:7670") + source = tmp_path / "doc.txt" + source.write_text("attempt id regression", encoding="utf-8") + submitted: list[tuple[str, str]] = [] + + async def create_job(*_args, **_kwargs): + return SimpleNamespace(job_id="JOB-1", trace_id=None) + + async def upload_one(*_args, **_kwargs): + return {"document_id": "stable-1", "attempt_id": "attempt-1"} + + async def consume_sse(_client, pending, uploads_done, tracker, **_kwargs): + await uploads_done.wait() + assert pending == {"attempt-1"} + tracker.mark_completed("attempt-1", {"id": "attempt-1", "status": "completed"}) + pending.clear() + + monkeypatch.setattr(rc, "_create_job", create_job) + monkeypatch.setattr(rc, "_upload_one", upload_one) + monkeypatch.setattr(rc, "_consume_sse", consume_sse) + + results = _run_async( + rc.ingest_documents( + [source], + show_progress=False, + on_file_submitted=lambda filename, doc_id: submitted.append((filename, doc_id)), + ) + ) + assert submitted == [("doc.txt", "stable-1")] + assert results == [{"id": "attempt-1", "status": "completed"}] + + # ---------------------------------------------------------------------- # _upload_one — guards the mid-rollout case (new gateway, stale worker) # ---------------------------------------------------------------------- +def test_upload_one_sends_inline_text_from_memory_with_classification_metadata() -> None: + requests: list[httpx.Request] = [] + + def _handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(202, json={"document_id": "doc-inline", "job_id": "JOB-1"}) + + rc = RetrieverServiceClient(base_url="http://nrl:7670") + source = InMemoryUpload( + filename="inline://00000003", + content="café".encode(), + content_type="text/plain; charset=utf-8", + classification_filename="inline-00000003.txt", + ) + + async def _call() -> dict[str, object]: + async with httpx.AsyncClient(transport=_make_transport(_handler)) as client: + return await rc._upload_one( + client, + source, + job_id="JOB-1", + pipeline_spec={"extraction_mode": "text", "stage_order": ["extract"]}, + ) + + assert _run_async(_call())["document_id"] == "doc-inline" + body = requests[0].content + assert b'filename="inline://00000003"' in body + assert "café".encode() in body + assert b'"filename": "inline-00000003.txt"' in body + assert b'"extraction_mode": "text"' in body + + def test_upload_one_raises_compat_error_on_404(tmp_path: Path) -> None: """Uploads to a stale pod (404 on the document path) ⇒ compat error. diff --git a/nemo_retriever/tests/test_service_ingest_async.py b/nemo_retriever/tests/test_service_ingest_async.py index b7539440f6..18157ac35e 100644 --- a/nemo_retriever/tests/test_service_ingest_async.py +++ b/nemo_retriever/tests/test_service_ingest_async.py @@ -61,6 +61,7 @@ def _fake_materialize_completed_document( document_id: str, *, return_results: bool, + client: Any = None, ) -> list[dict[str, Any]] | None: if not return_results and self._save_to_disk_dir is None: return None @@ -280,3 +281,14 @@ def test_ingest_async_forwards_return_results(stub_ingestor: ServiceIngestor) -> out = future.result(timeout=5.0) assert isinstance(out, ServiceIngestResult) assert out.dataframe is None + + +def test_ingest_return_results_false_creates_no_result_client(stub_ingestor: ServiceIngestor) -> None: + with patch.object( + stub_ingestor, + "_new_result_fetch_client", + side_effect=AssertionError("result client must not be created"), + ): + result = stub_ingestor.ingest(return_results=False) + + assert result.dataframe is None diff --git a/nemo_retriever/tests/test_service_ingest_router.py b/nemo_retriever/tests/test_service_ingest_router.py index 11e565fa73..14adfbae38 100644 --- a/nemo_retriever/tests/test_service_ingest_router.py +++ b/nemo_retriever/tests/test_service_ingest_router.py @@ -18,22 +18,34 @@ from __future__ import annotations +import asyncio +import hashlib import json import re +from types import SimpleNamespace from typing import Any import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from nemo_retriever.common.schemas.collections import IngestOperation +from nemo_retriever.common.schemas.requests import IngestRequest from nemo_retriever.service.app import create_app from nemo_retriever.service.config import ( + AuthConfig, PipelineOverridesConfig, PipelinePoolConfig, ServiceConfig, + VectorDbConfig, ) from nemo_retriever.service import tracing +from nemo_retriever.service.routers.dashboard import VdbQueryRequest, vdb_query +from nemo_retriever.service.routers.ingest import _route_by_page_count +from nemo_retriever.service.services.job_tracker import get_job_tracker from nemo_retriever.service.services.pipeline_pool import PoolType, WorkItem +from nemo_retriever.service.utils.file_type import FileCategory from .conftest import create_test_job @@ -82,6 +94,7 @@ def _stub_batch(_config: ServiceConfig): cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), pipeline_overrides=PipelineOverridesConfig(), ) @@ -127,6 +140,7 @@ def _stub_batch(_config: ServiceConfig): cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), pipeline_overrides=PipelineOverridesConfig(), ) @@ -154,7 +168,10 @@ def traced_gateway_app(monkeypatch: pytest.MonkeyPatch): ) monkeypatch.setattr("nemo_retriever.service.tracing.BatchSpanProcessor", SimpleSpanProcessor) - cfg = ServiceConfig(mode="gateway") + cfg = ServiceConfig( + mode="gateway", + auth=AuthConfig(allow_unscoped_dev=True), + ) try: app = create_app(cfg) with TestClient(app) as client: @@ -168,6 +185,16 @@ def _make_pdf_bytes() -> bytes: return b"%PDF-1.4\n%stub\n" +@pytest.mark.parametrize("category", [FileCategory.TEXT, FileCategory.HTML, FileCategory.IMAGE]) +def test_non_pdf_categories_skip_pdf_page_count(monkeypatch: pytest.MonkeyPatch, category: FileCategory) -> None: + monkeypatch.setattr( + "nemo_retriever.service.routers.ingest._count_pdf_pages", + lambda _: pytest.fail("non-PDF uploads must not be parsed as PDFs"), + ) + + assert _route_by_page_count(b"not a PDF", IngestRequest(), category) is PoolType.REALTIME + + def _wait_for_items(captured_items: list[WorkItem], count: int) -> None: import time as _time @@ -188,8 +215,12 @@ def test_ingest_without_spec_falls_back_to_legacy_pipeline( assert resp.status_code == 202, resp.text body = resp.json() assert "document_id" in body + assert body["document_id"] == body["attempt_id"] assert body["job_id"] == job_id + status = app_with_stub_pool.get(f"/v1/ingest/job/{job_id}/document/{body['document_id']}") + assert status.status_code in (200, 202), status.text + # Wait briefly for the async worker loop to consume the queued item. import time as _time @@ -235,7 +266,9 @@ def test_ingest_with_valid_spec_attaches_to_work_item( assert item.pipeline_spec["stage_order"] == ["extract"] -def test_ingest_rejects_trust_sensitive_override(app_with_stub_pool: TestClient) -> None: +def test_ingest_rejects_trust_sensitive_override( + app_with_stub_pool: TestClient, +) -> None: job_id = create_test_job(app_with_stub_pool) metadata = {"pipeline": {"extract_params": {"page_elements_invoke_url": "http://attacker/"}}} resp = app_with_stub_pool.post( @@ -247,7 +280,9 @@ def test_ingest_rejects_trust_sensitive_override(app_with_stub_pool: TestClient) assert "trust-sensitive" in resp.json()["detail"] -def test_ingest_rejects_caption_when_endpoint_not_configured(app_with_stub_pool: TestClient) -> None: +def test_ingest_rejects_caption_when_endpoint_not_configured( + app_with_stub_pool: TestClient, +) -> None: """Without ``nim_endpoints.caption_invoke_url``, caption overrides are 403.""" job_id = create_test_job(app_with_stub_pool) metadata = {"pipeline": {"caption_params": {"prompt": "Describe"}}} @@ -260,10 +295,17 @@ def test_ingest_rejects_caption_when_endpoint_not_configured(app_with_stub_pool: assert "caption" in resp.json()["detail"].lower() -def test_ingest_rejects_webhook_when_sinks_disabled(app_with_stub_pool: TestClient) -> None: +def test_ingest_rejects_webhook_when_sinks_disabled( + app_with_stub_pool: TestClient, +) -> None: """Without ``sinks.webhook_url_prefixes`` set, the ``webhook`` stage is not allowed.""" job_id = create_test_job(app_with_stub_pool) - metadata = {"pipeline": {"webhook_params": {"endpoint_url": "http://x/"}, "stage_order": ["webhook"]}} + metadata = { + "pipeline": { + "webhook_params": {"endpoint_url": "http://x/"}, + "stage_order": ["webhook"], + } + } resp = app_with_stub_pool.post( f"/v1/ingest/job/{job_id}/document", files={"file": ("doc.pdf", _make_pdf_bytes(), "application/pdf")}, @@ -290,7 +332,9 @@ def test_ingest_rejects_webhook_params_without_stage_when_sinks_disabled( assert "disabled" in resp.json()["detail"].lower() -def test_create_job_returns_201_and_aggregate_fields(app_with_stub_pool: TestClient) -> None: +def test_create_job_returns_201_and_aggregate_fields( + app_with_stub_pool: TestClient, +) -> None: """POST /v1/ingest/job opens a fresh aggregate with status=pending.""" resp = app_with_stub_pool.post( "/v1/ingest/job", @@ -304,6 +348,34 @@ def test_create_job_returns_201_and_aggregate_fields(app_with_stub_pool: TestCli assert body["job_id"] +@pytest.mark.parametrize( + "payload", + [ + {"expected_documents": 1, "operation": "replace", "target_document_id": "doc"}, + { + "expected_documents": 2, + "collection_name": "research", + "operation": "replace", + "target_document_id": "doc", + }, + {"expected_documents": 1, "collection_name": "research", "operation": "replace"}, + {"expected_documents": 1, "operation": "append", "target_document_id": "doc"}, + ], +) +def test_invalid_collection_job_is_rejected_before_registration( + app_with_stub_pool: TestClient, + payload: dict[str, Any], +) -> None: + tracker = get_job_tracker() + assert tracker is not None + initial_jobs = len(tracker.all_jobs()) + + response = app_with_stub_pool.post("/v1/ingest/job", json=payload) + + assert response.status_code == 422 + assert len(tracker.all_jobs()) == initial_jobs + + def test_create_job_succeeds_when_tracing_span_setup_fails( app_with_stub_pool: TestClient, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -486,7 +558,11 @@ def test_job_upload_routes_emit_accept_spans( page_resp = client.post( f"/v1/ingest/job/{job_id}/page", files={"file": ("page.png", b"page", "image/png")}, - data={"document_id": "source-doc", "page_number": "1", "filename": "source.pdf"}, + data={ + "document_id": "source-doc", + "page_number": "1", + "filename": "source.pdf", + }, ) whole_resp = client.post( f"/v1/ingest/job/{job_id}/whole", @@ -497,6 +573,8 @@ def test_job_upload_routes_emit_accept_spans( assert document_resp.status_code == 202, document_resp.text assert page_resp.status_code == 202, page_resp.text assert whole_resp.status_code == 202, whole_resp.text + assert document_resp.json()["document_id"] == document_resp.json()["attempt_id"] + assert whole_resp.json()["document_id"] == whole_resp.json()["attempt_id"] _wait_for_items(captured_items, 3) names = {span.name for span in exported_spans} @@ -572,7 +650,9 @@ def test_job_upload_accept_span_prefers_inbound_traceparent_over_job_trace( assert captured_items[0].trace_context["traceparent"].split("-")[1] == inbound_trace_id -def test_dashboard_job_views_include_trace_id(traced_gateway_app: tuple[TestClient, list[Any]]) -> None: +def test_dashboard_job_views_include_trace_id( + traced_gateway_app: tuple[TestClient, list[Any]], +) -> None: client, exported_spans = traced_gateway_app resp = client.post( @@ -602,7 +682,56 @@ def test_dashboard_job_views_include_trace_id(traced_gateway_app: tuple[TestClie assert exported_spans -def test_create_job_retain_results_persisted_on_aggregate(app_with_stub_pool: TestClient) -> None: +def test_dashboard_vdb_query_forwards_internal_token(monkeypatch) -> None: + captured: dict[str, Any] = {} + + class FakeClient: + def __init__(self, **_kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def post(self, url: str, **kwargs): + captured.update(url=url, **kwargs) + return SimpleNamespace( + raise_for_status=lambda: None, + json=lambda: {"results": []}, + ) + + monkeypatch.setattr( + "nemo_retriever.service.routers.dashboard.httpx.AsyncClient", + FakeClient, + ) + request = SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + config=ServiceConfig( + vectordb=VectorDbConfig( + enabled=True, + vectordb_url="http://vectordb", + internal_api_token="internal-secret", + ) + ) + ) + ) + ) + asyncio.run( + vdb_query( + VdbQueryRequest(query="retrieval"), + request, + ) + ) + + assert captured["headers"] == {"X-NRL-Internal-Token": "internal-secret"} + + +def test_create_job_retain_results_persisted_on_aggregate( + app_with_stub_pool: TestClient, +) -> None: from nemo_retriever.service.services.job_tracker import get_job_tracker resp = app_with_stub_pool.post( @@ -618,6 +747,46 @@ def test_create_job_retain_results_persisted_on_aggregate(app_with_stub_pool: Te assert agg.retain_results is True +def test_create_job_idempotency_includes_behavior_but_not_label( + app_with_stub_pool: TestClient, +) -> None: + payload = { + "expected_documents": 1, + "idempotency_key": "request-key", + "retain_results": False, + "label": "first label", + } + created = app_with_stub_pool.post("/v1/ingest/job", json=payload) + assert created.status_code == 201, created.text + + replay = app_with_stub_pool.post( + "/v1/ingest/job", + json={**payload, "label": "updated label"}, + ) + assert replay.status_code == 200, replay.text + assert replay.json()["job_id"] == created.json()["job_id"] + + conflict = app_with_stub_pool.post( + "/v1/ingest/job", + json={**payload, "retain_results": True}, + ) + assert conflict.status_code == 409, conflict.text + + +def test_create_job_rejects_invalid_collection_name_before_registration( + app_with_stub_pool: TestClient, +) -> None: + tracker = get_job_tracker() + assert tracker is not None + response = app_with_stub_pool.post( + "/v1/ingest/job", + json={"expected_documents": 1, "collection_name": "foo/documents"}, + ) + + assert response.status_code == 422 + assert tracker.all_jobs() == [] + + def test_get_job_returns_aggregate_snapshot(app_with_stub_pool: TestClient) -> None: job_id = create_test_job(app_with_stub_pool, expected_documents=2) resp = app_with_stub_pool.get(f"/v1/ingest/job/{job_id}") @@ -673,6 +842,68 @@ async def test_gateway_enqueue_unregisters_pending_when_broker_unavailable(monke assert unregistered == ["work-id"] +def test_collection_page_is_rejected_before_registration( + app_with_stub_pool: TestClient, + captured_items: list[WorkItem], +) -> None: + from nemo_retriever.service.services.job_tracker import get_job_tracker + + tracker = get_job_tracker() + assert tracker is not None + tracker.register_job( + "collection-page", + expected_documents=1, + collection_name="research", + scope="default", + ) + response = app_with_stub_pool.post( + "/v1/ingest/job/collection-page/page", + files={"file": ("page.png", b"page", "image/png")}, + data={"document_id": "source", "page_number": "1", "filename": "source.pdf"}, + ) + assert response.status_code == 422 + assert tracker.job_documents("collection-page") == [] + assert captured_items == [] + + +@pytest.mark.parametrize("endpoint", ["document", "whole"]) +def test_collection_upload_propagates_server_storage_context( + app_with_stub_pool: TestClient, + captured_items: list[WorkItem], + endpoint: str, +) -> None: + """Both upload routes must emit the same collection work envelope.""" + from nemo_retriever.service.services.job_tracker import get_job_tracker + + tracker = get_job_tracker() + assert tracker is not None + job_id = f"collection-{endpoint}" + tracker.register_job( + job_id, + expected_documents=1, + collection_name="research", + scope="workspace", + operation="append", + ) + response = app_with_stub_pool.post( + f"/v1/ingest/job/{job_id}/{endpoint}", + headers={"X-NRL-Scope": "workspace"}, + files={"file": ("report.txt", b"finding", "text/plain")}, + data={"metadata": "{}"}, + ) + assert response.status_code == 202, response.text + _wait_for_items(captured_items, 1) + item = captured_items[0] + body = response.json() + assert item.write.scope == "workspace" + assert item.write.collection_name == "research" + assert item.write.operation is IngestOperation.APPEND + assert item.write.content_sha256 == hashlib.sha256(b"finding").hexdigest() + assert item.write.storage_document_id == body["document_id"] + assert item.id == body["attempt_id"] + assert body["document_id"] != body["attempt_id"] + + def test_upload_beyond_capacity_returns_409(app_with_stub_pool: TestClient, captured_items: list[WorkItem]) -> None: """The (expected_documents + 1)th upload must be rejected with 409.""" job_id = create_test_job(app_with_stub_pool, expected_documents=1) @@ -701,3 +932,21 @@ def test_pipeline_config_endpoint_reports_allowed_overrides( assert body["allowed_overrides"]["mode"] == "allow_list" assert "dpi" in body["allowed_overrides"]["allowed_extract_keys"] assert "ocr_invoke_url" in body["allowed_overrides"]["denied_key_substrings"] + + +def test_document_registration_returns_503_if_tracker_is_unavailable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from nemo_retriever.service.routers import ingest + + monkeypatch.setattr(ingest, "get_job_tracker", lambda: None) + + with pytest.raises(HTTPException) as exc_info: + ingest._register_document_under_job( + document_id="attempt", + job_id="job", + filename="document.pdf", + ) + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == "Job tracker not available" diff --git a/nemo_retriever/tests/test_service_job_callback_diagnostics.py b/nemo_retriever/tests/test_service_job_callback_diagnostics.py index b3eb3ba63c..a0bc144219 100644 --- a/nemo_retriever/tests/test_service_job_callback_diagnostics.py +++ b/nemo_retriever/tests/test_service_job_callback_diagnostics.py @@ -39,6 +39,7 @@ from nemo_retriever.service.app import create_app from nemo_retriever.service.config import ( + AuthConfig, PipelineOverridesConfig, PipelinePoolConfig, ServiceConfig, @@ -190,6 +191,7 @@ def _stub_factory(_config: ServiceConfig): cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), pipeline_overrides=PipelineOverridesConfig(), ) diff --git a/nemo_retriever/tests/test_service_job_tracker.py b/nemo_retriever/tests/test_service_job_tracker.py index 5f6609ac38..95b15ae976 100644 --- a/nemo_retriever/tests/test_service_job_tracker.py +++ b/nemo_retriever/tests/test_service_job_tracker.py @@ -39,7 +39,6 @@ MarkOutcome, ) - # ---------------------------------------------------------------------- # Helpers # ---------------------------------------------------------------------- @@ -161,6 +160,57 @@ def test_expired_job_id_can_be_reused() -> None: assert replacement.expected_documents == 2 +def test_expired_job_releases_idempotency_key() -> None: + tracker = JobTracker(max_jobs=1, stale_job_ttl_s=60.0) + tracker.register_job( + "expired", + expected_documents=1, + scope="workspace-a", + idempotency_key="request-key", + idempotency_fingerprint="old-request", + ) + _age_job(tracker, "expired", seconds=61) + + replacement = tracker.register_job( + "replacement", + expected_documents=2, + scope="workspace-a", + idempotency_key="request-key", + idempotency_fingerprint="new-request", + ) + + assert replacement.job_id == "replacement" + + +def test_dropping_a_keyed_job_leaves_the_current_key_owner_replayable() -> None: + tracker = JobTracker(stale_job_ttl_s=60.0) + tracker.register_job( + "unfingerprinted", + expected_documents=1, + scope="workspace-a", + idempotency_key="request-key", + ) + owner = tracker.register_job( + "owner", + expected_documents=1, + scope="workspace-a", + idempotency_key="request-key", + idempotency_fingerprint="request", + ) + _age_job(tracker, "unfingerprinted", seconds=61) + tracker.get_job("unfingerprinted") + + replay = tracker.register_job( + "replay", + expected_documents=1, + scope="workspace-a", + idempotency_key="request-key", + idempotency_fingerprint="request", + ) + + assert replay.job_id == owner.job_id + + @pytest.mark.parametrize( "reader", [ diff --git a/nemo_retriever/tests/test_service_local_models.py b/nemo_retriever/tests/test_service_local_models.py index 7b4e5eac74..58e5cc8dba 100644 --- a/nemo_retriever/tests/test_service_local_models.py +++ b/nemo_retriever/tests/test_service_local_models.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import pytest + from nemo_retriever.service.config import ( LocalEmbedConfig, LocalModelsConfig, @@ -84,6 +86,32 @@ def test_build_extract_params_nim_url_wins_over_local_flags() -> None: assert ep.use_table_structure is True +def test_build_extract_params_from_nemotron_parse_nim_config() -> None: + nim = NimEndpointsConfig( + nemotron_parse_invoke_url="http://parse-nim/v1/chat/completions", + nemotron_parse_model="nvidia/nemotron-parse-v1.2", + api_key="k", + ) + ep = build_extract_params(nim, LocalModelsConfig()) + assert ep.method == "nemotron_parse" + assert ep.nemotron_parse_invoke_url == "http://parse-nim/v1/chat/completions" + assert ep.nemotron_parse_model == "nvidia/nemotron-parse-v1.2" + assert ep.api_key == "k" + + +def test_nemotron_parse_model_requires_endpoint() -> None: + with pytest.raises(ValueError, match="nemotron_parse_model requires"): + NimEndpointsConfig(nemotron_parse_model="nvidia/nemotron-parse-v1.2") + + +def test_build_extract_params_accepts_nemotron_parse_endpoint_only() -> None: + nim = NimEndpointsConfig(nemotron_parse_invoke_url="https://integrate.api.nvidia.com/v1/chat/completions") + ep = build_extract_params(nim, LocalModelsConfig()) + assert ep.method == "nemotron_parse" + assert ep.nemotron_parse_invoke_url == "https://integrate.api.nvidia.com/v1/chat/completions" + assert ep.nemotron_parse_model is None + + def test_build_asr_params_local_when_enabled() -> None: local = LocalModelsConfig(enabled=True) asr = build_asr_params(NimEndpointsConfig(), local) diff --git a/nemo_retriever/tests/test_service_mcp.py b/nemo_retriever/tests/test_service_mcp.py index 59d54261ec..bd4d45d73b 100644 --- a/nemo_retriever/tests/test_service_mcp.py +++ b/nemo_retriever/tests/test_service_mcp.py @@ -14,6 +14,7 @@ from nemo_retriever.service.app import create_app from nemo_retriever.service.config import ( + AgenticConfig, AuthConfig, LoggingConfig, MCPConfig, @@ -24,6 +25,7 @@ MCPDocumentInput, ServiceMCPClient, ServiceMCPSettings, + build_mcp, settings_from_service_config, ) from nemo_retriever.service.services.pipeline_pool import WorkItem @@ -80,6 +82,129 @@ def _handler(request: httpx.Request) -> httpx.Response: assert result["results"][0]["hits"][0]["text"] == "match" +def test_query_tool_strips_agentic_flag_from_payload() -> None: + seen: dict[str, Any] = {} + + def _handler(request: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(request.content) + return httpx.Response(200, json={"results": [{"hits": []}]}) + + client = ServiceMCPClient( + ServiceMCPSettings(base_url="http://service:7670"), + transport=httpx.MockTransport(_handler), + ) + + _run(client.query("q", payload={"agentic": True, "filters": {"k": "v"}})) + + assert seen["body"] == { + "filters": {"k": "v"}, + "query": "q", + "top_k": 5, + "format": "hits", + } + assert "agentic" not in seen["body"] + + +def test_agentic_query_client_posts_agentic_flag_on_v1_query() -> None: + seen: dict[str, Any] = {} + + def _handler(request: httpx.Request) -> httpx.Response: + seen["path"] = request.url.path + seen["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "results": [ + { + "hits": [ + { + "text": None, + "metadata": { + "result_source": "selection_agent", + "rank": 1, + }, + "source": "report.pdf", + "source_id": None, + "path": None, + "page_number": None, + "pdf_basename": None, + "pdf_page": None, + } + ] + } + ] + }, + ) + + client = ServiceMCPClient( + ServiceMCPSettings( + base_url="http://service:7670", + query_methods="all", + agentic_request_timeout_s=300.0, + ), + transport=httpx.MockTransport(_handler), + ) + + result = _run(client.agentic_query("What is indexed?", top_k=2)) + + assert seen == { + "path": "/v1/query", + "body": { + "query": "What is indexed?", + "top_k": 2, + "format": "hits", + "agentic": True, + }, + } + assert result["results"][0]["hits"][0]["source"] == "report.pdf" + + +def test_query_methods_gate_mcp_retrieval_tools() -> None: + classic = {tool.name for tool in _run(build_mcp(ServiceMCPSettings(query_methods="classic")).list_tools())} + agentic = {tool.name for tool in _run(build_mcp(ServiceMCPSettings(query_methods="agentic")).list_tools())} + all_tools = {tool.name for tool in _run(build_mcp(ServiceMCPSettings(query_methods="all")).list_tools())} + + assert "query" in classic + assert "agentic_query" not in classic + assert "query" not in agentic + assert "agentic_query" in agentic + assert "query" in all_tools + assert "agentic_query" in all_tools + assert "answer" in classic and "answer" in agentic and "answer" in all_tools + + +def test_settings_from_service_config_maps_query_methods_when_agentic_enabled() -> None: + settings = settings_from_service_config( + ServiceConfig( + agentic=AgenticConfig( + enabled=True, + llm_model="model", + invoke_url="https://llm.example/v1/chat/completions", + request_timeout_s=321.0, + ), + mcp=MCPConfig(query_methods="all"), + ) + ) + + assert settings.query_methods == "all" + assert settings.enable_agentic_query is True + assert settings.enable_classic_query is True + assert settings.agentic_request_timeout_s == 321.0 + + +def test_settings_from_service_config_drops_agentic_tools_when_agentic_disabled() -> None: + settings = settings_from_service_config( + ServiceConfig( + agentic=AgenticConfig(enabled=False), + mcp=MCPConfig(query_methods="all"), + ) + ) + + assert settings.query_methods == "classic" + assert settings.enable_agentic_query is False + assert settings.enable_classic_query is True + + def test_ingest_documents_accepts_inline_base64_upload() -> None: calls: list[tuple[str, str]] = [] upload_body = b"" @@ -187,7 +312,7 @@ async def _stub_work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: mode="standalone", logging=LoggingConfig(file=str(tmp_path / "service.log")), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), - auth=AuthConfig(api_token="secret"), + auth=AuthConfig(enabled=True, api_token="secret"), ) app = create_app(cfg) @@ -197,3 +322,62 @@ async def _stub_work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: assert unauthorized.status_code == 401 assert authorized.status_code != 401 + + +def test_query_tool_client_posts_rerank_controls() -> None: + seen: dict[str, Any] = {} + + def _handler(request: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(request.content) + return httpx.Response(200, json={"results": [{"hits": []}]}) + + client = ServiceMCPClient( + ServiceMCPSettings(base_url="http://service:7670"), + transport=httpx.MockTransport(_handler), + ) + + _run(client.query("What is indexed?", top_k=3, rerank=True, rerank_top_k=20)) + + assert seen["body"] == { + "query": "What is indexed?", + "top_k": 3, + "format": "hits", + "rerank": True, + "rerank_top_k": 20, + } + + +def test_query_tool_explicit_controls_override_payload() -> None: + seen: dict[str, Any] = {} + + def _handler(request: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(request.content) + return httpx.Response(200, json={"results": [{"hits": []}]}) + + client = ServiceMCPClient( + ServiceMCPSettings(base_url="http://service:7670"), + transport=httpx.MockTransport(_handler), + ) + + _run( + client.query( + "authoritative query", + top_k=3, + payload={ + "query": "payload query", + "top_k": 99, + "format": "evidence", + "agentic": True, + "rerank": True, + "rerank_top_k": 50, + "filters": {"source": "a.pdf"}, + }, + ) + ) + + assert seen["body"] == { + "query": "authoritative query", + "top_k": 3, + "format": "hits", + "filters": {"source": "a.pdf"}, + } diff --git a/nemo_retriever/tests/test_service_media_dependency_gate.py b/nemo_retriever/tests/test_service_media_dependency_gate.py index 263150085a..b618022ae4 100644 --- a/nemo_retriever/tests/test_service_media_dependency_gate.py +++ b/nemo_retriever/tests/test_service_media_dependency_gate.py @@ -34,6 +34,7 @@ from nemo_retriever.service.app import _check_media_dependencies, create_app from nemo_retriever.service.config import ( + AuthConfig, PipelineOverridesConfig, PipelinePoolConfig, ServiceConfig, @@ -158,6 +159,7 @@ def _stub_batch(_config: ServiceConfig): cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), pipeline_overrides=PipelineOverridesConfig(), ) diff --git a/nemo_retriever/tests/test_service_metrics_lifecycle.py b/nemo_retriever/tests/test_service_metrics_lifecycle.py new file mode 100644 index 0000000000..a1133f8386 --- /dev/null +++ b/nemo_retriever/tests/test_service_metrics_lifecycle.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lifecycle projection tests for tracker-owned ingest metrics.""" + +from nemo_retriever.service.services.job_tracker import JobTracker, MarkOutcome +from nemo_retriever.service.services import metrics as metrics_module +from nemo_retriever.service.services.metrics import IngestMetrics + + +def _wired_tracker() -> tuple[JobTracker, IngestMetrics]: + tracker = JobTracker() + metrics = IngestMetrics() + tracker.add_terminal_observer(metrics.record_terminal_transition) + return tracker, metrics + + +def test_terminal_transitions_reconcile_document_and_job_metrics() -> None: + tracker, metrics = _wired_tracker() + tracker.register_job("job", expected_documents=2) + metrics.record_job_created("job") + for document_id in ("ok", "bad"): + tracker.register_document(document_id, job_id="job") + metrics.record_document_accepted(document_id=document_id, job_id="job") + tracker.mark_processing(document_id) + + assert tracker.mark_completed("ok", elapsed_s=1.25) is MarkOutcome.TRANSITIONED + assert tracker.mark_failed("bad", "pipeline failed", elapsed_s=2.5) is MarkOutcome.TRANSITIONED + + job = metrics.get_job("job") + assert job is not None + assert (job.documents_completed, job.documents_failed, job.status) == ( + 1, + 1, + "partial_success", + ) + assert job.completed_at is not None + assert job.wall_duration_s is not None + completed = metrics.get_document("ok") + failed = metrics.get_document("bad") + assert completed is not None and completed.status == "completed" + assert completed.completed_at is not None and completed.processing_duration_s == 1.25 + assert failed is not None and failed.status == "failed" + assert failed.error == "pipeline failed" and failed.processing_duration_s == 2.5 + + +def test_duplicate_and_unknown_terminal_transitions_do_not_change_metrics() -> None: + tracker, metrics = _wired_tracker() + tracker.register_job("job", expected_documents=1) + metrics.record_job_created("job") + tracker.register_document("doc", job_id="job") + metrics.record_document_accepted(document_id="doc", job_id="job") + + assert tracker.mark_completed("doc", elapsed_s=1) is MarkOutcome.TRANSITIONED + before = metrics.get_job("job") + assert tracker.mark_completed("doc", elapsed_s=99) is MarkOutcome.IDEMPOTENT + assert tracker.mark_failed("unknown", "missing") is MarkOutcome.UNKNOWN_DOCUMENT + assert metrics.get_job("job") == before + + +def test_explicit_page_terminal_transitions_reconcile_page_counts() -> None: + tracker, metrics = _wired_tracker() + tracker.register_job("job", expected_documents=2) + metrics.record_job_created("job") + for page_id in ("page-1", "page-2"): + tracker.register_document(page_id, job_id="job") + metrics.record_page_accepted(page_id=page_id, document_id="source", job_id="job") + tracker.mark_processing(page_id) + + tracker.mark_completed("page-1", elapsed_s=0.5) + tracker.mark_failed("page-2", "OCR failed", elapsed_s=0.75) + + job = metrics.get_job("job") + assert job is not None + assert (job.pages_total, job.pages_completed, job.pages_failed) == (2, 1, 1) + page = metrics.get_page("page-2") + assert page is not None + assert page.job_id == "job" and page.status == "failed" and page.error == "OCR failed" + + +def test_page_terminal_counts_survive_recent_page_eviction(monkeypatch) -> None: + monkeypatch.setattr(metrics_module, "MAX_RECENT_PAGES", 1) + tracker, metrics = _wired_tracker() + tracker.register_job("job", expected_documents=2) + metrics.record_job_created("job") + for page_id in ("page-1", "page-2"): + tracker.register_document(page_id, job_id="job") + metrics.record_page_accepted(page_id=page_id, document_id="source", job_id="job") + tracker.mark_processing(page_id) + + tracker.mark_completed("page-1") + tracker.mark_failed("page-2", "OCR failed") + + job = metrics.get_job("job") + assert job is not None + assert (job.pages_total, job.pages_completed, job.pages_failed) == (2, 1, 1) diff --git a/nemo_retriever/tests/test_service_packaging.py b/nemo_retriever/tests/test_service_packaging.py index 789bf31741..47de46076a 100644 --- a/nemo_retriever/tests/test_service_packaging.py +++ b/nemo_retriever/tests/test_service_packaging.py @@ -25,4 +25,13 @@ def test_service_extra_includes_litellm_for_answer_generation() -> None: litellm = next((req for req in requirements if req.name == "litellm"), None) assert litellm is not None - assert any(str(spec).startswith(">=") and "1.86.0" in str(spec) for spec in litellm.specifier) + assert any(str(spec).startswith(">=") and "1.95.0rc3" in str(spec) for spec in litellm.specifier) + + +def test_core_dependencies_include_tokenizer_stack() -> None: + pyproject = tomllib.loads((PROJECT_ROOT / "pyproject.toml").read_text(encoding="utf-8")) + requirements = [Requirement(dep) for dep in pyproject["project"]["dependencies"]] + names = {req.name for req in requirements} + + assert "tokenizers" in names + assert "huggingface-hub" in names diff --git a/nemo_retriever/tests/test_service_pipeline_spec.py b/nemo_retriever/tests/test_service_pipeline_spec.py index eca2567f2b..600a48cd86 100644 --- a/nemo_retriever/tests/test_service_pipeline_spec.py +++ b/nemo_retriever/tests/test_service_pipeline_spec.py @@ -16,6 +16,8 @@ from __future__ import annotations +import json + import pytest from nemo_retriever.common.params import DedupParams, EmbedParams, ExtractParams @@ -25,14 +27,19 @@ from nemo_retriever.service.services.pipeline_executor import ( _build_graph_ingestor_from_spec, _merge_server_owned, + _post_records_to_vectordb, _request_needs_asr_params, + _resolve_extract_params, _resolve_service_extraction_mode, _run_pipeline_in_process, _TRUST_OWNED_EMBED_KEYS, _TRUST_OWNED_EXTRACT_KEYS, ) +from nemo_retriever.common.schemas.collections import IngestOperation +from nemo_retriever.service.services.pipeline_pool import DocumentWriteContext from nemo_retriever.service.utils.file_type import infer_extraction_mode_from_filename from nemo_retriever.service.service_ingestor import ServiceIngestor +from nemo_retriever.service.client import InMemoryUpload class _TinyTokenizer: @@ -73,6 +80,122 @@ def test_compact_result_schema_populates_pipeline_payload() -> None: assert PipelineSpec.model_validate(payload).result_schema == "compact" +def test_service_inline_text_builds_in_memory_uploads(monkeypatch: pytest.MonkeyPatch) -> None: + ingestor = ServiceIngestor(base_url="http://retriever.example") + monkeypatch.setattr("tempfile.mkdtemp", lambda *args, **kwargs: pytest.fail("inline text must remain in memory")) + + ingestor.texts(["first", "first"]).extract(split_config={"text": {"max_tokens": 12}}) + + assert ingestor._collect_inputs() == [ + InMemoryUpload( + filename="inline://00000000", + content=b"first", + content_type="text/plain; charset=utf-8", + classification_filename="inline-00000000.txt", + ), + InMemoryUpload( + filename="inline://00000001", + content=b"first", + content_type="text/plain; charset=utf-8", + classification_filename="inline-00000001.txt", + ), + ] + assert ingestor._pipeline_payload()["extraction_mode"] == "auto" + assert ingestor._pipeline_payload()["split_config"] == {"text": {"max_tokens": 12}} + + +def test_service_inline_text_replaces_and_validates_inputs() -> None: + ingestor = ServiceIngestor(base_url="http://retriever.example").texts("first").texts(["second"]) + + assert [item.filename for item in ingestor._collect_inputs()] == ["inline://00000000"] + assert [item.content for item in ingestor._collect_inputs()] == [b"second"] + + with pytest.raises(TypeError, match=r"texts\[1\] must be a string"): + ServiceIngestor(base_url="http://retriever.example").texts(["valid", None]) + + +@pytest.mark.parametrize("files_first", [True, False]) +def test_service_inline_text_composes_with_files_and_uses_auto_routing(tmp_path, files_first: bool) -> None: + document = tmp_path / "document.txt" + document.write_text("document", encoding="utf-8") + + ingestor = ServiceIngestor(base_url="http://retriever.example") + if files_first: + ingestor.files(str(document)).texts(["inline"]) + else: + ingestor.texts(["inline"]).files(str(document)) + ingestor.extract(split_config={"text": {"max_tokens": 12}}) + + inputs = ingestor._collect_inputs() + assert inputs[0] == document + assert inputs[1] == InMemoryUpload( + filename="inline://00000000", + content=b"inline", + content_type="text/plain; charset=utf-8", + classification_filename="inline-00000000.txt", + ) + assert ingestor._pipeline_payload()["extraction_mode"] == "auto" + assert ingestor._pipeline_payload()["split_config"] == {"text": {"max_tokens": 12}} + + +@pytest.mark.parametrize("inline_texts", [[], ["", " \n"]]) +def test_service_empty_inline_text_does_not_hide_files(tmp_path, inline_texts: list[str]) -> None: + document = tmp_path / "document.txt" + document.write_text("document", encoding="utf-8") + + ingestor = ServiceIngestor(base_url="http://retriever.example").files(str(document)).texts(inline_texts) + + assert ingestor._collect_inputs()[0] == document + assert ingestor._pipeline_payload() is None + + +@pytest.mark.parametrize(("inline_texts", "expected_mode"), [([], "pdf"), ([""], "auto")]) +def test_service_empty_inline_list_preserves_explicit_extraction_mode( + tmp_path, inline_texts: list[str], expected_mode: str +) -> None: + document = tmp_path / "document.pdf" + document.write_bytes(b"%PDF-1.4 stub") + + ingestor = ( + ServiceIngestor(base_url="http://retriever.example") + .files(str(document)) + .texts(inline_texts) + .extract(extraction_mode="pdf") + ) + + assert ingestor._pipeline_payload()["extraction_mode"] == expected_mode + + +@pytest.mark.parametrize("values", [[], ["", " \n"]]) +@pytest.mark.parametrize( + ("result_schema", "expected_columns"), + [ + ("legacy", ["text", "content", "path", "page_number", "metadata"]), + ("compact", ["text", "source_id", "element_type", "page_number"]), + ], +) +def test_service_inline_empty_corpus_short_circuits_with_schema( + values: list[str], + result_schema: str, + expected_columns: list[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + ingestor = ServiceIngestor(base_url="http://retriever.example").texts(values).embed() + monkeypatch.setattr( + ingestor, + "ingest_stream", + lambda **kwargs: pytest.fail("empty inline corpus must not contact the service"), + ) + + result = ingestor.ingest(result_schema=result_schema) + + assert result.job_id is None + assert result.failures == [] + assert result.dataframe.empty + assert result.dataframe.columns.tolist() == expected_columns + assert ingestor._collect_inputs() == [] + + def test_legacy_pipeline_payload_disables_bulk_result_payloads() -> None: ing = ServiceIngestor(base_url="http://example:7670") ing.all_tasks() @@ -201,6 +324,13 @@ def test_client_rejects_server_owned_keys() -> None: ing.extract(ExtractParams(page_elements_invoke_url="http://attacker/")) +def test_policy_rejects_client_nemotron_parse_model_override() -> None: + policy = PipelineOverridesConfig().to_policy() + spec = PipelineSpec(extract_params={"method": "nemotron_parse", "nemotron_parse_model": "attacker/model"}) + with pytest.raises(PolicyError): + validate_pipeline_spec(spec, policy) + + def test_future_phase_methods_raise_informative_error() -> None: """Methods deferred to follow-up phases still produce a clear error. @@ -321,14 +451,48 @@ def test_merge_preserves_server_extract_endpoints() -> None: "page_elements_invoke_url": "http://server/page_elements", "ocr_invoke_url": "http://server/ocr", "api_key": "server-token", + "nemotron_parse_invoke_url": "http://server/parse", + "nemotron_parse_model": "nvidia/nemotron-parse-v1.2", "dpi": 150, } - override = {"dpi": 600, "page_elements_invoke_url": "http://attacker/"} + override = { + "dpi": 600, + "page_elements_invoke_url": "http://attacker/", + "nemotron_parse_model": "attacker/model", + } merged = _merge_server_owned(base, override, _TRUST_OWNED_EXTRACT_KEYS) assert merged["dpi"] == 600 assert merged["page_elements_invoke_url"] == "http://server/page_elements" assert merged["ocr_invoke_url"] == "http://server/ocr" assert merged["api_key"] == "server-token" + assert merged["nemotron_parse_model"] == "nvidia/nemotron-parse-v1.2" + + +@pytest.mark.parametrize("method", ["pdfium", "pdfium_hybrid", "ocr"]) +def test_resolve_extract_params_drops_parse_fields_for_other_methods(method: str) -> None: + base = { + "method": "nemotron_parse", + "nemotron_parse_invoke_url": "http://server/parse", + "nemotron_parse_model": "nvidia/nemotron-parse-v1.2", + "api_key": "server-token", + } + resolved = _resolve_extract_params(base, {"method": method}) + assert resolved.method == method + assert resolved.nemotron_parse_invoke_url is None + assert resolved.nemotron_parse_model is None + assert resolved.api_key == "server-token" + + +def test_resolve_extract_params_preserves_parse_fields_for_parse_method() -> None: + base = { + "method": "nemotron_parse", + "nemotron_parse_invoke_url": "http://server/parse", + "nemotron_parse_model": "nvidia/nemotron-parse-v1.2", + } + resolved = _resolve_extract_params(base, {"method": "nemotron_parse"}) + assert resolved.method == "nemotron_parse" + assert resolved.nemotron_parse_invoke_url == "http://server/parse" + assert resolved.nemotron_parse_model == "nvidia/nemotron-parse-v1.2" def test_merge_preserves_server_embed_endpoints() -> None: @@ -486,6 +650,7 @@ def test_build_graph_ingestor_attaches_asr_params_for_explicit_audio_mode() -> N ("README.md", "text"), ("payload.json", "text"), ("setup.sh", "text"), + ("inline://00000000", "text"), ("page.html", "html"), ("report.pdf", "pdf"), ("diagram.png", "image"), @@ -501,6 +666,7 @@ def test_infer_extraction_mode_from_filename(filename: str, expected: str | None ("extraction_mode", "filename", "resolved"), [ ("auto", "notes.txt", "text"), + ("auto", "inline://00000000", "text"), ("auto", "page.html", "html"), ("auto", "report.pdf", "pdf"), ("pdf", "notes.txt", "pdf"), @@ -511,7 +677,7 @@ def test_resolve_service_extraction_mode(extraction_mode: str, filename: str, re assert _resolve_service_extraction_mode(extraction_mode, filename) == resolved -def test_build_graph_ingestor_uses_typed_txt_html_shortcuts() -> None: +def test_build_graph_ingestor_routes_txt_and_html_inputs() -> None: base_extract: dict[str, object] = {} spec = {"extraction_mode": "auto", "stage_order": ["extract"]} @@ -524,7 +690,16 @@ def test_build_graph_ingestor_uses_typed_txt_html_shortcuts() -> None: ) assert txt_mode == "text" assert txt_ingestor._extraction_mode == "text" - assert txt_ingestor._text_params is not None + + inline_ingestor, inline_mode, _ = _build_graph_ingestor_from_spec( + "inline://00000000", + b"The quick brown fox", + base_extract, + None, + None, + ) + assert inline_mode == "text" + assert inline_ingestor._extraction_mode == "text" html_ingestor, html_mode, _ = _build_graph_ingestor_from_spec( "page.html", @@ -570,6 +745,24 @@ def test_run_pipeline_in_process_html_txt_produce_rows(monkeypatch: pytest.Monke assert txt_rows >= 1 +def test_run_pipeline_in_process_preserves_service_inline_identity(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("nemo_retriever.common.modality.txt.split._get_tokenizer", lambda *_, **__: _TinyTokenizer()) + + row_count, rows, _ = _run_pipeline_in_process( + "inline://00000003", + "café service".encode("utf-8"), + {}, + None, + None, + None, + ) + + assert row_count == 1 + assert rows[0]["text"] == "café service" + assert rows[0]["path"] == "inline://00000003" + assert rows[0]["metadata"]["source_path"] == "inline://00000003" + + def test_build_graph_ingestor_omits_asr_params_when_worker_unconfigured() -> None: """When the worker has no ASR endpoint, nothing should be attached regardless of filename or extraction mode. @@ -587,3 +780,168 @@ def test_build_graph_ingestor_omits_asr_params_when_worker_unconfigured() -> Non ) assert ingestor._asr_params is None + + +def test_run_pipeline_posts_canonical_pdf_table_image_provenance( + monkeypatch: pytest.MonkeyPatch, +) -> None: + posted: dict[str, object] = {} + graph_rows = [ + { + "text": "table content", + "text_embeddings_1b_v2": {"embedding": [0.1, 0.2]}, + "path": "/documents/report.pdf", + "page_number": 1, + "_page_number": 7, + "_content_type": "table_caption", + "_stored_image_uri": "s3://artifacts/table.png", + "_bbox_xyxy_norm": [0.1, 0.2, 0.8, 0.9], + "metadata": {"content_metadata": {"page_number": 1}}, + } + ] + + class _Ingestor: + def ingest(self): + return graph_rows + + class _Response: + status = 200 + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + def _urlopen(request, timeout): + posted["url"] = request.full_url + posted["timeout"] = timeout + posted["json"] = json.loads(request.data) + return _Response() + + monkeypatch.setattr( + "nemo_retriever.service.services.pipeline_executor._build_graph_ingestor_from_spec", + lambda *_args, **_kwargs: (_Ingestor(), "pdf", False), + ) + monkeypatch.setattr( + "nemo_retriever.service.services.pipeline_executor._sanitize_result_data", + lambda _result, **_kwargs: [], + ) + monkeypatch.setattr("urllib.request.urlopen", _urlopen) + + row_count, _, _ = _run_pipeline_in_process( + "report.pdf", + b"%PDF-1.4 stub", + {}, + None, + vectordb_url="http://vectordb:7671", + write_context=DocumentWriteContext( + scope="tenant-a", + collection_name="papers", + storage_document_id="document-1", + content_sha256="a" * 64, + document_version="version-2", + document_metadata={ + "category": "Finance_Investment", + "source_path": "Finance_Investment/report.pdf", + "source_filename": "report.pdf", + "page_number": 999, + }, + ), + job_id="job-1", + ) + + assert row_count == 1 + assert posted["url"] == "http://vectordb:7671/internal/vectordb/write" + payload = posted["json"] + assert isinstance(payload, dict) + record = payload["records"][0][0] + assert record["document_type"] == "text" + metadata = record["metadata"] + assert metadata["embedding"] == [0.1, 0.2] + assert metadata["content"] == "table content" + assert metadata["source_metadata"] == { + "source_id": "/documents/report.pdf", + "source_name": "report.pdf", + } + assert metadata["content_metadata"] == { + "page_number": 7, + "type": "table", + "fidelity": "ocr", + "stored_image_uri": "s3://artifacts/table.png", + "bbox_xyxy_norm": [0.1, 0.2, 0.8, 0.9], + "category": "Finance_Investment", + "source_path": "Finance_Investment/report.pdf", + "source_filename": "report.pdf", + } + assert payload["scope"] == "tenant-a" + assert payload["collection_name"] == "papers" + assert payload["document_id"] == "document-1" + assert "rows" not in payload + + +def test_post_records_to_vectordb_uses_canonical_internal_payload(monkeypatch: pytest.MonkeyPatch) -> None: + posted: dict[str, object] = {} + + class _Response: + status = 200 + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + def _urlopen(request, timeout): + posted["url"] = request.full_url + posted["headers"] = dict(request.headers) + posted["timeout"] = timeout + posted["json"] = json.loads(request.data) + return _Response() + + monkeypatch.setattr("urllib.request.urlopen", _urlopen) + records = [ + [ + { + "document_type": "text", + "metadata": { + "embedding": [0.1, 0.2], + "content": "canonical chunk", + "content_metadata": {"page_number": 7}, + }, + } + ] + ] + + _post_records_to_vectordb( + records, + "http://vectordb:7671/", + "document.pdf", + context=DocumentWriteContext( + scope="tenant-a", + collection_name="papers", + storage_document_id="document-1", + content_sha256="a" * 64, + document_version="version-2", + operation=IngestOperation.REPLACE, + ), + job_id="job-1", + internal_api_token="internal-token", + ) + + assert posted["url"] == "http://vectordb:7671/internal/vectordb/write" + assert posted["timeout"] == 30 + assert posted["headers"]["X-nrl-internal-token"] == "internal-token" + assert posted["json"] == { + "records": records, + "scope": "tenant-a", + "collection_name": "papers", + "document_id": "document-1", + "job_id": "job-1", + "filename": "document.pdf", + "content_sha256": "a" * 64, + "document_version": "version-2", + "operation": "replace", + } + assert "rows" not in posted["json"] + assert "artifact_prefix" not in posted["json"] diff --git a/nemo_retriever/tests/test_service_pipeline_tracing.py b/nemo_retriever/tests/test_service_pipeline_tracing.py index 87beb90dd4..43e2b2d70b 100644 --- a/nemo_retriever/tests/test_service_pipeline_tracing.py +++ b/nemo_retriever/tests/test_service_pipeline_tracing.py @@ -17,6 +17,7 @@ from nemo_retriever.service import tracing from nemo_retriever.service.services import pipeline_executor +from nemo_retriever.service.services.pipeline_pool import DocumentWriteContext class _CollectingExporter: @@ -137,6 +138,8 @@ class _WorkItem: filename = "contract.pdf" payload = b"%PDF-1.4\n" pipeline_spec = None + job_id = None + write = DocumentWriteContext() def _fake_process_pool_executor(*args: Any, **kwargs: Any) -> ThreadPoolExecutor: return ThreadPoolExecutor(max_workers=1) @@ -163,7 +166,11 @@ def _fake_run_pipeline_in_process(*args: Any, **kwargs: Any) -> tuple[int, list[ max_tasks_per_child=None, model_dump=lambda *args, **kwargs: {}, ), - vectordb=SimpleNamespace(enabled=False, vectordb_url=None), + vectordb=SimpleNamespace( + enabled=False, + vectordb_url=None, + internal_api_token=None, + ), pipeline=SimpleNamespace( realtime_workers=1, batch_workers=1, diff --git a/nemo_retriever/tests/test_service_query_client.py b/nemo_retriever/tests/test_service_query_client.py index 755b2db4d6..5232a68445 100644 --- a/nemo_retriever/tests/test_service_query_client.py +++ b/nemo_retriever/tests/test_service_query_client.py @@ -4,6 +4,8 @@ from __future__ import annotations +import asyncio +import inspect from typing import Any import pytest @@ -20,27 +22,31 @@ def _install_query_response( class FakeResponse: status_code = 200 text = "" + content = b"{}" def json(self) -> dict[str, Any]: return body + # The synchronous methods are facades over the async implementation, so + # even a blocking ``query()`` call goes out over ``httpx.AsyncClient``. class FakeHttpClient: def __init__(self, *, timeout: Any, headers: dict[str, str]) -> None: if calls is not None: calls.append({"timeout": timeout, "headers": headers}) - def __enter__(self) -> "FakeHttpClient": + async def __aenter__(self) -> "FakeHttpClient": return self - def __exit__(self, *_args: Any) -> None: + async def __aexit__(self, *_args: Any) -> None: return None - def post(self, url: str, *, json: dict[str, Any]) -> FakeResponse: + async def request(self, method: str, url: str, *, json: dict[str, Any]) -> FakeResponse: + assert method == "POST" if calls is not None: calls.append({"url": url, "json": json}) return FakeResponse() - monkeypatch.setattr(service_client_module.httpx, "Client", FakeHttpClient) + monkeypatch.setattr(service_client_module.httpx, "AsyncClient", FakeHttpClient) def test_service_client_query_posts_to_v1_query_with_auth(monkeypatch) -> None: @@ -89,3 +95,191 @@ def test_service_client_query_rejects_result_count_mismatch_for_multi_query_requ with pytest.raises(RuntimeError, match=r"expected 2 result set\(s\), got 1"): RetrieverServiceClient(base_url="http://svc:7670").query(["deployment?", "scaling?"], top_k=2) + + +def test_sync_and_async_query_agree_on_the_wire_payload(monkeypatch) -> None: + """The sync facade must send exactly what the async implementation sends.""" + sync_calls: list[dict[str, Any]] = [] + _install_query_response(monkeypatch, {"results": [{"hits": []}]}, sync_calls) + RetrieverServiceClient(base_url="http://svc:7670").query("deployment?", top_k=3) + + async_calls: list[dict[str, Any]] = [] + _install_query_response(monkeypatch, {"results": [{"hits": []}]}, async_calls) + asyncio.run(RetrieverServiceClient(base_url="http://svc:7670").aquery("deployment?", top_k=3)) + + assert sync_calls[1] == async_calls[1] + + +def test_query_requires_top_k_while_aquery_defaults_to_ten() -> None: + """A released signature asymmetry: keep it, do not unify the two. + + ``query`` has always required ``top_k``; ``aquery`` has always defaulted + it. Sharing one implementation must not quietly change either. + """ + sync_top_k = inspect.signature(RetrieverServiceClient.query).parameters["top_k"] + async_top_k = inspect.signature(RetrieverServiceClient.aquery).parameters["top_k"] + + assert sync_top_k.default is inspect.Parameter.empty + assert async_top_k.default == 10 + + +def test_sync_facade_works_from_inside_a_running_event_loop(monkeypatch) -> None: + """``asyncio.run`` cannot nest, so the facade must fall back to a thread.""" + _install_query_response(monkeypatch, {"results": [{"hits": [{"text": "passage"}]}]}) + client = RetrieverServiceClient(base_url="http://svc:7670") + + async def _call_sync_from_async() -> Any: + # Deliberately blocking: this is the call that used to be impossible. + return client.query("deployment?", top_k=1) + + assert asyncio.run(_call_sync_from_async()) == [[{"text": "passage"}]] + + +# ---------------------------------------------------------------------- +# Sync/async parity for the collection, document and job lifecycle +# +# Each operation is implemented once as a coroutine, with the synchronous +# method as a thin facade. These pin that both entry points issue the same +# request and parse the same response, which nothing covered before. +# ---------------------------------------------------------------------- + + +def _install_lifecycle_response( + monkeypatch: pytest.MonkeyPatch, + body: dict[str, Any], + calls: list[dict[str, Any]], +) -> None: + class FakeResponse: + status_code = 200 + text = "" + content = b"{}" + + def json(self) -> dict[str, Any]: + return body + + class FakeAsyncClient: + def __init__(self, *, timeout: Any, headers: dict[str, str]) -> None: + self._headers = headers + + async def __aenter__(self) -> "FakeAsyncClient": + return self + + async def __aexit__(self, *_args: Any) -> None: + return None + + async def request(self, method: str, url: str, **kwargs: Any) -> FakeResponse: + calls.append({"method": method, "url": url, "headers": self._headers, **kwargs}) + return FakeResponse() + + monkeypatch.setattr(service_client_module.httpx, "AsyncClient", FakeAsyncClient) + + +_TS = "2026-01-01T00:00:00+00:00" +_COLLECTION_BODY = { + "name": "research", + "scope": "workspace", + "status": "active", + "created_at": _TS, + "updated_at": _TS, +} +_DOCUMENT_BODY = { + "document_id": "document-1", + "collection_name": "research", + "scope": "workspace", + "filename": "paper.pdf", + "content_sha256": "a" * 64, + "document_version": "v1", + "status": "completed", + "chunk_count": 3, + "created_at": _TS, + "updated_at": _TS, +} +_JOB_BODY = { + "job_id": "job-1", + "expected_documents": 1, + "status": "completed", + "created_at": _TS, +} + +_LIFECYCLE_CASES = [ + ("get_collection", ("research",), {}, _COLLECTION_BODY, "GET", "/v1/collections/research"), + ("list_collections", (), {}, {"items": [], "next_token": None}, "GET", "/v1/collections"), + ("update_collection", ("research",), {"description": "d"}, _COLLECTION_BODY, "PATCH", "/v1/collections/research"), + ( + "delete_collection", + ("research",), + {}, + {"name": "research", "scope": "workspace", "existed": True, "deleted": False, "status": "deleting"}, + "DELETE", + "/v1/collections/research", + ), + ( + "list_documents", + ("research",), + {}, + {"items": [], "next_token": None}, + "GET", + "/v1/collections/research/documents", + ), + ( + "get_document", + ("research", "document-1"), + {}, + _DOCUMENT_BODY, + "GET", + "/v1/collections/research/documents/document-1", + ), + ( + "delete_document", + ("research", "document-1"), + {}, + { + "document_id": "document-1", + "collection_name": "research", + "scope": "workspace", + "existed": True, + "deleted": False, + "status": "deleting", + }, + "DELETE", + "/v1/collections/research/documents/document-1", + ), + ("get_job", ("job-1",), {}, _JOB_BODY, "GET", "/v1/ingest/job/job-1"), + ( + "list_job_documents", + ("job-1",), + {}, + {"job_id": "job-1", "total": 0, "total_filtered": 0, "offset": 0, "limit": 100, "items": []}, + "GET", + "/v1/ingest/job/job-1/documents", + ), +] + + +@pytest.mark.parametrize( + ("name", "args", "kwargs", "body", "method", "path"), + _LIFECYCLE_CASES, + ids=[case[0] for case in _LIFECYCLE_CASES], +) +def test_sync_and_async_lifecycle_methods_issue_identical_requests( + monkeypatch: pytest.MonkeyPatch, + name: str, + args: tuple[Any, ...], + kwargs: dict[str, Any], + body: dict[str, Any], + method: str, + path: str, +) -> None: + sync_calls: list[dict[str, Any]] = [] + _install_lifecycle_response(monkeypatch, body, sync_calls) + client = RetrieverServiceClient(base_url="http://svc:7670", scope="workspace") + sync_result = getattr(client, name)(*args, **kwargs) + + async_calls: list[dict[str, Any]] = [] + _install_lifecycle_response(monkeypatch, body, async_calls) + async_result = asyncio.run(getattr(client, f"a{name}")(*args, **kwargs)) + + assert sync_calls == async_calls + assert sync_calls[0]["method"] == method + assert sync_calls[0]["url"] == f"http://svc:7670{path}" + assert sync_result == async_result diff --git a/nemo_retriever/tests/test_service_query_rerank.py b/nemo_retriever/tests/test_service_query_rerank.py new file mode 100644 index 0000000000..7580926aeb --- /dev/null +++ b/nemo_retriever/tests/test_service_query_rerank.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json + +from fastapi.testclient import TestClient + +from nemo_retriever.service.app import create_app +from nemo_retriever.service.config import ( + AuthConfig, + LoggingConfig, + NimEndpointsConfig, + PipelinePoolConfig, + ServiceConfig, + VectorDbConfig, +) + + +def _configure_noop_workers(monkeypatch): + async def _stub_work(_item): + return 0, [] + + monkeypatch.setattr( + "nemo_retriever.service.services.pipeline_executor.create_realtime_work_fn", + lambda _config: _stub_work, + ) + monkeypatch.setattr( + "nemo_retriever.service.services.pipeline_executor.create_batch_work_fn", + lambda _config: _stub_work, + ) + + +def test_reranked_query_uses_main_service_orchestration(monkeypatch, tmp_path) -> None: + _configure_noop_workers(monkeypatch) + seen: dict[str, object] = {} + + class _Response: + status_code = 200 + content = json.dumps( + { + "results": [ + { + "hits": [ + {"text": "first", "source": "a"}, + {"text": "second", "source": "b"}, + {"text": "third", "source": "c"}, + ] + } + ] + } + ).encode() + + class _Client: + def __init__(self, *args, **kwargs) -> None: + seen["timeout"] = kwargs["timeout"] + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + async def post(self, url, **kwargs): + seen["url"] = url + seen["body"] = json.loads(kwargs["content"]) + return _Response() + + def _rerank(query, hits, **kwargs): + seen["rerank"] = {"query": query, "kwargs": kwargs} + return [{"_rerank_score": 0.9, **hits[2]}, {"_rerank_score": 0.8, **hits[1]}] + + monkeypatch.setattr("httpx.AsyncClient", _Client) + monkeypatch.setattr("nemo_retriever.operators.rerank.rerank_hits", _rerank) + config = ServiceConfig( + mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), + logging=LoggingConfig(file=str(tmp_path / "service.log")), + pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), + vectordb=VectorDbConfig(enabled=True, vectordb_url="http://vectordb:7671"), + nim_endpoints=NimEndpointsConfig( + rerank_invoke_url="http://reranker:8080", + rerank_model_name="rerank-model", + api_key="server-secret", + ), + ) + + with TestClient(create_app(config)) as client: + response = client.post( + "/v1/query", + json={"query": "revenue", "top_k": 2, "rerank": True, "rerank_top_k": 3}, + ) + + assert response.status_code == 200 + assert response.json()["results"][0]["hits"] == [ + {"_rerank_score": 0.9, "text": "third", "source": "c"}, + {"_rerank_score": 0.8, "text": "second", "source": "b"}, + ] + assert seen["url"] == "http://vectordb:7671/v1/query" + assert seen["body"] == {"query": "revenue", "top_k": 3} + assert seen["rerank"] == { + "query": "revenue", + "kwargs": { + "rerank_invoke_url": "http://reranker:8080", + "model_name": "rerank-model", + "api_key": "server-secret", + "top_n": 2, + }, + } + + +def test_reranked_query_requires_main_service_reranker(monkeypatch, tmp_path) -> None: + _configure_noop_workers(monkeypatch) + config = ServiceConfig( + mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), + logging=LoggingConfig(file=str(tmp_path / "service.log")), + pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), + vectordb=VectorDbConfig(enabled=True, vectordb_url="http://vectordb:7671"), + ) + + with TestClient(create_app(config)) as client: + response = client.post("/v1/query", json={"query": "revenue", "rerank": True}) + + assert response.status_code == 400 + assert "rerank_invoke_url" in response.json()["detail"] + + +def test_reranked_query_uses_lazy_local_main_service_model(monkeypatch, tmp_path) -> None: + _configure_noop_workers(monkeypatch) + seen: dict[str, object] = {} + local_model = object() + + class _Response: + status_code = 200 + content = json.dumps({"results": [{"hits": [{"text": "first"}, {"text": "second"}]}]}).encode() + + class _Client: + def __init__(self, *args, **kwargs) -> None: + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + async def post(self, url, **kwargs): + seen["body"] = json.loads(kwargs["content"]) + return _Response() + + def _create_local_reranker(model_name, **kwargs): + seen["factory"] = {"model_name": model_name, "kwargs": kwargs} + return local_model + + def _rerank(query, hits, **kwargs): + seen["rerank"] = {"query": query, "kwargs": kwargs} + return [{"_rerank_score": 0.9, **hits[1]}] + + monkeypatch.setattr("httpx.AsyncClient", _Client) + monkeypatch.setattr("nemo_retriever.models.create_local_reranker", _create_local_reranker) + monkeypatch.setattr("nemo_retriever.operators.rerank.rerank_hits", _rerank) + config = ServiceConfig( + mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), + logging=LoggingConfig(file=str(tmp_path / "service.log")), + pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), + vectordb=VectorDbConfig(enabled=True, vectordb_url="http://vectordb:7671"), + local_models={ + "enabled": False, + "device": "cuda:0", + "hf_cache_dir": "/models", + "rerank": { + "enabled": True, + "model_name": "local-reranker", + "backend": "hf", + "gpu_memory_utilization": 0.6, + "max_length": 256, + "batch_size": 4, + }, + }, + ) + + with TestClient(create_app(config)) as client: + response = client.post("/v1/query", json={"query": "revenue", "top_k": 1, "rerank": True}) + + assert response.status_code == 200 + assert seen["body"] == {"query": "revenue", "top_k": 50} + assert seen["factory"] == { + "model_name": "local-reranker", + "kwargs": { + "backend": "hf", + "device": "cuda:0", + "hf_cache_dir": "/models", + "gpu_memory_utilization": 0.6, + }, + } + assert seen["rerank"] == { + "query": "revenue", + "kwargs": { + "model": local_model, + "model_name": "local-reranker", + "max_length": 256, + "batch_size": 4, + "top_n": 1, + }, + } + + +def test_false_rerank_string_uses_normal_query_path(monkeypatch, tmp_path) -> None: + _configure_noop_workers(monkeypatch) + seen: dict[str, object] = {} + + class _Response: + status_code = 200 + content = json.dumps({"results": [{"hits": [{"text": "first"}]}]}).encode() + + class _Client: + def __init__(self, *args, **kwargs) -> None: + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + async def post(self, url, **kwargs): + seen["body"] = json.loads(kwargs["content"]) + return _Response() + + monkeypatch.setattr("httpx.AsyncClient", _Client) + config = ServiceConfig( + mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), + logging=LoggingConfig(file=str(tmp_path / "service.log")), + pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), + vectordb=VectorDbConfig(enabled=True, vectordb_url="http://vectordb:7671"), + ) + + with TestClient(create_app(config)) as client: + response = client.post("/v1/query", json={"query": "revenue", "rerank": "false"}) + + assert response.status_code == 200 + assert seen["body"] == {"query": "revenue", "rerank": "false"} diff --git a/nemo_retriever/tests/test_service_save_to_disk.py b/nemo_retriever/tests/test_service_save_to_disk.py index 00a1b48b51..82fc853717 100644 --- a/nemo_retriever/tests/test_service_save_to_disk.py +++ b/nemo_retriever/tests/test_service_save_to_disk.py @@ -19,6 +19,7 @@ from typing import Any from unittest.mock import patch +import httpx import pytest from nemo_retriever.service.service_ingestor import ServiceIngestor @@ -146,7 +147,12 @@ def test_materialize_fetches_once_when_return_results_and_save_to_disk(tmp_path: rows = [{"page": 1, "text": "shared"}] fetch_calls = 0 - def _counting_fetch(self: ServiceIngestor, document_id: str) -> list[dict[str, Any]]: + def _counting_fetch( + self: ServiceIngestor, + document_id: str, + *, + client: httpx.Client | None = None, + ) -> list[dict[str, Any]]: nonlocal fetch_calls fetch_calls += 1 assert document_id == "doc-1" @@ -168,3 +174,140 @@ def test_save_document_authorisation_header_sent_when_token_present(tmp_path: Pa ing._save_document_to_disk("doc-x") assert captured["kwargs"]["headers"] == {"Authorization": "Bearer sekret"} + + +# ---------------------------------------------------------------------- +# ingest-scoped result client reuse and retry +# ---------------------------------------------------------------------- + + +def _completion_events(*document_ids: str) -> list[dict[str, Any]]: + return [ + {"event": "job_created", "job_id": "job-1"}, + *[ + { + "event": "document_complete", + "document_id": document_id, + "status": "completed", + "result_rows": 1, + } + for document_id in document_ids + ], + {"event": "job_finalized", "job_id": "job-1"}, + ] + + +def _result_row(document_id: str) -> dict[str, Any]: + return { + "path": f"/uploads/{document_id}.pdf", + "page_number": 1, + "text": f"content-{document_id}", + "metadata": {"source_id": document_id}, + } + + +def test_ingest_reuses_one_result_client_across_documents(monkeypatch: pytest.MonkeyPatch) -> None: + ing = ServiceIngestor(base_url="http://example:7670") + monkeypatch.setattr(ing, "ingest_stream", lambda **_kwargs: iter(_completion_events("doc-a", "doc-b"))) + requests: list[httpx.Request] = [] + clients: list[httpx.Client] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + document_id = request.url.path.rsplit("/", 1)[-1] + return httpx.Response(200, json={"result_data": [_result_row(document_id)]}) + + def client_factory() -> httpx.Client: + client = httpx.Client(transport=httpx.MockTransport(handler)) + clients.append(client) + return client + + monkeypatch.setattr(ing, "_new_result_fetch_client", client_factory) + result = ing.ingest(result_schema="compact") + + assert len(clients) == 1 + assert clients[0].is_closed + assert [request.url.path for request in requests] == [ + "/v1/ingest/status/doc-a", + "/v1/ingest/status/doc-b", + ] + assert result.dataframe is not None + assert len(result.dataframe) == 2 + assert result.failures == [] + + +@pytest.mark.parametrize("error_type", [httpx.ConnectError, httpx.ReadError, httpx.RemoteProtocolError]) +def test_ingest_retries_transient_result_fetch_on_fresh_client( + monkeypatch: pytest.MonkeyPatch, + error_type: type[Exception], +) -> None: + ing = ServiceIngestor(base_url="http://example:7670") + monkeypatch.setattr(ing, "ingest_stream", lambda **_kwargs: iter(_completion_events("doc-a"))) + clients: list[httpx.Client] = [] + + def failing_handler(request: httpx.Request) -> httpx.Response: + raise error_type("transient result failure", request=request) + + def success_handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"result_data": [_result_row("doc-a")]}) + + handlers = [failing_handler, success_handler] + + def client_factory() -> httpx.Client: + client = httpx.Client(transport=httpx.MockTransport(handlers[len(clients)])) + clients.append(client) + return client + + monkeypatch.setattr(ing, "_new_result_fetch_client", client_factory) + result = ing.ingest(result_schema="compact") + + assert len(clients) == 2 + assert all(client.is_closed for client in clients) + assert result.dataframe is not None + assert len(result.dataframe) == 1 + assert result.failures == [] + + +def test_ingest_exhausted_result_retry_remains_visible(monkeypatch: pytest.MonkeyPatch) -> None: + ing = ServiceIngestor(base_url="http://example:7670") + monkeypatch.setattr(ing, "ingest_stream", lambda **_kwargs: iter(_completion_events("doc-a"))) + clients: list[httpx.Client] = [] + + def failing_handler(request: httpx.Request) -> httpx.Response: + raise httpx.ReadError("persistent result failure", request=request) + + def client_factory() -> httpx.Client: + client = httpx.Client(transport=httpx.MockTransport(failing_handler)) + clients.append(client) + return client + + monkeypatch.setattr(ing, "_new_result_fetch_client", client_factory) + result = ing.ingest(result_schema="compact") + + assert len(clients) == 2 + assert all(client.is_closed for client in clients) + assert len(result.failures) == 1 + assert result.failures[0][0] == "doc-a" + assert "persistent result failure" in result.failures[0][1] + + +def test_ingest_does_not_retry_http_status_failure(monkeypatch: pytest.MonkeyPatch) -> None: + ing = ServiceIngestor(base_url="http://example:7670") + monkeypatch.setattr(ing, "ingest_stream", lambda **_kwargs: iter(_completion_events("doc-a"))) + clients: list[httpx.Client] = [] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, json={"detail": "permanent"}) + + def client_factory() -> httpx.Client: + client = httpx.Client(transport=httpx.MockTransport(handler)) + clients.append(client) + return client + + monkeypatch.setattr(ing, "_new_result_fetch_client", client_factory) + result = ing.ingest(result_schema="compact") + + assert len(clients) == 1 + assert clients[0].is_closed + assert len(result.failures) == 1 + assert "500" in result.failures[0][1] diff --git a/nemo_retriever/tests/test_service_sidecar.py b/nemo_retriever/tests/test_service_sidecar.py index b9223a98d2..4b5eb7480b 100644 --- a/nemo_retriever/tests/test_service_sidecar.py +++ b/nemo_retriever/tests/test_service_sidecar.py @@ -25,6 +25,7 @@ from nemo_retriever.service.app import create_app from nemo_retriever.service.config import ( + AuthConfig, PipelineOverridesConfig, PipelinePoolConfig, ServiceConfig, @@ -266,6 +267,7 @@ async def _stub_work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), pipeline_overrides=PipelineOverridesConfig(sinks=SinksConfig(vdb_uri_schemes=["s3://"])), ) diff --git a/nemo_retriever/tests/test_service_sinks.py b/nemo_retriever/tests/test_service_sinks.py index e25b1da9d4..2214759816 100644 --- a/nemo_retriever/tests/test_service_sinks.py +++ b/nemo_retriever/tests/test_service_sinks.py @@ -25,6 +25,7 @@ from nemo_retriever.common.params import StoreParams, VdbUploadParams, WebhookParams from nemo_retriever.service.app import create_app from nemo_retriever.service.config import ( + AuthConfig, PipelineOverridesConfig, PipelinePoolConfig, ServiceConfig, @@ -364,6 +365,7 @@ async def _stub_work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), pipeline_overrides=PipelineOverridesConfig( sinks=SinksConfig( diff --git a/nemo_retriever/tests/test_service_sse.py b/nemo_retriever/tests/test_service_sse.py index 8595541663..1e017f7afd 100644 --- a/nemo_retriever/tests/test_service_sse.py +++ b/nemo_retriever/tests/test_service_sse.py @@ -41,6 +41,7 @@ from nemo_retriever.service.app import create_app from nemo_retriever.service.config import ( + AuthConfig, PipelineOverridesConfig, PipelinePoolConfig, ServiceConfig, @@ -72,6 +73,7 @@ async def _stub_work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: cfg = ServiceConfig( mode="standalone", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), pipeline_overrides=PipelineOverridesConfig(), ) diff --git a/nemo_retriever/tests/test_service_vectordb_app.py b/nemo_retriever/tests/test_service_vectordb_app.py index 3ae75b8b33..2d41f04a88 100644 --- a/nemo_retriever/tests/test_service_vectordb_app.py +++ b/nemo_retriever/tests/test_service_vectordb_app.py @@ -5,20 +5,277 @@ from __future__ import annotations import sys -from unittest.mock import MagicMock, PropertyMock, patch +from typing import Any +from unittest.mock import MagicMock, patch import pytest from fastapi.testclient import TestClient import nemo_retriever.service.vectordb_app as vectordb_module +from nemo_retriever.common.schemas.collections import ( + CollectionCreateRequest, + CollectionDeleteResult, + CollectionInfo, + CollectionPage, + CollectionUpdateRequest, + DocumentDeleteResult, + DocumentInfo, + DocumentPage, +) +from nemo_retriever.common.vdb.adt_vdb import ( + CollectionWriteContext, + CollectionWriteResult, + UnsupportedVDBOperation, + VDB, + VDBInvalidRequest, + VDBResourceConflict, + VDBResourceNotFound, +) +from nemo_retriever.common.vdb.lancedb import LanceDB +from nemo_retriever.common.vdb.records import RetrievalContractError from nemo_retriever.service.vectordb_app import ( VectorDBState, _embed_queries_remote, - _strategies_for_retrieval_mode, _tensor_to_embedding_rows, create_vectordb_app, ) +_NOW = "2026-07-27T00:00:00+00:00" + + +class FakeVDB(VDB): + """Backend-neutral in-memory fake for HTTP contract tests.""" + + def __init__(self, *, table_exists: bool = False) -> None: + self.collections: dict[tuple[str, str], CollectionInfo] = {} + self.documents: dict[tuple[str, str, str], DocumentInfo] = {} + self.table_exists = table_exists + self.legacy_rows = 1 if table_exists else 0 + self.last_write_context: CollectionWriteContext | None = None + self.last_retrieval: dict[str, Any] | None = None + self.health_calls = 0 + + def create_index(self, **kwargs): + return None + + def write_to_index(self, records: list, **kwargs): + self.legacy_rows += sum(len(batch) for batch in records) + self.table_exists = self.legacy_rows > 0 + + def retrieval(self, queries: list, **kwargs): + return [ + [ + { + "text": "legacy hit", + "source": "legacy.pdf", + "metadata": {}, + "_distance": 0.2, + } + ] + for _ in queries + ] + + def run(self, records): + self.write_to_index(records) + + def create_collection(self, *, scope: str, request: CollectionCreateRequest) -> CollectionInfo: + key = (scope, request.name) + if key in self.collections: + raise VDBResourceConflict("Collection already exists") + info = CollectionInfo( + name=request.name, + scope=scope, + description=request.description, + metadata=request.metadata, + created_at=_NOW, + updated_at=_NOW, + expires_at=request.expires_at, + ) + self.collections[key] = info + return info + + def get_collection(self, *, scope: str, collection_name: str) -> CollectionInfo: + try: + return self.collections[(scope, collection_name)] + except KeyError as exc: + raise VDBResourceNotFound("Collection not found") from exc + + def list_collections(self, *, scope: str, limit: int, continuation_token: str | None) -> CollectionPage: + if continuation_token == "invalid": + raise VDBInvalidRequest("Invalid continuation token") + items = [item for (item_scope, _), item in self.collections.items() if item_scope == scope] + return CollectionPage(items=items[:limit]) + + def update_collection( + self, + *, + scope: str, + collection_name: str, + request: CollectionUpdateRequest, + ) -> CollectionInfo: + current = self.get_collection(scope=scope, collection_name=collection_name) + updated = current.model_copy( + update={ + "description": request.description, + "metadata": (request.metadata if request.metadata is not None else current.metadata), + "expires_at": request.expires_at, + "updated_at": _NOW, + } + ) + self.collections[(scope, collection_name)] = updated + return updated + + def delete_collection(self, *, scope: str, collection_name: str, if_exists: bool) -> CollectionDeleteResult: + existed = self.collections.pop((scope, collection_name), None) is not None + if not existed and not if_exists: + raise VDBResourceNotFound("Collection not found") + return CollectionDeleteResult( + name=collection_name, + scope=scope, + existed=existed, + deleted=True, + status="deleted", + ) + + def get_document(self, *, scope: str, collection_name: str, document_id: str) -> DocumentInfo: + try: + return self.documents[(scope, collection_name, document_id)] + except KeyError as exc: + raise VDBResourceNotFound("Document not found") from exc + + def list_documents( + self, + *, + scope: str, + collection_name: str, + limit: int, + continuation_token: str | None, + ) -> DocumentPage: + self.get_collection(scope=scope, collection_name=collection_name) + items = [ + item + for (item_scope, item_collection, _), item in self.documents.items() + if (item_scope, item_collection) == (scope, collection_name) + ] + return DocumentPage(items=items[:limit]) + + def delete_document( + self, + *, + scope: str, + collection_name: str, + document_id: str, + if_exists: bool, + ) -> DocumentDeleteResult: + existed = self.documents.pop((scope, collection_name, document_id), None) is not None + if not existed and not if_exists: + raise VDBResourceNotFound("Document not found") + return DocumentDeleteResult( + document_id=document_id, + collection_name=collection_name, + scope=scope, + existed=existed, + deleted=True, + status="deleted", + ) + + def write_collection(self, records: list, *, context: CollectionWriteContext) -> CollectionWriteResult: + self.get_collection(scope=context.scope, collection_name=context.collection_name) + self.last_write_context = context + written = sum(len(batch) for batch in records) + self.documents[(context.scope, context.collection_name, context.document_id)] = DocumentInfo( + document_id=context.document_id, + collection_name=context.collection_name, + scope=context.scope, + filename=context.filename, + content_sha256=context.content_sha256, + document_version=context.document_version, + status="completed", + chunk_count=written, + job_id=context.job_id, + created_at=_NOW, + updated_at=_NOW, + ) + return CollectionWriteResult(written=written, total_rows=written) + + def retrieve_collection( + self, + vectors: list, + *, + scope: str, + collection_name: str, + query_texts: list[str], + top_k: int, + **kwargs: Any, + ) -> tuple[list[list[dict[str, Any]]], list[str]]: + self.get_collection(scope=scope, collection_name=collection_name) + self.last_retrieval = { + "scope": scope, + "collection_name": collection_name, + "query_texts": query_texts, + "top_k": top_k, + } + hit = { + "chunk_id": "chunk-1", + "document_id": "document-1", + "text": "collection hit", + "distance": 0.2, + "filename": "report.pdf", + "page_number": 1, + "content_type": "text", + "source": "report.pdf", + "source_id": "report.pdf", + "metadata": {}, + "physical_table": "private-table", + "_distance": 0.2, + } + return ([[hit] for _ in vectors], ["dense"]) + + def health(self) -> dict[str, Any]: + self.health_calls += 1 + return { + "total_rows": self.legacy_rows, + "table_exists": self.table_exists, + "effective_retrieval_mode": "dense" if self.table_exists else None, + "retrieval_strategies": ["dense"] if self.table_exists else [], + "collections": { + "active": len(self.collections), + "deleting": 0, + "expired": 0, + }, + "cleanup": {"pending": 0, "oldest_age_seconds": 0}, + "reconciliation": {"successes": 0, "failures": 0}, + "open_table_cache_count": 0, + } + + +class ContractFailureVDB(FakeVDB): + def retrieve_collection(self, *args: Any, **kwargs: Any): + raise RetrievalContractError("physical table secret must not be returned") + + +class ValueFailureVDB(FakeVDB): + def retrieve_collection(self, *args: Any, **kwargs: Any): + raise ValueError("backend parsing failed") + + +class HealthFailureVDB(FakeVDB): + def health(self) -> dict[str, Any]: + raise RuntimeError("backend unavailable") + + +class DefaultHealthVDB(FakeVDB): + health = VDB.health + + +def _app(vdb: VDB, **kwargs: Any): + return create_vectordb_app( + vdb=vdb, + embed_endpoint="http://embed.example/v1/embeddings", + reconciliation_interval_seconds=0, + **kwargs, + ) + @pytest.mark.parametrize( ("extra_args", "expected_key"), @@ -36,111 +293,323 @@ def test_main_resolves_remote_embed_api_key(monkeypatch, extra_args, expected_ke assert create_app.call_args.kwargs["embed_api_key"] == expected_key -def test_query_empty_index_returns_422(tmp_path) -> None: - app = create_vectordb_app( - lancedb_uri=str(tmp_path), - table_name="test_table", - embed_endpoint="http://embed.example/v1/embeddings", - embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", - ) - with TestClient(app) as client: - resp = client.post("/v1/query", json={"query": "hello", "top_k": 3}) +def test_fake_vdb_completes_collection_http_flow_without_backend_details() -> None: + backend = FakeVDB() + app = _app(backend) + record = { + "document_type": "text", + "metadata": { + "embedding": [1.0, 0.0], + "content": "hello", + "content_metadata": {"page_number": 1}, + "source_metadata": {"source_id": "report.pdf"}, + }, + } - assert resp.status_code == 422 - assert "No data has been ingested yet" in resp.json()["detail"] + with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0]]): + with TestClient(app) as client: + created = client.post( + "/v1/collections", + headers={"X-NRL-Scope": "tenant-a"}, + json={"name": "research", "description": "docs"}, + ) + listed = client.get("/v1/collections", headers={"X-NRL-Scope": "tenant-a"}) + updated = client.patch( + "/v1/collections/research", + headers={"X-NRL-Scope": "tenant-a"}, + json={"description": "updated"}, + ) + written = client.post( + "/internal/vectordb/write", + json={ + "records": [[record]], + "scope": "tenant-a", + "collection_name": "research", + "document_id": "document-1", + "job_id": "job-1", + "filename": "report.pdf", + "content_sha256": "sha256", + "document_version": "v1", + }, + ) + documents = client.get( + "/v1/collections/research/documents", + headers={"X-NRL-Scope": "tenant-a"}, + ) + queried = client.post( + "/v1/query", + headers={"X-NRL-Scope": "tenant-a"}, + json={"query": "hello", "collection_name": "research"}, + ) + deleted_document = client.delete( + "/v1/collections/research/documents/document-1", + headers={"X-NRL-Scope": "tenant-a"}, + ) + deleted_collection = client.delete( + "/v1/collections/research", + headers={"X-NRL-Scope": "tenant-a"}, + ) + + assert created.status_code == 201 + assert listed.json()["items"][0]["name"] == "research" + assert updated.json()["description"] == "updated" + assert written.json() == {"written": 1, "total_rows": 1} + assert documents.json()["items"][0]["document_id"] == "document-1" + assert queried.status_code == 200 + hit = queried.json()["results"][0]["hits"][0] + assert backend.health_calls == 0 + assert hit["text"] == "collection hit" + assert "physical_table" not in hit + assert "_distance" not in hit + assert backend.last_write_context is not None + assert backend.last_write_context.scope == "tenant-a" + assert backend.last_retrieval == { + "scope": "tenant-a", + "collection_name": "research", + "query_texts": ["hello"], + "top_k": 10, + } + assert deleted_document.status_code == 200 + assert deleted_collection.status_code == 200 -def test_query_without_embed_backend_returns_501(tmp_path) -> None: - app = create_vectordb_app(lancedb_uri=str(tmp_path)) - with TestClient(app) as client: - resp = client.post("/v1/query", json={"query": "hello", "top_k": 3}) +def test_unsupported_collection_retrieval_returns_501_without_legacy_fallback() -> None: + class UnsupportedCollectionRetrievalVDB(FakeVDB): + def __init__(self) -> None: + super().__init__() + self.legacy_retrieval_calls = 0 - assert resp.status_code == 501 - assert "No embedding backend configured" in resp.json()["detail"] + def retrieval(self, queries: list, **kwargs): + self.legacy_retrieval_calls += 1 + return super().retrieval(queries, **kwargs) + def retrieve_collection(self, *args: Any, **kwargs: Any): + raise UnsupportedVDBOperation("Collection retrieval mode is unsupported") -def test_health_reports_embed_mode(tmp_path) -> None: - app = create_vectordb_app( - lancedb_uri=str(tmp_path), - local_embed=True, - embed_model="nvidia/llama-nemotron-embed-1b-v2", + backend = UnsupportedCollectionRetrievalVDB() + backend.create_collection( + scope="tenant-a", + request=CollectionCreateRequest(name="research"), ) - with TestClient(app) as client: - resp = client.get("/v1/health") - assert resp.status_code == 200 - assert resp.json()["embed_mode"] == "local" + with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0]]): + with TestClient(_app(backend)) as client: + response = client.post( + "/v1/query", + headers={"X-NRL-Scope": "tenant-a"}, + json={"query": "hello", "collection_name": "research"}, + ) + assert response.status_code == 501 + assert response.json()["detail"] == ("The configured VectorDB backend does not support this operation.") + assert backend.legacy_retrieval_calls == 0 -def test_health_reports_effective_retrieval_mode_none_without_table(tmp_path) -> None: - app = create_vectordb_app( + +def test_production_vdb_preserves_legacy_service_write_without_index_rebuild( + tmp_path, + monkeypatch, +) -> None: + backend = vectordb_module._production_vdb( lancedb_uri=str(tmp_path), - embed_endpoint="http://embed.example/v1/embeddings", + table_name="legacy", + expiration_cleanup_enabled=True, + ) + assert isinstance(backend, LanceDB) + assert backend.build_index is False + + index_writes = [] + monkeypatch.setattr( + backend, + "create_index", + lambda records, table_name: object(), + ) + monkeypatch.setattr( + backend, + "write_to_index", + lambda *args, **kwargs: index_writes.append((args, kwargs)), ) - with TestClient(app) as client: - resp = client.get("/v1/health") - - assert resp.status_code == 200 - body = resp.json() - assert "retrieval_mode" not in body - assert body["effective_retrieval_mode"] is None + assert backend.run([]) == [] + assert index_writes == [] + + +def test_legacy_write_and_query_keep_existing_vdb_path() -> None: + backend = FakeVDB() + app = _app(backend) + record = { + "document_type": "text", + "metadata": { + "embedding": [1.0, 0.0], + "content": "legacy", + "content_metadata": {}, + "source_metadata": {}, + }, + } -def test_health_stays_ok_when_mode_resolution_errors(tmp_path) -> None: - app = create_vectordb_app( - lancedb_uri=str(tmp_path), - embed_endpoint="http://embed.example/v1/embeddings", + with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0]]): + with TestClient(app) as client: + written = client.post( + "/internal/vectordb/write", + json={ + "records": [[record]], + "scope": "default", + "filename": "legacy.pdf", + "document_version": "1", + }, + ) + queried = client.post("/v1/query", json={"query": "legacy"}) + + assert written.json() == {"written": 1, "total_rows": 1} + hit = queried.json()["results"][0]["hits"][0] + assert hit["text"] == "legacy hit" + assert hit["_distance"] == 0.2 + + +def test_service_managed_lancedb_preserves_multimodal_fields(tmp_path) -> None: + backend = LanceDB( + uri=str(tmp_path), + table_name="legacy", + vector_dim=None, + overwrite=False, + build_index=False, + _service_table_schema=True, ) + app = _app(backend) + record = { + "document_type": "text", + "metadata": { + "embedding": [1.0, 0.0, 0.0, 0.0], + "content": "table content", + "content_metadata": { + "page_number": 7, + "type": "table_caption", + "stored_image_uri": "s3://artifacts/table.png", + "bbox_xyxy_norm": [0.1, 0.2, 0.8, 0.9], + }, + "source_metadata": { + "source_id": "/documents/report.pdf", + "source_name": "report.pdf", + }, + }, + } + + with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0, 0.0, 0.0]]): + with TestClient(app) as client: + written = client.post( + "/internal/vectordb/write", + json={ + "records": [[record]], + "filename": "report.pdf", + "document_version": "1", + }, + ) + queried = client.post("/v1/query", json={"query": "table", "top_k": 1}) + + assert written.status_code == 200 + assert queried.status_code == 200 + hit = queried.json()["results"][0]["hits"][0] + assert hit["content_type"] == "table" + assert hit["stored_image_uri"] == "s3://artifacts/table.png" + assert hit["bbox_xyxy_norm"] == "[0.1, 0.2, 0.8, 0.9]" + assert hit["page_number"] == 7 + assert hit["source_id"] == "/documents/report.pdf" + + +def test_typed_collection_errors_map_to_http_contract() -> None: + backend = FakeVDB() + with TestClient(_app(backend)) as client: + missing = client.get("/v1/collections/missing") + assert missing.status_code == 404 + + assert client.post("/v1/collections", json={"name": "duplicate"}).status_code == 201 + conflict = client.post("/v1/collections", json={"name": "duplicate"}) + invalid = client.get("/v1/collections?continuation_token=invalid") + + assert conflict.status_code == 409 + assert invalid.status_code == 422 + + +@pytest.mark.parametrize("backend_cls", [ContractFailureVDB, ValueFailureVDB]) +def test_backend_query_failures_return_safe_500(backend_cls) -> None: + backend = backend_cls() + backend.create_collection(scope="default", request=CollectionCreateRequest(name="research")) + with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0]]): + with TestClient(_app(backend), raise_server_exceptions=False) as client: + response = client.post( + "/v1/query", + json={"query": "hello", "collection_name": "research"}, + ) + + assert response.status_code == 500 + assert "physical table secret" not in response.text + if backend_cls is ContractFailureVDB: + assert response.json()["detail"] == "VectorDB retrieval contract violation." + else: + assert response.json()["detail"] == "VectorDB backend operation failed." + + +def test_query_empty_legacy_index_returns_422() -> None: + with TestClient(_app(FakeVDB())) as client: + response = client.post("/v1/query", json={"query": "hello", "top_k": 3}) + + assert response.status_code == 422 + assert "No data has been ingested yet" in response.json()["detail"] + + +def test_legacy_query_allows_backend_with_default_empty_health() -> None: + app = _app(DefaultHealthVDB()) + with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0]]): + with TestClient(app) as client: + response = client.post("/v1/query", json={"query": "hello", "top_k": 3}) + + assert response.status_code == 200 + assert response.json()["results"][0]["hits"][0]["text"] == "legacy hit" + + +def test_query_without_embed_backend_returns_501() -> None: + app = create_vectordb_app(vdb=FakeVDB(), reconciliation_interval_seconds=0) with TestClient(app) as client: - with patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), patch.object( - VectorDBState, - "resolve_effective_retrieval_mode", - side_effect=OSError("transient I/O error"), - ): - resp = client.get("/v1/health") + response = client.post("/v1/query", json={"query": "hello", "top_k": 3}) - # Health backs k8s probes; a mode-resolution failure must not 500. - assert resp.status_code == 200 - assert resp.json()["effective_retrieval_mode"] == "unknown" + assert response.status_code == 501 + assert "No embedding backend configured" in response.json()["detail"] -def test_strategies_for_retrieval_mode() -> None: - assert _strategies_for_retrieval_mode("dense") == ["dense"] - assert _strategies_for_retrieval_mode("hybrid") == ["hybrid"] +def test_internal_auth_is_optional_and_can_be_enabled() -> None: + app = _app(FakeVDB(), internal_api_token="internal-secret") + with TestClient(app) as client: + assert client.get("/v1/health").status_code == 200 + assert client.get("/v1/collections").status_code == 401 + assert ( + client.get( + "/v1/collections", + headers={"X-NRL-Internal-Token": "internal-secret"}, + ).status_code + == 200 + ) + + +def test_health_and_metrics_use_backend_neutral_health() -> None: + backend = FakeVDB(table_exists=True) + app = _app(backend) + with TestClient(app) as client: + health = client.get("/v1/health") + metrics = client.get("/metrics") + assert health.status_code == 200 + assert health.json()["table_exists"] is True + assert health.json()["effective_retrieval_mode"] == "dense" + assert metrics.status_code == 200 + assert "nrl_vectordb_collections" in metrics.text -def test_write_rows_creates_then_appends_table(tmp_path) -> None: - state = VectorDBState( - lancedb_uri=str(tmp_path), - table_name="nemo_retriever", - embed_endpoint="http://embed.example/v1/embeddings", - embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", - embed_api_key="", - ) - assert state.table_exists is False - - row = { - "vector": [1.0, 0.0, 0.0, 0.0], - "text": "seed", - "pdf_page": "p1", - "filename": "f.pdf", - "pdf_basename": "f.pdf", - "page_number": 1, - "source": "f.pdf", - "source_id": "f.pdf", - "path": "/f.pdf", - "metadata": "{}", - "stored_image_uri": "", - "content_type": "text", - "bbox_xyxy_norm": "", - } - assert state.write_rows([row]) == 1 - assert state.table_exists is True - assert state.total_rows() == 1 - # A second write appends rather than overwriting the existing table. - assert state.write_rows([dict(row, text="second")]) == 1 - assert state.total_rows() == 2 +def test_health_returns_503_when_backend_is_unavailable() -> None: + app = _app(HealthFailureVDB()) + with TestClient(app) as client: + health = client.get("/v1/health") + + assert health.status_code == 503 + assert health.json() == {"detail": "VectorDB backend is unavailable"} def test_tensor_to_embedding_rows_handles_batch() -> None: @@ -158,10 +627,8 @@ def test_vector_db_state_local_embed_queries() -> None: tensor.cpu.return_value = tensor tensor.tolist.return_value = [[1.0, 2.0]] mock_embedder.embed_queries.return_value = tensor - state = VectorDBState( - lancedb_uri="/tmp/unused", - table_name="t", + vdb=FakeVDB(), embed_endpoint="", embed_model="nvidia/llama-nemotron-embed-1b-v2", embed_api_key="", @@ -185,7 +652,6 @@ def fake_infer_microservice(data, **kwargs): return [[0.1, 0.2]] monkeypatch.setattr("nemo_retriever.models.nim.util.infer_microservice", fake_infer_microservice) - vectors = _embed_queries_remote( ["hello"], embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", @@ -199,120 +665,3 @@ def fake_infer_microservice(data, **kwargs): assert calls["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" assert calls["model_provider_prefix"] == "nvidia" assert calls["embedding_endpoint"] == "https://litellm.example.com/v1/embeddings" - - -_CANNED_HITS = [ - { - "text": "Revenue grew 12% year over year.", - "pdf_basename": "10k_2023.pdf", - "page_number": 12, - "content_type": "text", - "_score": 0.91, - "metadata": {}, - } -] - - -def _query_app(tmp_path): - return create_vectordb_app( - lancedb_uri=str(tmp_path), - table_name="t", - embed_endpoint="http://embed.example/v1/embeddings", - embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", - ) - - -def test_query_evidence_format_returns_evidence_coverage(tmp_path) -> None: - app = _query_app(tmp_path) - with patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), patch.object( - VectorDBState, "embed_queries", return_value=[[0.1, 0.2]] - ), patch.object(VectorDBState, "search", return_value=([_CANNED_HITS], ["dense"])): - with TestClient(app) as client: - resp = client.post("/v1/query", json={"query": "revenue", "top_k": 5, "format": "evidence"}) - - assert resp.status_code == 200 - body = resp.json() - assert list(body) == ["results"] - assert len(body["results"]) == 1 - item = body["results"][0] - assert set(item) == {"evidence", "coverage"} - - ev = item["evidence"][0] - assert ev["source"] == "10k_2023" - assert ev["citation"] == "10k_2023 p.12" - assert ev["locator"] == {"kind": "page", "value": 12} - assert ev["modality"] == "text" - assert ev["fidelity"] == "verbatim" - assert ev["score"] == 0.91 - - coverage = item["coverage"] - assert coverage["strategies_used"] == ["dense"] - assert coverage["n_docs_seen"] == 1 - assert coverage["thin_spots"] == ["single source"] - - -def test_query_hybrid_evidence_reports_hybrid_strategy(tmp_path) -> None: - # A table whose capabilities resolve to hybrid (vector + FTS) reports hybrid. - app = _query_app(tmp_path) - with patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), patch.object( - VectorDBState, "embed_queries", return_value=[[0.1, 0.2]] - ), patch.object(VectorDBState, "search", return_value=([_CANNED_HITS], ["hybrid"])): - with TestClient(app) as client: - resp = client.post("/v1/query", json={"query": "revenue", "top_k": 5, "format": "evidence"}) - - assert resp.status_code == 200 - assert resp.json()["results"][0]["coverage"]["strategies_used"] == ["hybrid"] - - -def test_query_unqueryable_table_returns_422(tmp_path) -> None: - # An unqueryable table (e.g. FTS-only / no vector column) surfaces as 422. - app = _query_app(tmp_path) - with patch.object(VectorDBState, "table_exists", new_callable=PropertyMock, return_value=True), patch.object( - VectorDBState, "embed_queries", return_value=[[0.1, 0.2]] - ), patch.object( - VectorDBState, - "search", - side_effect=ValueError( - "LanceDB table 't' at '" + str(tmp_path) + "' has an FTS index but no vector " - "column; sparse-only retrieval is not supported by the VectorDB service." - ), - ): - with TestClient(app) as client: - resp = client.post("/v1/query", json={"query": "revenue", "top_k": 5}) - - assert resp.status_code == 422 - assert "sparse-only retrieval is not supported" in resp.json()["detail"] - - -def test_search_hybrid_delegates_to_lancedb_wrapper(tmp_path) -> None: - state = VectorDBState( - lancedb_uri=str(tmp_path), - table_name="docs", - embed_endpoint="http://embed.example/v1/embeddings", - embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", - embed_api_key="", - ) - state._table_exists = True - - mock_caps = MagicMock() - mock_caps.has_vector = True - mock_caps.has_fts = True - mock_caps.retrieval_mode = "hybrid" - mock_caps.vector_column = "vector" - - mock_vdb = MagicMock() - mock_vdb.retrieval.return_value = [[{"text": "hit", "_score": 0.5}]] - - with patch.object(state, "_table_capabilities", return_value=mock_caps), patch( - "nemo_retriever.common.vdb.lancedb.LanceDB", - return_value=mock_vdb, - ): - hits, strategies = state.search([[0.1, 0.2]], ["revenue"], top_k=3) - - assert strategies == ["hybrid"] - assert hits[0][0]["text"] == "hit" - mock_vdb.retrieval.assert_called_once() - call_kwargs = mock_vdb.retrieval.call_args.kwargs - assert call_kwargs["top_k"] == 3 - assert call_kwargs["hybrid"] is True - assert call_kwargs["query_texts"] == ["revenue"] diff --git a/nemo_retriever/tests/test_service_vectordb_evidence_integration.py b/nemo_retriever/tests/test_service_vectordb_evidence_integration.py index 9008c77e92..db2760872b 100644 --- a/nemo_retriever/tests/test_service_vectordb_evidence_integration.py +++ b/nemo_retriever/tests/test_service_vectordb_evidence_integration.py @@ -14,36 +14,38 @@ from __future__ import annotations -import json from unittest.mock import patch from fastapi.testclient import TestClient +from nemo_retriever.common.vdb.lancedb import LanceDB from nemo_retriever.service.vectordb_app import VectorDBState, create_vectordb_app _DIM = 4 -# One stored chunk, in the real LanceDB schema (metadata is a JSON *string* column). -_ROW = { - "vector": [1.0, 0.0, 0.0, 0.0], - "pdf_page": "10k_2023_12", - "filename": "10k_2023.pdf", - "pdf_basename": "10k_2023.pdf", - "page_number": 12, - "source": "10k_2023.pdf", - "source_id": "10k_2023.pdf", - "path": "/data/10k_2023.pdf", - "text": "Revenue grew 12% year over year.", - "metadata": json.dumps({"page_number": 12, "type": "text"}), - "stored_image_uri": "", - "content_type": "text", - "bbox_xyxy_norm": "", +_RECORD = { + "document_type": "text", + "metadata": { + "embedding": [1.0, 0.0, 0.0, 0.0], + "content": "Revenue grew 12% year over year.", + "content_metadata": {"page_number": 12, "type": "text"}, + "source_metadata": { + "source_id": "/data/10k_2023.pdf", + "source_name": "10k_2023.pdf", + }, + }, } def test_query_evidence_format_end_to_end_over_real_lancedb(tmp_path) -> None: - app = create_vectordb_app( - lancedb_uri=str(tmp_path), + backend = LanceDB( + uri=str(tmp_path), table_name="nemo_retriever", + overwrite=False, + build_index=False, + vector_dim=_DIM, + ) + app = create_vectordb_app( + vdb=backend, embed_endpoint="http://embed.example/v1/embeddings", # -> embed_mode="remote" embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", ) @@ -51,7 +53,7 @@ def test_query_evidence_format_end_to_end_over_real_lancedb(tmp_path) -> None: # Stub ONLY the embedding model; real LanceDB does the rest. with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0, 0.0, 0.0]]): with TestClient(app) as client: - write = client.post("/internal/vectordb/write", json={"rows": [_ROW]}) + write = client.post("/internal/vectordb/write", json={"records": [[_RECORD]]}) assert write.status_code == 200, write.text assert write.json()["total_rows"] == 1 @@ -62,7 +64,8 @@ def test_query_evidence_format_end_to_end_over_real_lancedb(tmp_path) -> None: assert resp.status_code == 200, resp.text body = resp.json() - assert list(body) == ["results"] + assert list(body) == ["results", "query_mode"] + assert body["query_mode"] == "classic" assert len(body["results"]) == 1 item = body["results"][0] @@ -75,7 +78,7 @@ def test_query_evidence_format_end_to_end_over_real_lancedb(tmp_path) -> None: assert ev["locator"] == {"kind": "page", "value": 12} assert ev["modality"] == "text" assert ev["fidelity"] == "verbatim" - # Score is the real LanceDB distance/relevance — present and numeric, value not asserted. + # Preserve the real LanceDB distance/relevance on the legacy evidence path. assert isinstance(ev["score"], (int, float)) coverage = item["coverage"] diff --git a/nemo_retriever/tests/test_service_vectordb_hybrid_integration.py b/nemo_retriever/tests/test_service_vectordb_hybrid_integration.py index 1708cc6dc7..a70df982bc 100644 --- a/nemo_retriever/tests/test_service_vectordb_hybrid_integration.py +++ b/nemo_retriever/tests/test_service_vectordb_hybrid_integration.py @@ -4,7 +4,6 @@ from __future__ import annotations -import json from datetime import timedelta from unittest.mock import patch @@ -12,36 +11,57 @@ import pytest from fastapi.testclient import TestClient +from nemo_retriever.common.vdb.lancedb import LanceDB +from nemo_retriever.common.vdb.lancedb_capabilities import inspect_lancedb_table_object from nemo_retriever.service.vectordb_app import VectorDBState, create_vectordb_app _DIM = 4 -_ROW = { - "vector": [1.0, 0.0, 0.0, 0.0], - "pdf_page": "10k_2023_12", - "filename": "10k_2023.pdf", - "pdf_basename": "10k_2023.pdf", - "page_number": 12, - "source": "10k_2023.pdf", - "source_id": "10k_2023.pdf", - "path": "/data/10k_2023.pdf", - "text": "Revenue grew 12% year over year.", - "metadata": json.dumps({"page_number": 12, "type": "text"}), - "stored_image_uri": "", - "content_type": "text", - "bbox_xyxy_norm": "", -} - - -def _state(tmp_path) -> VectorDBState: - return VectorDBState( - lancedb_uri=str(tmp_path), + + +def _record( + *, + vector: list[float] | None = None, + text: str = "Revenue grew 12% year over year.", +) -> dict: + return { + "document_type": "text", + "metadata": { + "embedding": vector or [1.0, 0.0, 0.0, 0.0], + "content": text, + "content_metadata": {"page_number": 12, "type": "text"}, + "source_metadata": { + "source_id": "/data/10k_2023.pdf", + "source_name": "10k_2023.pdf", + }, + }, + } + + +_RECORD = _record() + + +def _backend(tmp_path, *, hybrid: bool = False) -> LanceDB: + return LanceDB( + uri=str(tmp_path), table_name="nemo_retriever", - embed_endpoint="http://embed.example/v1/embeddings", - embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", - embed_api_key="", + vector_dim=_DIM, + overwrite=False, + build_index=False, + hybrid=hybrid, ) +def _write(backend: LanceDB, *records: dict) -> int: + backend.run([list(records)]) + table = lancedb.connect(backend.uri).open_table(backend.table_name) + return int(table.count_rows()) + + +def _capabilities(backend: LanceDB): + table = lancedb.connect(backend.uri).open_table(backend.table_name) + return inspect_lancedb_table_object(table) + + def _prebuild_fts_index(uri: str, table_name: str) -> None: """Simulate an ingestion pipeline that wrote a table with a BM25/FTS index. @@ -57,72 +77,73 @@ def _prebuild_fts_index(uri: str, table_name: str) -> None: @pytest.mark.integration def test_write_rows_persists_rows_without_building_fts(tmp_path) -> None: - state = _state(tmp_path) - assert state.write_rows([_ROW]) == 1 + backend = _backend(tmp_path) + assert _write(backend, _RECORD) == 1 - caps = state._table_capabilities() + caps = _capabilities(backend) assert caps is not None assert caps.has_vector # The service must not build an FTS index on write; the table stays dense. assert not caps.has_fts - assert state.resolve_effective_retrieval_mode() == "dense" + assert backend.health()["effective_retrieval_mode"] == "dense" @pytest.mark.integration def test_append_does_not_build_or_mutate_fts(tmp_path) -> None: - state = _state(tmp_path) - assert state.write_rows([_ROW]) == 1 + backend = _backend(tmp_path) + assert _write(backend, _RECORD) == 1 - appended = dict(_ROW) - appended["vector"] = [0.0, 1.0, 0.0, 0.0] - appended["text"] = "Zephyr quarterly guidance mentions unicorn synergy." - assert state.write_rows([appended]) == 1 + appended = _record( + vector=[0.0, 1.0, 0.0, 0.0], + text="Zephyr quarterly guidance mentions unicorn synergy.", + ) + assert _write(backend, appended) == 2 - table = state._db.open_table("nemo_retriever") + table = lancedb.connect(str(tmp_path)).open_table("nemo_retriever") assert table.count_rows() == 2 # Still no FTS index — appends only persist rows. - caps = state._table_capabilities() + caps = _capabilities(backend) assert not caps.has_fts @pytest.mark.integration def test_auto_resolves_hybrid_when_fts_prebuilt(tmp_path) -> None: # Ingestion built the table with both a vector column and an FTS index. - seed = _state(tmp_path) - seed.write_rows([_ROW]) + seed = _backend(tmp_path, hybrid=True) + _write(seed, _RECORD) _prebuild_fts_index(str(tmp_path), "nemo_retriever") - state = _state(tmp_path) - caps = state._table_capabilities() + backend = _backend(tmp_path, hybrid=True) + caps = _capabilities(backend) assert caps.has_vector assert caps.has_fts - assert state.resolve_effective_retrieval_mode() == "hybrid" + assert backend.health()["effective_retrieval_mode"] == "hybrid" @pytest.mark.integration def test_auto_resolves_dense_when_no_fts(tmp_path) -> None: - seed = _state(tmp_path) - seed.write_rows([_ROW]) + seed = _backend(tmp_path) + _write(seed, _RECORD) - state = _state(tmp_path) - caps = state._table_capabilities() + backend = _backend(tmp_path) + caps = _capabilities(backend) assert caps.has_vector assert not caps.has_fts - assert state.resolve_effective_retrieval_mode() == "dense" + assert backend.health()["effective_retrieval_mode"] == "dense" @pytest.mark.integration def test_query_auto_selects_hybrid_when_fts_prebuilt(tmp_path) -> None: + backend = _backend(tmp_path, hybrid=True) app = create_vectordb_app( - lancedb_uri=str(tmp_path), - table_name="nemo_retriever", + vdb=backend, embed_endpoint="http://embed.example/v1/embeddings", embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", ) with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0, 0.0, 0.0]]): with TestClient(app) as client: - write = client.post("/internal/vectordb/write", json={"rows": [_ROW]}) + write = client.post("/internal/vectordb/write", json={"records": [[_RECORD]]}) assert write.status_code == 200, write.text # Ingestion builds the FTS index; the query path detects it. @@ -140,16 +161,16 @@ def test_query_auto_selects_hybrid_when_fts_prebuilt(tmp_path) -> None: @pytest.mark.integration def test_query_auto_selects_dense_when_no_fts(tmp_path) -> None: + backend = _backend(tmp_path) app = create_vectordb_app( - lancedb_uri=str(tmp_path), - table_name="nemo_retriever", + vdb=backend, embed_endpoint="http://embed.example/v1/embeddings", embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", ) with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0, 0.0, 0.0]]): with TestClient(app) as client: - write = client.post("/internal/vectordb/write", json={"rows": [_ROW]}) + write = client.post("/internal/vectordb/write", json={"records": [[_RECORD]]}) assert write.status_code == 200, write.text resp = client.post( diff --git a/nemo_retriever/tests/test_service_work_queue.py b/nemo_retriever/tests/test_service_work_queue.py index 592c47ffdf..a0814d6919 100644 --- a/nemo_retriever/tests/test_service_work_queue.py +++ b/nemo_retriever/tests/test_service_work_queue.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio +import hashlib import logging import threading import time @@ -27,7 +28,7 @@ init_job_tracker, shutdown_job_tracker, ) -from nemo_retriever.service.services.pipeline_pool import PoolType, WorkItem, _Pool +from nemo_retriever.service.services.pipeline_pool import PipelinePool, PoolType, WorkItem, _Pool from nemo_retriever.service.services.prometheus import POOL_ACTIVE_SLOTS, WORK_QUEUE_CLAIMS from nemo_retriever.service.services.work_queue import ( GatewayWorkClient, @@ -340,11 +341,16 @@ def test_gateway_upload_claim_payload_and_callback_lifecycle(tmp_path, monkeypat monkeypatch.setenv("NEMO_RETRIEVER_RESULTS_DIR", str(results_dir)) config = ServiceConfig( mode="gateway", + auth=AuthConfig(allow_unscoped_dev=True), logging=LoggingConfig(file=str(tmp_path / "service.log")), mcp=MCPConfig(enabled=False), pipeline=PipelinePoolConfig(realtime_queue_size=2, batch_queue_size=2), work_queue=_config(tmp_path / "spool", gateway_url="http://testserver"), ) + # This test covers the claim/payload/callback lifecycle, not lease expiry. + # Without this the reaper can retire the lease mid-test and the callback + # fails with 409 purely because the host was slow. + monkeypatch.setattr(WorkBroker, "_expire_locked", lambda self, pool: None) with TestClient(create_app(config)) as client: created = client.post("/v1/ingest/job", json={"expected_documents": 1}) @@ -354,7 +360,11 @@ def test_gateway_upload_claim_payload_and_callback_lifecycle(tmp_path, monkeypat accepted = client.post( f"/v1/ingest/job/{job_id}/whole", files={"file": ("document.txt", b"hello gateway", "text/plain")}, - data={"metadata": "{}"}, + data={ + "metadata": ( + '{"metadata":{"category":"Finance_Investment",' '"source_path":"Finance_Investment/document.txt"}}' + ) + }, ) assert accepted.status_code == 202 document_id = accepted.json()["document_id"] @@ -371,6 +381,20 @@ def test_gateway_upload_claim_payload_and_callback_lifecycle(tmp_path, monkeypat claim = claim_response.json() assert claim["work_id"] == document_id assert claim["delivery_attempt"] == 1 + assert claim["extra"] == { + "write": { + "scope": "default", + "collection_name": None, + "operation": "append", + "content_sha256": hashlib.sha256(b"hello gateway").hexdigest(), + "document_version": None, + "storage_document_id": document_id, + "document_metadata": { + "category": "Finance_Investment", + "source_path": "Finance_Investment/document.txt", + }, + } + } processing = client.get(f"/v1/ingest/job/{job_id}/document/{document_id}") assert processing.json()["status"] == "processing" @@ -415,6 +439,7 @@ def test_gateway_callback_treats_stale_acknowledge_as_idempotent(tmp_path, monke monkeypatch.setenv("NEMO_RETRIEVER_RESULTS_DIR", str(results_dir)) config = ServiceConfig( mode="gateway", + auth=AuthConfig(allow_unscoped_dev=True), logging=LoggingConfig(file=str(tmp_path / "service.log")), mcp=MCPConfig(enabled=False), pipeline=PipelinePoolConfig(realtime_queue_size=2, batch_queue_size=2), @@ -466,7 +491,7 @@ def test_internal_work_endpoints_require_configured_service_auth(tmp_path): mode="gateway", logging=LoggingConfig(file=str(tmp_path / "service.log")), mcp=MCPConfig(enabled=False), - auth=AuthConfig(api_token="secret"), + auth=AuthConfig(enabled=True, api_token="secret"), work_queue=_config(tmp_path / "spool", gateway_url="http://testserver"), ) with TestClient(create_app(config)) as client: @@ -482,9 +507,32 @@ def test_internal_work_endpoints_require_configured_service_auth(tmp_path): ) +def test_split_worker_uses_internal_gateway_credential_when_configured(tmp_path): + config = _config(tmp_path, gateway_url="http://gateway") + internal_pool = PipelinePool( + PipelinePoolConfig(), + mode="batch", + work_queue_config=config, + auth_config=AuthConfig(api_token="public-secret"), + internal_api_token="internal-secret", + ) + assert internal_pool._batch is not None + assert internal_pool._batch._pull_client.headers == {"X-NRL-Internal-Token": "internal-secret"} + + compatibility_pool = PipelinePool( + PipelinePoolConfig(), + mode="batch", + work_queue_config=config, + auth_config=AuthConfig(api_token="public-secret"), + ) + assert compatibility_pool._batch is not None + assert compatibility_pool._batch._pull_client.headers == {"Authorization": "Bearer public-secret"} + + def test_gateway_dry_run_does_not_register_or_enqueue_work(tmp_path): config = ServiceConfig( mode="gateway", + auth=AuthConfig(allow_unscoped_dev=True), logging=LoggingConfig(file=str(tmp_path / "service.log")), mcp=MCPConfig(enabled=False), pipeline=PipelinePoolConfig(realtime_queue_size=2, batch_queue_size=2), @@ -540,6 +588,7 @@ def test_gateway_restart_is_explicit_loss_boundary(tmp_path, monkeypatch): spool = tmp_path / "spool" config = ServiceConfig( mode="gateway", + auth=AuthConfig(allow_unscoped_dev=True), logging=LoggingConfig(file=str(tmp_path / "service.log")), mcp=MCPConfig(enabled=False), pipeline=PipelinePoolConfig(realtime_queue_size=2, batch_queue_size=2), @@ -594,7 +643,10 @@ def test_gateway_restart_is_explicit_loss_boundary(tmp_path, monkeypatch): @pytest.mark.anyio -async def test_shutdown_removes_queued_and_leased_payloads_and_records(tmp_path): +async def test_shutdown_removes_queued_and_leased_payloads_and_records(tmp_path, monkeypatch): + # Shutdown, not expiry, is what must invalidate the lease here; the reaper + # would otherwise clear it first and mask the behaviour under test. + monkeypatch.setattr(WorkBroker, "_expire_locked", lambda self, pool: None) broker = WorkBroker(_config(tmp_path), PipelinePoolConfig(batch_queue_size=2)) await broker.start() leased = await _enqueue(broker, "leased") @@ -609,8 +661,6 @@ async def test_shutdown_removes_queued_and_leased_payloads_and_records(tmp_path) assert not broker._records assert all(not queue for queue in broker._queues.values()) assert broker._spool_bytes == 0 - with pytest.raises(StaleLease): - broker.validate_callback("leased", claim.lease.lease_id, claim.lease.generation) @pytest.mark.anyio diff --git a/nemo_retriever/tests/test_service_worker_callback.py b/nemo_retriever/tests/test_service_worker_callback.py index 2b9c10be60..af0e5eff85 100644 --- a/nemo_retriever/tests/test_service_worker_callback.py +++ b/nemo_retriever/tests/test_service_worker_callback.py @@ -90,6 +90,7 @@ def test_worker_document_result_endpoint_is_idempotent() -> None: from nemo_retriever.service.app import create_app from nemo_retriever.service.config import ( + AuthConfig, PipelineOverridesConfig, PipelinePoolConfig, ServiceConfig, @@ -97,6 +98,7 @@ def test_worker_document_result_endpoint_is_idempotent() -> None: cfg = ServiceConfig( mode="batch", + auth=AuthConfig(allow_unscoped_dev=True), pipeline=PipelinePoolConfig(realtime_workers=1, batch_workers=1), pipeline_overrides=PipelineOverridesConfig(), ) @@ -323,7 +325,7 @@ def test_worker_result_endpoint_returns_retryable_503( from fastapi.testclient import TestClient from nemo_retriever.service.app import create_app - from nemo_retriever.service.config import ServiceConfig + from nemo_retriever.service.config import AuthConfig, ServiceConfig from nemo_retriever.service.routers import ingest def unavailable(_: str) -> None: @@ -331,7 +333,8 @@ def unavailable(_: str) -> None: monkeypatch.setattr(ingest, "get_result_data", unavailable) - with TestClient(create_app(ServiceConfig(mode="batch"))) as client: + config = ServiceConfig(mode="batch", auth=AuthConfig(allow_unscoped_dev=True)) + with TestClient(create_app(config)) as client: response = client.get("/v1/internal/document-result/doc-unavailable") assert response.status_code == 503 @@ -358,7 +361,9 @@ def unavailable(_: str) -> None: assert error.value.headers == {"Retry-After": "60"} -def test_gateway_fetch_reads_shared_result_off_event_loop(monkeypatch: pytest.MonkeyPatch) -> None: +def test_gateway_fetch_reads_shared_result_off_event_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: from nemo_retriever.service.routers import ingest event_loop_thread = threading.get_ident() @@ -407,13 +412,14 @@ def test_gateway_status_routes_read_shared_results_idempotently( from fastapi.testclient import TestClient from nemo_retriever.service.app import create_app - from nemo_retriever.service.config import ServiceConfig + from nemo_retriever.service.config import AuthConfig, ServiceConfig from nemo_retriever.service.services.job_tracker import get_job_tracker monkeypatch.setenv("NEMO_RETRIEVER_RESULTS_DIR", str(tmp_path)) rows = [{"text": "shared route"}] - with TestClient(create_app(ServiceConfig(mode="gateway"))) as client: + config = ServiceConfig(mode="gateway", auth=AuthConfig(allow_unscoped_dev=True)) + with TestClient(create_app(config)) as client: tracker = get_job_tracker() assert tracker is not None tracker.register_job("job-shared", expected_documents=2, retain_results=True) @@ -527,7 +533,7 @@ def test_gateway_callback_copies_result_before_completing( from fastapi.testclient import TestClient from nemo_retriever.service.app import create_app - from nemo_retriever.service.config import ServiceConfig + from nemo_retriever.service.config import AuthConfig, ServiceConfig from nemo_retriever.service.routers import ingest from nemo_retriever.service.services.job_tracker import ( DocumentStatus, @@ -564,7 +570,8 @@ async def get(self, url: str) -> _Resp: return _Resp() monkeypatch.setenv("NEMO_RETRIEVER_RESULTS_DIR", str(tmp_path)) - with TestClient(create_app(ServiceConfig(mode="gateway"))) as client: + config = ServiceConfig(mode="gateway", auth=AuthConfig(allow_unscoped_dev=True)) + with TestClient(create_app(config)) as client: tracker = get_job_tracker() assert tracker is not None tracker.register_job("handoff-job", expected_documents=1, retain_results=True) @@ -603,7 +610,7 @@ def test_gateway_callback_does_not_complete_when_result_handoff_fails( from fastapi.testclient import TestClient from nemo_retriever.service.app import create_app - from nemo_retriever.service.config import ServiceConfig + from nemo_retriever.service.config import AuthConfig, ServiceConfig from nemo_retriever.service.routers import ingest from nemo_retriever.service.services.job_tracker import ( DocumentStatus, @@ -627,7 +634,8 @@ async def get(self, url: str) -> _Resp: return _Resp() monkeypatch.setenv("NEMO_RETRIEVER_RESULTS_DIR", str(tmp_path)) - with TestClient(create_app(ServiceConfig(mode="gateway"))) as client: + config = ServiceConfig(mode="gateway", auth=AuthConfig(allow_unscoped_dev=True)) + with TestClient(create_app(config)) as client: tracker = get_job_tracker() assert tracker is not None tracker.register_job("failed-handoff-job", expected_documents=1, retain_results=True) @@ -660,11 +668,12 @@ def test_gateway_callback_permanently_rejects_retained_result_for_unknown_docume from fastapi.testclient import TestClient from nemo_retriever.service.app import create_app - from nemo_retriever.service.config import ServiceConfig + from nemo_retriever.service.config import AuthConfig, ServiceConfig from nemo_retriever.service.services.job_tracker import get_job_tracker monkeypatch.setenv("NEMO_RETRIEVER_RESULTS_DIR", str(tmp_path)) - with TestClient(create_app(ServiceConfig(mode="gateway"))) as client: + config = ServiceConfig(mode="gateway", auth=AuthConfig(allow_unscoped_dev=True)) + with TestClient(create_app(config)) as client: response = client.post( "/v1/internal/job-callback", json={ @@ -722,10 +731,18 @@ def make_request(peer: str) -> Request: assert missing.value.status_code == 503 -def test_internal_auth_headers_support_default_and_custom_header_names() -> None: +def test_internal_auth_headers_use_only_the_dedicated_internal_credential() -> None: + from nemo_retriever.service.auth import internal_auth_headers + + assert internal_auth_headers(None) == {} + assert internal_auth_headers("internal-secret") == {"X-NRL-Internal-Token": "internal-secret"} + + +def test_service_auth_headers_preserve_worker_pull_credentials() -> None: from nemo_retriever.service.auth import auth_headers from nemo_retriever.service.config import AuthConfig + assert auth_headers(AuthConfig()) == {} assert auth_headers(AuthConfig(api_token="secret")) == {"Authorization": "Bearer secret"} assert auth_headers(AuthConfig(api_token="secret", header_name="X-Service-Token")) == {"X-Service-Token": "secret"} @@ -942,7 +959,7 @@ def test_gateway_result_pull_sends_configured_internal_auth( from starlette.requests import Request - from nemo_retriever.service.config import AuthConfig, ServiceConfig + from nemo_retriever.service.config import AuthConfig, ServiceConfig, VectorDbConfig from nemo_retriever.service.routers import ingest client_kwargs: dict[str, Any] = {} @@ -971,6 +988,7 @@ async def get(self, url: str) -> _Resp: config=ServiceConfig( mode="gateway", auth=AuthConfig(api_token="secret", header_name="X-Service-Token"), + vectordb=VectorDbConfig(internal_api_token="internal-secret"), ) ) ) @@ -991,7 +1009,7 @@ async def get(self, url: str) -> _Resp: monkeypatch.setattr(ingest.httpx, "AsyncClient", _Client) asyncio.run(ingest._pull_and_store_worker_result(request, "auth-pull-doc", "10.1.2.3")) - assert client_kwargs["headers"] == {"X-Service-Token": "secret"} + assert client_kwargs["headers"] == {"X-NRL-Internal-Token": "internal-secret"} assert get_result_data("auth-pull-doc") == [{"text": "authenticated handoff"}] diff --git a/nemo_retriever/tests/test_tokenizer_provider.py b/nemo_retriever/tests/test_tokenizer_provider.py new file mode 100644 index 0000000000..d98085d3c6 --- /dev/null +++ b/nemo_retriever/tests/test_tokenizer_provider.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for dependency-light, revision-pinned chunk tokenizers.""" + +from __future__ import annotations + +from pathlib import Path + +import pandas as pd +import pytest +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace + +import nemo_retriever.common.modality.txt.tokenizer_provider as provider +import nemo_retriever.operators.extract.html.ray_data as html_actor_module +import nemo_retriever.operators.extract.txt.ray_data as txt_actor_module +from nemo_retriever.common.modality.html.convert import html_bytes_to_chunks_df +from nemo_retriever.common.modality.txt.split import ( + DEFAULT_TOKENIZER_MODEL_ID, + txt_bytes_to_chunks_df, +) +from nemo_retriever.models import resolve_embed_model + + +def _write_tokenizer(path: Path) -> None: + tokenizer = Tokenizer( + WordLevel( + {"[UNK]": 0, "hello": 1, "world": 2}, + unk_token="[UNK]", + ) + ) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.save(str(path)) + + +def test_default_chunk_tokenizer_matches_default_embedding_model() -> None: + assert DEFAULT_TOKENIZER_MODEL_ID == resolve_embed_model(None) + + +def test_load_chunk_tokenizer_uses_pinned_artifact( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + tokenizer_path = tmp_path / "tokenizer.json" + _write_tokenizer(tokenizer_path) + provider.load_chunk_tokenizer.cache_clear() + monkeypatch.setattr( + provider, + "hf_hub_download_with_pinned_revision", + lambda **_kwargs: str(tokenizer_path), + ) + + tokenizer = provider.load_chunk_tokenizer("nvidia/llama-nemotron-embed-vl-1b-v2") + + token_ids = tokenizer.encode("hello world") + assert tokenizer.decode(token_ids) == "hello world" + + +def test_txt_and_html_chunk_without_transformers( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + tokenizer_path = tmp_path / "tokenizer.json" + _write_tokenizer(tokenizer_path) + provider.load_chunk_tokenizer.cache_clear() + monkeypatch.setattr( + provider, + "hf_hub_download_with_pinned_revision", + lambda **_kwargs: str(tokenizer_path), + ) + + txt = txt_bytes_to_chunks_df(b"hello world", "document.txt") + html = html_bytes_to_chunks_df( + b"

hello world

", + "document.html", + ) + + assert txt["text"].tolist() == ["hello world"] + assert html["text"].tolist() == ["hello world"] + + +def test_load_chunk_tokenizer_surfaces_actionable_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider.load_chunk_tokenizer.cache_clear() + + def fail_download(**_kwargs: object) -> str: + raise RuntimeError("offline cache miss") + + monkeypatch.setattr( + provider, + "hf_hub_download_with_pinned_revision", + fail_download, + ) + + with pytest.raises( + provider.TokenizerUnavailableError, + match="Pre-cache tokenizer.json", + ): + provider.load_chunk_tokenizer("nvidia/llama-nemotron-embed-vl-1b-v2") + + +def test_load_chunk_tokenizer_wraps_unregistered_model_error() -> None: + provider.load_chunk_tokenizer.cache_clear() + + with pytest.raises( + provider.TokenizerUnavailableError, + match="unregistered/model", + ) as exc_info: + provider.load_chunk_tokenizer("unregistered/model") + + assert isinstance(exc_info.value.__cause__, ValueError) + + +@pytest.mark.parametrize( + ("module", "actor_type", "splitter_name"), + [ + (txt_actor_module, txt_actor_module.TxtSplitCPUActor, "txt_bytes_to_chunks_df"), + ( + html_actor_module, + html_actor_module.HtmlSplitCPUActor, + "html_bytes_to_chunks_df", + ), + ], +) +def test_text_actors_propagate_tokenizer_failures( + module: object, + actor_type: type, + splitter_name: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail_split(*_args: object, **_kwargs: object) -> pd.DataFrame: + raise provider.TokenizerUnavailableError("tokenizer unavailable") + + monkeypatch.setattr(module, splitter_name, fail_split) + actor = actor_type() + batch = pd.DataFrame([{"bytes": b"content", "path": "document.txt"}]) + + with pytest.raises(provider.TokenizerUnavailableError, match="unavailable"): + actor.process(batch) + + +@pytest.mark.parametrize( + ("module", "actor_type", "splitter_name"), + [ + (txt_actor_module, txt_actor_module.TxtSplitCPUActor, "txt_bytes_to_chunks_df"), + ( + html_actor_module, + html_actor_module.HtmlSplitCPUActor, + "html_bytes_to_chunks_df", + ), + ], +) +def test_text_actors_isolate_non_tokenizer_document_failures( + module: object, + actor_type: type, + splitter_name: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def split_document(_payload: object, path: str, **_kwargs: object) -> pd.DataFrame: + if path == "bad.txt": + raise ValueError("malformed document") + return pd.DataFrame([{"text": "good", "path": path, "page_number": 1, "metadata": {}}]) + + monkeypatch.setattr(module, splitter_name, split_document) + actor = actor_type() + batch = pd.DataFrame( + [ + {"bytes": b"bad", "path": "bad.txt"}, + {"bytes": b"good", "path": "good.txt"}, + ] + ) + + result = actor.process(batch) + + assert result["path"].tolist() == ["good.txt"] diff --git a/nemo_retriever/tests/test_txt_split.py b/nemo_retriever/tests/test_txt_split.py index 83265960be..1c844f747f 100644 --- a/nemo_retriever/tests/test_txt_split.py +++ b/nemo_retriever/tests/test_txt_split.py @@ -12,7 +12,13 @@ import pandas as pd import pytest -from nemo_retriever.common.modality.txt.split import split_text_by_tokens, txt_file_to_chunks_df, TextChunkParams +from nemo_retriever.common.modality.txt.split import ( + TextChunkParams, + split_text_by_tokens, + text_to_chunks_df, + txt_bytes_to_chunks_df, + txt_file_to_chunks_df, +) class _MockTokenizer: @@ -58,7 +64,6 @@ def test_split_text_by_tokens_max_tokens_positive(): def test_txt_file_to_chunks_df(tmp_path: Path, monkeypatch): - pytest.importorskip("transformers") monkeypatch.setattr( "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda model_id, cache_dir=None: _MockTokenizer() ) @@ -78,10 +83,63 @@ def test_txt_file_to_chunks_df(tmp_path: Path, monkeypatch): def test_txt_file_to_chunks_df_empty_file(tmp_path: Path): - pytest.importorskip("transformers") f = tmp_path / "empty.txt" f.write_text("", encoding="utf-8") df = txt_file_to_chunks_df(str(f), params=TextChunkParams(max_tokens=512)) assert isinstance(df, pd.DataFrame) - assert list(df.columns) == ["text", "path", "page_number", "metadata"] + assert list(df.columns) == ["text", "content", "path", "page_number", "metadata"] assert len(df) == 0 + + +def test_text_to_chunks_df_preserves_logical_source_id(monkeypatch): + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda model_id, cache_dir=None: _MockTokenizer() + ) + + df = text_to_chunks_df( + "one two three four", + "inline://00000000", + params=TextChunkParams(max_tokens=2), + ) + + assert df["text"].tolist() == ["one two", "three four"] + assert df["path"].tolist() == ["inline://00000000", "inline://00000000"] + assert [metadata["source_path"] for metadata in df["metadata"]] == [ + "inline://00000000", + "inline://00000000", + ] + + +def test_txt_bytes_preserves_service_inline_identity_and_utf8_transport(monkeypatch): + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda model_id, cache_dir=None: _MockTokenizer() + ) + + df = txt_bytes_to_chunks_df( + "café document".encode("utf-8"), + "inline://00000007", + params=TextChunkParams(max_tokens=10, encoding="utf-16"), + ) + + assert df["text"].tolist() == ["café document"] + assert df["path"].tolist() == ["inline://00000007"] + assert df["metadata"].iloc[0]["source_path"] == "inline://00000007" + + +def test_file_and_decoded_text_helpers_produce_equivalent_chunks(tmp_path: Path, monkeypatch): + monkeypatch.setattr( + "nemo_retriever.common.modality.txt.split._get_tokenizer", lambda model_id, cache_dir=None: _MockTokenizer() + ) + text = "one two three four five" + path = tmp_path / "document.txt" + path.write_text(text, encoding="utf-8") + params = TextChunkParams(max_tokens=2) + + file_df = txt_file_to_chunks_df(str(path), params=params) + inline_df = text_to_chunks_df(text, "inline://00000000", params=params) + + assert file_df["text"].tolist() == inline_df["text"].tolist() + assert file_df["page_number"].tolist() == inline_df["page_number"].tolist() + assert [metadata["chunk_index"] for metadata in file_df["metadata"]] == [ + metadata["chunk_index"] for metadata in inline_df["metadata"] + ] diff --git a/nemo_retriever/tests/test_vdb_collection_contract.py b/nemo_retriever/tests/test_vdb_collection_contract.py new file mode 100644 index 0000000000..7a69547290 --- /dev/null +++ b/nemo_retriever/tests/test_vdb_collection_contract.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +from dataclasses import FrozenInstanceError, replace +from typing import Any + +import pytest +from pydantic import ValidationError + +from nemo_retriever.common.schemas.collections import IngestOperation +from nemo_retriever.common.schemas.requests import JobCreateRequest +from nemo_retriever.common.vdb.adt_vdb import ( + CollectionWriteContext, + CollectionWriteResult, + VDB, +) +from nemo_retriever.service.services.pipeline_pool import DocumentWriteContext, WorkItem + + +class IncompleteVDB(VDB): + """VDB missing the required collection-management contract.""" + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + + def create_index(self, **kwargs: Any) -> None: + return None + + def write_to_index(self, records: list, **kwargs: Any) -> None: + return None + + def retrieval(self, queries: list, **kwargs: Any) -> list[list[dict[str, Any]]]: + return [[] for _ in queries] + + def run(self, records: list) -> None: + return None + + +def test_vdb_requires_collection_management_implementations() -> None: + with pytest.raises(TypeError) as exc_info: + IncompleteVDB() + + for method in ( + "create_collection", + "get_collection", + "list_collections", + "update_collection", + "delete_collection", + "get_document", + "list_documents", + "delete_document", + "write_collection", + "retrieve_collection", + ): + assert method in str(exc_info.value) + + +def test_optional_collection_maintenance_has_safe_defaults() -> None: + assert VDB.reconcile_collections(None) == {"successes": 0, "failures": 0} # type: ignore[arg-type] + assert VDB.health(None) == {} # type: ignore[arg-type] + + +def test_collection_write_contract_is_immutable_and_reports_counts() -> None: + context = CollectionWriteContext( + scope="workspace-a", + collection_name="collection-a", + document_id="document-a", + document_version="version-a", + content_sha256="sha256-a", + filename="document.txt", + job_id=None, + operation="replace", + ) + + assert context.operation == "replace" + assert CollectionWriteResult(written=3, total_rows=7) == CollectionWriteResult( + written=3, + total_rows=7, + ) + with pytest.raises(FrozenInstanceError): + context.collection_name = "other" # type: ignore[misc] + + +def _write_context(**overrides: Any) -> CollectionWriteContext: + fields: dict[str, Any] = { + "scope": "workspace-a", + "collection_name": "collection-a", + "document_id": "document-a", + "document_version": "version-a", + "content_sha256": "sha256-a", + "filename": "document.txt", + } + fields.update(overrides) + return CollectionWriteContext(**fields) + + +@pytest.mark.parametrize( + ("supplied", "expected"), + [ + ("append", IngestOperation.APPEND), + ("replace", IngestOperation.REPLACE), + (IngestOperation.REPLACE, IngestOperation.REPLACE), + ], +) +def test_collection_write_context_coerces_operation_to_enum(supplied: Any, expected: IngestOperation) -> None: + """Backends compare with ``is``, so a wire string must become the member.""" + context = _write_context(operation=supplied) + + assert context.operation is expected + assert replace(context, content_sha256="other").operation is expected + + +def test_collection_write_context_rejects_unknown_operation() -> None: + with pytest.raises(ValueError): + _write_context(operation="upsert") + + +def test_collection_write_context_defaults_to_append() -> None: + assert _write_context().operation is IngestOperation.APPEND + + +@pytest.mark.parametrize( + ("wire_value", "target_document_id"), + [("append", None), ("replace", "document-1")], +) +def test_ingest_operation_round_trips_as_a_plain_wire_string( + wire_value: str, + target_document_id: str | None, +) -> None: + """The enum is an internal type only; the REST contract stays strings.""" + request = JobCreateRequest( + expected_documents=1, + collection_name="c", + operation=wire_value, + target_document_id=target_document_id, + ) + + assert isinstance(request.operation, IngestOperation) + assert request.model_dump(mode="json")["operation"] == wire_value + assert json.dumps({"operation": request.operation}) == json.dumps({"operation": wire_value}) + + +def test_ingest_operation_rejects_unknown_wire_values() -> None: + with pytest.raises(ValidationError): + JobCreateRequest(expected_documents=1, collection_name="c", operation="upsert") + + +def test_job_idempotency_fingerprint_is_unchanged_by_the_enum() -> None: + """The stored fingerprint must not shift, or retries would stop matching.""" + fingerprint_input = {"operation": IngestOperation.APPEND, "expected_documents": 1} + literal_input = {"operation": "append", "expected_documents": 1} + + def _digest(payload: dict[str, Any]) -> str: + return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + assert _digest(fingerprint_input) == _digest(literal_input) + + +def test_work_item_rebuilds_the_write_context_from_a_broker_claim() -> None: + """``RichModel`` ignores unknown keys, so the nested key must survive intact.""" + original = DocumentWriteContext( + scope="workspace", + collection_name="research", + operation=IngestOperation.REPLACE, + content_sha256="a" * 64, + storage_document_id="document-1", + ) + + claim_extra = {"write": original.model_dump(mode="json")} + rebuilt = WorkItem(id="attempt-1", **claim_extra) + + assert rebuilt.write == original + assert rebuilt.write.operation is IngestOperation.REPLACE + assert rebuilt.write.storage_document_id == "document-1" + + +def test_work_item_write_context_falls_back_to_the_attempt_id() -> None: + item = WorkItem(id="attempt-1") + + assert item.write.resolved(fallback_document_id=item.id).storage_document_id == "attempt-1" + assert item.write.resolved(fallback_document_id=item.id).operation is IngestOperation.APPEND diff --git a/nemo_retriever/tests/test_vdb_records.py b/nemo_retriever/tests/test_vdb_records.py new file mode 100644 index 0000000000..7497960924 --- /dev/null +++ b/nemo_retriever/tests/test_vdb_records.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from nemo_retriever.common.schemas.collections import QueryHit +from nemo_retriever.common.vdb.records import ( + normalize_retrieval_results, + to_client_vdb_records, +) + + +def _normalize_one(hit: dict) -> dict: + return normalize_retrieval_results([[hit]])[0][0] + + +def test_legacy_entity_is_flattened_once_with_top_level_precedence() -> None: + hit = _normalize_one( + { + "entity": { + "text": "nested text", + "source": {"source_id": "nested.pdf"}, + "content_metadata": {"page_number": 2}, + "chunk_id": "nested-chunk", + "document_id": "nested-document", + }, + "text": "flat text", + "source": {"source_id": "flat.pdf"}, + "content_metadata": {"page_number": 3}, + "chunk_id": "flat-chunk", + "document_id": "flat-document", + "filename": "flat.pdf", + } + ) + + assert hit["text"] == "flat text" + assert hit["source_id"] == "flat.pdf" + assert hit["page_number"] == 3 + assert hit["chunk_id"] == "flat-chunk" + assert hit["document_id"] == "flat-document" + assert "entity" not in hit + + +def test_legacy_entity_accepts_only_pre_collection_fields() -> None: + hit = _normalize_one( + { + "entity": { + "text": "legacy text", + "source": {"source_id": "legacy.pdf"}, + "content_metadata": {"page_number": 4}, + "chunk_id": "nested-chunk", + "document_id": "nested-document", + "document_version": "nested-version", + } + } + ) + + assert hit["text"] == "legacy text" + assert hit["source_id"] == "legacy.pdf" + assert hit["page_number"] == 4 + assert "chunk_id" not in hit + assert "document_id" not in hit + assert "document_version" not in hit + + +def test_flat_hit_is_canonicalized_without_entity() -> None: + hit = _normalize_one( + { + "text": "flat", + "source": {"source_id": "flat.pdf"}, + "content_metadata": {"page_number": "5"}, + } + ) + + assert hit["text"] == "flat" + assert hit["source_id"] == "flat.pdf" + assert hit["page_number"] == 5 + assert hit["pdf_page"] == "flat_5" + + +@pytest.mark.parametrize("content_type", ["audio", "video", "video_frame"]) +def test_legacy_media_page_values_remain_unchanged(content_type: str) -> None: + metadata = { + "type": content_type, + "page_number": 3, + "chunk_index": 3, + "frame_timestamp_seconds": 9.0, + } + + hit = _normalize_one({"text": "media", "content_metadata": metadata}) + + assert hit["page_number"] == 3 + assert hit["metadata"] == metadata + + +def test_query_hit_validates_canonical_page_instead_of_repairing_it() -> None: + payload = { + "chunk_id": "chunk", + "document_id": "document", + "text": "text", + "distance": 0.2, + "filename": "document.pdf", + } + + assert QueryHit(**payload, page_number=None).page_number is None + with pytest.raises(ValidationError): + QueryHit(**payload, page_number=0) + + +def test_canonical_record_batches_pass_through_without_reconversion() -> None: + records = [ + [ + { + "document_type": "text", + "metadata": { + "embedding": [0.1, 0.2], + "content": "already canonical", + }, + } + ] + ] + + assert to_client_vdb_records(records) is records + + +def test_graph_record_conversion_preserves_service_provenance() -> None: + records = to_client_vdb_records( + [ + { + "text": "table content", + "text_embeddings_1b_v2": {"embedding": [0.1, 0.2]}, + "path": "/tmp/source.pdf", + "page_number": 1, + "_page_number": 7, + "_content_type": "table_caption", + "_stored_image_uri": "s3://artifacts/table.png", + "_bbox_xyxy_norm": [0.1, 0.2, 0.8, 0.9], + "page_elements_v3_num_detections": 3, + "page_elements_v3_counts_by_label": {"table": 2, "chart": 1}, + "table": [{"content": "a"}, {"content": "b"}], + "metadata": { + "chunk_index": 4, + "chunk_count": 9, + "segment_start_seconds": 1.5, + "frame_timestamp_seconds": 2.5, + "content_metadata": {"page_number": 1}, + }, + } + ] + ) + + metadata = records[0][0]["metadata"] + assert metadata["embedding"] == [0.1, 0.2] + assert metadata["content"] == "table content" + assert metadata["source_metadata"] == { + "source_id": "/tmp/source.pdf", + "source_name": "source.pdf", + } + assert metadata["content_metadata"] == { + "page_number": 7, + "type": "table", + "fidelity": "ocr", + "stored_image_uri": "s3://artifacts/table.png", + "bbox_xyxy_norm": [0.1, 0.2, 0.8, 0.9], + "page_elements_v3_num_detections": 3, + "page_elements_v3_counts_by_label": {"table": 2, "chart": 1}, + "ocr_table_detections": 2, + "chunk_index": 4, + "chunk_count": 9, + "segment_start_seconds": 1.5, + "frame_timestamp_seconds": 2.5, + } + + +def test_narrow_lancedb_hit_promotes_canonical_multimodal_metadata() -> None: + hit = _normalize_one( + { + "text": "table content", + "metadata": { + "page_number": 7, + "type": "table_caption", + "stored_image_uri": "s3://artifacts/table.png", + "bbox_xyxy_norm": [0.1, 0.2, 0.8, 0.9], + }, + "source": {"source_id": "/tmp/source.pdf"}, + } + ) + + assert hit["content_type"] == "table" + assert hit["stored_image_uri"] == "s3://artifacts/table.png" + assert hit["bbox_xyxy_norm"] == [0.1, 0.2, 0.8, 0.9] + assert hit["page_number"] == 7 + assert hit["source_id"] == "/tmp/source.pdf" diff --git a/nemo_retriever/uv.lock b/nemo_retriever/uv.lock index 05d48fc5fa..1e2d8fb0e8 100644 --- a/nemo_retriever/uv.lock +++ b/nemo_retriever/uv.lock @@ -15,8 +15,21 @@ required-markers = [ [manifest] overrides = [ + { name = "aiohttp", specifier = ">=3.14.3" }, + { name = "cryptography", specifier = ">=48.0.1" }, { name = "fastparquet", specifier = ">=2024.11.0,<2026" }, + { name = "idna", specifier = ">=3.15" }, + { name = "langsmith", specifier = ">=0.8.18" }, + { name = "litellm", specifier = ">=1.95.0rc3" }, + { name = "msgpack", specifier = ">=1.2.1" }, + { name = "nltk", specifier = ">=3.10.1" }, { name = "opencv-python", marker = "sys_platform == 'never'" }, + { name = "pillow", specifier = ">=12.3.0" }, + { name = "pydantic-settings", specifier = ">=2.14.2" }, + { name = "ray", specifier = ">=2.56.1" }, + { name = "requests", specifier = ">=2.33.0" }, + { name = "transformers", specifier = ">=5.14.1" }, + { name = "vllm", specifier = "==0.25.1" }, ] [[package]] @@ -97,7 +110,7 @@ wheels = [ [[package]] name = "aiohttp" -version = "3.14.1" +version = "3.14.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohappyeyeballs" }, @@ -109,38 +122,26 @@ dependencies = [ { name = "typing-extensions" }, { name = "yarl" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/82/78/8ea7308cac6934de8c74a14f3d5f65d1c89287426688be79538d0e5c013d/aiohttp-3.14.1.tar.gz", hash = "sha256:307f2cff90a764d329e77040603fa032db89c5c24fdad50c4c15334cba744035", size = 7955794, upload-time = "2026-06-07T21:09:35.529Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/21/151624b51cd92553d95424daf4bf19f19ce9be9002d19253e7e7ce67197b/aiohttp-3.14.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d35143e27778b4bb0fb189562d7f275bff79c62ab8e98459717c0ea617ff2480", size = 757402, upload-time = "2026-06-07T21:06:40.311Z" }, - { url = "https://files.pythonhosted.org/packages/c2/82/280619e0bd7bf2454987e19282616e84762255dd9c8468f62382e8c191f1/aiohttp-3.14.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:bcfb80a2cc36fba2534e5e5b5264dc7ae6fcd9bf15256da3e53d2f499e6fa29d", size = 512310, upload-time = "2026-06-07T21:06:42.207Z" }, - { url = "https://files.pythonhosted.org/packages/55/b2/2aac325583aaa1353045f96dffa586d8a34e8322e14a7ba49cffeb103ab4/aiohttp-3.14.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:27fd7c91e51729b4f7e1577865fa6d34c9adccbc39aabe9000285b48af9f0ec2", size = 512448, upload-time = "2026-06-07T21:06:43.813Z" }, - { url = "https://files.pythonhosted.org/packages/8a/72/a60607cb849faa8af8a356c9329ea2eb6f395d49e82cc82ccba1fd8deb8f/aiohttp-3.14.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:64c567bf9eaf664280116a8688f63016e6b32db2505908e2bdaca1b6438142f2", size = 1766854, upload-time = "2026-06-07T21:06:45.391Z" }, - { url = "https://files.pythonhosted.org/packages/b5/d3/d9fe1c9ec7557ab4d0d82bebaa728c6418f0b93295ec2f4ab015f7710cc7/aiohttp-3.14.1-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:f5e6ff2bdbb8f4cd3fbe41f99e25bbcd58e3bf9f13d3dd31a11e7917251cc77a", size = 1740884, upload-time = "2026-06-07T21:06:47.413Z" }, - { url = "https://files.pythonhosted.org/packages/c1/dc/f2cecfaf9337ba3e63f181500814ff502aa3d00d9c7ec93a9d23d10a27b2/aiohttp-3.14.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2f73e01dc37122325caf079982621262f96d74823c179038a82fddfc50359264", size = 1810034, upload-time = "2026-06-07T21:06:50.165Z" }, - { url = "https://files.pythonhosted.org/packages/66/d7/2ff65c5e65c0d7476daf7e15c032e0805e36811185b9623e3238ad6c763e/aiohttp-3.14.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bb2c0c80d431c0d03f2c7dbf125150fedd4f0de17366a7ca33f7ccb822391842", size = 1904054, upload-time = "2026-06-07T21:06:52.035Z" }, - { url = "https://files.pythonhosted.org/packages/20/9c/d445818389df371f56d141d881153ba23183c4735a03f7356ffb43f7757d/aiohttp-3.14.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3e6fc1a85fa7194a1a7d19f44e8609180f4a8eb5fa4c7ed8b4355f080fad235c", size = 1790278, upload-time = "2026-06-07T21:06:54.049Z" }, - { url = "https://files.pythonhosted.org/packages/4d/aa/bf04cb4d865fc6101c2229a294ad744973b72e513fdc5a6b791e6983d72a/aiohttp-3.14.1-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:686b6c0d3911ec387b444ddf5dc62fb7f7c0a7d5186a7861626496a5ab4aff95", size = 1591795, upload-time = "2026-06-07T21:06:55.911Z" }, - { url = "https://files.pythonhosted.org/packages/dc/b4/4dac0038960427ba832f6609dfb4ea5437d7fd80c72001b9e48f834f428b/aiohttp-3.14.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c6fa4dc7ad6f8109c70bb1499e589f76b0b792baf39f9b017eb92c8a81d0a199", size = 1728397, upload-time = "2026-06-07T21:06:57.777Z" }, - { url = "https://files.pythonhosted.org/packages/2b/f9/7cd4e8ad7aa3b75f17d56bb5498dd604a93d4e6eece822ba0568c413fff0/aiohttp-3.14.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:87a5eea1b2a5e21e1ebdbb33ad4165359189327e63fc4e4894693e7f821ac817", size = 1766504, upload-time = "2026-06-07T21:07:00.009Z" }, - { url = "https://files.pythonhosted.org/packages/f9/df/fc01d9fcad0f73fed3f3d361f1f94f975947b50dff82919f6dc2bf4316cc/aiohttp-3.14.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:1c1421eb01d4fd608d88cc8290211d177a58532b55ad94076fb349c5bf467f0a", size = 1777806, upload-time = "2026-06-07T21:07:02.064Z" }, - { url = "https://files.pythonhosted.org/packages/41/09/47e2d090bddcc8fb4ccb4c314aadc32d7c5d9bb55f50f6ad1c92fc15d501/aiohttp-3.14.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:34b257ec41345c1e8f2df68fa908a7952f5de932723871eb633ecbbff396c9a4", size = 1580707, upload-time = "2026-06-07T21:07:03.942Z" }, - { url = "https://files.pythonhosted.org/packages/3d/36/f1a4ce904ae0b6930cfe9afc96d0896f7ec1a620c400405d63783bb95a9c/aiohttp-3.14.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:de538791a80e5d862addbc183f70f0158ac9b9bb872bb147f1fd2a683691e087", size = 1798121, upload-time = "2026-06-07T21:07:05.987Z" }, - { url = "https://files.pythonhosted.org/packages/70/0a/e0075ce9ca0279ee1d4f0c0b85f54fea02ebc83c3007651a72bece658fec/aiohttp-3.14.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6f71173be42d3241d428f760122febb748de0623f44308a6f120d0dd9ec572e3", size = 1767580, upload-time = "2026-06-07T21:07:07.873Z" }, - { url = "https://files.pythonhosted.org/packages/3e/61/a0c0a8f327a9c52095cdd8e312391b00d3ed64ab6c72bb5c33d8ec251cf7/aiohttp-3.14.1-cp312-cp312-win32.whl", hash = "sha256:ec8dc383ee57ea3e883477dcca3f11b65d58199f1080acaf4cd6ad9a99698be4", size = 452771, upload-time = "2026-06-07T21:07:09.669Z" }, - { url = "https://files.pythonhosted.org/packages/df/d9/ea367c75f16ac9c6cdc8febb25e8318fa21a2b1bc8d6514d4b2d890bface/aiohttp-3.14.1-cp312-cp312-win_amd64.whl", hash = "sha256:2aa92c87868cd13674989f9ee83e5f9f7ea4237589b728048e1f0c8f6caa3271", size = 479873, upload-time = "2026-06-07T21:07:11.538Z" }, - { url = "https://files.pythonhosted.org/packages/03/64/8d96784a7851156db8a4c6c3f6f91042fdf39fb15a4cc38c8b3c14833c45/aiohttp-3.14.1-cp312-cp312-win_arm64.whl", hash = "sha256:2c840c90759922cb5e6dda94596e079a30fb5a5ba548e7e0dc00574703940847", size = 448073, upload-time = "2026-06-07T21:07:13.637Z" }, -] - -[[package]] -name = "aiohttp-cors" -version = "0.8.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiohttp" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/d89e846a5444b3d5eb8985a6ddb0daef3774928e1bfbce8e84ec97b0ffa7/aiohttp_cors-0.8.1.tar.gz", hash = "sha256:ccacf9cb84b64939ea15f859a146af1f662a6b1d68175754a07315e305fb1403", size = 38626, upload-time = "2025-03-31T14:16:20.048Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/98/3b/40a68de458904bcc143622015fff2352b6461cd92fd66d3527bf1c6f5716/aiohttp_cors-0.8.1-py3-none-any.whl", hash = "sha256:3180cf304c5c712d626b9162b195b1db7ddf976a2a25172b35bb2448b890a80d", size = 25231, upload-time = "2025-03-31T14:16:18.478Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/58/d9/22ce5786ac0c1653ae8b6c23bded02c1686d11f0dbb45b31ce128e0df985/aiohttp-3.14.3.tar.gz", hash = "sha256:9491196535a88924a60afd5b5f434b5b203b6cc616250878dbdb223a8f7844bc", size = 7971213, upload-time = "2026-07-23T01:57:27.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/d4/eb96299230e20acf2efae207cb8d69051f1f68e357e5ea5e479bf6fb097a/aiohttp-3.14.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:39aded8c7f3b935b54aab1d8d73c70ec0ee2d3ec3b943e0e86611bc150ba47f5", size = 754690, upload-time = "2026-07-23T01:53:47.332Z" }, + { url = "https://files.pythonhosted.org/packages/88/11/e7a70a209eb9a067c0d3212b518a0134e3484f5178c7533878b6b514d469/aiohttp-3.14.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5bcb6ff3fdab1258a192679ff1a05d44f59626430aa05cd1a9d2447423599228", size = 509484, upload-time = "2026-07-23T01:53:51.159Z" }, + { url = "https://files.pythonhosted.org/packages/30/07/4bbc222cc8dbe31d4c3e8a5baad2286e4d42026ac0c570027b89afce6344/aiohttp-3.14.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:617105e2c3018ee38d0c8ce5ee3c84f621a6d8b9f723202aacaff28449ca91ee", size = 511949, upload-time = "2026-07-23T01:53:55.083Z" }, + { url = "https://files.pythonhosted.org/packages/54/b9/42e74c46b7b7c794b995bbc1f573fb48950c38b19d8600c62a6804ee2d67/aiohttp-3.14.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f631fe87a6f30df5fbe6d79640b25e4cffb38c31c7fb6f10871517b84b0f8c1a", size = 1765282, upload-time = "2026-07-23T01:53:59.662Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ed/62bc4d74363ad346d518e0720363a949f63e2e23439a79eb5813d4d29bb3/aiohttp-3.14.3-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:a94dbaae5ae27bd849c93570669bff91e0510f33a80805738e3de72a7be0447b", size = 1741511, upload-time = "2026-07-23T01:54:04.063Z" }, + { url = "https://files.pythonhosted.org/packages/d0/9f/181e8a8bc79e47d13c7fc4540bd7a3b729d9505609c61f392a8dd2fbfe55/aiohttp-3.14.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8f2f1c4c032c7cedd7d8da6f54c97b70266c6570c3108d3fdffee7188bb70529", size = 1810680, upload-time = "2026-07-23T01:54:09.882Z" }, + { url = "https://files.pythonhosted.org/packages/5c/9a/dec94d6ad694552fe3424e3f1928d7a606a5d9d9433a04e7ecdd9d38ae7f/aiohttp-3.14.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ea05e1f97ceea523942d9b2a7d7c0359d781d683d6b043f5943a602b14da4787", size = 1905646, upload-time = "2026-07-23T01:54:13.475Z" }, + { url = "https://files.pythonhosted.org/packages/52/b7/7cd31f29d6055bd711ae6e669367fba6f5ae9de463910a793e30556a8db7/aiohttp-3.14.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:543906c127fb1d929b95076db19b83fa2d46751006ff1e23b093aa5ac4d8db42", size = 1792122, upload-time = "2026-07-23T01:54:15.752Z" }, + { url = "https://files.pythonhosted.org/packages/66/73/10b1ef93afa61f4963c746257b70ced619cf31a4798671de5fdb2608501d/aiohttp-3.14.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0a5ff2dfbb9ce645fa5b8ef3e02c6c0b9cc3f6030ff863d0c51fffc50cb5541b", size = 1591127, upload-time = "2026-07-23T01:54:19.489Z" }, + { url = "https://files.pythonhosted.org/packages/49/ed/3b203fa6de1b338c14acdc06bf6ca9b043b7944f005966958c2ced932cde/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:041badb8f84396357c4d3ad26de6afd7a32b112f43d3c63045c0c8278cfd2043", size = 1725210, upload-time = "2026-07-23T01:54:24.129Z" }, + { url = "https://files.pythonhosted.org/packages/28/b7/1c2aab8c706436dcc28598452488ac9cd7c409da815237c28c27d58993e6/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:530125ee1163c4219af35dc3aa1206e541e7b31b6efc1a3f93b70a136f65d427", size = 1764848, upload-time = "2026-07-23T01:54:27.973Z" }, + { url = "https://files.pythonhosted.org/packages/54/50/94c28f08b131c4bf10984ea2c7a536c9920608bb2d6e7f95642c30cc87b7/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:c8653fd547c93a61aadc612007790f5555cdd18946fa48cf45e26d8ea4ea473d", size = 1777102, upload-time = "2026-07-23T01:54:31.775Z" }, + { url = "https://files.pythonhosted.org/packages/13/d4/e7d09ba7d345fb2d74440fd2fa033c5e079fac05552927705986f41a364f/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:89176250f686cb9853c0fb7ead90e639e915b84a6f43eedc2a4e7ec21f1037f0", size = 1580205, upload-time = "2026-07-23T01:54:34.518Z" }, + { url = "https://files.pythonhosted.org/packages/a3/84/072a91d68e1e1eb587985b54baab94221277f877e8ef274fc213a0ceae28/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:3a26434dafe408229ff3403458ca58de24fb51936504decac49ce6755f77e59d", size = 1797219, upload-time = "2026-07-23T01:54:36.995Z" }, + { url = "https://files.pythonhosted.org/packages/e0/eb/aad34e897e668424d6e995da5dff8a4a09af93363d3392488772957a63aa/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d1558173930a5a8d3069cee5c92fc91c87c4dbcb099debbb3622053717145a19", size = 1768629, upload-time = "2026-07-23T01:54:40.103Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2b/6bb88ddba0fecd9122aa3ebcad25996cf6c083a4a7040dbb3a4f97972af6/aiohttp-3.14.3-cp312-cp312-win32.whl", hash = "sha256:16100ad3ab8d649fdfbee87602d9d2dcdca9df0b9eda8a1b5fdc0d41f96da559", size = 451481, upload-time = "2026-07-23T01:54:42.547Z" }, + { url = "https://files.pythonhosted.org/packages/76/9b/f2f8f108da17ecef2cc3efc424e8b7ad3782b1a8360f7b8eae8ced84f6ea/aiohttp-3.14.3-cp312-cp312-win_amd64.whl", hash = "sha256:33a2d7c28d33797a2e99923dffa63f83d908a19b6bf26cfe80fa790aa5e1a75a", size = 476845, upload-time = "2026-07-23T01:54:44.853Z" }, + { url = "https://files.pythonhosted.org/packages/3e/44/28dac80a8941b604f4da10ce21097614ca1bf905ce93dca28d8d7de9c1e7/aiohttp-3.14.3-cp312-cp312-win_arm64.whl", hash = "sha256:362a3fd481769cac1a824514bcd86fda51c65e8fe6e051099e008fddde6db17c", size = 448050, upload-time = "2026-07-23T01:54:47.087Z" }, ] [[package]] @@ -256,10 +257,12 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/6f/60/1e787a0b5ebf318483235be2a689ee367173983067e441b8379564f667c0/apache_tvm_ffi-0.1.9.tar.gz", hash = "sha256:d2d402587e8906de0a07f4746aa78f3d452c7efe3625d4bb39ac2ad693bce530", size = 2513731, upload-time = "2026-02-27T19:28:06.602Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/df/f2/b8c4b151169f6d7ba8773c8af68b2e0c1013d7fb3f1bdf87573f47157ce9/apache_tvm_ffi-0.1.9-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:49e52350b0470654847de752e65603b604a4d3323e7e9f5e8a982f44acc4c143", size = 2041756, upload-time = "2026-02-27T19:27:23.931Z" }, { url = "https://files.pythonhosted.org/packages/a7/c0/6d3d54f50012255b41bc3e24944c086f63c4707c8686c7c6780e9283eb96/apache_tvm_ffi-0.1.9-cp312-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7d503029e66c43b1a1cb1a42a1e9bb428c8a28dcbdec31c28e705472ca648a3a", size = 2203712, upload-time = "2026-02-27T19:27:25.867Z" }, { url = "https://files.pythonhosted.org/packages/c6/dd/2bab4c6cd86257dbf99e93452a1af833113f8dc3e25a25579f6e4e4c8a94/apache_tvm_ffi-0.1.9-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28241371934ea8af10d5067087ba1229ebddded7b2c02d33a258ec2a96df8c46", size = 2299704, upload-time = "2026-02-27T19:27:27.477Z" }, { url = "https://files.pythonhosted.org/packages/7a/4a/b469bcb2e1014cb84d336d2a59f42958a058251c577a4c2680cacad346e2/apache_tvm_ffi-0.1.9-cp312-abi3-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:87cacce81df55685fc6a76e1e3c5db1200e85e87bf5974b692c59d131b7bc622", size = 2130865, upload-time = "2026-02-27T19:27:29.092Z" }, { url = "https://files.pythonhosted.org/packages/70/ef/5402da5d37f5270fd88ea0348acca78dba9be8bdbf6c2bcae0935eb03ef1/apache_tvm_ffi-0.1.9-cp312-abi3-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f45eb43499acac45ff6c93564f0ff2d3ca27b69656d540fd56ce59d51c0b4c65", size = 2278991, upload-time = "2026-02-27T19:27:30.729Z" }, + { url = "https://files.pythonhosted.org/packages/b5/23/1b7dc5f0807f83098183a57db6ee85b2c93b646d74a6e03781c9208aaeb0/apache_tvm_ffi-0.1.9-cp312-abi3-win_amd64.whl", hash = "sha256:d1dcf4c041d5ec05e3da1d545800c33cdbb95c113baa7705085ff79fa262752b", size = 1973200, upload-time = "2026-02-27T19:27:32.367Z" }, ] [[package]] @@ -351,6 +354,8 @@ version = "1.0.9" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/26/6a/4cc5a9dd40fd8a6d283fd3761e5f59c490109571ef8e3c73245417e5a305/blake3-1.0.9.tar.gz", hash = "sha256:5fa374fa5070ca084368776c19b420157eb0f2d3f091343d6bc59189929d62e2", size = 116872, upload-time = "2026-06-22T18:02:25.366Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/d2/9bdf8345c70993aaef635398f52edfb915d6e8ad2c000c801204e387c456/blake3-1.0.9-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:a70c20542d5e7960983a0ff32999049a2b0e5ef1f22dbbbdfb51cf04828a4156", size = 344587, upload-time = "2026-06-22T18:00:34.244Z" }, + { url = "https://files.pythonhosted.org/packages/36/9d/be8b1f7f85b12bb45a0fade6ca7bdbf83a507d23d0b6141ba29fe69c8cea/blake3-1.0.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:72cdecf088a9d25e6ec79948a578995649b0dbee407e7a46c543a9ecc0f6f281", size = 328864, upload-time = "2026-06-22T18:00:35.59Z" }, { url = "https://files.pythonhosted.org/packages/f2/78/66580635d744c826671fd219938caffb16281a26f62c4f856695d4233677/blake3-1.0.9-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42fa57bf462285ef16400601b0fd32214c248ba92505bbb94b1221ab9af5a092", size = 373795, upload-time = "2026-06-22T18:00:36.887Z" }, { url = "https://files.pythonhosted.org/packages/b1/79/b5b17d3004bb81a5732c0b176c812703d200ed8c652b3b7713b9633bbe10/blake3-1.0.9-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b25ccde5a64be070f20e5c7a81da70292db40b164b6c77588cbd6230856badbb", size = 374183, upload-time = "2026-06-22T18:00:38.205Z" }, { url = "https://files.pythonhosted.org/packages/3c/63/0d209c44b2041bbe130ced12a23c92dd995fbfe5bce7ee77fffea16f5cb0/blake3-1.0.9-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2a800b87433955f37691b5f361ad29c7dd3ee089c9cd109adc5aea8e24bc4c1f", size = 446783, upload-time = "2026-06-22T18:00:39.493Z" }, @@ -360,6 +365,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/34/cf/c7863a185550706a9624f6aa7b6d46470aaed0bb46a827c5cda2a7d03151/blake3-1.0.9-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:a288664d08dee154cc496e06e62517fc9e655ecec12b0d7db538d244ac79edf1", size = 380067, upload-time = "2026-06-22T18:00:45.249Z" }, { url = "https://files.pythonhosted.org/packages/54/0a/e7af679c719368b400c9ba9c3460072aac2ba077ddbd4bc806fef28cda03/blake3-1.0.9-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:91db52a809b68b5bebe7c413ddcd230e1f759398e7fa7a873104595a4fa648b6", size = 549471, upload-time = "2026-06-22T18:00:46.793Z" }, { url = "https://files.pythonhosted.org/packages/2c/3c/37c1dd3539b7bd9b6d2eef019802aacdb4a3d48ab484b140603bbf9c5b5a/blake3-1.0.9-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:cfaa671b07eb73883162ca940442193868358b0b904cfa266e4b74131ce966da", size = 591396, upload-time = "2026-06-22T18:00:48.122Z" }, + { url = "https://files.pythonhosted.org/packages/ae/55/4f0a23b72795292e74084834130900ea778c0583004519c86698dfffe1a5/blake3-1.0.9-cp312-cp312-win32.whl", hash = "sha256:ae47c3d5729ff89baa6ddf6de47fcfcc915985d39eb1bfcd6db653331f3c6fcc", size = 229271, upload-time = "2026-06-22T18:00:49.377Z" }, + { url = "https://files.pythonhosted.org/packages/12/91/7db93e4689f0f145bcb954dc62936e5f5090548a9fa20c6bbebfaeaa648a/blake3-1.0.9-cp312-cp312-win_amd64.whl", hash = "sha256:15566065ff90ab3da46ec0be1417406f00507af902b6fb0fbc6563e77f02fc42", size = 218220, upload-time = "2026-06-22T18:00:50.659Z" }, ] [[package]] @@ -446,10 +453,14 @@ version = "6.1.3" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/3a/6f/07b4af8da8bd27f640362b1ac8271d80895407f2ede0c2bcc9433c06e1ca/cbor2-6.1.3.tar.gz", hash = "sha256:8d70680acb55c04ea5b5ad86da094f9612b53d5a8a65d0f5b3aafc3ce917ecbb", size = 89503, upload-time = "2026-07-04T10:36:48.793Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/31/16/cff14259c3d19a7f0ae88b6996fe4c85f6ff1764dad889ac8a39e843e39c/cbor2-6.1.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3d939f55097c21e032f5a2d67592fcc57298986281f219356e2f519e4466f4ea", size = 412779, upload-time = "2026-07-04T10:36:04.975Z" }, { url = "https://files.pythonhosted.org/packages/50/6c/f3641d19b7b85a63cb2756c10164131489c2cb46b379ec51ae22283fefb9/cbor2-6.1.3-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b025009478d644dab407164fd60e3ef4381af284f5af6966df94c663756d949e", size = 457781, upload-time = "2026-07-04T10:36:06.349Z" }, { url = "https://files.pythonhosted.org/packages/55/85/0c55a66f3037056bfb8e1c7184168085fdea67ae5830404498bcf466233b/cbor2-6.1.3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:2226d32e102e375737656ad5d141ad8c6ae3e705e04e263f24756f0eb379c6c1", size = 468373, upload-time = "2026-07-04T10:36:07.769Z" }, { url = "https://files.pythonhosted.org/packages/46/74/40f7db3e0d880560193916a5c9b744fcf299558bed7113f77c28237c7c29/cbor2-6.1.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e61d465244d66ffed36492eef3b44d43795d76a2bba0663a2f15c186af7f7513", size = 523844, upload-time = "2026-07-04T10:36:09.404Z" }, { url = "https://files.pythonhosted.org/packages/b5/a1/b5e07d6a08441c3a552fe2ae48ccb7e9dfc5065b9f6a3bae9879b4f0fbc0/cbor2-6.1.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:87fe7be8fab6ec4796aa127c1a52e09e79dbafd2aa31caf809cf04b8080a5975", size = 536238, upload-time = "2026-07-04T10:36:10.914Z" }, + { url = "https://files.pythonhosted.org/packages/c9/99/e166be0fd74bf3a91f5a0d103e34883efbc438d970f72cc8200e274787e5/cbor2-6.1.3-cp312-cp312-win32.whl", hash = "sha256:da25d345f01e6a40b2e5c57ef96b4dcff7be69394fb62f0f70e07f437f2376a9", size = 279858, upload-time = "2026-07-04T10:36:12.247Z" }, + { url = "https://files.pythonhosted.org/packages/d2/1b/90b4a121e40aba189c55a5822dd3c698eaf487e1d4a780ab18c804a5ef1c/cbor2-6.1.3-cp312-cp312-win_amd64.whl", hash = "sha256:d5514f693db6fa6f433b4096e9b604e6a7bf151c9ef1d2db86d0858e4c5e768f", size = 300929, upload-time = "2026-07-04T10:36:13.564Z" }, + { url = "https://files.pythonhosted.org/packages/19/db/52c58a8d33464927389dde8103997b3fa51b081ce29b347ac2cc4fd0dfbf/cbor2-6.1.3-cp312-cp312-win_arm64.whl", hash = "sha256:3d43183d7beb3d3cd198d69b31bd2ee487ed704a1150c75cb0a66d6ad63d8c1a", size = 290908, upload-time = "2026-07-04T10:36:14.857Z" }, ] [[package]] @@ -548,31 +559,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a7/06/3d6badcf13db419e25b07041d9c7b4a2c331d3f4e7134445ec5df57714cd/coloredlogs-15.0.1-py2.py3-none-any.whl", hash = "sha256:612ee75c546f53e92e70049c9dbfcc18c935a2b9a53b66085ce9ef6a6e5c0934", size = 46018, upload-time = "2021-06-11T10:22:42.561Z" }, ] -[[package]] -name = "colorful" -version = "0.5.8" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/82/31/109ef4bedeb32b4202e02ddb133162457adc4eb890a9ed9c05c9dd126ed0/colorful-0.5.8.tar.gz", hash = "sha256:bb16502b198be2f1c42ba3c52c703d5f651d826076817185f0294c1a549a7445", size = 209361, upload-time = "2025-10-29T11:53:21.663Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/11/25cdf9d5fc21efd30134fc74c43702c6f7ef09ebae8ed927f1283403ad8d/colorful-0.5.8-py2.py3-none-any.whl", hash = "sha256:a9381fdda3337fbaba5771991020abc69676afa102646650b759927892875992", size = 201334, upload-time = "2025-10-29T11:53:20.251Z" }, -] - [[package]] name = "compressed-tensors" -version = "0.15.0.1" +version = "0.17.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "loguru" }, { name = "pydantic" }, - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "transformers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/41/1b/c3c4a98ec5f2727656336f07a0c35862195c310d8eb0b2fa5b4be6848680/compressed_tensors-0.15.0.1.tar.gz", hash = "sha256:a8e93054e8a5ec49c980b09ed36c4c1249b4a8ee167920a8e461c4da26e78d99", size = 229412, upload-time = "2026-04-10T14:23:54.708Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/9e/d7f18bd9a0354088abc11a0c1f2c7698f7c49e5a709faedf6a46e388f693/compressed_tensors-0.17.0.tar.gz", hash = "sha256:15c20d06bdbcf35b51fc99fd125e7b9be1e1855567c33b7a46dfac26ad6fb126", size = 257091, upload-time = "2026-06-03T16:49:17.208Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/52/93833dc1610e017ac5b7dcd59b8304d8ef67d1114c2d124e728a2cbbea12/compressed_tensors-0.15.0.1-py3-none-any.whl", hash = "sha256:e1b1f322e82e475715e242bad46925a304ea8e5c98b5055a15b8eb22fb6bfea9", size = 194260, upload-time = "2026-04-10T14:23:53.098Z" }, + { url = "https://files.pythonhosted.org/packages/35/63/6edf0415b072fff0bf8b546074dea3f0f9b148e49b601ac98bdc60a76c68/compressed_tensors-0.17.0-py3-none-any.whl", hash = "sha256:4a1b89b508f7efb8ffb4eee8a6e69e0452d9b080cae130146025c64fbe9fa9aa", size = 211714, upload-time = "2026-06-03T16:49:15.672Z" }, ] [[package]] @@ -657,6 +657,7 @@ dependencies = [ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/67/5e7dba1ba576dd73da5dee894ca076ca5e959450dfff66d6d510a255d1f7/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49", size = 6025351, upload-time = "2026-05-29T23:11:49.685Z" }, { url = "https://files.pythonhosted.org/packages/39/2a/6d2e9047d1fb243dbaa364b01e0297534b9ed7fd27dba1c9f361519cf69b/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e32d08f71ebcdf00f0f41eab2eb37e8da94c8ed411cc9f7f7a019ce6b34abe3a", size = 6657965, upload-time = "2026-05-29T23:11:52.227Z" }, + { url = "https://files.pythonhosted.org/packages/7c/95/872a0392122f1fb43fcb06869790ef3171f37beee9f7db8f441739113570/cuda_bindings-13.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:b134dd8c5c66ae4c4ad814f7aee88fd215353c077010cbc47e3b55ed35ec9eff", size = 5875099, upload-time = "2026-05-29T23:11:54.635Z" }, ] [[package]] @@ -670,6 +671,7 @@ dependencies = [ wheels = [ { url = "https://files.pythonhosted.org/packages/0d/a0/1daeae599cadd612689dbbf70d7da1c01883964fc2fbc7386f3c630a68cf/cuda_core-1.0.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6816dc020aee6103d8071bc02d8e4e1d91f2b49596f666896d608d92224d79d1", size = 4789856, upload-time = "2026-05-12T20:11:30.862Z" }, { url = "https://files.pythonhosted.org/packages/a1/4d/603557ab3cb171cc2a61d3678a39cb4dae3fd21275078bfbd1c0b0b5230b/cuda_core-1.0.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:be7b65311bf78964b7905adbf3c0f8f717d432f2854dc45169277729bf60f1e2", size = 5106023, upload-time = "2026-05-12T20:11:33.509Z" }, + { url = "https://files.pythonhosted.org/packages/c2/1a/ae079963c9df7f4274227eb63cf8f6083a532a6443adb340d951fd21c626/cuda_core-1.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:1a5c1aa3b738a7599ea289498d038fe625d259fd7ab795394541eee58a8e29bc", size = 4663076, upload-time = "2026-05-12T20:11:35.784Z" }, ] [[package]] @@ -693,22 +695,58 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/31/7ff3f7768eded7535c621abc2fecb9d181a34ea4cae3afe682feb796f242/cuda_python-13.3.1-py3-none-any.whl", hash = "sha256:280b014139ab447b6dd70a377db1596f310d6e887d9d342e6651b919ec145fb3", size = 8295, upload-time = "2026-05-29T23:28:47.012Z" }, ] +[[package]] +name = "cuda-tile" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine != 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f3/49/4592bc94ca05a07c7947ea114fd12734c8497f2daffee9faa79a03e39fb5/cuda_tile-1.3.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:375316b64c51ee7cfadb2f170a30c1547bc41eb39f1e233a6556713857d2e81f", size = 245744, upload-time = "2026-04-20T15:52:09.621Z" }, + { url = "https://files.pythonhosted.org/packages/40/76/84cb68be463c827bf79da9fa0aa5140838de6455ef6f438bbe0ffa75d378/cuda_tile-1.3.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:e4865acbff1172aaee304bf9c550586088d8b4545a384423597a590899386709", size = 247301, upload-time = "2026-04-20T15:51:04.042Z" }, +] + +[package.optional-dependencies] +tileiras = [ + { name = "nvidia-cuda-nvcc", version = "13.2.86", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-cuda-tileiras", version = "13.2.86", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-nvvm", version = "13.2.86", source = { registry = "https://pypi.org/simple" } }, +] + [[package]] name = "cuda-tile" version = "1.5.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform == 'win32'", + "sys_platform == 'darwin'", + "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] dependencies = [ { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/7c/6d/cc2fb5a25689a501564a2eced4acf654f307e801a2c1506be97c0d100491/cuda_tile-1.5.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:87652483baa9c81a9a24e4450f016e4ee78fd205d8422dad8996571bd1f2622e", size = 322641, upload-time = "2026-07-08T01:49:23.388Z" }, - { url = "https://files.pythonhosted.org/packages/1c/f5/b4ba9d0fc71198d939ebf9a090228179995d8411ee9def8f638a0e3ccdc5/cuda_tile-1.5.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:cef6d30acc37557643ece0de3770fc4c33497c4af40209e424f72fbfcbe6ea5a", size = 324990, upload-time = "2026-07-08T01:49:17.739Z" }, + { url = "https://files.pythonhosted.org/packages/a5/6e/7a60f317c503580ab7946dbb7fd080438fe953d0ddfdc81904beb9a1fab7/cuda_tile-1.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:16d97a60ed1d33388abbca85ea08cdae6325cc700476b3135f190d0fb50329f4", size = 304817, upload-time = "2026-07-08T01:49:38.853Z" }, +] + +[package.optional-dependencies] +tileiras = [ + { name = "cuda-toolkit", version = "13.3.1", source = { registry = "https://pypi.org/simple" }, extra = ["nvcc", "nvvm", "tileiras"] }, ] [[package]] name = "cuda-toolkit" version = "13.0.2" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine != 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] wheels = [ { url = "https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb", size = 2364, upload-time = "2025-12-19T23:24:07.328Z" }, ] @@ -718,7 +756,7 @@ cublas = [ { name = "nvidia-cublas" }, ] cudart = [ - { name = "nvidia-cuda-runtime" }, + { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" } }, ] cufft = [ { name = "nvidia-cufft" }, @@ -739,7 +777,7 @@ cusparse = [ { name = "nvidia-cusparse" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" } }, ] nvrtc = [ { name = "nvidia-cuda-nvrtc" }, @@ -748,6 +786,36 @@ nvtx = [ { name = "nvidia-nvtx" }, ] +[[package]] +name = "cuda-toolkit" +version = "13.3.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform == 'win32'", + "sys_platform == 'darwin'", + "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/a1/54c1e9498ba0df91ca15a46f41af6320cb9faed6ec2dbb30b6cbff8887c4/cuda_toolkit-13.3.1-py2.py3-none-any.whl", hash = "sha256:2ceda460a540323d52469bcfde48b48c1861f6482e4b5ea3cb5bdac00a1b11bd", size = 2656, upload-time = "2026-06-29T17:23:23.848Z" }, +] + +[package.optional-dependencies] +nvcc = [ + { name = "nvidia-cuda-crt", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "nvidia-cuda-nvcc", version = "13.3.73", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "nvidia-cuda-runtime", version = "13.3.29", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "nvidia-nvvm", version = "13.3.73", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, +] +nvvm = [ + { name = "nvidia-nvvm", version = "13.3.73", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, +] +tileiras = [ + { name = "nvidia-cuda-nvcc", version = "13.3.73", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "nvidia-cuda-tileiras", version = "13.3.36", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "nvidia-nvjitlink", version = "13.3.33", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "nvidia-nvvm", version = "13.3.73", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, +] + [[package]] name = "cycler" version = "0.12.1" @@ -880,15 +948,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl", hash = "sha256:5e31b2d5fbad117cc363ebaf6b689474db18a1f6438bc82358b024abd4c2ca19", size = 45550, upload-time = "2023-08-31T06:11:58.822Z" }, ] -[[package]] -name = "distlib" -version = "0.4.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c9/02/bd72be9134d25ed783ecbbc38a539ffaefbf90c78418c7fb7229600dbac7/distlib-0.4.3.tar.gz", hash = "sha256:f152097224a0ae24be5a0f6bae1b9359af82133bce63f98a95f86cae1aede9ed", size = 615141, upload-time = "2026-06-12T08:04:52.847Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/02/08/9c41fb51ab5b43eb21674aff13df270e8ba6c4b29c8624e328dc7a9482af/distlib-0.4.3-py2.py3-none-any.whl", hash = "sha256:4b0ce306c966eb73bc3a7b6abad017c556dadd92c44701562cd528ac7fde4d5b", size = 470628, upload-time = "2026-06-12T08:04:50.506Z" }, -] - [[package]] name = "distro" version = "1.9.0" @@ -990,7 +1049,7 @@ wheels = [ [[package]] name = "fastapi" -version = "0.139.2" +version = "0.136.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-doc" }, @@ -999,9 +1058,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/cd/95/d3f0ae10836324a2eab98a52b61210ac609f08200bf4bb0dc8132d32f78a/fastapi-0.139.2.tar.gz", hash = "sha256:333145a6891e9b5b3cfceb69baf817e8240cde4d4588ae5a10bf56ffacb6255e", size = 423428, upload-time = "2026-07-16T15:06:17.912Z" } +sdist = { url = "https://files.pythonhosted.org/packages/81/2d/ff8d91d7b564d464629a0fd50a4489c97fcb836ac230bf3a7269232a9b1f/fastapi-0.136.3.tar.gz", hash = "sha256:e487fae93ad408e6f47641ee4dfe389864fd7bec92e547ea8498fc13f43e83ab", size = 396410, upload-time = "2026-05-23T18:53:15.192Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5f/c7/cb03251d9dfb177246a9809a76f189d21df32dbd4a845951881d11323b7f/fastapi-0.139.2-py3-none-any.whl", hash = "sha256:b9ad015a835173d59865e2f5d8296fbc2b317bf56a2ba1a5bfbdd03de2fd4b1c", size = 130234, upload-time = "2026-07-16T15:06:19.557Z" }, + { url = "https://files.pythonhosted.org/packages/e0/82/45359b62a067409bd929ae8a56b8ed13e5a8c8a61194b3c236920999ab83/fastapi-0.136.3-py3-none-any.whl", hash = "sha256:3d2a69bdf04b7e9f3afa292c3bc7a98816bbfafa10bc9b45f3f3700d2f761620", size = 117481, upload-time = "2026-05-23T18:53:16.924Z" }, ] [package.optional-dependencies] @@ -1063,6 +1122,8 @@ version = "0.11.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/03/0f/0aeb3fc50046617702acc0078b277b58367fd62eb727b9ec733ae0e8bbcc/fastar-0.11.0.tar.gz", hash = "sha256:aa7f100f7313c03fdb20f1385927ba95671071ba308ad0c1763fef295e1895ce", size = 70238, upload-time = "2026-04-13T17:11:17.143Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/06/a5773706afc8bd496769786590bbc56d2d0ee419a299cc12ea3f5717fcf3/fastar-0.11.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3c51f1c2cdddbd1420d2897ace7738e36c65e17f6ae84e0bfe763f8d1068bb97", size = 708394, upload-time = "2026-04-13T17:09:57.269Z" }, + { url = "https://files.pythonhosted.org/packages/cc/a6/d5e2a4e48495616440a21eed07558219ca90243ad00b0502586f95bd4833/fastar-0.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0d9d6b052baf5380baea866675dab6ccd04ec2460d12b1c46f10ce3f4ee6a820", size = 628417, upload-time = "2026-04-13T17:09:42.145Z" }, { url = "https://files.pythonhosted.org/packages/ab/69/9816d69ac8265c9e50456637a487ccfb7a9c566efd9dbcd673df9c2558c2/fastar-0.11.0-cp312-cp312-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:bd2f05666d4df7e14885b5c38fefd92a785917387513d33d837ff42ec143a22f", size = 863950, upload-time = "2026-04-13T17:09:11.506Z" }, { url = "https://files.pythonhosted.org/packages/5b/0d/f88daad53aff2e754b6b5ff2a7113f72447a34f6ef17cc23ca99988117b7/fastar-0.11.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c1e6e74aba1ae77ca4aedcaf1697cd413319f4c88a5ccbe5b42c709517c5097e", size = 760737, upload-time = "2026-04-13T17:07:55.958Z" }, { url = "https://files.pythonhosted.org/packages/2f/a6/82ef4ecd969d50d92ed3ed9dbd8fe77faa24be5e5736f716edc9f4ce8d62/fastar-0.11.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:38ef77fe940bbc9b37a98bd838727f844b11731cd39358a2640ff864fb385086", size = 757603, upload-time = "2026-04-13T17:08:10.623Z" }, @@ -1074,6 +1135,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a5/37/e8bb24f506ba2b08fbaf36c5800e843bd4d542954e9331f00418e2d23349/fastar-0.11.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:4bb4dc0fc8f7a6807febcebce8a2f3626ba4955a9263d81ecc630aad83be84c0", size = 1035185, upload-time = "2026-04-13T17:10:30.207Z" }, { url = "https://files.pythonhosted.org/packages/9a/bf/be753736296338149ee4cb3e92e2b5423d6ba17c7b951d15218fd7e99bbf/fastar-0.11.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4ec95af56aa173f6e320e1183001bf108ba59beaf13edd1fc8200648db203588", size = 1072191, upload-time = "2026-04-13T17:10:47.072Z" }, { url = "https://files.pythonhosted.org/packages/d2/cd/a81c1aaafb5a22ce57c98ae22f39c89413ed53e4ee6e1b1444b0bd666a6c/fastar-0.11.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:136cf342735464091c39dc3708168f9fdeb9ebea40b1ead937c61afaf46143d9", size = 1028054, upload-time = "2026-04-13T17:11:04.293Z" }, + { url = "https://files.pythonhosted.org/packages/ec/88/1ce4eed3d70627c95f49ca017f6bbbf2ddcc4b0c601d293259de7689bc20/fastar-0.11.0-cp312-cp312-win32.whl", hash = "sha256:35f23c11b556cc4d3704587faacbc0037f7bdf6c4525cd1d09c70bda4b1c6809", size = 454198, upload-time = "2026-04-13T17:11:45.168Z" }, + { url = "https://files.pythonhosted.org/packages/8f/1d/26ce92f4331cd61a69840db9ca6115829805eec24f285481a854f578e917/fastar-0.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:920bc56c3c0b8a8ca492904941d1883c1c947c858cd93343356c29122a38f44c", size = 486697, upload-time = "2026-04-13T17:11:31.084Z" }, + { url = "https://files.pythonhosted.org/packages/ed/96/e6eda4480559c69b05d466e7b5ea9170e81fef3795a73e059959a3258319/fastar-0.11.0-cp312-cp312-win_arm64.whl", hash = "sha256:395248faf89e8a6bd5dc1fd544c8465113b627cb6d7c8b296796b60ebea33593", size = 462591, upload-time = "2026-04-13T17:11:20.577Z" }, ] [[package]] @@ -1148,6 +1212,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/07/4c/f17bd54c933fd2364 wheels = [ { url = "https://files.pythonhosted.org/packages/d1/3a/f95f7fc099ac1fc4c22aa46257d159eac88e29dd0765d21c4fc91caedb01/fastsafetensors-0.3.3-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2c2f788a936ffd17938484360339645812e14cf1b2bdf4c18c035c713218e5a9", size = 1887326, upload-time = "2026-07-07T07:21:34.624Z" }, { url = "https://files.pythonhosted.org/packages/92/8c/e3347b2a44a8ab9aced94fa450df4f309baa21f7f2981a8a7bd6a977f4d3/fastsafetensors-0.3.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3587bc66b8dec560ad903becf9540889013d4d47f0e10cf45f19bedc7b7bffa7", size = 1915478, upload-time = "2026-07-07T07:21:35.901Z" }, + { url = "https://files.pythonhosted.org/packages/80/65/388a55e6b2b3023fb732843335de14803f16a6d92f0cd47f1125d28b77ac/fastsafetensors-0.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:39c252a5528fa8653366f979d2ab8fa81159db4456b7c72cb5c288adb7699079", size = 424934, upload-time = "2026-07-07T07:21:37.216Z" }, ] [[package]] @@ -1192,20 +1257,21 @@ wheels = [ [[package]] name = "flashinfer-cubin" -version = "0.6.8.post1" +version = "0.6.13" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/11/b7/5e3b1a8c67031b421a8bd29c2bc29b900a550bb3392e8bda18bb15b5e476/flashinfer_cubin-0.6.8.post1-py3-none-any.whl", hash = "sha256:43636d4cd39e694a83d76a89f87fefcdf4cecb4c4f7dd22dac25ec368c1e901f", size = 295154113, upload-time = "2026-04-18T18:28:21.738Z" }, + { url = "https://files.pythonhosted.org/packages/19/43/ce916b4cdec4705173e222ca29c68e09004b47526888746094c5ffb29fca/flashinfer_cubin-0.6.13-py3-none-any.whl", hash = "sha256:41e4848c2d09d220e8394489b2fb6cfec6b6ad09f897b5ab8b39fc23055f6c24", size = 457984995, upload-time = "2026-06-25T00:29:26.08Z" }, ] [[package]] name = "flashinfer-python" -version = "0.6.8.post1" +version = "0.6.13" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "apache-tvm-ffi" }, { name = "click" }, - { name = "cuda-tile" }, + { name = "cuda-tile", version = "1.3.0", source = { registry = "https://pypi.org/simple" }, extra = ["tileiras"], marker = "sys_platform == 'linux'" }, + { name = "cuda-tile", version = "1.5.0", source = { registry = "https://pypi.org/simple" }, extra = ["tileiras"], marker = "sys_platform != 'linux'" }, { name = "einops" }, { name = "ninja" }, { name = "numpy" }, @@ -1215,12 +1281,13 @@ dependencies = [ { name = "packaging" }, { name = "requests" }, { name = "tabulate" }, - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "tqdm" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/53/1e/2760fef9e74abc4480961048e5790b4c9e955872fb4d7d97900cfddced5a/flashinfer_python-0.6.8.post1.tar.gz", hash = "sha256:b18e4121baf9b93fa9a9f368ba9b981a0342895f50ab9dddc224aeb964ed346f", size = 6675885, upload-time = "2026-04-18T18:28:13.299Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/f7/7f6dd2b03f4277509dfd1e5c7a8ec1de2662fd245d2e663f44a3493882b1/flashinfer_python-0.6.13.tar.gz", hash = "sha256:8a6d7d3708c7c87952390ec4e3aabe6e1c356defa8c7211b26bccaa355a61c59", size = 9638085, upload-time = "2026-06-24T22:46:29.391Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/73/6d/1e8a8533913e33a50a486332ce0673f4fdb860f6eb9ed450327c5c1762cb/flashinfer_python-0.6.8.post1-py3-none-any.whl", hash = "sha256:818f9b8cc2fe66c42a1f6264be4841ac8821ada703685a02cfccb2b5124a710b", size = 9385316, upload-time = "2026-04-18T18:28:10.285Z" }, + { url = "https://files.pythonhosted.org/packages/e4/50/e8920ed7f68e0116a385e3ab814ac2f0010579852fc483bfb48819d11976/flashinfer_python-0.6.13-py3-none-any.whl", hash = "sha256:239e6ddc3cbbaf0bee251861a8c7c69438b1171830d69ddfa133ddea4494850d", size = 14191198, upload-time = "2026-06-24T22:46:26.565Z" }, ] [[package]] @@ -1308,50 +1375,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/71/ae30dadffc90b9006d77af76b393cb9dfbfc9629f339fc1574a1c52e6806/future-1.0.0-py3-none-any.whl", hash = "sha256:929292d34f5872e70396626ef385ec22355a1fae8ad29e1a734c3e43f9fbc216", size = 491326, upload-time = "2024-02-21T11:52:35.956Z" }, ] -[[package]] -name = "gguf" -version = "0.19.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "tqdm" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/48/ae/17f1308ae45cd7b08ebb521747d5b23f4efc4d172038a4e228dd5106c3ff/gguf-0.19.0.tar.gz", hash = "sha256:dbadcd6cc7ccd44256f2229fe7c2dff5e8aa5cf0612ab987fd2b1a57e428923f", size = 111220, upload-time = "2026-05-06T13:04:03.667Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/bb/d71d6da82763528c2c2ed6b59a9d6142c6595545a4c448e2085d155e88c2/gguf-0.19.0-py3-none-any.whl", hash = "sha256:70bcd10edfe697fb2dad6e40af2234b9d8ece9a41a99761405121ebda1c3c1cd", size = 118475, upload-time = "2026-05-06T13:04:02.588Z" }, -] - -[[package]] -name = "google-api-core" -version = "2.32.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-auth" }, - { name = "googleapis-common-protos" }, - { name = "proto-plus" }, - { name = "protobuf" }, - { name = "requests" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/03/33/00277be1305fd68355d08197f05e22db259c0cff49a10c8590a1869ade9b/google_api_core-2.32.0.tar.gz", hash = "sha256:2b33aad226b19272458c46abfe5c5a38d9531ece0c44502129a1463ce83674ac", size = 177659, upload-time = "2026-07-16T20:36:07.717Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/81/44/5018c5ac1526c98169db98d87a6ff7d5508f5246621c3ee1a046fdd5e0a6/google_api_core-2.32.0-py3-none-any.whl", hash = "sha256:ae1f0d58a6c8869350bf469f8eb3092e7f8c494a942d9525494afb6c162b0904", size = 174198, upload-time = "2026-07-16T20:35:41.865Z" }, -] - -[[package]] -name = "google-auth" -version = "2.56.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cryptography" }, - { name = "pyasn1-modules" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/58/66/b4ba60005743e01933e22b4f62313e063f7460458b7d8a358427b4930013/google_auth-2.56.0.tar.gz", hash = "sha256:f90fa030b569a92654b9d690665a073841df33d57487be53db583a9a0867a553", size = 364629, upload-time = "2026-07-13T19:09:57.143Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/7d/cd3e187f14ce832e419e70709bfcc40cb0dc11517d5d03c9d3919bcc3101/google_auth-2.56.0-py3-none-any.whl", hash = "sha256:6e88c10217e07a92bfd01cac8ee99e32ccfb08414c3102e6c5b8d58f37a0d1e0", size = 257976, upload-time = "2026-07-13T19:09:42.685Z" }, -] - [[package]] name = "googleapis-common-protos" version = "1.75.0" @@ -1512,21 +1535,22 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "0.36.2" +version = "1.27.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "click" }, { name = "filelock" }, { name = "fsspec" }, - { name = "hf-xet", marker = "platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, + { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, + { name = "httpx" }, { name = "packaging" }, { name = "pyyaml" }, - { name = "requests" }, { name = "tqdm" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7c/b7/8cb61d2eece5fb05a83271da168186721c450eb74e3c31f7ef3169fa475b/huggingface_hub-0.36.2.tar.gz", hash = "sha256:1934304d2fb224f8afa3b87007d58501acfda9215b334eed53072dd5e815ff7a", size = 649782, upload-time = "2026-02-06T09:24:13.098Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/9b/ddf3d02a8681f1b9ce52fda03d755dad6b74c4f8172304c4c8d2975450f9/huggingface_hub-1.27.0.tar.gz", hash = "sha256:c1fed40ea82a6b41b477f5243546549b792ae0a93abcea608cff66089bf8f8df", size = 942668, upload-time = "2026-08-07T12:48:05.161Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/af/48ac8483240de756d2438c380746e7130d1c6f75802ef22f3c6d49982787/huggingface_hub-0.36.2-py3-none-any.whl", hash = "sha256:48f0c8eac16145dfce371e9d2d7772854a4f591bcb56c9cf548accf531d54270", size = 566395, upload-time = "2026-02-06T09:24:11.133Z" }, + { url = "https://files.pythonhosted.org/packages/de/d8/95b735e183957c1f26d94c52977f09d466d55119cbbc1558ea4975e4c216/huggingface_hub-1.27.0-py3-none-any.whl", hash = "sha256:7df6827c2f956c60fbaa64646e979e566db76f619dd0a9729dfb8c5a3eb4f68d", size = 784926, upload-time = "2026-08-07T12:48:02.905Z" }, ] [[package]] @@ -1541,6 +1565,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f0/0f/310fb31e39e2d734ccaa2c0fb981ee41f7bd5056ce9bc29b2248bd569169/humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477", size = 86794, upload-time = "2021-09-17T21:40:39.897Z" }, ] +[[package]] +name = "humming-kernels" +version = "0.1.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-bindings" }, + { name = "jinja2" }, + { name = "numpy" }, + { name = "nvidia-ml-py" }, + { name = "pyelftools" }, + { name = "safetensors" }, + { name = "tabulate" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "tqdm" }, + { name = "triton" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/51/4f/6977a31451c3f7aa1deaa76506d6cdeb74ef418ad8bcba2e98f7510b26ec/humming_kernels-0.1.10.tar.gz", hash = "sha256:da3e46fb9fc9eba2a9327c2e8135ead68e390c955acd7449f97ee7c71666c8b1", size = 220110, upload-time = "2026-07-02T10:22:57.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/ba/869bc24591d2b4fb0d8da821528072971052a934af077a11f77a0f2b3e79/humming_kernels-0.1.10-py3-none-any.whl", hash = "sha256:4ded0998ff085afeddde70baf93f97c2929969ec3d4a63a52cfec5072bc972b4", size = 184889, upload-time = "2026-07-02T10:22:56.031Z" }, +] + +[package.optional-dependencies] +cu13 = [ + { name = "nvidia-cuda-cccl" }, + { name = "nvidia-cuda-nvcc", version = "13.2.86", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvcc", version = "13.3.73", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, + { name = "nvidia-cuda-nvrtc" }, + { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "nvidia-cuda-runtime", version = "13.3.29", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'win32'" }, +] + [[package]] name = "idna" version = "3.18" @@ -1556,12 +1612,18 @@ version = "3.5.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/3a/06/b31f040a8764336a11152e474a7abcb3782fedb0d1cdf78f442b82878c56/ijson-3.5.1.tar.gz", hash = "sha256:af40bd1a85f55db0b8b30715c858761306bd92d5590148636f75c3309e6e76bd", size = 69913, upload-time = "2026-07-06T17:37:42.923Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/6e/f3ded1ebb85ccc89a30f7b10a0076f30db70ae1d1e0b6423ff93c57b7539/ijson-3.5.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ee60c7741012671867678eae71c51872cac938b76f3d4ca40a778e6c361774d2", size = 88643, upload-time = "2026-07-06T17:36:28.529Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f2/18f14a1d79ef4898e746b4f50dcdbe60abab317cc2bd8390f043b9553c4e/ijson-3.5.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:11c1d7d36a13054b5872ecd5d745dc4009d9abdbcba2312de69e66c2f92a46d2", size = 60611, upload-time = "2026-07-06T17:36:29.597Z" }, + { url = "https://files.pythonhosted.org/packages/30/c7/6e3e591324fd4c7a7a9e1bc23548bacbd84c0d91766b71f09f13e945e7e9/ijson-3.5.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b9517efbe6604bce16f3e50d49b0cd1bdc58917f98cf2eab026599c5c0422991", size = 60447, upload-time = "2026-07-06T17:36:30.747Z" }, { url = "https://files.pythonhosted.org/packages/4d/a5/9af7be670381ddac26dd55107ed0110b50f5161673b053311db67f510dcc/ijson-3.5.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:ea4fd7bec203a600b1cc88a492dfe6b75ce4b1b87488a66adcd5406022213f64", size = 139092, upload-time = "2026-07-06T17:36:31.749Z" }, { url = "https://files.pythonhosted.org/packages/41/fb/f9c1664d75467453e6bd4e5f9cd2211b730b09e049445ab64cbac68cc6a3/ijson-3.5.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:350caea815e53151994b597abc80cf669454276b5ac6aadcec69ef6d48f7e90b", size = 149921, upload-time = "2026-07-06T17:36:32.912Z" }, { url = "https://files.pythonhosted.org/packages/43/80/d20b1c49c4aa7cc6644131e2e57192b45346ef4816566ed1cd9fd05bae38/ijson-3.5.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e4fcebfe1685bb7ba06a8255a5d428ea6b4b895d7acf979cb637d8bbc9db2f47", size = 149848, upload-time = "2026-07-06T17:36:34.032Z" }, { url = "https://files.pythonhosted.org/packages/fd/fc/5baa710869f5ab939e6233583ced1546889b55c35f35b844c518ac10abc3/ijson-3.5.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d78f362f51c8691798758a9e6ac3c9d385ee1228cb82987c91562a2fae235cd3", size = 150810, upload-time = "2026-07-06T17:36:35.19Z" }, { url = "https://files.pythonhosted.org/packages/54/16/a12b3d987a5c1677b04557c6f9b9feb7e04b7d4171e9a344856cb9136e9b/ijson-3.5.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:0b184180d45f85fd4479659582749b109e49f4a29c21ac700ccc9c2280fe015e", size = 142989, upload-time = "2026-07-06T17:36:36.23Z" }, { url = "https://files.pythonhosted.org/packages/ed/63/1026c535671fc334fc85aeb78f0945c825e7a338575edc753c0f455459ae/ijson-3.5.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e353891d33a2e6aa5caf72c2a5fbadd7a46f5f9b32dcfd0c84113b2444c255b8", size = 151702, upload-time = "2026-07-06T17:36:37.296Z" }, + { url = "https://files.pythonhosted.org/packages/cb/af/b58aa3a2bf4d31c388ea78b49826605f60932891ce97e404d196766b4ea3/ijson-3.5.1-cp312-cp312-win32.whl", hash = "sha256:936f28671f018f8ac4d3f003ae9fa01d0467ab4ef4cfd0c97f23beda485b61c6", size = 52613, upload-time = "2026-07-06T17:36:38.345Z" }, + { url = "https://files.pythonhosted.org/packages/04/66/ce70a92949c2a753dad91fdd5761dc14f3a44517e80cfc3c26612982ed61/ijson-3.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:322c783f3ee0c6b383bbd4db88370b10172168808cc2a0bf811f1253f7435602", size = 54729, upload-time = "2026-07-06T17:36:39.337Z" }, + { url = "https://files.pythonhosted.org/packages/a5/ff/e17784240c9cf1d58de2f2853ebaf9cc54f6bce117a1f12a6150bbb4a5aa/ijson-3.5.1-cp312-cp312-win_arm64.whl", hash = "sha256:e2ac204b59f09e38e16d277f906240e9fd38780e42076599419265af183dc4b4", size = 53714, upload-time = "2026-07-06T17:36:40.308Z" }, ] [[package]] @@ -2044,7 +2106,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.93.0" +version = "1.95.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -2060,20 +2122,24 @@ dependencies = [ { name = "tiktoken" }, { name = "tokenizers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/93/e1/4f05ca4cbb4efb739c9e66a182ecd5c816bc05bf3665ec8e0fb4ab408379/litellm-1.93.0.tar.gz", hash = "sha256:140bf215e264c71601bca9c06d2436c5451bb59e1e195ea23fc2d3d87b6929ec", size = 15948866, upload-time = "2026-07-19T03:01:24.389Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0e/96/8cdfb9aaf584b57af35a0423c111a1c1264a78b548cebbb5ed96defacdab/litellm-1.95.0.tar.gz", hash = "sha256:0ef126d52c7a559f8353e50d60fd0d5e7e6c8767ad54df25ddaf79b9edca1afc", size = 17513577, upload-time = "2026-08-02T02:52:49.465Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/be/69/cabe7e747fea4c744752bd7ff8f7f208723151a63a89bb7c2437212523ff/litellm-1.93.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:3daf5c5aceb07f5d68871071e0ecdc678caccebadca9143cc00bb76f5a8c54e8", size = 20164234, upload-time = "2026-07-19T03:01:05.713Z" }, - { url = "https://files.pythonhosted.org/packages/c5/db/6af798603c6e2cf21ad7f2edf7e95019bd859dda82284b94014d608fcd85/litellm-1.93.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:0784172435de48f66ef7ad89d421604db1ad0db1321ef7f39dbcfe6b20111417", size = 20156724, upload-time = "2026-07-19T03:01:09.037Z" }, + { url = "https://files.pythonhosted.org/packages/33/d0/ad0272853cc450f8bb4a40a93d206e767e18ac2ff3f91870374e1d9fc090/litellm-1.95.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:cb667f84f08520f32b076e03c7a3fa51bf3f7e8b641dade34ab046bf00314d6b", size = 26421359, upload-time = "2026-08-02T02:52:16.048Z" }, + { url = "https://files.pythonhosted.org/packages/02/c1/4301aa8ef6d2fb0e4a2b8dec973d7c4499f680b26d2e1b77643864235a4b/litellm-1.95.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:1bdf7153557cc0851fa9477b137fde476c56d5de92a5778ecfc6c3a75439a4e1", size = 26300401, upload-time = "2026-08-02T02:52:19.501Z" }, + { url = "https://files.pythonhosted.org/packages/7c/3d/6cd087bd541f18d924f17bd8e1bb68a7f9d73d03274c5339f9de563bc992/litellm-1.95.0-cp312-cp312-win_amd64.whl", hash = "sha256:62cc5d834e8223dbd16c9ad0b46c73354b6d67cc7fa0eba2764ce65b3b8c474f", size = 24917446, upload-time = "2026-08-02T02:52:23.119Z" }, ] [[package]] name = "llguidance" -version = "1.3.0" +version = "1.7.6" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/95/48/3f7a9d3ff1b36bba92b5107a3a21286821227afe9ea464736133994d61fb/llguidance-1.3.0.tar.gz", hash = "sha256:861249afd51dc325646834462ea827e57a5c2b2042e108e6aae7059fdad9104d", size = 1070460, upload-time = "2025-10-20T19:58:44.164Z" } +sdist = { url = "https://files.pythonhosted.org/packages/da/91/6bc8bb503dc259e46d253b5424385a54fe06c38a4c7a12befe69a3c2455a/llguidance-1.7.6.tar.gz", hash = "sha256:db7febbe412ed2015501904646750071d7e00e6df7f85c4b956ad4f206fd2df7", size = 1156574, upload-time = "2026-06-03T20:13:25.316Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/aa/11/44389d3d1526d7a5c38ffd587a5ebc61d7bee443ac1dea95f2089ad58f5f/llguidance-1.3.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6f6caca5d78db7f76e1fbb0fff8607b861c32d47fa3d5dee2fc49de27ee269df", size = 2835242, upload-time = "2025-10-20T19:58:34.518Z" }, - { url = "https://files.pythonhosted.org/packages/83/a8/1ff2bedb8f9acb46a2d2d603415d272bb622c142ea86f5b95445cc6e366c/llguidance-1.3.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc17e9dd602c3879bf91664a64bf72f54c74dbfbeb24ccfab6a5fe435b12f7aa", size = 3033133, upload-time = "2025-10-20T19:58:38.721Z" }, + { url = "https://files.pythonhosted.org/packages/fa/1d/5a9a13421b1f3f1c1acf82beb63ed72fa4d302e65099b72f4a4fe5a098ab/llguidance-1.7.6-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:eabf4572c8731734c0444c353b9ea06bc5c156986d2ff0a4ec0499159271381f", size = 3227892, upload-time = "2026-06-03T20:13:09.533Z" }, + { url = "https://files.pythonhosted.org/packages/46/fe/bb185f11bad82f2637e3cd8cbf6b200cbb6ed56ac395de47ea05a60d4649/llguidance-1.7.6-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:9c54c899db8cb4b4fba128a7d844730066576c70d806c95ada92b2bd2d6ab498", size = 3138127, upload-time = "2026-06-03T20:13:11.649Z" }, + { url = "https://files.pythonhosted.org/packages/51/b9/dc76d7716e04dc7b3427cae52eaa32bd20771382d4d1dd9f4538a9dd2086/llguidance-1.7.6-cp39-abi3-manylinux_2_31_aarch64.whl", hash = "sha256:e70fa25ed550c2b50c2fd70baa9e2808b4ecb859d01e453bd5459aff62ba38c3", size = 2899993, upload-time = "2026-06-03T20:13:13.563Z" }, + { url = "https://files.pythonhosted.org/packages/1a/64/d74336f22242ef94356a456057d4ff1be7c1bc9c7dbc867171c6982a5512/llguidance-1.7.6-cp39-abi3-manylinux_2_31_x86_64.whl", hash = "sha256:ceec951d29a74309984e3be0fe7f5f56c1362434cd937abd517b259a60908b1e", size = 3074809, upload-time = "2026-06-03T20:13:15.498Z" }, + { url = "https://files.pythonhosted.org/packages/49/37/99d700f0e2c83acf25a8d8946b2bee9f5eac47bc530bfbd53ba3126c667f/llguidance-1.7.6-cp39-abi3-win_amd64.whl", hash = "sha256:ace7e81cd31950a87186356ab24bd7f75fbc10a05ca9d9f7f8748f931963f763", size = 2879207, upload-time = "2026-06-03T20:13:23.341Z" }, ] [[package]] @@ -2366,10 +2432,14 @@ version = "0.21.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e3/60/f79b9b013a16fa3a58350c9295ddc6789f2e335f36ea61ed10a21b215364/msgspec-0.21.1.tar.gz", hash = "sha256:2313508e394b0d208f8f56892ca9b2799e2561329de9763b19619595a6c0f72c", size = 319193, upload-time = "2026-04-12T21:44:50.394Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/6e/cf/317224852c00248c620a9bcf4b26e2e4ab8afd752f18d2a6ef73ebd423b6/msgspec-0.21.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d4248cf0b6129b7d230eacd493c17cc2d4f3989f3bb7f633a928a85b7dcfa251", size = 196188, upload-time = "2026-04-12T21:44:07.181Z" }, + { url = "https://files.pythonhosted.org/packages/6d/81/074612945c0666078f7366f40000013de9f6ba687491d450df699bceebc9/msgspec-0.21.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5102c7e9b3acff82178449b85006d96310e690291bb1ea0142f1b24bcb8aabcb", size = 188473, upload-time = "2026-04-12T21:44:08.736Z" }, { url = "https://files.pythonhosted.org/packages/8a/37/655101799590bcc5fddb2bd3fe0e6194e816c2d1da7c361725f5eb89a910/msgspec-0.21.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:846758412e9518252b2ac9bffd6f0e54d9ff614f5f9488df7749f81ff5c80920", size = 218871, upload-time = "2026-04-12T21:44:09.917Z" }, { url = "https://files.pythonhosted.org/packages/b5/d1/d4cd9fe89c7d400d7a18f86ccc94daa3f0927f53558846fcb60791dce5d6/msgspec-0.21.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:21995e74b5c598c2e004110ad66ec7f1b8c20bf2bcf3b2de8fd9a3094422d3ff", size = 225025, upload-time = "2026-04-12T21:44:11.191Z" }, { url = "https://files.pythonhosted.org/packages/24/bf/e20549e602b9edccadeeff98760345a416f9cce846a657e8b18e3396b212/msgspec-0.21.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6129f0cca52992e898fd5344187f7c8127b63d810b2fd73e36fca73b4c6475ee", size = 222672, upload-time = "2026-04-12T21:44:12.481Z" }, { url = "https://files.pythonhosted.org/packages/b4/68/04d7a8f0f786545cf9b8c280c57aa6befb5977af6e884b8b54191cbe44b3/msgspec-0.21.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ef3ec2296248d1f8b9231acb051b6d471dfde8f21819e86c9adaaa9f42918521", size = 227303, upload-time = "2026-04-12T21:44:13.709Z" }, + { url = "https://files.pythonhosted.org/packages/cc/4d/619866af2840875be408047bf9e70ceafbae6ab50660de7134ed1b25eb86/msgspec-0.21.1-cp312-cp312-win_amd64.whl", hash = "sha256:d4ab834a054c6f0cbeef6df9e7e1b33d5f1bc7b86dea1d2fd7cad003873e783d", size = 190017, upload-time = "2026-04-12T21:44:14.977Z" }, + { url = "https://files.pythonhosted.org/packages/5e/2e/a8f9eca8fd00e097d7a9e99ba8a4685db994494448e3d4f0b7f6e9a3c0f7/msgspec-0.21.1-cp312-cp312-win_arm64.whl", hash = "sha256:628aaa35c74950a8c59da330d7e98917e1c7188f983745782027748ee4ca573e", size = 175345, upload-time = "2026-04-12T21:44:16.431Z" }, ] [[package]] @@ -2435,6 +2505,8 @@ dependencies = [ { name = "ffmpeg-python" }, { name = "fsspec" }, { name = "httpx" }, + { name = "huggingface-hub" }, + { name = "jinja2" }, { name = "lancedb" }, { name = "langchain-nvidia-ai-endpoints" }, { name = "markitdown" }, @@ -2451,11 +2523,12 @@ dependencies = [ { name = "pypdfium2" }, { name = "python-multipart" }, { name = "pyyaml" }, - { name = "ray", extra = ["data", "serve"] }, + { name = "ray" }, { name = "requests" }, { name = "rich" }, { name = "s3fs" }, { name = "sqlglot" }, + { name = "tokenizers" }, { name = "tqdm" }, { name = "typer" }, { name = "universal-pathlib" }, @@ -2494,14 +2567,13 @@ all = [ { name = "scipy" }, { name = "soundfile" }, { name = "timm" }, - { name = "tokenizers" }, { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" }, { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" }, { name = "torchvision", version = "0.26.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "transformers" }, { name = "tritonclient" }, - { name = "vllm", marker = "sys_platform == 'linux'" }, + { name = "vllm" }, ] benchmarks = [ { name = "datasets" }, @@ -2531,14 +2603,13 @@ local = [ { name = "psutil" }, { name = "scikit-learn" }, { name = "timm" }, - { name = "tokenizers" }, { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" }, { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" }, { name = "torchvision", version = "0.26.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "transformers" }, { name = "tritonclient" }, - { name = "vllm", marker = "sys_platform == 'linux'" }, + { name = "vllm" }, ] multimedia = [ { name = "cairosvg" }, @@ -2599,24 +2670,26 @@ requires-dist = [ { name = "fastapi", specifier = ">=0.114.0" }, { name = "fastmcp", specifier = ">=2.0.0" }, { name = "ffmpeg-python" }, - { name = "flashinfer-cubin", marker = "sys_platform == 'linux' and extra == 'local'", specifier = "==0.6.8.post1" }, - { name = "flashinfer-python", marker = "sys_platform == 'linux' and extra == 'local'", specifier = "==0.6.8.post1" }, + { name = "flashinfer-cubin", marker = "sys_platform == 'linux' and extra == 'local'", specifier = "==0.6.13" }, + { name = "flashinfer-python", marker = "sys_platform == 'linux' and extra == 'local'", specifier = "==0.6.13" }, { name = "fsspec", specifier = ">=2025.5.1" }, { name = "httpx", specifier = ">=0.27.0" }, + { name = "huggingface-hub", specifier = ">=0.34.0" }, + { name = "jinja2", specifier = ">=3.1" }, { name = "lancedb" }, { name = "langchain-nvidia-ai-endpoints", specifier = ">=1.4.0" }, { name = "langgraph", marker = "extra == 'tabular'", specifier = ">=1.2.0" }, { name = "librosa", marker = "extra == 'multimedia'", specifier = ">=0.10.2" }, { name = "librosa", marker = "extra == 'service'", specifier = ">=0.10.2" }, - { name = "litellm", marker = "extra == 'llm'", specifier = ">=1.86.0,<2" }, - { name = "litellm", marker = "extra == 'service'", specifier = ">=1.86.0,<2" }, + { name = "litellm", marker = "extra == 'llm'", specifier = ">=1.95.0rc3,<2" }, + { name = "litellm", marker = "extra == 'service'", specifier = ">=1.95.0rc3,<2" }, { name = "markitdown" }, { name = "nemo-retriever", extras = ["benchmarks", "llm", "local", "multimedia", "nemotron-parse", "service", "tabular"], marker = "extra == 'all'" }, { name = "nemotron-ocr", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'local') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'local')", specifier = ">=2.0.0,<3" }, { name = "nemotron-page-elements-v3", marker = "extra == 'local'", specifier = ">=3.0.1,<4" }, { name = "nemotron-table-structure-v1", marker = "extra == 'local'", specifier = ">=1.0.0,<2" }, { name = "neo4j", marker = "extra == 'tabular'", specifier = ">=5.0" }, - { name = "nltk", specifier = "==3.9.4" }, + { name = "nltk", specifier = ">=3.10.1" }, { name = "numba", marker = "extra == 'multimedia'", specifier = ">=0.59" }, { name = "numba", marker = "extra == 'service'", specifier = ">=0.59" }, { name = "numpy", specifier = ">=1.26.0" }, @@ -2629,8 +2702,8 @@ requires-dist = [ { name = "opentelemetry-exporter-otlp-proto-grpc", specifier = ">=1.41.1" }, { name = "opentelemetry-sdk", specifier = ">=1.41.1" }, { name = "pandas", specifier = ">=2.0,<3" }, - { name = "pillow", specifier = "==12.2.0" }, - { name = "prometheus-fastapi-instrumentator", specifier = ">=7.0,<8" }, + { name = "pillow", specifier = ">=12.3.0" }, + { name = "prometheus-fastapi-instrumentator", specifier = ">=8.0,<9" }, { name = "psutil", marker = "extra == 'local'", specifier = ">=5.9.0" }, { name = "psutil", marker = "extra == 'service'", specifier = ">=5.9.0" }, { name = "pydantic", specifier = ">=2.8.0" }, @@ -2638,8 +2711,8 @@ requires-dist = [ { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.2" }, { name = "python-multipart", specifier = ">=0.0.9" }, { name = "pyyaml", specifier = ">=6.0" }, - { name = "ray", extras = ["data", "serve"], specifier = ">=2.49.0" }, - { name = "requests", specifier = ">=2.32.5" }, + { name = "ray", extras = ["data", "serve"], specifier = ">=2.56.1" }, + { name = "requests", specifier = ">=2.33.0" }, { name = "rich", specifier = ">=13.7.0" }, { name = "s3fs", specifier = ">=2025.5.1" }, { name = "scikit-learn", marker = "extra == 'local'", specifier = ">=1.6.0" }, @@ -2648,7 +2721,7 @@ requires-dist = [ { name = "soundfile", marker = "extra == 'multimedia'", specifier = ">=0.12.0" }, { name = "sqlglot", specifier = ">=30.0.0" }, { name = "timm", marker = "extra == 'local'", specifier = "==1.0.22" }, - { name = "tokenizers", marker = "extra == 'local'", specifier = ">=0.21.1" }, + { name = "tokenizers", specifier = ">=0.21.1" }, { name = "torch", marker = "sys_platform == 'darwin' and extra == 'local'", specifier = "==2.11.0" }, { name = "torch", marker = "sys_platform == 'linux' and extra == 'local'", specifier = "==2.11.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "torch", marker = "sys_platform == 'win32' and extra == 'local'", specifier = "==2.11.0", index = "https://download.pytorch.org/whl/cu130" }, @@ -2656,13 +2729,13 @@ requires-dist = [ { name = "torchvision", marker = "sys_platform == 'linux' and extra == 'local'", specifier = "==0.26.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "torchvision", marker = "sys_platform == 'win32' and extra == 'local'", specifier = "==0.26.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "tqdm", specifier = ">=4.66.0" }, - { name = "transformers", marker = "extra == 'local'", specifier = ">=4.57.6,<5" }, + { name = "transformers", marker = "extra == 'local'", specifier = ">=5.14.1,<6" }, { name = "tritonclient", marker = "extra == 'local'" }, { name = "typer", specifier = ">=0.12.0" }, { name = "universal-pathlib", specifier = ">=0.2.0" }, { name = "urllib3", specifier = "==2.7.0" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.30.0" }, - { name = "vllm", marker = "sys_platform == 'linux' and extra == 'local'", specifier = "==0.20.0" }, + { name = "vllm", marker = "sys_platform == 'linux' and extra == 'local'", specifier = "==0.25.1" }, ] provides-extras = ["service", "local", "multimedia", "nemotron-parse", "tabular", "benchmarks", "llm", "dev", "all"] @@ -2787,17 +2860,18 @@ wheels = [ [[package]] name = "nltk" -version = "3.9.4" +version = "3.10.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, + { name = "defusedxml" }, { name = "joblib" }, { name = "regex" }, { name = "tqdm" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/74/a1/b3b4adf15585a5bc4c357adde150c01ebeeb642173ded4d871e89468767c/nltk-3.9.4.tar.gz", hash = "sha256:ed03bc098a40481310320808b2db712d95d13ca65b27372f8a403949c8b523d0", size = 2946864, upload-time = "2026-03-24T06:13:40.641Z" } +sdist = { url = "https://files.pythonhosted.org/packages/72/16/24d639531e73cbc6884fb251d116dfe469df9c595e0dcf24668c54d0e8d3/nltk-3.10.2.tar.gz", hash = "sha256:fcfd80fb77931868cea8357573c79838b8abc609942ef9914d1c9f6070d4645c", size = 3101716, upload-time = "2026-08-05T09:56:20.657Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/91/04e965f8e717ba0ab4bdca5c112deeab11c9e750d94c4d4602f050295d39/nltk-3.9.4-py3-none-any.whl", hash = "sha256:f2fa301c3a12718ce4a0e9305c5675299da5ad9e26068218b69d692fda84828f", size = 1552087, upload-time = "2026-03-24T06:13:38.47Z" }, + { url = "https://files.pythonhosted.org/packages/1f/2b/bf677eb32ca6684b270c0d19ab133c2271e9f3375997e5a8dd2b08e3152d/nltk-3.10.2-py3-none-any.whl", hash = "sha256:2c7ccacb765c5e26b0cb60fb1b57080af522c6924d12a714a243305ba3637412", size = 1725815, upload-time = "2026-08-05T09:56:09.657Z" }, ] [[package]] @@ -2844,6 +2918,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:ee8722c1f0145ab246bccb9e452153b5e0515fd094c3678df50b2a0888b8b171", size = 423133236, upload-time = "2025-10-09T08:59:32.536Z" }, ] +[[package]] +name = "nvidia-cuda-cccl" +version = "13.3.3.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/bd/572971ffc14bd36676c821fc15d991b08fe6179cb09368250147475f954d/nvidia_cuda_cccl-13.3.3.4.1-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:067d19b4b3c9d0f2ebec9f29a311b2863db96bf98e058bbc331597d51ce818cf", size = 3454030, upload-time = "2026-06-29T16:41:49.092Z" }, + { url = "https://files.pythonhosted.org/packages/f8/ab/049726d90147865a3ea53bae6cb7c35b98bf1fdf96cdb967101329625f83/nvidia_cuda_cccl-13.3.3.4.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc0adc188d570b09f4d606c7dc05a42aa3d8aa082e0d60f7bbfc5b6435f627c6", size = 3454034, upload-time = "2026-06-29T16:42:07.435Z" }, + { url = "https://files.pythonhosted.org/packages/24/d3/b1afcd9c40ceca72022579215fcaf5318cd747fd896cb928d4a1de924ff8/nvidia_cuda_cccl-13.3.3.4.1-py3-none-win_amd64.whl", hash = "sha256:d7c92cc03047031fa7af30866636d35ce4af409c28fc7dd8f69cb17053741399", size = 3454014, upload-time = "2026-06-29T17:09:09.012Z" }, +] + +[[package]] +name = "nvidia-cuda-crt" +version = "13.3.73" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/41/2089e411507d66458d67208bdd1bc562d492bb6458c3d2aea4603072a219/nvidia_cuda_crt-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:60aacc0b5e1e8b40c62abe4d1ab16440add91b99bd2f17f62dd091586b73d166", size = 157353, upload-time = "2026-06-29T16:42:38.163Z" }, + { url = "https://files.pythonhosted.org/packages/7e/ce/16d76f4b5b3f7460f5ebd17516685495c149c66651ffdd381f90e4d4e65c/nvidia_cuda_crt-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:df14a17ae1c5c3171265411212246654d780f89344ea85344466c6b955247543", size = 157352, upload-time = "2026-06-29T16:43:09.209Z" }, + { url = "https://files.pythonhosted.org/packages/49/b3/6791ffba6f4b8e0d3ed875285aad8078ee407afa464ecd934ae298c205b1/nvidia_cuda_crt-13.3.73-py3-none-win_amd64.whl", hash = "sha256:af04e75148db1f0eea30958f33a9ec5a5a2dc2afa99ca4323f9a93b840602ca5", size = 158286, upload-time = "2026-06-29T17:09:28.621Z" }, +] + [[package]] name = "nvidia-cuda-cupti" version = "13.0.85" @@ -2854,12 +2948,40 @@ wheels = [ ] [[package]] -name = "nvidia-cuda-nvdisasm" +name = "nvidia-cuda-nvcc" +version = "13.2.86" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine != 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "nvidia-cuda-crt" }, + { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-nvvm", version = "13.2.86", source = { registry = "https://pypi.org/simple" } }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/22/a3/403638f80960e677ae8ecc78059d8c85dc2519b85ed8c0229840dc6f3b54/nvidia_cuda_nvcc-13.2.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:909140a1f942b943982b2eff120e618c94e29d75d9e33f5cd074f0e64eb411e8", size = 38716270, upload-time = "2026-07-16T09:37:12.754Z" }, + { url = "https://files.pythonhosted.org/packages/38/7b/ad9b5f84e8820af24afa8db66b5b2785b9bdc3b8724dffcd5e6ef2d40e53/nvidia_cuda_nvcc-13.2.86-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7d3f56c8d705bad35bbba69eb0470d90856fd5d4dc48c6a6173aaf1e9f887cf5", size = 44042846, upload-time = "2026-07-16T09:37:47.98Z" }, +] + +[[package]] +name = "nvidia-cuda-nvcc" version = "13.3.73" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform == 'win32'", + "sys_platform == 'darwin'", + "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "nvidia-cuda-crt" }, + { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "nvidia-cuda-runtime", version = "13.3.29", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'win32'" }, + { name = "nvidia-nvvm", version = "13.3.73", source = { registry = "https://pypi.org/simple" } }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/92/be/e9de501cb71b10f7654381a485fa4ebf470ea25c3dce018cccaecf8a8f9a/nvidia_cuda_nvdisasm-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:dd4751884f9016b9b6dbf007abdeb5681d0a2edc731dd3d2fda9d6d878e88f73", size = 4744517, upload-time = "2026-06-29T16:48:33.527Z" }, - { url = "https://files.pythonhosted.org/packages/86/3e/88460ebd737e559e8e9843db7a63f8ced9ec7be1882344438819dd13aebc/nvidia_cuda_nvdisasm-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fa17084b07c0dca68a42892f771b4b1b40fbe9b91660209623e61cea611cae8c", size = 4782824, upload-time = "2026-06-29T16:49:05.109Z" }, + { url = "https://files.pythonhosted.org/packages/79/89/97eb797bb8bdee1d4e74069d072c24b79ae90c012fa3b539f2a7ccecf6cf/nvidia_cuda_nvcc-13.3.73-py3-none-win_amd64.whl", hash = "sha256:3d9da631bcac3dee49d1357b84cd05abe56aa3ccf76b05a7df8a80ef78addcb5", size = 32536529, upload-time = "2026-06-29T17:11:25.455Z" }, ] [[package]] @@ -2869,17 +2991,69 @@ source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, + { url = "https://files.pythonhosted.org/packages/4a/af/345fedb9f4c76c84ab4fa445b36bd4048a4d9db60e6bc76b4f913ff4b852/nvidia_cuda_nvrtc-13.0.88-py3-none-win_amd64.whl", hash = "sha256:6bcd4e7f8e205cbe644f5a98f2f799bef9556fefc89dd786e79a16312ce49872", size = 76807835, upload-time = "2025-09-04T08:39:15.274Z" }, ] [[package]] name = "nvidia-cuda-runtime" version = "13.0.96" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine != 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "sys_platform == 'darwin'", + "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] wheels = [ { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, ] +[[package]] +name = "nvidia-cuda-runtime" +version = "13.3.29" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform == 'win32'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/27/b53a5e0397842a5c11f0e1a39d4e5b2f22638a4126e83b3c4e196f62c969/nvidia_cuda_runtime-13.3.29-py3-none-win_amd64.whl", hash = "sha256:0667ec61c3d897388efa305ed4f7609ace88849a753ba9c6311d06dca55fff4f", size = 2630354, upload-time = "2026-05-26T17:00:05.389Z" }, +] + +[[package]] +name = "nvidia-cuda-tileiras" +version = "13.2.86" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine != 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "nvidia-cuda-nvcc", version = "13.2.86", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-nvvm", version = "13.2.86", source = { registry = "https://pypi.org/simple" } }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/a2/ade26f7fb55a5bb87815f7650660463d6601db411d5a9228f414b3ed5dfe/nvidia_cuda_tileiras-13.2.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5f79c6a9bf8583dae105cd67b985555f39f4576416664a86383ba892e8c346c9", size = 36418795, upload-time = "2026-07-16T09:43:35.006Z" }, + { url = "https://files.pythonhosted.org/packages/a4/ec/e01bcfe4f48fabd8fd1af139061a74f92f3b71fe3cac8a64e943157c2585/nvidia_cuda_tileiras-13.2.86-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:832f360aa8ce478ff878c4e6630cd767349a597ea54382e7c07a835b0daead96", size = 36970477, upload-time = "2026-07-16T09:44:08.421Z" }, +] + +[[package]] +name = "nvidia-cuda-tileiras" +version = "13.3.36" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform == 'win32'", +] +dependencies = [ + { name = "nvidia-cuda-nvcc", version = "13.3.73", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-nvjitlink", version = "13.3.33", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-nvvm", version = "13.3.73", source = { registry = "https://pypi.org/simple" } }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/b1/9dcc1aa140205c8f9ecdb671b488e1189462d870e5ed0eac02522722cb9a/nvidia_cuda_tileiras-13.3.36-py3-none-win_amd64.whl", hash = "sha256:0dd286086c6d273826218d02c076ddea8e285b50ee223bd1429756b706b7bbc7", size = 29715011, upload-time = "2026-05-26T17:05:29.974Z" }, +] + [[package]] name = "nvidia-cudnn-cu13" version = "9.19.0.56" @@ -2894,11 +3068,12 @@ wheels = [ [[package]] name = "nvidia-cudnn-frontend" -version = "1.18.0" +version = "1.27.0" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e3/b4/604e230378680ee117849a4e1045baca092f93161a829291a84d5acce70c/nvidia_cudnn_frontend-1.18.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:310b417f2848a83d1437203fcaeea320a74fb7f28af20bf42bf5afc9c01f1c12", size = 2027408, upload-time = "2026-01-27T23:32:46.576Z" }, - { url = "https://files.pythonhosted.org/packages/c6/52/08f98262e77b1cbcc834cc1a5db494d0661ea1dbdea58c2e2d51a57fdaca/nvidia_cudnn_frontend-1.18.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6c023539ca6de99234cf5102c3ec0d6af817f5396fc93028a22ba5b834a35b8a", size = 2159245, upload-time = "2026-01-27T23:07:32.664Z" }, + { url = "https://files.pythonhosted.org/packages/df/cd/d6b6910b79389955d9c33596c03380db63d76f1bcc6cdd24efc3ced68a3b/nvidia_cudnn_frontend-1.27.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:649b9f5a20bded17bc917122bcabd83826b82cb9bbd5b74573b769a9f4930798", size = 4589494, upload-time = "2026-08-06T22:42:10.993Z" }, + { url = "https://files.pythonhosted.org/packages/cb/b5/3a998fa2ba1aa527b35136d2c675ee3f8394c6a7f30605c63c3d9b64023b/nvidia_cudnn_frontend-1.27.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ef1f1b4927f2f9e76a5ba83ac4bdc01a6a84750032429ea9c132599ec60c8947", size = 4749917, upload-time = "2026-08-06T22:42:36.857Z" }, + { url = "https://files.pythonhosted.org/packages/de/c9/934518aa93cd19fe2dd19c9fd7572a69adb4dbf8160edcd57301c22e3ea3/nvidia_cudnn_frontend-1.27.0-cp312-cp312-win_amd64.whl", hash = "sha256:bd917aa2f83af6f72ccbe3e999a7175c467e1004aeb45435fac67e45be3b5a5d", size = 4120145, upload-time = "2026-08-06T22:42:57.674Z" }, ] [[package]] @@ -2906,7 +3081,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" } }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -2938,7 +3113,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas" }, { name = "nvidia-cusparse" }, - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" } }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -2950,7 +3125,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" } }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -2968,63 +3143,47 @@ wheels = [ [[package]] name = "nvidia-cutlass-dsl" -version = "4.6.0" +version = "4.5.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cutlass-dsl-libs-base" }, - { name = "nvidia-cutlass-dsl-libs-cu12" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/8b/1c/fbddb760a0228df87a9e9d1e60b76ecbe6e18035f5853efe0b4563651b2b/nvidia_cutlass_dsl-4.6.0-py3-none-any.whl", hash = "sha256:e3e0e4d8df20d82c8401fa013f4d82021f41daa5fca3d24b55d4a677f2308ca8", size = 10459, upload-time = "2026-07-02T03:23:18.43Z" }, + { url = "https://files.pythonhosted.org/packages/f0/15/575d7df4fe2f3406f1cfc68be72aeff2834f8a696daf1cd5bee8017e4507/nvidia_cutlass_dsl-4.5.2-py3-none-any.whl", hash = "sha256:68ed1b63ca74aae87955012da9dfd7fdaae471329d0028b229b841c7192ccf52", size = 10179, upload-time = "2026-05-25T03:38:56.364Z" }, ] -[[package]] -name = "nvidia-cutlass-dsl-libs-base" -version = "4.6.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cuda-python" }, - { name = "numpy" }, - { name = "nvidia-cuda-nvdisasm" }, - { name = "nvidia-cutlass-dsl-libs-core" }, - { name = "protobuf" }, - { name = "typing-extensions" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/2c/f8/22653971fcab2a7ed581934f7a2708c9873fa6a8e8eb285422c8eed4ae01/nvidia_cutlass_dsl_libs_base-4.6.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:6412572899b1c6d182e516b20f2b0a21874ec88d25234e6040fb2a4381de7a1a", size = 3321728, upload-time = "2026-07-02T03:25:28.888Z" }, - { url = "https://files.pythonhosted.org/packages/ce/38/e91f66739d2f8711d1a2457e68cd86d6fbae307ce66ce270a405d4dc6dc7/nvidia_cutlass_dsl_libs_base-4.6.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:e41cd5db4de4b535c30ae9ca4412b957800a62560019ae91fa51cf3ea89bf254", size = 2824817, upload-time = "2026-07-02T03:25:53.814Z" }, +[package.optional-dependencies] +cu13 = [ + { name = "nvidia-cutlass-dsl-libs-cu13" }, ] [[package]] -name = "nvidia-cutlass-dsl-libs-core" -version = "4.6.0" +name = "nvidia-cutlass-dsl-libs-base" +version = "4.5.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cuda-python" }, { name = "numpy" }, - { name = "nvidia-cuda-nvdisasm" }, - { name = "protobuf" }, { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/84/94/e4e2404ac06a477096ccf8127bf5d391510d36cafb4be86c8c15b4873b0d/nvidia_cutlass_dsl_libs_core-4.6.0-py3-none-any.whl", hash = "sha256:f9ea6d313a03cb11fa177da32e8747ad0cac51358850810f36aa6c4736192c27", size = 767713, upload-time = "2026-07-02T03:23:39.876Z" }, + { url = "https://files.pythonhosted.org/packages/b1/ef/e827e3c67d72adbf4e8f680bdf03b1b67723d9e1ae7c3d0a1751f39f69ce/nvidia_cutlass_dsl_libs_base-4.5.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:d2a3c412287e356fbe48fe9f845d6d33cd35dea5e20d7e4f628c20957967cacd", size = 75643473, upload-time = "2026-05-25T03:49:15.857Z" }, + { url = "https://files.pythonhosted.org/packages/97/68/c1247ab848f26c4ab56e562eea0e3f31fc14c9aaf0d883afaa92d8f05592/nvidia_cutlass_dsl_libs_base-4.5.2-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:15ef6a59193667e663934ef4873f8ccad37455e9b7c3c419c3072113b8aedf61", size = 74513226, upload-time = "2026-05-25T03:51:32.496Z" }, ] [[package]] -name = "nvidia-cutlass-dsl-libs-cu12" -version = "4.6.0" +name = "nvidia-cutlass-dsl-libs-cu13" +version = "4.5.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cuda-python" }, { name = "numpy" }, - { name = "nvidia-cuda-nvdisasm" }, { name = "nvidia-cutlass-dsl-libs-base" }, - { name = "protobuf" }, { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/11/38/62def848b65bf067f434df7680c7e8c48519b25bbd3f03f9cdff3606353b/nvidia_cutlass_dsl_libs_cu12-4.6.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:87f132ccc30946949868989f3b1b1adaa714ccdf5c636e5379b54909cc29576c", size = 86992102, upload-time = "2026-07-02T03:29:41.47Z" }, - { url = "https://files.pythonhosted.org/packages/bf/64/f3f8962a9b91dd9368b90e23b2ac81614d6e9df72b55365ec0c216c3f8f9/nvidia_cutlass_dsl_libs_cu12-4.6.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:abc341ff0fce40ed0bdadf160f6afac07fb9d01768d4daebd1628c330b3e4210", size = 88436835, upload-time = "2026-07-02T03:30:13.676Z" }, + { url = "https://files.pythonhosted.org/packages/21/e5/aeb570713a7bd6c2cb08102c2ebe6de234ef1bbc276d1af4643266cd71a8/nvidia_cutlass_dsl_libs_cu13-4.5.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:3032405dff28892340f96b467e744a822079cae454dce534fc17b77e85190e42", size = 79084280, upload-time = "2026-05-25T03:40:57.547Z" }, + { url = "https://files.pythonhosted.org/packages/03/60/443e559139da15ab544761ac14f4206dffb981af48cc9856cd5b5b7cf0e7/nvidia_cutlass_dsl_libs_cu13-4.5.2-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:80f0cd402e0f1d1571e5aed33bfa17dbc9cb90cc5b1352f0f806b4788558e80e", size = 78759198, upload-time = "2026-05-25T03:45:59.297Z" }, ] [[package]] @@ -3049,11 +3208,26 @@ wheels = [ name = "nvidia-nvjitlink" version = "13.0.88" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine != 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] wheels = [ { url = "https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b", size = 40713933, upload-time = "2025-09-04T08:35:43.553Z" }, { url = "https://files.pythonhosted.org/packages/ab/2c/93c5250e64df4f894f1cbb397c6fd71f79813f9fd79d7cd61de3f97b3c2d/nvidia_nvjitlink-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e931536ccc7d467a98ba1d8b89ff7fa7f1fa3b13f2b0069118cd7f47bff07d0c", size = 38768748, upload-time = "2025-09-04T08:35:20.008Z" }, ] +[[package]] +name = "nvidia-nvjitlink" +version = "13.3.33" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform == 'win32'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/f2/ec9c05a108095828dfc58840978c627b3c313fdf2a567c6de9ffbbb46901/nvidia_nvjitlink-13.3.33-py3-none-win_amd64.whl", hash = "sha256:4297ee49639b4f2e07255a1d69b3acc7ab2d011bb892b403e91ac98368962e3b", size = 37766359, upload-time = "2026-05-26T17:11:28.96Z" }, +] + [[package]] name = "nvidia-nvshmem-cu13" version = "3.4.5" @@ -3072,6 +3246,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, ] +[[package]] +name = "nvidia-nvvm" +version = "13.2.86" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine != 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/1d/b50325516e259d4de9fca78b069840d39b6a172c5ba012b88d9985e01fa3/nvidia_nvvm-13.2.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a9244f3209922d655612c11ebb4117ea1f80983cb2b215c9cccded127282cebc", size = 64280112, upload-time = "2026-07-16T09:58:27.625Z" }, + { url = "https://files.pythonhosted.org/packages/e6/64/039ad70d68355634581d3b66f99b7aa8f75ca6078e1a6a2c9223677bbee3/nvidia_nvvm-13.2.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a502dcc2859f17a925adba495c222d44a221b9eb10e7d111a7046dc2cc883b69", size = 61886756, upload-time = "2026-07-16T09:57:51.818Z" }, +] + +[[package]] +name = "nvidia-nvvm" +version = "13.3.73" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform == 'win32'", + "sys_platform == 'darwin'", + "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ad/6b/d5756f485012b920475cbc01457c1b9a7d0485bfb04b92598c5e1ef3e9ab/nvidia_nvvm-13.3.73-py3-none-win_amd64.whl", hash = "sha256:b5c91dfa59ee4cee90b2dfb19c6203f31c914b9c9b5ca10726c2da7cf8ed401d", size = 59981103, upload-time = "2026-06-29T17:21:43.334Z" }, +] + [[package]] name = "nvidia-riva-client" version = "2.26.0" @@ -3086,6 +3286,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/21/53/0988b79ee818cb2b8792abeda5a75156946b64da7a2ef2719b9d2068a34f/nvidia_riva_client-2.26.0-py3-none-any.whl", hash = "sha256:16ffc98266fa7be7261e0675de6b7028e7f973c2ac3dfd679668148ff497cc0c", size = 57539, upload-time = "2026-05-28T06:11:43.205Z" }, ] +[[package]] +name = "nvtx" +version = "0.2.15" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/dd/692765e87de30bae1522cdffaa0f2b52949658a92a0fa6d96b1a01eae9d2/nvtx-0.2.15.tar.gz", hash = "sha256:2287d3be05b85661deb386f878d1f536c2e532774aa9ec7a50c434942ed81ae5", size = 121230, upload-time = "2026-03-18T10:01:25.547Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/07/698355285a03a366ef63ea9762fc1feef3f9f25483e1655408f72d827090/nvtx-0.2.15-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2cc530cd0f1a2c14a3a7e683833db509888ac5ed4ead94e5c9e2c7317c6937a7", size = 807159, upload-time = "2026-03-18T10:09:49.232Z" }, + { url = "https://files.pythonhosted.org/packages/c0/d1/08f22448d83481408d663065764ba583df091a7de629ed38fc97e522f1af/nvtx-0.2.15-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3ca8030a6d197952318013dd1c12c22da1d4b9feb76ba72e0fcd449961183c2c", size = 806187, upload-time = "2026-03-18T10:13:32.972Z" }, + { url = "https://files.pythonhosted.org/packages/54/23/c97c39e3b7ba256aa343cb828ca0d1c8421f705ca84795658ecd14ca95ed/nvtx-0.2.15-cp312-cp312-win_amd64.whl", hash = "sha256:70a1e768964e0520b68ccabc4df391cc227537c45936a7eba6507bc65e617e00", size = 129178, upload-time = "2026-03-18T10:02:55.299Z" }, +] + [[package]] name = "onnx" version = "1.22.0" @@ -3198,6 +3409,7 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/3e/92/2d038d096f29179c7c9571b431f9e739f87a487121901725e23fe338dd9d/openai_harmony-0.0.8.tar.gz", hash = "sha256:6e43f98e6c242fa2de6f8ea12eab24af63fa2ed3e89c06341fb9d92632c5cbdf", size = 284777, upload-time = "2025-11-05T19:07:06.727Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/45/c6/2502f416d46be3ec08bb66d696cccffb57781a499e3ff2e4d7c174af4e8f/openai_harmony-0.0.8-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:029ec25ca74abe48fdb58eb9fdd2a8c1618581fc33ce8e5653f8a1ffbfbd9326", size = 2627806, upload-time = "2025-11-05T19:06:57.063Z" }, { url = "https://files.pythonhosted.org/packages/d3/d2/ce6953ca87db9cae3e775024184da7d1c5cb88cead19a2d75b42f00a959c/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e4f709815924ec325b9a890e6ab2bbb0ceec8e319a4e257328eb752cf36b2efc", size = 2948463, upload-time = "2025-11-05T19:06:48.17Z" }, { url = "https://files.pythonhosted.org/packages/fa/4c/b553c9651662d6ce102ca7f3629d268b23df1abe5841e24bed81e8a8e949/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5cfcfd963b50a41fc656c84d3440ca6eecdccd6c552158ce790b8f2e33dfb5a9", size = 2704083, upload-time = "2025-11-05T19:06:50.205Z" }, { url = "https://files.pythonhosted.org/packages/9b/af/4eec8f9ab9c27bcdb444460c72cf43011d176fc44c79d6e113094ca1e152/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0a3a16972aa1cee38ea958470cd04ac9a2d5ac38fdcf77ab686611246220c158", size = 2959765, upload-time = "2025-11-05T19:06:53.62Z" }, @@ -3207,6 +3419,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1d/10/4327dbf87f75ae813405fd9a9b4a5cde63d506ffed0a096a440a4cabd89c/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:cbaa3bda75ef0d8836e1f8cc84af62f971b1d756d740efc95c38c3e04c0bfde2", size = 2932931, upload-time = "2025-11-05T19:07:01.437Z" }, { url = "https://files.pythonhosted.org/packages/8a/c8/1774eec4f6f360ef57618fb8f52e3d3af245b2491bd0297513aa09eec04b/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:772922a9bd24e133950fad71eb1550836f415a88e8c77870e12d0c3bd688ddc2", size = 2996140, upload-time = "2025-11-05T19:07:03.438Z" }, { url = "https://files.pythonhosted.org/packages/60/c3/3d1e01e2dba517a91760e4a03e4f20ffc75039a6fe584d0e6f9b5c78fd15/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:007b0476a1f331f8130783f901f1da6f5a7057af1a4891f1b6a31dec364189b5", size = 3205080, upload-time = "2025-11-05T19:07:05.078Z" }, + { url = "https://files.pythonhosted.org/packages/14/63/119de431572d7c70a7bf1037034a9be6ed0a7502a7498ba7302bca5b3242/openai_harmony-0.0.8-cp38-abi3-win32.whl", hash = "sha256:a9b5f893326b28d9e935ade14b4f655f5a840942473bc89b201c25f7a15af9cf", size = 2082457, upload-time = "2025-11-05T19:07:09.631Z" }, + { url = "https://files.pythonhosted.org/packages/40/1f/c83cf5a206c263ee70448a5ae4264682555f4d0b5bed0d2cc6ca1108103d/openai_harmony-0.0.8-cp38-abi3-win_amd64.whl", hash = "sha256:39d44f0d8f466bd56698e7ead708bead3141e27b9b87e3ab7d5a6d0e4a869ee5", size = 2438369, upload-time = "2025-11-05T19:07:08.1Z" }, ] [[package]] @@ -3221,29 +3435,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/12/cf/03675d8bd8ecbf4445504d8071adab19f5f993676795708e36402ab38263/openapi_pydantic-0.5.1-py3-none-any.whl", hash = "sha256:a3a09ef4586f5bd760a8df7f43028b60cafb6d9f61de2acba9574766255ab146", size = 96381, upload-time = "2025-01-08T19:29:25.275Z" }, ] -[[package]] -name = "opencensus" -version = "0.11.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-api-core" }, - { name = "opencensus-context" }, - { name = "six" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/15/a7/a46dcffa1b63084f9f17fe3c8cb20724c4c8f91009fd0b2cfdb27d5d2b35/opencensus-0.11.4.tar.gz", hash = "sha256:cbef87d8b8773064ab60e5c2a1ced58bbaa38a6d052c41aec224958ce544eff2", size = 64966, upload-time = "2024-01-03T18:04:07.085Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b5/ed/9fbdeb23a09e430d87b7d72d430484b88184633dc50f6bfb792354b6f661/opencensus-0.11.4-py2.py3-none-any.whl", hash = "sha256:a18487ce68bc19900336e0ff4655c5a116daf10c1b3685ece8d971bddad6a864", size = 128225, upload-time = "2024-01-03T18:04:05.127Z" }, -] - -[[package]] -name = "opencensus-context" -version = "0.1.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4c/96/3b6f638f6275a8abbd45e582448723bffa29c1fb426721dedb5c72f7d056/opencensus-context-0.1.3.tar.gz", hash = "sha256:a03108c3c10d8c80bb5ddf5c8a1f033161fa61972a9917f9b9b3a18517f0088c", size = 4066, upload-time = "2022-08-03T22:20:22.359Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/10/68/162c97ea78c957d68ecf78a5c5041d2e25bd5562bdf5d89a6cbf7f8429bf/opencensus_context-0.1.3-py2.py3-none-any.whl", hash = "sha256:073bb0590007af276853009fac7e4bab1d523c3f03baf4cb4511ca38967c6039", size = 5060, upload-time = "2022-08-03T22:20:20.352Z" }, -] - [[package]] name = "opencv-python" version = "5.0.0.93" @@ -3345,20 +3536,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cd/d0/fdeb1a98d8d3a6205f5f297c51b4a9bfe65126ab60339669bbe3dd54c2e2/opentelemetry_exporter_otlp_proto_http-1.44.0-py3-none-any.whl", hash = "sha256:838592fce774c1c8bb7b9a0a7facbfa82e17be5a8a4e94cef10cb84ae026bae3", size = 21850, upload-time = "2026-07-16T15:25:20.006Z" }, ] -[[package]] -name = "opentelemetry-exporter-prometheus" -version = "0.65b0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-sdk" }, - { name = "prometheus-client" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f4/85/3a1af7b90a76d6c069bcbbc86cc641c93d15057f2674cc8226f47c5260e8/opentelemetry_exporter_prometheus-0.65b0.tar.gz", hash = "sha256:2777cbf41c403c119e10f418fce5d645c956b47f673a2ee120285d1d0c6df2d5", size = 16411, upload-time = "2026-07-16T15:25:39.971Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/82/0d/d7cdc030edeed0fea03448446b88f1a1f8c4a565bad30dac0f5477ebe290/opentelemetry_exporter_prometheus-0.65b0-py3-none-any.whl", hash = "sha256:3b3d24b586d0ad9712c7b52b7d19c8a9dfbb318b9b284121b5f95e90ed019367", size = 13031, upload-time = "2026-07-16T15:25:20.906Z" }, -] - [[package]] name = "opentelemetry-proto" version = "1.44.0" @@ -3457,8 +3634,14 @@ version = "0.2.14" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/6a/04/4a0812eb27c086cfd2e66e7ec9150f33e105912a9b7f8b335e3479f03a06/outlines_core-0.2.14.tar.gz", hash = "sha256:64808deed1591ca3029ff64346ceb974cd5d780c916ea82504951fe83523039e", size = 191539, upload-time = "2026-01-09T15:59:10.016Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/66/93/30b9188648a479b32be429a24166db47a7bfdb0f9a8aac4c6dcf569e0a52/outlines_core-0.2.14-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:95e6476d9702d2fcc4e85370dbbfb6933a46c816e9c90107f6ce36eb68b5d64a", size = 2049651, upload-time = "2026-01-09T15:58:28.549Z" }, + { url = "https://files.pythonhosted.org/packages/0d/06/f3557daa8e87d5b95f64de269a301d73ec3c2202ab897c3e1f1cb93eb1db/outlines_core-0.2.14-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:f04731a5e29a190e2cc9f692a1f3fb2414a645355ca7d01b83df43439c38bea8", size = 2201046, upload-time = "2026-01-09T15:58:29.958Z" }, + { url = "https://files.pythonhosted.org/packages/0c/67/d8acf778990964c951080d568284e858d466f27dfd6f2674781927faba1c/outlines_core-0.2.14-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:0e4c69f0a8565edb56464c4c9b6c291a10805f3a96dff84182980e90ae1a5e2f", size = 2049558, upload-time = "2026-01-09T15:58:31.003Z" }, + { url = "https://files.pythonhosted.org/packages/17/e1/0320b14b49b8379ced1ab195ecf5875dbd2267b90148847541f43bfde6c1/outlines_core-0.2.14-cp312-cp312-macosx_15_0_x86_64.whl", hash = "sha256:63f53cfd9614e754499ae86dd699f3abcecf42d6a4e58d80fd80347881d85960", size = 2197854, upload-time = "2026-01-09T15:58:32.39Z" }, { url = "https://files.pythonhosted.org/packages/29/29/3a04944407207a5d214879ca5ca33c2bd3e65199a4e927051c1bdaaa4d50/outlines_core-0.2.14-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3bb2060c240c4507f334965a8948dbeeb22007560d797f6debd92346c0b620cb", size = 2341426, upload-time = "2026-01-09T15:58:33.553Z" }, { url = "https://files.pythonhosted.org/packages/b2/a7/a77f746272504bac3f628047d56ea1731b61549a3e1d9bbfd226f2968246/outlines_core-0.2.14-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:1de34681c7e0e7e1551fc9036e4fa3c57986336c905a10536591ceb6d869c258", size = 2236941, upload-time = "2026-01-09T15:58:35.118Z" }, + { url = "https://files.pythonhosted.org/packages/99/0d/9f599d938923ab8ceeff26fdf2f9ea53bea3c962085c4927a08338a32349/outlines_core-0.2.14-cp312-cp312-win32.whl", hash = "sha256:870e8e038853818cb202ccc8cde92251f300f96805bfcc3be1c883adda7b5297", size = 1842940, upload-time = "2026-01-09T15:58:36.544Z" }, + { url = "https://files.pythonhosted.org/packages/f8/df/0f145c52ebd156d80273e2f5278227ea57e0275b2aa863bed33f44f77923/outlines_core-0.2.14-cp312-cp312-win_amd64.whl", hash = "sha256:87b42440478764cce1353a87d8560ef82f3b39b9d753bfe93195ea3584f369e3", size = 2137266, upload-time = "2026-01-09T15:58:37.831Z" }, ] [[package]] @@ -3520,21 +3703,19 @@ wheels = [ [[package]] name = "pillow" -version = "12.2.0" +version = "12.3.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/8c/21/c2bcdd5906101a30244eaffc1b6e6ce71a31bd0742a01eb89e660ebfac2d/pillow-12.2.0.tar.gz", hash = "sha256:a830b1a40919539d07806aa58e1b114df53ddd43213d9c8b75847eee6c0182b5", size = 46987819, upload-time = "2026-04-01T14:46:17.687Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/be/7482c8a5ebebbc6470b3eb791812fff7d5e0216c2be3827b30b8bb6603ed/pillow-12.2.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2d192a155bbcec180f8564f693e6fd9bccff5a7af9b32e2e4bf8c9c69dbad6b5", size = 5308279, upload-time = "2026-04-01T14:43:13.246Z" }, - { url = "https://files.pythonhosted.org/packages/d8/95/0a351b9289c2b5cbde0bacd4a83ebc44023e835490a727b2a3bd60ddc0f4/pillow-12.2.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f3f40b3c5a968281fd507d519e444c35f0ff171237f4fdde090dd60699458421", size = 4695490, upload-time = "2026-04-01T14:43:15.584Z" }, - { url = "https://files.pythonhosted.org/packages/de/af/4e8e6869cbed569d43c416fad3dc4ecb944cb5d9492defaed89ddd6fe871/pillow-12.2.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:03e7e372d5240cc23e9f07deca4d775c0817bffc641b01e9c3af208dbd300987", size = 6284462, upload-time = "2026-04-01T14:43:18.268Z" }, - { url = "https://files.pythonhosted.org/packages/e9/9e/c05e19657fd57841e476be1ab46c4d501bffbadbafdc31a6d665f8b737b6/pillow-12.2.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b86024e52a1b269467a802258c25521e6d742349d760728092e1bc2d135b4d76", size = 8094744, upload-time = "2026-04-01T14:43:20.716Z" }, - { url = "https://files.pythonhosted.org/packages/2b/54/1789c455ed10176066b6e7e6da1b01e50e36f94ba584dc68d9eebfe9156d/pillow-12.2.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7371b48c4fa448d20d2714c9a1f775a81155050d383333e0a6c15b1123dda005", size = 6398371, upload-time = "2026-04-01T14:43:23.443Z" }, - { url = "https://files.pythonhosted.org/packages/43/e3/fdc657359e919462369869f1c9f0e973f353f9a9ee295a39b1fea8ee1a77/pillow-12.2.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:62f5409336adb0663b7caa0da5c7d9e7bdbaae9ce761d34669420c2a801b2780", size = 7087215, upload-time = "2026-04-01T14:43:26.758Z" }, - { url = "https://files.pythonhosted.org/packages/8b/f8/2f6825e441d5b1959d2ca5adec984210f1ec086435b0ed5f52c19b3b8a6e/pillow-12.2.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:01afa7cf67f74f09523699b4e88c73fb55c13346d212a59a2db1f86b0a63e8c5", size = 6509783, upload-time = "2026-04-01T14:43:29.56Z" }, - { url = "https://files.pythonhosted.org/packages/67/f9/029a27095ad20f854f9dba026b3ea6428548316e057e6fc3545409e86651/pillow-12.2.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fc3d34d4a8fbec3e88a79b92e5465e0f9b842b628675850d860b8bd300b159f5", size = 7212112, upload-time = "2026-04-01T14:43:32.091Z" }, - { url = "https://files.pythonhosted.org/packages/be/42/025cfe05d1be22dbfdb4f264fe9de1ccda83f66e4fc3aac94748e784af04/pillow-12.2.0-cp312-cp312-win32.whl", hash = "sha256:58f62cc0f00fd29e64b29f4fd923ffdb3859c9f9e6105bfc37ba1d08994e8940", size = 6378489, upload-time = "2026-04-01T14:43:34.601Z" }, - { url = "https://files.pythonhosted.org/packages/5d/7b/25a221d2c761c6a8ae21bfa3874988ff2583e19cf8a27bf2fee358df7942/pillow-12.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:7f84204dee22a783350679a0333981df803dac21a0190d706a50475e361c93f5", size = 7084129, upload-time = "2026-04-01T14:43:37.213Z" }, - { url = "https://files.pythonhosted.org/packages/10/e1/542a474affab20fd4a0f1836cb234e8493519da6b76899e30bcc5d990b8b/pillow-12.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:af73337013e0b3b46f175e79492d96845b16126ddf79c438d7ea7ff27783a414", size = 2463612, upload-time = "2026-04-01T14:43:39.421Z" }, + { url = "https://files.pythonhosted.org/packages/37/bf/fb3ebff8ddcb76aac5a01389251bbbb9519922a9b520d8247c1ca864a25d/pillow-12.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ba09209fbe443b4acccebe845d8a138b89a8f4fbaeedd44953490b5315d5e965", size = 5345969, upload-time = "2026-07-01T11:54:06.397Z" }, + { url = "https://files.pythonhosted.org/packages/d8/66/9a386a92561f402389a4fc70c18838bf6d35eb5eb5c6850b4b2dc64f5048/pillow-12.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ffd0c5368496f41b0944be820fcb7a838aa6e623d250b01acf2643939c3f99d7", size = 4780323, upload-time = "2026-07-01T11:54:09.351Z" }, + { url = "https://files.pythonhosted.org/packages/25/27/ac8f99618ffd3dde21db0f4d4b1d2ab00c0880595bfd17df103f7f39fd0c/pillow-12.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9c7f76c0673154f044e9d78c8655fb4213f6ca31a836df48b40fe5d187717b9", size = 6266838, upload-time = "2026-07-01T11:54:11.71Z" }, + { url = "https://files.pythonhosted.org/packages/84/21/a35af28dcc61f37ed850a2d64c65c701321dfbf25085e469d5559360cbbf/pillow-12.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91", size = 6940830, upload-time = "2026-07-01T11:54:13.732Z" }, + { url = "https://files.pythonhosted.org/packages/eb/51/8b08617af3ad95e33ce6d7dd2c99ed6c8298f7fb131636303956be022e25/pillow-12.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e491916b378fba47242221bb9ead245211b70d504f495d105d17b14a24b4907c", size = 6344383, upload-time = "2026-07-01T11:54:15.756Z" }, + { url = "https://files.pythonhosted.org/packages/1d/72/cf78ac9780bb93c28328f408973845a309d4d145041665f734572ced1b52/pillow-12.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df", size = 7052934, upload-time = "2026-07-01T11:54:17.721Z" }, + { url = "https://files.pythonhosted.org/packages/20/20/25e0f4dc178a6bc0696793720055519a0de89e7661dae886992decbd2f81/pillow-12.3.0-cp312-cp312-win32.whl", hash = "sha256:dbce0b29841537a2fa4a214c2bbf14de3587c9680caa9b4e217568472490b28f", size = 6472684, upload-time = "2026-07-01T11:54:19.839Z" }, + { url = "https://files.pythonhosted.org/packages/45/89/da2f7971a317f83d807fdd4065c0af40208e59e692cc43d315a71a0e96d1/pillow-12.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:a2b55dd6b2a4c4b7d87ffa56bdb33fdc5fdb9a462173861a7bc097f17d91cb09", size = 7227137, upload-time = "2026-07-01T11:54:22.025Z" }, + { url = "https://files.pythonhosted.org/packages/de/47/4845a0a6c0dbf1db8456bd9fc791f13c5ced7ced20606d08a0aacfd25b49/pillow-12.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:331b624368d4f1d069149002f25f44bc61c8919ce8ddb3c45bdad8f6e2d89510", size = 2568267, upload-time = "2026-07-01T11:54:24.051Z" }, ] [[package]] @@ -3580,15 +3761,15 @@ wheels = [ [[package]] name = "prometheus-fastapi-instrumentator" -version = "7.1.0" +version = "8.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "prometheus-client" }, { name = "starlette" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/69/6d/24d53033cf93826aa7857699a4450c1c67e5b9c710e925b1ed2b320c04df/prometheus_fastapi_instrumentator-7.1.0.tar.gz", hash = "sha256:be7cd61eeea4e5912aeccb4261c6631b3f227d8924542d79eaf5af3f439cbe5e", size = 20220, upload-time = "2025-03-19T19:35:05.351Z" } +sdist = { url = "https://files.pythonhosted.org/packages/95/f4/cdcebf7094b03b99fba71ac8f56bd6f227973642662f49d272332d8419b3/prometheus_fastapi_instrumentator-8.1.0.tar.gz", hash = "sha256:b77f3043665e8d28e2bbd21017506195a43d9adf1d402d01bf95b494b7e560e1", size = 20492, upload-time = "2026-07-26T11:12:44.202Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/27/72/0824c18f3bc75810f55dacc2dd933f6ec829771180245ae3cc976195dec0/prometheus_fastapi_instrumentator-7.1.0-py3-none-any.whl", hash = "sha256:978130f3c0bb7b8ebcc90d35516a6fe13e02d2eb358c8f83887cdef7020c31e9", size = 19296, upload-time = "2025-03-19T19:35:04.323Z" }, + { url = "https://files.pythonhosted.org/packages/44/b9/91a2246e6cf01b7ccb14479803c8a50f9c258ae5c6a0f16ba3294820632b/prometheus_fastapi_instrumentator-8.1.0-py3-none-any.whl", hash = "sha256:b9f40b2cff3f7891ca0610b3ae4fc6ec723fd326b04bb659819aaeb821a0fc7d", size = 19649, upload-time = "2026-07-26T11:12:45.168Z" }, ] [[package]] @@ -3617,18 +3798,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/ed/1cdcab6ba3d6ab7feca11fc14f0eeea80755bb53ef4e892079f31b10a25f/propcache-0.5.2-py3-none-any.whl", hash = "sha256:be1ddfcbb376e3de5d2e2db1d58d6d67463e6b4f9f040c000de8e300295465fe", size = 14036, upload-time = "2026-05-08T21:02:10.673Z" }, ] -[[package]] -name = "proto-plus" -version = "1.28.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "protobuf" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/87/44/767757fd2cdd4a60d7e4440d9f7b491d6131103d313638d2c03e06c268fb/proto_plus-1.28.1.tar.gz", hash = "sha256:832e68e7fe064cf90ab153b6e5eb935b27891bb89aaeb68b115e9b702f6cb168", size = 57166, upload-time = "2026-07-08T17:04:02.367Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6a/34/2f2b57dbfd145b995a29847a16b0903fce5ef6ad3c7aad740a609c5d3678/proto_plus-1.28.1-py3-none-any.whl", hash = "sha256:6660f5f1970874bdcfc3088b435188a36a37bd3596668f7d726417c4ae8cfbed", size = 50408, upload-time = "2026-07-08T17:03:34.532Z" }, -] - [[package]] name = "protobuf" version = "6.33.5" @@ -3694,21 +3863,6 @@ memory = [ { name = "cachetools" }, ] -[[package]] -name = "py-spy" -version = "0.4.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/93/d8/5b71371f50cf153b1307e5a11ac8a4ce4d85651dae946bd7e9a064146545/py_spy-0.4.2.tar.gz", hash = "sha256:90e600b27bb6bb40479637baca5a5b4bc2ba3395c93d889e672315d93042c4ae", size = 286374, upload-time = "2026-04-24T22:08:54.906Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ef/21/ec030145a0c7992bd4b9eafb2f06f56358b3a5339eab4a16534baf3c69aa/py_spy-0.4.2-py2.py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:1ccf688393105111684435f035bc14ec3f22117dd2b85b2414612cf27a22755a", size = 3743992, upload-time = "2026-04-24T22:08:45.438Z" }, - { url = "https://files.pythonhosted.org/packages/50/80/de5fd27243c2be03692ecd317bf0dbe24b4c6f78f689ce111e7277a7cb09/py_spy-0.4.2-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:a0e6f6810ccf0fc5e64e85e0182a5b626c4496eec01b14fb8755154b363a4831", size = 1859057, upload-time = "2026-04-24T22:08:46.946Z" }, - { url = "https://files.pythonhosted.org/packages/89/23/3eb4c23c684ebd667674ce1d076ae855e0621d1d9bd5e052aa3f7982f757/py_spy-0.4.2-py2.py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:142887e984a4e541071c99a4401ff8c3770f255d329dbd0f64e8c1dd51882cce", size = 2828136, upload-time = "2026-04-24T22:08:48.519Z" }, - { url = "https://files.pythonhosted.org/packages/ca/01/6314152cf9ad3310ebacbf2c47b5ed858086530f8e12b1a665725ca5e0f4/py_spy-0.4.2-py2.py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7f1c6d9b0e2379ead5bf792df43f4cf36153aa79e6dda4fb8ac7740cf8017110", size = 2857707, upload-time = "2026-04-24T22:08:49.677Z" }, - { url = "https://files.pythonhosted.org/packages/cc/1f/0960a129d504728d28a51dbd5a04ce94031eb75bac676341da7aefdd8232/py_spy-0.4.2-py2.py3-none-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:24720573f95230653b457671a1dcc3c5a381fcf4e92677761e328a430ad251b2", size = 2301852, upload-time = "2026-04-24T22:08:51.152Z" }, - { url = "https://files.pythonhosted.org/packages/f9/34/dd7d3c763a00b7b965e25a5eab0acd1a345dbaf0f45fffe595278873a1c0/py_spy-0.4.2-py2.py3-none-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:aeb0323409199c785f730645e9f4bb7a7b9ca2c481f2c331a55642b5d13fa52f", size = 2936518, upload-time = "2026-04-24T22:08:52.264Z" }, - { url = "https://files.pythonhosted.org/packages/6f/ed/1409cdb557e558a6c98003ab12fdd4284699e158c167c187cb0f124eea4c/py_spy-0.4.2-py2.py3-none-win_amd64.whl", hash = "sha256:8b06a353c177677e4e1701b288d8c58e2f8d4208ee81a8048d9f72ba800918f8", size = 1894002, upload-time = "2026-04-24T22:08:53.811Z" }, -] - [[package]] name = "pyarrow" version = "25.0.0" @@ -3724,33 +3878,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/fe/81d1e5f8beed15c01e98649d5c6e2167b67fd395884a2488f18bf1cf0dba/pyarrow-25.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:3f356afe61186395c861d5cd63dc21ff7d5fa335012a4668d979257df7fea0f5", size = 27945954, upload-time = "2026-07-10T08:27:32.903Z" }, ] -[[package]] -name = "pyasn1" -version = "0.6.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a4/9a/23310166d960def5897e91fe20e5b724601b02a22e84ba1f94232c0b7f67/pyasn1-0.6.4.tar.gz", hash = "sha256:9c447d8431c947fe4c8febc4ed9e760bc29011a5b01e5c74b67025bd9fb8ce81", size = 151262, upload-time = "2026-07-09T01:12:33.988Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9a/3b/6163796d69c3977d1e4287bea4a6979161cbbdd170ebb430511e8e1999ce/pyasn1-0.6.4-py3-none-any.whl", hash = "sha256:deda9277cfd454080ec40b207fb6df82206a3a2688735233cdcd8d3d565f088b", size = 84410, upload-time = "2026-07-09T01:12:32.92Z" }, -] - -[[package]] -name = "pyasn1-modules" -version = "0.4.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pyasn1" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892, upload-time = "2025-03-28T02:41:22.17Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259, upload-time = "2025-03-28T02:41:19.028Z" }, -] - [[package]] name = "pybase64" version = "1.4.3" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/aa/b8/4ed5c7ad5ec15b08d35cc79ace6145d5c1ae426e46435f4987379439dfea/pybase64-1.4.3.tar.gz", hash = "sha256:c2ed274c9e0ba9c8f9c4083cfe265e66dd679126cd9c2027965d807352f3f053", size = 137272, upload-time = "2025-12-06T13:27:04.013Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/86/a7/efcaa564f091a2af7f18a83c1c4875b1437db56ba39540451dc85d56f653/pybase64-1.4.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:18d85e5ab8b986bb32d8446aca6258ed80d1bafe3603c437690b352c648f5967", size = 38167, upload-time = "2025-12-06T13:23:16.821Z" }, + { url = "https://files.pythonhosted.org/packages/db/c7/c7ad35adff2d272bf2930132db2b3eea8c44bb1b1f64eb9b2b8e57cde7b4/pybase64-1.4.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3f5791a3491d116d0deaf4d83268f48792998519698f8751efb191eac84320e9", size = 31673, upload-time = "2025-12-06T13:23:17.835Z" }, { url = "https://files.pythonhosted.org/packages/43/1b/9a8cab0042b464e9a876d5c65fe5127445a2436da36fda64899b119b1a1b/pybase64-1.4.3-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:f0b3f200c3e06316f6bebabd458b4e4bcd4c2ca26af7c0c766614d91968dee27", size = 68210, upload-time = "2025-12-06T13:23:18.813Z" }, { url = "https://files.pythonhosted.org/packages/62/f7/965b79ff391ad208b50e412b5d3205ccce372a2d27b7218ae86d5295b105/pybase64-1.4.3-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:bb632edfd132b3eaf90c39c89aa314beec4e946e210099b57d40311f704e11d4", size = 71599, upload-time = "2025-12-06T13:23:20.195Z" }, { url = "https://files.pythonhosted.org/packages/03/4b/a3b5175130b3810bbb8ccfa1edaadbd3afddb9992d877c8a1e2f274b476e/pybase64-1.4.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:356ef1d74648ce997f5a777cf8f1aefecc1c0b4fe6201e0ef3ec8a08170e1b54", size = 59922, upload-time = "2025-12-06T13:23:21.487Z" }, @@ -3765,8 +3900,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fd/7d/931c2539b31a7b375e7d595b88401eeb5bd6c5ce1059c9123f9b608aaa14/pybase64-1.4.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:66e3791f2ed725a46593f8bd2761ff37d01e2cdad065b1dceb89066f476e50c6", size = 54333, upload-time = "2025-12-06T13:23:32.422Z" }, { url = "https://files.pythonhosted.org/packages/de/5e/537601e02cc01f27e9d75f440f1a6095b8df44fc28b1eef2cd739aea8cec/pybase64-1.4.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:72bb0b6bddadab26e1b069bb78e83092711a111a80a0d6b9edcb08199ad7299b", size = 56492, upload-time = "2025-12-06T13:23:33.515Z" }, { url = "https://files.pythonhosted.org/packages/96/97/2a2e57acf8f5c9258d22aba52e71f8050e167b29ed2ee1113677c1b600c1/pybase64-1.4.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5b3365dbcbcdb0a294f0f50af0c0a16b27a232eddeeb0bceeefd844ef30d2a23", size = 70974, upload-time = "2025-12-06T13:23:36.27Z" }, + { url = "https://files.pythonhosted.org/packages/75/2e/a9e28941c6dab6f06e6d3f6783d3373044be9b0f9a9d3492c3d8d2260ac0/pybase64-1.4.3-cp312-cp312-win32.whl", hash = "sha256:7bca1ed3a5df53305c629ca94276966272eda33c0d71f862d2d3d043f1e1b91a", size = 33686, upload-time = "2025-12-06T13:23:37.848Z" }, + { url = "https://files.pythonhosted.org/packages/83/e3/507ab649d8c3512c258819c51d25c45d6e29d9ca33992593059e7b646a33/pybase64-1.4.3-cp312-cp312-win_amd64.whl", hash = "sha256:9f2da8f56d9b891b18b4daf463a0640eae45a80af548ce435be86aa6eff3603b", size = 35833, upload-time = "2025-12-06T13:23:38.877Z" }, + { url = "https://files.pythonhosted.org/packages/bc/8a/6eba66cd549a2fc74bb4425fd61b839ba0ab3022d3c401b8a8dc2cc00c7a/pybase64-1.4.3-cp312-cp312-win_arm64.whl", hash = "sha256:0631d8a2d035de03aa9bded029b9513e1fee8ed80b7ddef6b8e9389ffc445da0", size = 31185, upload-time = "2025-12-06T13:23:39.908Z" }, + { url = "https://files.pythonhosted.org/packages/17/45/92322aec1b6979e789b5710f73c59f2172bc37c8ce835305434796824b7b/pybase64-1.4.3-graalpy312-graalpy250_312_native-macosx_10_13_x86_64.whl", hash = "sha256:2baaa092f3475f3a9c87ac5198023918ea8b6c125f4c930752ab2cbe3cd1d520", size = 38746, upload-time = "2025-12-06T13:26:25.869Z" }, + { url = "https://files.pythonhosted.org/packages/11/94/f1a07402870388fdfc2ecec0c718111189732f7d0f2d7fe1386e19e8fad0/pybase64-1.4.3-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:cde13c0764b1af07a631729f26df019070dad759981d6975527b7e8ecb465b6c", size = 32573, upload-time = "2025-12-06T13:26:27.792Z" }, { url = "https://files.pythonhosted.org/packages/fa/8f/43c3bb11ca9bacf81cb0b7a71500bb65b2eda6d5fe07433c09b543de97f3/pybase64-1.4.3-graalpy312-graalpy250_312_native-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:5c29a582b0ea3936d02bd6fe9bf674ab6059e6e45ab71c78404ab2c913224414", size = 43461, upload-time = "2025-12-06T13:26:28.906Z" }, { url = "https://files.pythonhosted.org/packages/2d/4c/2a5258329200be57497d3972b5308558c6de42e3749c6cc2aa1cbe34b25a/pybase64-1.4.3-graalpy312-graalpy250_312_native-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b6b664758c804fa919b4f1257aa8cf68e95db76fc331de5f70bfc3a34655afe1", size = 36058, upload-time = "2025-12-06T13:26:30.092Z" }, + { url = "https://files.pythonhosted.org/packages/ea/6d/41faa414cde66ec023b0ca8402a8f11cb61731c3dc27c082909cbbd1f929/pybase64-1.4.3-graalpy312-graalpy250_312_native-win_amd64.whl", hash = "sha256:f7537fa22ae56a0bf51e4b0ffc075926ad91c618e1416330939f7ef366b58e3b", size = 36231, upload-time = "2025-12-06T13:26:31.656Z" }, ] [[package]] @@ -3887,6 +4028,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/c1/6e422f34e569cf8e18df68d1939c81c099d2b61e4f7d9621c8a77560799c/pydantic_settings-2.14.2-py3-none-any.whl", hash = "sha256:a20c97b37910b6550d5ea50fbcc2d4187defe58cd57070b73863d069419c9440", size = 61715, upload-time = "2026-06-19T13:44:55.02Z" }, ] +[[package]] +name = "pyelftools" +version = "0.33" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/11/767522582afab1b884d277de0e6e011640cb9d7292a38694b4b1a1df1ae8/pyelftools-0.33.tar.gz", hash = "sha256:660d82dcbeb8e83d1702bd97f223f761625da06111c0cc988eac6b8ab0c1b61f", size = 15068655, upload-time = "2026-05-29T12:56:22.553Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/2a/f9697576603dae937727827505a6126a066affb227034e77e6f9068910da/pyelftools-0.33-py3-none-any.whl", hash = "sha256:f215ad5f47d3f1373a21496a6c9e0707c622840d0622f23ff7ce08678b020036", size = 201178, upload-time = "2026-05-29T12:56:20.587Z" }, +] + [[package]] name = "pygments" version = "2.20.0" @@ -3910,6 +4060,17 @@ crypto = [ { name = "cryptography" }, ] +[[package]] +name = "pynvvideocodec" +version = "2.0.4" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/1c/78f6fdf85133157a6a3405eab5ef4c2bc8048194dbda1c91bb9b8645bb36/pynvvideocodec-2.0.4-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:bad9e25f494abdcfa8f9dffa33a840509eda3ffcdf6e7cf6465d73be307c0c82", size = 28630316, upload-time = "2026-07-08T04:25:54.596Z" }, + { url = "https://files.pythonhosted.org/packages/ca/49/98da271686e00676f41b1197ba5431ddc341b96d8efb68ea9d68e2b0d870/pynvvideocodec-2.0.4-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b59cec7a1a3f78fad13fead78cad8b6d9686827f9ff4477080245457675a01d0", size = 43176147, upload-time = "2026-05-27T04:04:08.297Z" }, + { url = "https://files.pythonhosted.org/packages/42/80/7b13c12fd5f3243b01190130ce098a44ddf62e030e6ed712911cbfe40311/pynvvideocodec-2.0.4-cp312-cp312-manylinux_2_34_aarch64.whl", hash = "sha256:a0daa28b09705806c8c6b26326df217c45e60c0a12a673ea3ea6ee5e2e7193b0", size = 35754893, upload-time = "2026-05-27T04:04:43.866Z" }, + { url = "https://files.pythonhosted.org/packages/76/50/eb1571ab1cee8ebb8a7bdfc355078beebe4b2bb2e5c6ad5d0e18ab8585db/pynvvideocodec-2.0.4-cp312-cp312-win_amd64.whl", hash = "sha256:46e2adb82dc6ac333d3535cc76e4e25c7e8d80dd272b1aba0c28702b861d5261", size = 25692590, upload-time = "2026-05-27T04:05:17.164Z" }, +] + [[package]] name = "pyparsing" version = "3.3.2" @@ -3994,19 +4155,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] -[[package]] -name = "python-discovery" -version = "1.4.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "filelock" }, - { name = "platformdirs" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/4c/81/58c70036dffeccb7fe7d79d6260c69f7a28272bbd3909c29a01ea9422744/python_discovery-1.4.4.tar.gz", hash = "sha256:5cad33982d412c1f3ffb8f9ca4ea292c9680bca3942451d30b69c37fce53a4a3", size = 72212, upload-time = "2026-07-08T23:06:50.691Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/ae/84bc0d2440c95772272bb6f4b3d09ccf08b2898fce89b3d4f969a9fc74e9/python_discovery-1.4.4-py3-none-any.whl", hash = "sha256:abebe9120b43453b68c908acfb1e72a19d1a959ed2cb620ad38fc57d08056dbe", size = 34181, upload-time = "2026-07-08T23:06:49.402Z" }, -] - [[package]] name = "python-dotenv" version = "1.2.2" @@ -4107,33 +4255,38 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/04/0b/3c9baedbdf613ecaa7aa07027780b8867f57b6293b6ee50de316c9f3222b/pyzmq-27.1.0.tar.gz", hash = "sha256:ac0765e3d44455adb6ddbf4417dcce460fc40a05978c08efdf2948072f6db540", size = 281750, upload-time = "2025-09-08T23:10:18.157Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/92/e7/038aab64a946d535901103da16b953c8c9cc9c961dadcbf3609ed6428d23/pyzmq-27.1.0-cp312-abi3-macosx_10_15_universal2.whl", hash = "sha256:452631b640340c928fa343801b0d07eb0c3789a5ffa843f6e1a9cee0ba4eb4fc", size = 1306279, upload-time = "2025-09-08T23:08:03.807Z" }, { url = "https://files.pythonhosted.org/packages/e8/5e/c3c49fdd0f535ef45eefcc16934648e9e59dace4a37ee88fc53f6cd8e641/pyzmq-27.1.0-cp312-abi3-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:1c179799b118e554b66da67d88ed66cd37a169f1f23b5d9f0a231b4e8d44a113", size = 895645, upload-time = "2025-09-08T23:08:05.301Z" }, { url = "https://files.pythonhosted.org/packages/f8/e5/b0b2504cb4e903a74dcf1ebae157f9e20ebb6ea76095f6cfffea28c42ecd/pyzmq-27.1.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3837439b7f99e60312f0c926a6ad437b067356dc2bc2ec96eb395fd0fe804233", size = 652574, upload-time = "2025-09-08T23:08:06.828Z" }, { url = "https://files.pythonhosted.org/packages/f8/9b/c108cdb55560eaf253f0cbdb61b29971e9fb34d9c3499b0e96e4e60ed8a5/pyzmq-27.1.0-cp312-abi3-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:43ad9a73e3da1fab5b0e7e13402f0b2fb934ae1c876c51d0afff0e7c052eca31", size = 840995, upload-time = "2025-09-08T23:08:08.396Z" }, { url = "https://files.pythonhosted.org/packages/c2/bb/b79798ca177b9eb0825b4c9998c6af8cd2a7f15a6a1a4272c1d1a21d382f/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0de3028d69d4cdc475bfe47a6128eb38d8bc0e8f4d69646adfbcd840facbac28", size = 1642070, upload-time = "2025-09-08T23:08:09.989Z" }, { url = "https://files.pythonhosted.org/packages/9c/80/2df2e7977c4ede24c79ae39dcef3899bfc5f34d1ca7a5b24f182c9b7a9ca/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_i686.whl", hash = "sha256:cf44a7763aea9298c0aa7dbf859f87ed7012de8bda0f3977b6fb1d96745df856", size = 2021121, upload-time = "2025-09-08T23:08:11.907Z" }, { url = "https://files.pythonhosted.org/packages/46/bd/2d45ad24f5f5ae7e8d01525eb76786fa7557136555cac7d929880519e33a/pyzmq-27.1.0-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:f30f395a9e6fbca195400ce833c731e7b64c3919aa481af4d88c3759e0cb7496", size = 1878550, upload-time = "2025-09-08T23:08:13.513Z" }, + { url = "https://files.pythonhosted.org/packages/e6/2f/104c0a3c778d7c2ab8190e9db4f62f0b6957b53c9d87db77c284b69f33ea/pyzmq-27.1.0-cp312-abi3-win32.whl", hash = "sha256:250e5436a4ba13885494412b3da5d518cd0d3a278a1ae640e113c073a5f88edd", size = 559184, upload-time = "2025-09-08T23:08:15.163Z" }, + { url = "https://files.pythonhosted.org/packages/fc/7f/a21b20d577e4100c6a41795842028235998a643b1ad406a6d4163ea8f53e/pyzmq-27.1.0-cp312-abi3-win_amd64.whl", hash = "sha256:9ce490cf1d2ca2ad84733aa1d69ce6855372cb5ce9223802450c9b2a7cba0ccf", size = 619480, upload-time = "2025-09-08T23:08:17.192Z" }, + { url = "https://files.pythonhosted.org/packages/78/c2/c012beae5f76b72f007a9e91ee9401cb88c51d0f83c6257a03e785c81cc2/pyzmq-27.1.0-cp312-abi3-win_arm64.whl", hash = "sha256:75a2f36223f0d535a0c919e23615fc85a1e23b71f40c7eb43d7b1dedb4d8f15f", size = 552993, upload-time = "2025-09-08T23:08:18.926Z" }, ] [[package]] name = "quack-kernels" -version = "0.6.1" +version = "0.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "apache-tvm-ffi" }, { name = "einops" }, { name = "nvidia-cutlass-dsl" }, - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "torch-c-dlpack-ext" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/01/17/890875f88f4d7da28faec9e6cf0a0cc565715b01474e50501fafc5bc71b4/quack_kernels-0.6.1.tar.gz", hash = "sha256:a694f89c91d137478de523c0227365a331ac9cb66790cfb08baa3dbfaafc71e7", size = 387353, upload-time = "2026-07-05T11:50:19.807Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e1/94/ee76e3a3dc74d986b7b24c5928f1d14b01bd5152375688c2ede369f6d19b/quack_kernels-0.5.0.tar.gz", hash = "sha256:c7c7338b67243397b6ca166e648bba161076e99f3858b532e1c877dcc6eaa03d", size = 366426, upload-time = "2026-05-29T05:00:25.985Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2b/65/a38a30a6ac96a757363a5be9d09cef799640bb143a64ba5a2f4d400d95d9/quack_kernels-0.6.1-py3-none-any.whl", hash = "sha256:266705ea82117e9b1c8a9e44d68a458519f2498d966c0efffd6812120c3995ad", size = 358439, upload-time = "2026-07-05T11:50:18.502Z" }, + { url = "https://files.pythonhosted.org/packages/2d/2b/a8f171d5e172880885571bf89e93204aaf231a0e92c4c84714eaf18c271a/quack_kernels-0.5.0-py3-none-any.whl", hash = "sha256:08821ebfb8e638cc20308d5c59410c6dbb3b637ccc7b07bd57c7a9261a06af74", size = 327709, upload-time = "2026-05-29T05:00:24.679Z" }, ] [[package]] name = "ray" -version = "2.55.1" +version = "2.56.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -4146,38 +4299,10 @@ dependencies = [ { name = "requests" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/ac/3a/4d34f471a68b958b7f94c974c19ad6836a61a2dc16393df4294169a2e4b0/ray-2.55.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:137f9006eee28caab8260803cca314f37bbda3fc94fdfa31c770b5d019626ad8", size = 65822379, upload-time = "2026-04-22T20:09:58.064Z" }, - { url = "https://files.pythonhosted.org/packages/f1/13/0db535102d0256b350ca116d8987588aca1a1f9ebb4638e1e1ff88bbcef8/ray-2.55.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:26541f69bb55607ef8335baac75b2ed12ff2ce02d56313219b29eda003039221", size = 72910802, upload-time = "2026-04-22T20:10:04.382Z" }, - { url = "https://files.pythonhosted.org/packages/4c/f8/fffadf3f4285eebd460e4d7f2ed1c0cd641ed89613c3f49eb881ee9fa7e2/ray-2.55.1-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:263705f6bab29e7622a94f82da25fd7f9cead76cdf89a07aab28f79cdf8f9d95", size = 73765203, upload-time = "2026-04-22T20:10:10.495Z" }, - { url = "https://files.pythonhosted.org/packages/10/f7/5acb86fc9625a0e6bbc40e1c7d42c60770e78585439a921c32738b6d675a/ray-2.55.1-cp312-cp312-win_amd64.whl", hash = "sha256:9ad56704c8bd7e92130162f9c58e4ef473609515637673d5a36e761f95335206", size = 27865547, upload-time = "2026-04-22T20:10:15.364Z" }, -] - -[package.optional-dependencies] -data = [ - { name = "fsspec" }, - { name = "numpy" }, - { name = "pandas" }, - { name = "pyarrow" }, -] -serve = [ - { name = "aiohttp" }, - { name = "aiohttp-cors" }, - { name = "colorful" }, - { name = "fastapi" }, - { name = "grpcio" }, - { name = "opencensus" }, - { name = "opentelemetry-exporter-prometheus" }, - { name = "opentelemetry-proto" }, - { name = "opentelemetry-sdk" }, - { name = "prometheus-client" }, - { name = "py-spy" }, - { name = "pydantic" }, - { name = "requests" }, - { name = "smart-open" }, - { name = "starlette" }, - { name = "uvicorn", extra = ["standard"] }, - { name = "virtualenv" }, - { name = "watchfiles" }, + { url = "https://files.pythonhosted.org/packages/50/d9/a17feef16a123f5d32d4c5fa7de853c59ba702f8404cf452a7ce20faca13/ray-2.56.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:44bc0000c5bfad85b2ff6e0ef91e95f901d1a2d2fdd72f94f08a046eb494cd61", size = 66346989, upload-time = "2026-07-17T21:28:39.4Z" }, + { url = "https://files.pythonhosted.org/packages/05/19/c3b1bcccd09decaf2a2e3370041ae67070bb1a8638f2665d36edfbb0261d/ray-2.56.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:8fdd6b096215906cf1f9acdc7898c9d6140606f2d27245778b8385a9f19e6cb0", size = 73319289, upload-time = "2026-07-17T21:28:44.502Z" }, + { url = "https://files.pythonhosted.org/packages/d2/7f/577a61bf2c8eff26e942afe53a6b03f4dbf5b4f233b832c228417d0c954e/ray-2.56.1-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:e5d3173696831134c76bd09451dfe95c32d72c271253b0d6b09d2df9994aa660", size = 74194147, upload-time = "2026-07-17T21:28:50.567Z" }, + { url = "https://files.pythonhosted.org/packages/00/e9/0fd1223597f9ca98ec496ec726043918a5301547789b1be95a635ec82649/ray-2.56.1-cp312-cp312-win_amd64.whl", hash = "sha256:8052573ee5ef8c4fdd7aeb6a257c80542e69c48f3f6d117101f95c970ffdc7e2", size = 28373294, upload-time = "2026-07-17T21:28:55.122Z" }, ] [[package]] @@ -4291,6 +4416,8 @@ version = "0.8.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/c9/77/6ba90ab4a538d3ec244329c57c4d26a78c8313ea6fa72c8768d46f11c1c9/rignore-0.8.0.tar.gz", hash = "sha256:2e5ad6b19834f04a877d26fe863fd77ed851ed4019fdca097fb1b744311e3562", size = 55358, upload-time = "2026-07-17T19:01:21.257Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/05/7e/0d270c1ed723b82bea8ccd504185acb5d71830975260e8de02af7acff728/rignore-0.8.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:a5e285ec58b3b66284a7f48805e4db0ea948370deb484a4935b147187ecf1e25", size = 847184, upload-time = "2026-07-17T18:58:31.348Z" }, + { url = "https://files.pythonhosted.org/packages/ff/dc/841941f8b0883a8038f9d540607238c456df20e98243a61142fd15699806/rignore-0.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:367d3cf401b477a5ba7eb05b5b94c491b0d704507d9eaf80378a8d843fe00674", size = 816590, upload-time = "2026-07-17T18:58:32.571Z" }, { url = "https://files.pythonhosted.org/packages/57/53/9e047a6cd95b553519703350f7a3530ddd31d309c9ade1ec913db5882f74/rignore-0.8.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d49870c032abcc28db2210ad92b3bd5706ec792ab25c9295a6c536a4b92226ce", size = 884130, upload-time = "2026-07-17T18:58:33.921Z" }, { url = "https://files.pythonhosted.org/packages/07/5a/ae444f30dfa47716ccadf9ed9512736494c0bf5222f330845a637f1c569c/rignore-0.8.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d50632823c273ee19212fd939b04e9b55d5a6e0c9b998fadef6dc9a0c2f6aecc", size = 857230, upload-time = "2026-07-17T18:58:35.32Z" }, { url = "https://files.pythonhosted.org/packages/ba/67/b2ddfbf5a42a8ea89e6a0330189f7f76f9113e0d8c6a4c68b68ff68f139c/rignore-0.8.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b3337e69856ba18c9d079e0ce8b342dc256f48585ae3e1c7aac67682b672e83a", size = 1133331, upload-time = "2026-07-17T18:58:36.895Z" }, @@ -4302,6 +4429,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a9/0e/d769dfa933861dd469c9158ffc9bc2cb3bc6d46f88000fb6f41cf2905a3a/rignore-0.8.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:e3e863b2cb1db481384bd43a25730ef00f7eccdd49c82a0e1d481a6412dc2653", size = 1132462, upload-time = "2026-07-17T18:58:45.583Z" }, { url = "https://files.pythonhosted.org/packages/c1/3d/57fc6264ebf8d9b95850589899e6192d9f92b67ea48e56fc32969b75474a/rignore-0.8.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:90cc363aedc7a93b4b15933f2104bfcae1c6c1635a9bc479665cae83044577d9", size = 1139679, upload-time = "2026-07-17T18:58:46.928Z" }, { url = "https://files.pythonhosted.org/packages/fa/be/24b12a8464e19d348aeb388267f897ede0fae6618052c3148b747acae214/rignore-0.8.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:303a9fd02d3612d15e6dc3474ae7d3eae1d07b668b7b4943fd4df305be4d2f76", size = 1138125, upload-time = "2026-07-17T18:58:48.42Z" }, + { url = "https://files.pythonhosted.org/packages/d6/8d/89f7cba3491164c04f8a64056ca973494536ffdf48bfba41565f54dfa122/rignore-0.8.0-cp312-cp312-win32.whl", hash = "sha256:e17a0914378fa15d1e29effce4b39fd2031f253e52d307ff4b71c443d1d5d30e", size = 637663, upload-time = "2026-07-17T18:58:49.76Z" }, + { url = "https://files.pythonhosted.org/packages/3f/ff/823a5ead8bba0a2054cbf2dbef99989204557904ad8578f514bec1b9b4e7/rignore-0.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:d5c8f1dff84c4d4114f1553564b5d909a0c4c09dbe1b5916a0506d719f1fc3b5", size = 728265, upload-time = "2026-07-17T18:58:51.028Z" }, + { url = "https://files.pythonhosted.org/packages/8f/ce/73c919505f9f270ee3e34205ff2dbe0b0d73e945f8c569922acff148bcf4/rignore-0.8.0-cp312-cp312-win_arm64.whl", hash = "sha256:6cdea3f85de8286a38ae75a0f9092cd3afc4d33ce6ee2e3f6005f97f7da9d249", size = 665216, upload-time = "2026-07-17T18:58:52.308Z" }, ] [[package]] @@ -4426,8 +4556,13 @@ version = "0.2.2" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/cc/33/ea3cb3839607eb175da835244a798f797f478c5ddf0e8ecdf57ea85a4c70/sentencepiece-0.2.2.tar.gz", hash = "sha256:3d2b5e824b5622038dc7b490897efe05ebbbb9e7350fc142f3ecc8789ef9bdf6", size = 8218435, upload-time = "2026-07-12T08:39:34.701Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/13/7a562289c8d5b49ebdf3f9c1e8ab67cf14a8743b1d90c8f406bfdec36b72/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:1edb10e520e4bddf74d85b0f5ae74cc2d60c2b448885080bfb618bc2b3a49f6b", size = 2188384, upload-time = "2026-07-12T08:38:28.486Z" }, + { url = "https://files.pythonhosted.org/packages/85/d1/912f14fd5eae168aba726ffb6a9a2dc1c71fe7676c53da6f5c442b886d4a/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f7c06c751c19d923435a54bff4f7e66e728fad160e8da28254f133abc9725820", size = 1441553, upload-time = "2026-07-12T08:38:30.552Z" }, + { url = "https://files.pythonhosted.org/packages/bd/44/caa9cab5f261a019e2808bc5046152775dc57352ba9cbae7525e9e7a1ed4/sentencepiece-0.2.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:38111ed1f79268f399c505028023d5eaaf0ab4e5eafceb709468b0d3323e7838", size = 1347176, upload-time = "2026-07-12T08:38:32.211Z" }, { url = "https://files.pythonhosted.org/packages/19/90/cd798935668cff71d309d8ff10385844ecf216b1fe454f1993ed8bf2cb91/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cbce24284f51f71d10a42b7b9c964dcb9048b28f1c8e5db40bcbcb6f428cba6a", size = 1325200, upload-time = "2026-07-12T08:38:33.689Z" }, { url = "https://files.pythonhosted.org/packages/b6/2d/37e3da037318a70066ded0d51bc2a7f35491ae6338dd993d5eb1503fc3b5/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c8a168b040bc61681293f79a949b5d911c8e25086f4260285b8d97ab5f1195da", size = 1397736, upload-time = "2026-07-12T08:38:35.771Z" }, + { url = "https://files.pythonhosted.org/packages/8d/11/753fca2e6b109be3ab7867abf357dfe48677fe726ae5a5363d0b54ca9450/sentencepiece-0.2.2-cp312-cp312-win_amd64.whl", hash = "sha256:7c6e7bf684dc12145bfa685d3060beaea55139134ba848289bee514ed42e7383", size = 1248030, upload-time = "2026-07-12T08:38:37.604Z" }, + { url = "https://files.pythonhosted.org/packages/e2/0a/70efbe861ca182d7d4b6e1a20f58e043400848fa9f2915229f082e221648/sentencepiece-0.2.2-cp312-cp312-win_arm64.whl", hash = "sha256:76ff5814db72e7462dece042d7593cdf102b8ec82c2b1cc201a2add34ee3050d", size = 1187325, upload-time = "2026-07-12T08:38:39.348Z" }, ] [[package]] @@ -4449,12 +4584,16 @@ version = "1.3.7" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/8d/48/49393a96a2eef1ab418b17475fb92b8fcfad83d099e678751b05472e69de/setproctitle-1.3.7.tar.gz", hash = "sha256:bc2bc917691c1537d5b9bca1468437176809c7e11e5694ca79a9ca12345dcb9e", size = 27002, upload-time = "2025-09-05T12:51:25.278Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/f0/2dc88e842077719d7384d86cc47403e5102810492b33680e7dadcee64cd8/setproctitle-1.3.7-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2dc99aec591ab6126e636b11035a70991bc1ab7a261da428491a40b84376654e", size = 18049, upload-time = "2025-09-05T12:49:36.241Z" }, + { url = "https://files.pythonhosted.org/packages/f0/b4/50940504466689cda65680c9e9a1e518e5750c10490639fa687489ac7013/setproctitle-1.3.7-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:cdd8aa571b7aa39840fdbea620e308a19691ff595c3a10231e9ee830339dd798", size = 13079, upload-time = "2025-09-05T12:49:38.088Z" }, { url = "https://files.pythonhosted.org/packages/d0/99/71630546b9395b095f4082be41165d1078204d1696c2d9baade3de3202d0/setproctitle-1.3.7-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2906b6c7959cdb75f46159bf0acd8cc9906cf1361c9e1ded0d065fe8f9039629", size = 32932, upload-time = "2025-09-05T12:49:39.271Z" }, { url = "https://files.pythonhosted.org/packages/50/22/cee06af4ffcfb0e8aba047bd44f5262e644199ae7527ae2c1f672b86495c/setproctitle-1.3.7-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6915964a6dda07920a1159321dcd6d94fc7fc526f815ca08a8063aeca3c204f1", size = 33736, upload-time = "2025-09-05T12:49:40.565Z" }, { url = "https://files.pythonhosted.org/packages/5c/00/a5949a8bb06ef5e7df214fc393bb2fb6aedf0479b17214e57750dfdd0f24/setproctitle-1.3.7-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cff72899861c765bd4021d1ff1c68d60edc129711a2fdba77f9cb69ef726a8b6", size = 35605, upload-time = "2025-09-05T12:49:42.362Z" }, { url = "https://files.pythonhosted.org/packages/b0/3a/50caca532a9343828e3bf5778c7a84d6c737a249b1796d50dd680290594d/setproctitle-1.3.7-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b7cb05bd446687ff816a3aaaf831047fc4c364feff7ada94a66024f1367b448c", size = 33143, upload-time = "2025-09-05T12:49:43.515Z" }, { url = "https://files.pythonhosted.org/packages/ca/14/b843a251296ce55e2e17c017d6b9f11ce0d3d070e9265de4ecad948b913d/setproctitle-1.3.7-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3a57b9a00de8cae7e2a1f7b9f0c2ac7b69372159e16a7708aa2f38f9e5cc987a", size = 34434, upload-time = "2025-09-05T12:49:45.31Z" }, { url = "https://files.pythonhosted.org/packages/c8/b7/06145c238c0a6d2c4bc881f8be230bb9f36d2bf51aff7bddcb796d5eed67/setproctitle-1.3.7-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d8828b356114f6b308b04afe398ed93803d7fca4a955dd3abe84430e28d33739", size = 32795, upload-time = "2025-09-05T12:49:46.419Z" }, + { url = "https://files.pythonhosted.org/packages/ef/dc/ef76a81fac9bf27b84ed23df19c1f67391a753eed6e3c2254ebcb5133f56/setproctitle-1.3.7-cp312-cp312-win32.whl", hash = "sha256:b0304f905efc845829ac2bc791ddebb976db2885f6171f4a3de678d7ee3f7c9f", size = 12552, upload-time = "2025-09-05T12:49:47.635Z" }, + { url = "https://files.pythonhosted.org/packages/e2/5b/a9fe517912cd6e28cf43a212b80cb679ff179a91b623138a99796d7d18a0/setproctitle-1.3.7-cp312-cp312-win_amd64.whl", hash = "sha256:9888ceb4faea3116cf02a920ff00bfbc8cc899743e4b4ac914b03625bdc3c300", size = 13247, upload-time = "2025-09-05T12:49:49.16Z" }, ] [[package]] @@ -4515,18 +4654,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] -[[package]] -name = "smart-open" -version = "8.0.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/76/53/9c513747547fd595d5c143259129ea8b9c3ea2f6b7bb9dcea2b1966ded3c/smart_open-8.0.1.tar.gz", hash = "sha256:18b1c4496003c6902be17c15f032b5c319f307c89c6ae9e6b028b508bed8b2cf", size = 61882, upload-time = "2026-07-15T13:56:10.492Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/96/325b8c507ccecc50421fecc0345a502ee6e4a44785af3c4e6ecbadad624a/smart_open-8.0.1-py3-none-any.whl", hash = "sha256:3e97f90e92a952cb57863dfe132082c400a52eeeb27c067692fb51dbcc5b0089", size = 73504, upload-time = "2026-07-15T13:56:09.033Z" }, -] - [[package]] name = "sniffio" version = "1.3.1" @@ -4626,15 +4753,15 @@ wheels = [ [[package]] name = "starlette" -version = "0.52.1" +version = "1.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c4/68/79977123bb7be889ad680d79a40f339082c1978b5cfcf62c2d8d196873ac/starlette-0.52.1.tar.gz", hash = "sha256:834edd1b0a23167694292e94f597773bc3f89f362be6effee198165a35d62933", size = 2653702, upload-time = "2026-01-18T13:34:11.062Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/3c/76d2fd1f1357ed0f0108d8a5aa233dcf16e2946a8559c84912fe08e01ac7/starlette-1.4.1.tar.gz", hash = "sha256:b7332de6e9375593a29ba9eee1e6ecfeb3eb2043e2e19a13b4b71da73ff35540", size = 2709041, upload-time = "2026-08-05T15:17:26.23Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/81/0d/13d1d239a25cbfb19e740db83143e95c772a1fe10202dda4b76792b114dd/starlette-0.52.1-py3-none-any.whl", hash = "sha256:0029d43eb3d273bc4f83a08720b4912ea4b071087a3b48db01b7c839f7954d74", size = 74272, upload-time = "2026-01-18T13:34:09.188Z" }, + { url = "https://files.pythonhosted.org/packages/75/0a/67e95f21498de41433babf7b1db0eeab449eb58872dfb831b27747a70fd0/starlette-1.4.1-py3-none-any.whl", hash = "sha256:7d078e0fbefae0d2cecfb80a799d6fb84b1c0c6acd4f14ac79d17d0e7ec27f19", size = 74019, upload-time = "2026-08-05T15:17:24.357Z" }, ] [[package]] @@ -4783,7 +4910,9 @@ dependencies = [ { name = "ml-dtypes" }, { name = "numpy" }, { name = "psutil" }, - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, + { name = "setuptools", marker = "sys_platform == 'darwin'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "torch-c-dlpack-ext" }, { name = "tqdm" }, { name = "typing-extensions" }, @@ -4791,6 +4920,7 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/56/70/5051f65821baa30a3d61fc48f8ba10c776490315e8c90f82559b92089756/tilelang-0.1.9.tar.gz", hash = "sha256:287f727c913bb648fcf6c1968809ba3390e55eeed257a5c6bb9a80bc05966af4", size = 93395292, upload-time = "2026-04-22T09:19:11.988Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/90/db/4dd76da8c8585c605639a21bc098d504e317fe324a72f01ce3c7370250b4/tilelang-0.1.9-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:00ed594fdeb229c5505b9ffa895c3c5daeb28641c78f783fa1f724cf1e08cecd", size = 36599020, upload-time = "2026-04-22T09:14:39.366Z" }, { url = "https://files.pythonhosted.org/packages/f7/8a/1cbeee79d62abaa02441c2d00621554e41aa62dbf3b94a4feb3867184b01/tilelang-0.1.9-cp38-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4bbccfe9035aed775ffafb6dc25a5994504b24e2c5d95d0f39643edfafa7bf12", size = 45419374, upload-time = "2026-04-22T09:15:56.014Z" }, { url = "https://files.pythonhosted.org/packages/c6/a7/f4bfb86f87e107703146e703204cec2c0eae2492b633e0052b0ace3febb6/tilelang-0.1.9-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:77ab0ee2f40f66ea015b6b21426d482751e28cbc635ef9d1198cbd6502454a7c", size = 42110365, upload-time = "2026-04-22T09:17:18.292Z" }, ] @@ -4851,6 +4981,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/72/f4/0de46cfa12cdcbcd464cc59fde36912af405696f687e53a091fb432f694c/tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc", size = 2612133, upload-time = "2026-01-05T10:45:17.232Z" }, ] +[[package]] +name = "tokenspeed-mla" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "apache-tvm-ffi" }, + { name = "nvidia-cutlass-dsl" }, + { name = "tokenspeed-triton" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/42/20/4110d624d81d63f0bee2f19dba7ea0e1d8a31ea50147e6c1db82223c88a4/tokenspeed_mla-0.1.2-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:592590f36d85e624ecdc5e357ff35e29e761e6d879900dce8b67a6785c8ce75c", size = 743769, upload-time = "2026-05-13T03:30:54.486Z" }, + { url = "https://files.pythonhosted.org/packages/84/01/4bf8b74ead3e8e7c1c809435396254c067a33fde48acc20f602aae622d97/tokenspeed_mla-0.1.2-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:c9466a351fe039792e56cf49f3e79744c1dc28c7af10306a02e62b8e92fa5985", size = 748681, upload-time = "2026-05-13T03:30:56.718Z" }, +] + +[[package]] +name = "tokenspeed-triton" +version = "3.8.10.post20260721" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1d/44/89740db8951918c9acd8731243eef8b44d0eb92ea423552639265c46018e/tokenspeed_triton-3.8.10.post20260721-cp312-abi3-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d793ad0eaebb1d08272c97a2b8f2c31304231748b03de9a08e70a362de92a6e0", size = 82966664, upload-time = "2026-07-21T17:14:38.568Z" }, + { url = "https://files.pythonhosted.org/packages/91/53/f46b401e8ec8998f5b9c39cff0614b796bf49113a09f588cfdfa342789a3/tokenspeed_triton-3.8.10.post20260721-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:66cba8d32a1539afd0ff3eec1782b082d01b4db6824d68017d1d789a03d0be37", size = 87210295, upload-time = "2026-07-21T17:14:42.173Z" }, +] + [[package]] name = "torch" version = "2.11.0" @@ -4883,7 +5037,7 @@ resolution-markers = [ ] dependencies = [ { name = "cuda-bindings", marker = "sys_platform != 'win32'" }, - { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform != 'win32'" }, + { name = "cuda-toolkit", version = "13.0.2", source = { registry = "https://pypi.org/simple" }, extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform != 'win32'" }, { name = "filelock" }, { name = "fsspec" }, { name = "jinja2" }, @@ -4908,12 +5062,15 @@ name = "torch-c-dlpack-ext" version = "0.1.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/37/de/921b6491efce5c389a5ef9bbed3d2d6660005840dae488124173180859ab/torch_c_dlpack_ext-0.1.5.tar.gz", hash = "sha256:d06f0357d575d22a168cc77acb9020fc4bae30968ceb6718a055dcbe92bacabe", size = 12913, upload-time = "2026-01-12T11:25:08.484Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/b1/67/10d236698525d7b7db4d74ec0a4b01f5b2db33968995fdd9ac6b4635e327/torch_c_dlpack_ext-0.1.5-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:c0f2bd51fcd99c0e5b50314e1985f2728c4941bfa821f065e6c30951d1f995ca", size = 5291237, upload-time = "2026-01-12T11:24:44.011Z" }, { url = "https://files.pythonhosted.org/packages/87/06/8d760997307a5c3be4384424667bf31aae0a42060838c532c7d846516175/torch_c_dlpack_ext-0.1.5-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3562ee411258676f9c38b8ad39306d1c8d027b6a86f6a87c920d2d009a9d1510", size = 443069, upload-time = "2026-01-12T11:24:45.451Z" }, { url = "https://files.pythonhosted.org/packages/e2/79/a914539b4785f3e44f891aa012a886edb8bc10fe081c440981c57543ce21/torch_c_dlpack_ext-0.1.5-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e6f9da4bb9af70e27facc777458be62e10dbbbddda7672d16138db0553c5a524", size = 897846, upload-time = "2026-01-12T11:24:48.168Z" }, + { url = "https://files.pythonhosted.org/packages/3a/e6/7d7a97a3953208d6d6ce749180c34d1dab48464ded9a76cecabe9d021ce6/torch_c_dlpack_ext-0.1.5-cp312-cp312-win_amd64.whl", hash = "sha256:670fbbab70123cc228bed41693a3720757af57a0ad22669063c9db25321e8f55", size = 1482855, upload-time = "2026-01-12T11:24:49.581Z" }, ] [[package]] @@ -4921,8 +5078,21 @@ name = "torchaudio" version = "2.11.0" source = { registry = "https://pypi.org/simple" } wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/b1/77658817acacd01a72b714440c62f419efc4d90170e704e8e7a2c0918988/torchaudio-2.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a1cf1acc883bee9cb906a933572fed6a8a933f86ef34e9ea7d803f72317e8c1b", size = 684226, upload-time = "2026-03-23T18:13:40.023Z" }, { url = "https://files.pythonhosted.org/packages/78/28/c7adc053039f286c2aca0038b766cbe3294e66fec6b29a820e95128f9ede/torchaudio-2.11.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:bc653defca1c16154398517a1adc98d0fb7f1dd08e58ced217558d213c2c6e29", size = 1626670, upload-time = "2026-03-23T18:13:42.162Z" }, { url = "https://files.pythonhosted.org/packages/88/d8/d6d0f896e064aa67377484efef4911cdcc07bce2929474e1417cc0af18c2/torchaudio-2.11.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:6503c0bdb29daf2e6281bb70ea2dfe2c3553b782b619eb5d73bdadd8a3f7cecf", size = 1771992, upload-time = "2026-03-23T18:13:33.188Z" }, + { url = "https://files.pythonhosted.org/packages/23/a8/941277ecc39f7a0a169d554302a1f1afd87c1d94a8aec828891916cea59a/torchaudio-2.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:478110f981e5d40a8d82221732c57a56c85a1d5895fb8fe646e86ee15eded3bd", size = 328663, upload-time = "2026-03-23T18:13:19.218Z" }, +] + +[[package]] +name = "torchcodec" +version = "0.15.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/5e/ba71ff29b3f957a7e05cfb5c1d189f34c4224166b5bbe900ec8320f506f7/torchcodec-0.15.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:3a4b24012f7a7fe962dfee8f06d9c91e9e3fd1f4b6302fdb5b8884a02aca3f37", size = 4576065, upload-time = "2026-07-15T10:14:06.554Z" }, + { url = "https://files.pythonhosted.org/packages/9d/de/c00b8d13e3e28de9c76f05b4c25fc4d882b4a3d1451b8d2073d089895684/torchcodec-0.15.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5c62f4257b49c6473b0a1006519274b7daef9ef9c1d66b1a6a025dba9df5daac", size = 2727846, upload-time = "2026-07-15T10:14:08.066Z" }, + { url = "https://files.pythonhosted.org/packages/d0/05/b7ba7ae04db4afeb1fd32d30ec6290d511c374adc464afe191c8fc8d4e22/torchcodec-0.15.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:fa31e33884829332cc55b301aa9d23ba90bf164aa8576a8c68aed6c0061c2d8c", size = 2988620, upload-time = "2026-07-15T10:14:09.504Z" }, + { url = "https://files.pythonhosted.org/packages/e0/0f/fa432d8c8b523f5891a66483f607ec80e28ae025d99ce1d1c50667d8446b/torchcodec-0.15.0-cp312-cp312-win_amd64.whl", hash = "sha256:589e127778870c691d8977c08311bf57c4fecb9eb56fa52cf29d9671fe78eb72", size = 3242793, upload-time = "2026-07-15T10:14:10.84Z" }, ] [[package]] @@ -4976,23 +5146,22 @@ wheels = [ [[package]] name = "transformers" -version = "4.57.6" +version = "5.14.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "filelock" }, { name = "huggingface-hub" }, { name = "numpy" }, { name = "packaging" }, { name = "pyyaml" }, { name = "regex" }, - { name = "requests" }, { name = "safetensors" }, { name = "tokenizers" }, { name = "tqdm" }, + { name = "typer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c4/35/67252acc1b929dc88b6602e8c4a982e64f31e733b804c14bc24b47da35e6/transformers-4.57.6.tar.gz", hash = "sha256:55e44126ece9dc0a291521b7e5492b572e6ef2766338a610b9ab5afbb70689d3", size = 10134912, upload-time = "2026-01-16T10:38:39.284Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/fb/2a2ba88f325e68a921d8b69ff63b477830b2e73ade9a3c8c8cab2f06d741/transformers-5.14.1.tar.gz", hash = "sha256:60d196c27781eacf8637e2b533f517582907ad6f9ae142046d6b69431a5b2173", size = 9295927, upload-time = "2026-07-16T09:41:57.773Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl", hash = "sha256:4c9e9de11333ddfe5114bc872c9f370509198acf0b87a832a0ab9458e2bd0550", size = 11993498, upload-time = "2026-01-16T10:38:31.289Z" }, + { url = "https://files.pythonhosted.org/packages/6f/67/8d85ca2323233ae3c0365a659c4e52ee1f587b440e4bc577e7d8e4416d0f/transformers-5.14.1-py3-none-any.whl", hash = "sha256:9db974c4079ede2d1a3ea7ca5a240df33f2cc26fc2b36ba64c5f2a4f43b6e725", size = 11625234, upload-time = "2026-07-16T09:41:54.143Z" }, ] [[package]] @@ -5165,24 +5334,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/99/39/6b3f7d234ba3964c428a6e40006340f53ba37993f46ed6e111c6e9141d18/uvloop-0.22.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:512fec6815e2dd45161054592441ef76c830eddaad55c8aa30952e6fe1ed07c0", size = 4296343, upload-time = "2025-10-16T22:16:35.149Z" }, ] -[[package]] -name = "virtualenv" -version = "21.6.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "distlib" }, - { name = "filelock" }, - { name = "platformdirs" }, - { name = "python-discovery" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/34/d9/b477fddb68840b570af8b22afe9b035cbc277b5fb7b33dea390617a8b10f/virtualenv-21.6.1.tar.gz", hash = "sha256:15f978b7cd329f24855ff4a0c4b4899cc7678589f49adbdcbbb4d3232e641128", size = 5526620, upload-time = "2026-07-10T19:33:53.312Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c1/7c/4e7225d46d634a0d8d534dd8a6ce0c319d09b4d0cf0337eb314ca4789d8c/virtualenv-21.6.1-py3-none-any.whl", hash = "sha256:afe991df855715a2b2f60edfcc0107ef95a79fdfd8cb4cdaa71603d1c12e463b", size = 5506392, upload-time = "2026-07-10T19:33:51.629Z" }, -] - [[package]] name = "vllm" -version = "0.20.0" +version = "0.25.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -5201,8 +5355,9 @@ dependencies = [ { name = "filelock" }, { name = "flashinfer-cubin" }, { name = "flashinfer-python" }, - { name = "gguf" }, + { name = "humming-kernels", extra = ["cu13"] }, { name = "ijson" }, + { name = "jsonschema" }, { name = "lark" }, { name = "llguidance", marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 'x86_64'" }, { name = "lm-format-enforcer" }, @@ -5214,7 +5369,8 @@ dependencies = [ { name = "numba" }, { name = "numpy" }, { name = "nvidia-cudnn-frontend" }, - { name = "nvidia-cutlass-dsl" }, + { name = "nvidia-cutlass-dsl", extra = ["cu13"] }, + { name = "nvtx" }, { name = "openai" }, { name = "openai-harmony" }, { name = "opencv-python-headless" }, @@ -5232,33 +5388,40 @@ dependencies = [ { name = "py-cpuinfo" }, { name = "pybase64" }, { name = "pydantic" }, + { name = "pynvvideocodec" }, { name = "python-json-logger" }, { name = "pyyaml" }, { name = "pyzmq" }, { name = "quack-kernels" }, { name = "regex" }, { name = "requests" }, + { name = "safetensors" }, { name = "sentencepiece" }, { name = "setproctitle" }, { name = "setuptools" }, { name = "six" }, + { name = "starlette" }, { name = "tiktoken" }, { name = "tilelang" }, { name = "tokenizers" }, - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, + { name = "tokenspeed-mla", marker = "sys_platform == 'linux'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "torchaudio" }, - { name = "torchvision", version = "0.26.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, + { name = "torchcodec" }, + { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torchvision", version = "0.26.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "tqdm" }, { name = "transformers" }, { name = "typing-extensions" }, { name = "watchfiles" }, - { name = "xgrammar", version = "0.2.3", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64'" }, - { name = "xgrammar", version = "0.2.4", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x'" }, + { name = "xgrammar", version = "0.2.3", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "xgrammar", version = "0.2.4", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'ppc64le' or platform_machine == 's390x' or (platform_machine == 'x86_64' and sys_platform != 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e7/80/9798ce5e16af5754183ef33a63dc27017e2b51c87f51cc741832ce47a2d5/vllm-0.20.0.tar.gz", hash = "sha256:a6d50152936ee292455af3ffbe359f7a284ac43bf3b68caccf29f368e196cc72", size = 33508260, upload-time = "2026-04-27T11:08:04.666Z" } +sdist = { url = "https://files.pythonhosted.org/packages/bf/ce/542315fe84e23a73b919d64794350233d4f1c70b2544ba0afd5348005aff/vllm-0.25.1.tar.gz", hash = "sha256:ddbdec3f1c0f21afa70b7eb6ddf3faa29d26b1302ee4b5d5e00ec3af41b0c2e4", size = 37731826, upload-time = "2026-07-14T08:58:23.252Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/63/5b/26379d3c522379373e50b9f77adf55eb94f4a0f62a6c8e3e7fe3f0bf0d39/vllm-0.20.0-cp38-abi3-manylinux_2_35_aarch64.whl", hash = "sha256:29a135ca0d70650f057f15c7c0b560d24659524c771f70fbddc24597c861c118", size = 235776358, upload-time = "2026-04-27T11:07:22.058Z" }, - { url = "https://files.pythonhosted.org/packages/47/bb/cb02d1e9679fce892a674f86caee25acc9ddd64d7dafa4cfe29e899993a8/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl", hash = "sha256:24d28892e210200f6e1bd13f699c42a74cd2bb7364c11248e2348f677c7f6dfb", size = 244415937, upload-time = "2026-04-27T11:07:48.135Z" }, + { url = "https://files.pythonhosted.org/packages/fd/36/40055709d6625a4f96fce5b82625860625d66723979ac092c10bef03d8c1/vllm-0.25.1-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:902be760af4c5ebfad8af5b8ea07a53ae14e5a6c839c8ab56da30581abb75ad2", size = 244036150, upload-time = "2026-07-14T08:58:44.737Z" }, + { url = "https://files.pythonhosted.org/packages/35/9d/c379618ce0abfc2679607d403c0f586b07e9c9c33d08c5bdd6196cb524e0/vllm-0.25.1-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:16fc7a28df1576eb6f7ca0455026551b8f9adb674c19c66059359ef3e964bd1e", size = 250100306, upload-time = "2026-07-14T08:59:06.476Z" }, ] [[package]] @@ -5392,16 +5555,22 @@ version = "0.2.4" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "platform_machine != 'x86_64' and sys_platform == 'linux'", + "sys_platform == 'win32'", + "sys_platform == 'darwin'", + "sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "apache-tvm-ffi" }, { name = "numpy" }, { name = "pydantic" }, - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "transformers" }, { name = "typing-extensions" }, ] wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/17/de9c2d0f2f37d78d4898ded27525bdb9d86817529c1c0dbd5b5ca9cc853f/xgrammar-0.2.4-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:2502403e179a73daccc7ea480d63f3a1a636d4b59738bb4dff0469135cc6e523", size = 23664403, upload-time = "2026-07-18T14:52:55.316Z" }, + { url = "https://files.pythonhosted.org/packages/99/c4/61c56f87eea9d49df76cf563099bc1136da10ba899507d2c32a02450d59c/xgrammar-0.2.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:bb13bda414b04caffb852a5203aa2519665668c7bb4868f1221612a7b9152945", size = 23614964, upload-time = "2026-07-18T14:52:57.841Z" }, { url = "https://files.pythonhosted.org/packages/18/10/baaac6c9ab3633bc5870d7cad893e6353b8345338610491679d1b61b5cfe/xgrammar-0.2.4-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b1d2bef13aec8122dfa6b96e3ece2cc0a42a732685f1d44f6248c305274e8eae", size = 44784353, upload-time = "2026-07-18T14:53:00.593Z" }, ] @@ -5469,8 +5638,12 @@ version = "4.15.4.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/8a/8e/0c8f17309549d2e5cde9a3ccefa6365437f1e7bafe71878eaf9478e47b18/z3_solver-4.15.4.0.tar.gz", hash = "sha256:928c29b58c4eb62106da51c1914f6a4a55d0441f8f48a81b9da07950434a8946", size = 5018600, upload-time = "2025-10-29T18:12:03.062Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/63/33/a3d5d2eaeb0f7b3174d57d405437eabb2075d4d50bd9ea0957696c435c7b/z3_solver-4.15.4.0-py3-none-macosx_13_0_arm64.whl", hash = "sha256:407e825cc9211f95ef46bdc8d151bf630e7ab2d62a21d24cd74c09cc5b73f3aa", size = 37052538, upload-time = "2025-10-29T18:11:46.233Z" }, + { url = "https://files.pythonhosted.org/packages/47/84/fd7ffac1551cd9f8d44fe41358f738be670fc4c24dfd514fab503f2cf3e7/z3_solver-4.15.4.0-py3-none-macosx_13_0_x86_64.whl", hash = "sha256:00bd10c5a6a5f6112d3a9a810d0799227e52f76caa860dafa5e00966bb47eb13", size = 39807925, upload-time = "2025-10-29T18:11:49.81Z" }, { url = "https://files.pythonhosted.org/packages/21/c9/bb51a96af0091324c81b803f16c49f719f9f6ea0b0bb52200f5c97ec4892/z3_solver-4.15.4.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e103a6f203f505b8b8b8e5c931cc407c95b61556512d4921c1ddc0b3f41b08e", size = 29268352, upload-time = "2025-10-29T18:11:53.032Z" }, { url = "https://files.pythonhosted.org/packages/bf/2e/0b49f7e4e53817cfb09a0f6585012b782dfe0b666e8abefcb4fac0570606/z3_solver-4.15.4.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:62c7e9cbdd711932301f29919ad9158de9b2f58b4d281dd259bbcd0a2f408ba1", size = 27226534, upload-time = "2025-10-29T18:11:55.59Z" }, + { url = "https://files.pythonhosted.org/packages/26/91/33de49538444d4aafbe47415c450c2f9abab1733e1226f276b496672f46c/z3_solver-4.15.4.0-py3-none-win32.whl", hash = "sha256:be3bc916545c96ffbf89e00d07104ff14f78336e55db069177a1bfbcc01b269d", size = 13191672, upload-time = "2025-10-29T18:11:58.424Z" }, + { url = "https://files.pythonhosted.org/packages/03/d6/a0b135e4419df475177ae78fc93c422430b0fd8875649486f9a5989772e6/z3_solver-4.15.4.0-py3-none-win_amd64.whl", hash = "sha256:00e35b02632ed085ea8199fb230f6015e6fc40554a6680c097bd5f060e827431", size = 16259597, upload-time = "2025-10-29T18:12:01.14Z" }, ] [[package]] diff --git a/ops/retriever-nightly/README.md b/ops/retriever-nightly/README.md index 41b4bdf8f5..aaff56ed96 100644 --- a/ops/retriever-nightly/README.md +++ b/ops/retriever-nightly/README.md @@ -3,261 +3,152 @@ # Retriever Nightly Launcher -This directory provides one public launcher for the library and ViDoRe v3 -benchmark suite: +`run-nightly.sh` runs the checked-in library and ViDoRe benchmark suite once. +It adds host locking, Git selection, standard dataset paths, artifact placement, +and optional Slack reporting around +[`retriever harness run-files`](../../nemo_retriever/harness/README.md). -| Workflow | Command | Code that runs | -| --- | --- | --- | -| Test the current checkout | `run-nightly.sh` | The current branch, including local changes. | -| Regular latest-main nightly | `run-nightly.sh --ref upstream/main` | Freshly fetched `upstream/main` in a clean worktree. | -| Reproduce an exact revision | `run-nightly.sh --ref ` | A clean worktree at an available local commit. | - -With no `--ref`, the launcher runs the checkout that contains the script and -does not fetch, switch branches, or reject local changes. With `--ref`, it -resolves one commit and creates or reuses an immutable detached worktree. It -exits after one terminal session summary. Recurrence is deliberately kept -outside its interface; the [daily `tmux` workflow](#daily-runs-with-tmux) is a -small shell loop rather than an installed scheduler. The launcher never merges -into or moves the invoking checkout, and it does not distribute datasets. +It does not install a scheduler or system service. -## Quick Start On A Standard Host +## Choose the Source -A workstation with `/datasets/nv-ingest` and writable `/raid/$USER` needs only -a Hugging Face token and, optionally, a Slack webhook to run its current -checkout: - -```bash -export HF_TOKEN=... -export SLACK_WEBHOOK_URL=https://hooks.slack.com/services/... -./ops/retriever-nightly/run-nightly.sh -``` - -That foreground command runs the current branch exactly as it exists, all -twelve benchmarks, and one terminal Slack summary. `SLACK_WEBHOOK_URL` is -optional; omit it to run without posting. Before GPU work, it checks ViDoRe -access. It then uses the checked-in `/datasets/nv-ingest` map, writes artifacts -under `/raid/$USER/retriever-nightly-artifacts`, and prints the terminal session -directory. No model-provider API key, `nightly.env`, `sudo`, systemd service, -or timer is required. - -The only required operator input on that host is `HF_TOKEN`. Set -`SLACK_WEBHOOK_URL` when the terminal summary should post to Slack. A custom -dataset map or artifact root is needed only when the host does not have the -standard paths. Use `--ref upstream/main` only when the operator deliberately -wants the newest clean upstream commit instead of the checked-out code. - -## Validate The Current PR Checkout - -### Prerequisites - -The supported v1 host is a Linux NVIDIA workstation with: - -- a NeMo Retriever Git checkout; -- `git`, Bash, `uv`, `flock`, `realpath`, and NVIDIA drivers available; -- access to the twelve benchmark datasets through `/datasets` or other local - paths; -- a Hugging Face read token in `HF_TOKEN` plus outbound HTTPS access to `huggingface.co`, - `cas-server.xethub.hf.co`, and `cas-bridge.xethub.hf.co` for ViDoRe - queries, qrels, and corpus metadata; and -- enough system RAM, local model cache, and artifact storage for the selected - runfiles. The complete batch suite is not validated on 128 GiB hosts. - -`uv` may be set with `RETRIEVER_UV_BIN`, discovered from `PATH`, or installed at -`$HOME/.local/bin/uv`. The locked `nemo_retriever` project selects Python 3.12 -and the repository dependencies. - -Batch mode starts its models locally, so it needs no model-provider API keys. -On a host with the standard `/datasets/nv-ingest` layout, the environment or -optional `nightly.env` accepts two secrets and optional path overrides: - -| Setting | Required | Purpose | +| Goal | Command | Source used | | --- | --- | --- | -| `HF_TOKEN` | every real launcher run | Read-only access for the automatic ViDoRe preflight. | -| `SLACK_WEBHOOK_URL` | no | Enables one terminal Slack post for real runs. | -| `RETRIEVER_DATASET_PATHS` | nonstandard hosts only | Path to a YAML file that replaces the checked-in `/datasets/nv-ingest` map. | -| `RETRIEVER_HARNESS_REFERENCE_FILE` | no | Current release snapshot shown beside nightly results in Slack without assigning a verdict. | +| Test the current checkout | `./ops/retriever-nightly/run-nightly.sh` | Current branch and local changes | +| Run the newest upstream code | `./ops/retriever-nightly/run-nightly.sh --ref upstream/main` | Freshly fetched `upstream/main` | +| Reproduce one commit | `./ops/retriever-nightly/run-nightly.sh --ref ` | Clean detached worktree | + +With no positional runfiles, the launcher runs twelve benchmarks: JP20, BO767, +Earnings, FinanceBench, and all eight ViDoRe v3 domains. Pass one or more +runfiles to run a smaller selection. -On hosts with a writable `/raid/$USER`, the launcher automatically keeps its -private configuration, artifacts, and managed Git checkouts there. -Other hosts use `$HOME`. +## First Run -Direct exports are the smallest configuration interface. A private file is -optional for operators who do not want to export the same values in every -shell. Already-exported supported settings take precedence over values in that -file. To create it: +On a standard host with `/datasets/nv-ingest` and writable `/raid/$USER`: ```bash -if [[ -d /raid/$USER && -w /raid/$USER ]]; then - RETRIEVER_NIGHTLY_ROOT=/raid/$USER -else - RETRIEVER_NIGHTLY_ROOT=$HOME -fi -RETRIEVER_NIGHTLY_CONFIG_DIR="$RETRIEVER_NIGHTLY_ROOT/.config/nemo-retriever/nightly" -mkdir -p "$RETRIEVER_NIGHTLY_CONFIG_DIR" -chmod 700 "$RETRIEVER_NIGHTLY_CONFIG_DIR" -test -e "$RETRIEVER_NIGHTLY_CONFIG_DIR/nightly.env" || \ - cp ops/retriever-nightly/nightly.env.example "$RETRIEVER_NIGHTLY_CONFIG_DIR/nightly.env" -chmod 600 "$RETRIEVER_NIGHTLY_CONFIG_DIR/nightly.env" -${EDITOR:-vi} "$RETRIEVER_NIGHTLY_CONFIG_DIR/nightly.env" +export HF_TOKEN=... +export SLACK_WEBHOOK_URL=https://hooks.slack.com/services/... # optional ``` -### Dataset Paths On A Nonstandard Host - -The checked-in `dataset_paths.datasets.yaml` describes the standard -`/datasets/nv-ingest` layout. Do not pass `--dataset-paths` on a host with that -layout. - -On any other host, `--dataset-paths` takes the path to a YAML configuration -file, not a dataset directory. Copy the complete twelve-dataset template -outside the checkout, replace its paths, and pass the resulting file: +Validate access and configuration before starting GPU work: ```bash -cp nemo_retriever/harness/dataset_paths.example.yaml \ - /raid/$USER/retriever-dataset-paths.yaml -${EDITOR:-vi} /raid/$USER/retriever-dataset-paths.yaml +./ops/retriever-nightly/run-nightly.sh --check-vidore-access +./ops/retriever-nightly/run-nightly.sh --dry-run ./ops/retriever-nightly/run-nightly.sh \ - --dataset-paths /raid/$USER/retriever-dataset-paths.yaml \ - --dry-run + --no-slack \ + nemo_retriever/harness/runfiles/jp20_beir.json ``` -For a JP20-only canary, the YAML file may contain only its local dataset and -ground-truth query paths: +All three commands should exit zero. Then run the full current checkout: -```yaml -schema_version: 1 -datasets: - jp20: - path: /raid/data/jp20 - query_file: /raid/data/jp20_query_gt.csv +```bash +./ops/retriever-nightly/run-nightly.sh ``` -Run that one benchmark by also supplying its runfile: +Use `--ref upstream/main` instead when the result must represent the newest +clean upstream commit. -```bash -./ops/retriever-nightly/run-nightly.sh \ - --dataset-paths /raid/data/retriever-dataset-paths.yaml \ - --no-slack \ - nemo_retriever/harness/runfiles/jp20_beir.json -``` +## Host Requirements -Keep the machine-local YAML outside the repository. For repeated runs, export -`RETRIEVER_DATASET_PATHS=/raid/$USER/retriever-dataset-paths.yaml` or set the -same value in the optional `nightly.env`; the command-line flag is simplest for -a one-off run. +The launcher expects: -The launcher loads the detected file only when it exists and uses its values as -defaults for settings that were not already exported. An existing secrets file -must be owned by the invoking user with mode `600`. The launcher does not -discover a repository `.env` file. `RETRIEVER_CONFIG_FILE` remains an optional -advanced path override. +- Linux with NVIDIA drivers +- Git, Bash, `uv`, `flock`, and `realpath` +- local access to the selected datasets +- `HF_TOKEN` and outbound Hugging Face access for real ViDoRe runs +- enough RAM, model cache, and artifact storage for the selected runfiles -Verify the token and read one byte from one remote parquet object in every -ViDoRe evaluation partition before starting GPU work: +The full batch suite is not validated on 128 GiB hosts. It starts models locally +and does not require a model-provider API key. -```bash -./ops/retriever-nightly/run-nightly.sh --check-vidore-access -``` +Common settings: -The access check does not download full parquet objects. A redirect failure -such as `302 -> 403 at cas-bridge.xethub.hf.co` is a Hugging Face/CAS delivery -failure; do not start the full suite until the check exits zero. +| Setting | When needed | +| --- | --- | +| `HF_TOKEN` | Every real launcher run; a read token is sufficient | +| `SLACK_WEBHOOK_URL` | Only when the terminal result should post to Slack | +| `RETRIEVER_HARNESS_REFERENCE_FILE` | When Slack should show the current RC beside matching results | +| `RETRIEVER_DATASET_PATHS` | Hosts without the standard dataset layout | +| `RETRIEVER_NIGHTLY_ROOT` | Hosts that need a nondefault config, artifact, and checkout root | -Then preflight the complete twelve-benchmark suite without starting ingest or -query: +Direct exports are the smallest setup. Persistent values may be copied from +[`nightly.env.example`](nightly.env.example) into: -```bash -./ops/retriever-nightly/run-nightly.sh --dry-run +```text +/.config/nemo-retriever/nightly/nightly.env ``` -Inspect the resulting `session_summary.json` and child plans. Dry-runs never -post to Slack. A real run posts when `SLACK_WEBHOOK_URL` is configured; use -`--no-slack` for a real functional test that must not post. The launcher prints -the timestamped session directory on success or terminal harness failure. +The launcher uses `/raid/$USER` as the nightly root when it is writable and +`$HOME` otherwise. The optional file must be owned by the current user with mode +`600`. Existing exported values take precedence. The launcher never loads a +repository `.env`. -Use one positional runfile for a smaller real canary before the full run: +## Use Nonstandard Dataset Paths + +The default map, +[`dataset_paths.datasets.yaml`](dataset_paths.datasets.yaml), uses +`/datasets/nv-ingest` for benchmark corpora. JP20 evaluation uses the checked-in +`data/jp20_query_gt.csv`. On another host, copy the portable template outside +the repository: ```bash -./ops/retriever-nightly/run-nightly.sh \ - --no-slack \ - nemo_retriever/harness/runfiles/jp20_beir.json +cp nemo_retriever/harness/dataset_paths.example.yaml \ + /local/path/to/dataset_paths.yaml +${EDITOR:-vi} /local/path/to/dataset_paths.yaml ``` -The launcher performs the ViDoRe access preflight before every real invocation, -including a JP20-only canary, so `HF_TOKEN` is still required for this command. - -Run the complete suite from the current checkout with no positional runfiles: +Pass the YAML file—not a dataset directory: ```bash -./ops/retriever-nightly/run-nightly.sh +./ops/retriever-nightly/run-nightly.sh \ + --dataset-paths /local/path/to/dataset_paths.yaml \ + --dry-run ``` -If `SLACK_WEBHOOK_URL` is configured, that real run posts its terminal summary. -Add `--no-slack` only when the full run is itself a functional test that must -not post. +For repeated runs, export the same path as `RETRIEVER_DATASET_PATHS`. ## Git Selection -With no `--ref`, `run-nightly.sh`: - -1. selects the Git checkout that contains the launcher; -2. runs its current branch and working tree without fetching or switching; -3. permits tracked, staged, and untracked changes; and -4. runs the ViDoRe access check before real GPU work. +Without `--ref`, the launcher runs the checkout containing the script exactly +as it exists. It records the commit and dirty state in the session. Dirty Slack +reports are prefixed with `[LOCAL CHANGES]`. -Every session records the checkout's HEAD in `run_commit` and whether it had -local changes in `working_tree_dirty`. Dirty runs also write -`source_worktree_status.txt` in the session directory and prefix the Slack -title with `[LOCAL CHANGES]`. The status artifact records paths and Git state, -not file contents, so a dirty run is intentionally identifiable but not fully -reproducible. +With `--ref`: -`--ref REF` requests a clean committed run. A local branch, tag, or SHA is -resolved without fetching. A remote branch such as `upstream/main` is fetched -first, then resolved fail-closed; a fetch failure never falls back to a stale -remote-tracking commit. The selected commit runs from an immutable detached -worktree named `commit-`. The launcher never runs `git pull`, merges -into the invoking checkout, or moves its current branch. +1. Local branches, tags, and SHAs resolve without fetching. +2. Remote branches such as `upstream/main` are fetched first. +3. Fetch failure stops the run instead of using a stale remote-tracking commit. +4. The selected commit runs in a clean detached worktree. -Immutable worktrees and one shared `uv` project environment live under the -detected nightly root at `retriever-nightly-checkouts`; on `/raid` hosts this is -`/raid/$USER/retriever-nightly-checkouts`. The seven most recently used SHA -worktrees are retained. Modified managed worktrees are never deleted -automatically. +Managed worktrees and a shared `uv` environment live under the nightly root. +The launcher retains the seven most recently used clean worktrees and never +moves or merges the invoking checkout. -The one-time latest-main setup must provide an `upstream` remote. Request a -clean latest-main preflight explicitly: +Configure `upstream` once if needed: ```bash git remote get-url upstream >/dev/null 2>&1 || \ git remote add upstream https://github.com/NVIDIA/NeMo-Retriever.git -./ops/retriever-nightly/run-nightly.sh --ref upstream/main --dry-run ``` -The dry-run fetches and selects the latest commit but skips remote access and -GPU execution. Use `run-nightly.sh --ref upstream/main --check-vidore-access` -to validate that commit and the machine credentials without starting a -session. The complete latest-main suite is: - -```bash -./ops/retriever-nightly/run-nightly.sh --ref upstream/main -``` +Use `--ref HEAD` to ignore local changes and run only the current commit. -Use `--ref HEAD` when local changes should be ignored and only the current -commit should run, or `--ref ` to reproduce an earlier run. The -invoking checkout may itself be dirty because the selected ref always runs in -a separate clean worktree. Ignored cache files do not mark a run dirty. +## Run Daily with `tmux` -## Daily Runs With `tmux` - -The launcher remains a one-shot developer tool. Use a transparent shell loop -inside `tmux` when a workstation should start the latest `upstream/main` -nightly approximately every 24 hours: +Start a session: ```bash tmux new -s retriever-nightly +``` +Inside it, export the environment and run a serial 24-hour loop: + +```bash export HF_TOKEN=... export SLACK_WEBHOOK_URL=https://hooks.slack.com/services/... +export RETRIEVER_HARNESS_REFERENCE_FILE=/path/to/current-release.json interval=86400 while true; do @@ -270,99 +161,54 @@ while true; do done ``` -Enter the exports inside the new `tmux` session so they do not depend on an -older tmux server's saved environment. Detach with `Ctrl-b d`, inspect it with -`tmux attach -t retriever-nightly`, and stop it with `tmux kill-session -t -retriever-nightly`. +Each iteration fetches the newest `upstream/main`. Runs do not overlap. If one +exceeds 24 hours, the next starts after it finishes. -The loop is serial: runs never overlap. It targets a 24-hour start-to-start -interval; if one run exceeds 24 hours, the next begins only after it finishes. -The loop survives an SSH disconnect but not a workstation reboot. This is an -operator-owned development workflow, not an installed service or timer. +Detach with `Ctrl-b d`, reconnect with `tmux attach -t retriever-nightly`, and +stop the loop with `tmux kill-session -t retriever-nightly`. The session +survives an SSH disconnect but not a host reboot. -While testing an unmerged branch, omit `--ref upstream/main` to run that -checkout on every iteration. Keep `--ref upstream/main` for the production -loop so every iteration fetches and runs the newest clean upstream commit. +To exercise an unmerged branch repeatedly, omit `--ref upstream/main`. -## Slack Report +## Slack and the Current Release -To enable Slack for real runs, export the incoming-webhook URL or place it in -the optional mode-`600` `nightly.env`: +When `SLACK_WEBHOOK_URL` is set, a real run posts once after +`session_summary.json` exists. Dry-runs and access checks never post. Use +`--no-slack` for a real canary. -```bash -export SLACK_WEBHOOK_URL=https://hooks.slack.com/services/... -``` +Set `RETRIEVER_HARNESS_REFERENCE_FILE` to show the current release beside +matching nightly results. The +[harness reporting guide](../../nemo_retriever/harness/README.md#report-completed-results) +defines the small external JSON format. + +The report presents observed values with their GPU context. It does not enforce +a score, assign a verdict, append history, or modify the release file. To move +to a new RC, replace the external file's label and values. + +## Runtime Behavior + +- A nonblocking host-local lock prevents overlapping launcher processes. +- Every run uses batch mode; each benchmark runs in a fresh child process. +- A failed child is recorded and later benchmarks still run. +- Each child has a six-hour wall-time limit. +- A configured Slack report is attempted once after a terminal session exists. +- `VLLM_DEEP_GEMM_WARMUP` defaults to `skip` unless the caller sets it. + +The command returns the harness status. If the harness succeeds but Slack +posting fails, it returns the Slack command's nonzero status. + +## Troubleshooting + +**ViDoRe access fails:** rerun `--check-vidore-access`. A final `403` from a +Hugging Face CAS host can indicate proxy, firewall, or egress policy rather than +an invalid token. Compare from another network before rotating credentials. + +**A large run stalls with high system memory:** batch ingest can materialize +page payloads in Python. Capture the child `run.log`, `status.json`, process RSS, +and Ray task summary, then reproduce only that runfile. Do not classify it as a +GPU OOM without GPU or kernel allocation evidence. -The URL itself is the Slack switch. If it is unset or empty, real runs complete -without posting. If it is set, the launcher validates it before expensive work -and posts once after `session_summary.json` exists. An invalid configured URL -fails preflight. Pass `--no-slack` for canaries or other real functional tests; -that flag suppresses webhook validation and posting. Dry-runs and access checks -never post. The launcher removes the URL from the benchmark child environment -and exposes it only to the final Slack command. - -To show the latest RC beside each matching nightly result, set -`RETRIEVER_HARNESS_REFERENCE_FILE` to the external snapshot documented in the -[harness README](../../nemo_retriever/harness/README.md#post-results-to-slack). -The harness reads only that configured snapshot; it does not append history or -apply pass/fail policy. When the next RC is ready, replace the snapshot's label -and observed values. - -Configuration precedence is, from highest to lowest: command-line flags, -supported variables already exported when the launcher starts, values loaded -from `RETRIEVER_CONFIG_FILE`, and launcher defaults. Run the launcher with -`--help` for its supported interface. - -## Runtime Contract - -The launcher takes a nonblocking host-local lock, forces batch mode, and runs -12 checked-in runfiles as one session: JP20, BO767, Earnings, FinanceBench, and -all eight public ViDoRe v3 domains. Real sessions execute each child in a fresh -spawned process so Ray and materialized dataframe memory are released before -the next benchmark; the parent still writes one terminal session summary. -Dry-runs stay in the parent process because they do not materialize datasets. A -configured Slack report runs once after a terminal session summary exists. A -`.slack_post_attempted` marker prevents a second attempt for the same session. -Incoming webhooks do not provide an idempotency key, so ambiguous transport -failures require human inspection. - -The Slack report keeps the library benchmarks detailed and collapses the full -ViDoRe v3 suite into total ingest time, aggregate pages/sec, macro-average -Recall@5 and nDCG@10 for the English and complete suites, and one accuracy row -per domain. Per-domain throughput and timing remain in the session artifacts. - -If one runtime child fails, `run-files` continues the remaining datasets and -writes a failed session summary. When Slack is configured, the launcher still -attempts one report and returns the harness status. If the harness succeeds but -Slack fails, it returns the Slack command's nonzero status. Process-isolated -children also have a six-hour wall-time limit; a child that exceeds it is -terminated, recorded as failed, and does not prevent later datasets from running. - -The launcher defaults `VLLM_DEEP_GEMM_WARMUP=skip` unless the caller explicitly -sets another vLLM-supported mode. This skips the optional compatibility-sensitive -warmup without disabling DeepGEMM kernels. It intentionally does not set -`VLLM_USE_DEEP_GEMM=0` or `VLLM_MOE_USE_DEEP_GEMM=0`. Set the warmup variable -explicitly, for example to `full`, only when validating another mode. This -matches the reliability direction under discussion in -[NVIDIA/NeMo-Retriever PR #2292](https://github.com/NVIDIA/NeMo-Retriever/pull/2292). - -## Troubleshooting Preflight And Host Memory - -`--check-vidore-access` validates the configured token, reads repository -metadata, and follows the same Hugging Face redirects used by `datasets` while -reading one byte from one parquet object in each of the queries, qrels, and -corpus partitions. If the token is valid but the final CAS host returns `403`, compare -the same check -from another network before rotating credentials. Success elsewhere points to -host proxy, firewall, or egress policy; failure from multiple networks should -be escalated with the named dataset object to Hugging Face or ViDoRe. - -Batch ingest currently materializes each terminal Ray dataset in Python. -High-resolution page payloads can therefore consume substantially more system -RAM than the final LanceDB table. The nightly's per-run process boundary -prevents that memory from accumulating across the twelve children, but an -individual large benchmark must still fit on the host. If Ray reports the -dataset and VDB write complete while a child remains idle at high RSS, capture -`run.log`, `status.json`, process RSS, and the Ray task summary. Retry only that -runfile as a focused reproduction; do not classify the symptom as GPU OOM -unless the GPU process or kernel logs show an actual allocation failure. +**Configuration is unclear:** run +`./ops/retriever-nightly/run-nightly.sh --help`. CLI flags override exported +values, which override the optional config file, which overrides launcher +defaults. diff --git a/ops/retriever-nightly/SECOND_HOST_VALIDATION.md b/ops/retriever-nightly/SECOND_HOST_VALIDATION.md deleted file mode 100644 index e225f19206..0000000000 --- a/ops/retriever-nightly/SECOND_HOST_VALIDATION.md +++ /dev/null @@ -1,170 +0,0 @@ - - - -# Portable Nightly Second-Host Validation - -Use this checklist on a separate Linux NVIDIA workstation before handing the -launcher to additional teammates. The review uses the pushed feature branch, -keeps datasets and artifacts outside the repository, validates the functional -path first, and finishes with one foreground full-suite command and terminal -Slack post. It does not install or start an operating-system service. - -## 1. Create a Local Review Branch - -In an existing clone whose `origin` points to the contributor fork: - -```bash -git fetch origin jioffe502/retriever-nightly-vidore-v3 -git switch --create review/portable-nightly \ - --track origin/jioffe502/retriever-nightly-vidore-v3 -git status --short -git rev-parse HEAD -``` - -`git status --short` must be empty. If this is a new clone, add the NVIDIA -repository as `upstream` for later comparisons: - -```bash -git remote add upstream https://github.com/NVIDIA/NeMo-Retriever.git -``` - -Every validation command below omits `--ref`, so the launcher runs exactly the -checked-out review branch without fetching or moving it. Keeping the review -checkout clean makes these validation results attributable to its HEAD. The -production latest-main workflow explicitly passes `--ref upstream/main`. - -## 2. Prepare the Host - -Confirm that `uv` and the NVIDIA driver are available: - -```bash -uv --version -nvidia-smi -``` - -This host has the standard `/datasets/nv-ingest` layout and `/raid/$USER`, so -the checked-in dataset map and launcher path defaults apply. Export the two -secrets; a read Hugging Face token is sufficient: - -```bash -export HF_TOKEN=... -export SLACK_WEBHOOK_URL=https://hooks.slack.com/services/... -export RETRIEVER_HARNESS_REFERENCE_FILE=/datasets/nv-ingest/nrl_rc_baselines/nrl-26.05-bo767.json -``` - -The reference file contains the currently selected RC observations shown beside -matching nightly results. Replace the path or file contents when advancing to -the next RC; the harness does not maintain reference history or apply a verdict. - -The launcher writes artifacts under `/raid/$USER/retriever-nightly-artifacts`. -No configuration file is required. On a host without the standard dataset -layout, copy and edit `nemo_retriever/harness/dataset_paths.example.yaml`, then -pass the resulting YAML file to `--dataset-paths`. The option takes the YAML -file itself, not a dataset directory: - -```bash -cp nemo_retriever/harness/dataset_paths.example.yaml \ - /raid/$USER/retriever-dataset-paths.yaml -${EDITOR:-vi} /raid/$USER/retriever-dataset-paths.yaml -./ops/retriever-nightly/run-nightly.sh \ - --dataset-paths /raid/$USER/retriever-dataset-paths.yaml \ - --dry-run -``` - -## 3. Verify ViDoRe Evaluation Access - -Before starting GPU work, validate the configured token and read one byte from -one remote parquet object in each of the queries, qrels, and corpus partitions: - -```bash -./ops/retriever-nightly/run-nightly.sh --check-vidore-access -``` - -The command should exit zero and report access for all eight ViDoRe v3 -datasets. It does not download the full objects. Do not start the complete -suite if this check reports a Hugging Face or CAS redirect failure. - -## 4. Preflight All Twelve Benchmarks - -```bash -./ops/retriever-nightly/run-nightly.sh --dry-run -``` - -The command should exit zero, report a new timestamped session directory, and -write a `session_summary.json` with `dry_run: true`, twelve runs, and a -`run_commit` matching `git rev-parse HEAD`. `isolate_runs` is `false` because -the dry-run does not materialize batch data. - -## 5. Run the JP20 Canary - -```bash -./ops/retriever-nightly/run-nightly.sh \ - --no-slack \ - nemo_retriever/harness/runfiles/jp20_beir.json -``` - -Confirm that the command exits zero and the session summary contains one -successful run with `isolate_runs: true`. Real `run-files` sessions isolate -each sequential child automatically. The launcher defaults the optional -DeepGEMM warmup to `skip`; no host setting is needed. - -## 6. Capture the Handoff Evidence - -For the dry-run and JP20 sessions, record: - -- the launcher exit code and printed session directory; -- `success`, `exit_code`, `dry_run`, `isolate_runs`, `run_commit`, and the - number of `runs` in `session_summary.json`; -- any failed child name and its artifact directory; and -- the GPU model and driver from `nvidia-smi`. - -The functional validation is complete when the ViDoRe access check, twelve-run -dry-run, and real JP20 canary succeed with terminal summaries attributed to the -review branch commit. These steps do not post to Slack. - -## 7. Run the Complete Nightly and Slack Report - -From the clean draft checkout, start all twelve benchmarks with one command: - -```bash -./ops/retriever-nightly/run-nightly.sh -``` - -The launcher runs the four library benchmarks and all eight ViDoRe v3 domains. -Each child runs in a fresh process; failures do not prevent later children from -running, and the parent writes one terminal `session_summary.json`. Because -`SLACK_WEBHOOK_URL` is exported, that terminal summary posts once. Confirm the -full `run_commit`, twelve child results, final command exit status, and Slack -message. The process remains attached to the invoking shell; use the host's -normal session manager if it must survive a disconnected terminal. - -## 8. Start The Post-Merge Daily Workflow - -After the launcher is merged to `upstream/main`, keep the one-shot launcher in -a transparent 24-hour loop inside `tmux`: - -```bash -tmux new -s retriever-nightly - -export HF_TOKEN=... -export SLACK_WEBHOOK_URL=https://hooks.slack.com/services/... -export RETRIEVER_HARNESS_REFERENCE_FILE=/datasets/nv-ingest/nrl_rc_baselines/nrl-26.05-bo767.json - -interval=86400 -while true; do - started="$(date +%s)" - ./ops/retriever-nightly/run-nightly.sh --ref upstream/main - elapsed=$(( $(date +%s) - started )) - if (( elapsed < interval )); then - sleep "$(( interval - elapsed ))" - fi -done -``` - -The explicit remote ref fetches the newest `upstream/main` on each iteration, -preflights ViDoRe access, runs the suite, and posts once when -`SLACK_WEBHOOK_URL` is set. Enter the exports inside the new tmux session, -detach with `Ctrl-b d`, and inspect it later with `tmux attach -t -retriever-nightly`. The serial loop does not overlap runs. It survives SSH -disconnects but must be restarted after a workstation reboot; no service or -timer is installed. diff --git a/ops/retriever-nightly/dataset_paths.datasets.yaml b/ops/retriever-nightly/dataset_paths.datasets.yaml index f62fce312a..395a6f6fbc 100644 --- a/ops/retriever-nightly/dataset_paths.datasets.yaml +++ b/ops/retriever-nightly/dataset_paths.datasets.yaml @@ -5,7 +5,7 @@ schema_version: 1 datasets: jp20: path: /datasets/nv-ingest/jp20 - query_file: /datasets/nv-ingest/ground_truths/jp20_query_gt.csv + query_file: ../../data/jp20_query_gt.csv bo767: path: /datasets/nv-ingest/bo767 query_file: /datasets/nv-ingest/ground_truths/bo767_query_gt.csv