Incremental Docs-to-Chroma RAG Sync with Deletions and a Retrieval Smoke Test

id: chroma-incremental-docs-rag-sync
namespace: company.ai

description: |
  Keep a Chroma collection in sync with a Markdown documentation repository
  without re-embedding everything. Chunks get content-addressed ids, so only
  new or edited sections are embedded and upserted, chunks of deleted or
  edited sections are removed, and a retrieval smoke test fails the run if
  a known question stops finding its page. Embeddings are computed locally
  with FastEmbed, so no API key is needed.

labels:
  team: ai-platform
  domain: rag

concurrency:
  limit: 1
  behavior: QUEUE

inputs:
  - id: repository_url
    type: STRING
    displayName: Docs repository
    description: Git repository that holds the documentation.
    defaults: https://github.com/encode/httpx

  - id: git_ref
    type: STRING
    displayName: Branch or tag
    description: Branch or tag to index.
    defaults: master

  - id: docs_glob
    type: STRING
    displayName: Docs glob
    description: Glob, relative to the repository root, selecting the Markdown files
      to index.
    defaults: "docs/**/*.md"

  - id: collection
    type: STRING
    displayName: Chroma collection
    defaults: httpx_docs

  - id: smoke_question
    type: STRING
    displayName: Smoke test question
    description: A question the docs must keep answering. Its nearest chunks are
      retrieved after every sync.
    defaults: How do I configure timeouts?

  - id: smoke_expected_source
    type: STRING
    displayName: Expected source
    description: Part of the file path that must appear among the top results for
      the smoke test question, otherwise the run fails.
    defaults: timeouts

  - id: full_resync
    type: BOOL
    displayName: Full resync
    description: Ignore the saved manifest and re-embed every chunk, for example
      after the collection was wiped.
    defaults: false

variables:
  embedding_model: BAAI/bge-small-en-v1.5
  manifest_key: "chroma_sync_{{ inputs.collection }}"
  nl: "\n"

triggers:
  - id: nightly
    type: io.kestra.plugin.core.trigger.Schedule
    description: Nightly catch-up sync of the default branch.
    cron: "0 2 * * *"

  - id: docs_changed
    type: io.kestra.plugin.core.trigger.Webhook
    description: Call from CI after a docs merge to sync within minutes. Pass
      git_ref in the body to index a release tag.
    key: replace-me-with-a-long-random-key

tasks:
  - id: manifest
    type: io.kestra.plugin.core.kv.Get
    description: Chunk ids and embedding model stored by the last successful sync.
    key: "{{ render(vars.manifest_key) }}"
    errorOnMissing: false

  - id: workspace
    type: io.kestra.plugin.core.flow.WorkingDirectory
    description: Clone and plan share one working directory. The retry sits here
      rather than on the inner tasks, so a network blip on the clone or the pip
      install reruns both, and an exhausted retry fails the run instead of
      leaving it RUNNING (kestra-io/kestra#21109).
    retry:
      type: constant
      interval: PT30S
      maxAttempts: 3
    tasks:
      - id: clone
        type: io.kestra.plugin.git.Clone
        url: "{{ inputs.repository_url }}"
        branch: "{{ inputs.git_ref }}"
        depth: 1
        directory: repo

      - id: plan
        type: io.kestra.plugin.scripts.python.Script
        description: Chunk the docs by heading, diff the content-addressed chunk ids
          against the manifest, embed only what is new, and embed the smoke test
          question.
        taskRunner:
          type: io.kestra.plugin.scripts.runner.docker.Docker
        containerImage: python:3.12-slim
        beforeCommands:
          - pip install --quiet --disable-pip-version-check fastembed==0.7.3
        outputFiles:
          - records.ion
        env:
          DOCS_GLOB: "{{ inputs.docs_glob }}"
          REPOSITORY_URL: "{{ inputs.repository_url }}"
          GIT_REF: "{{ inputs.git_ref }}"
          MODEL: "{{ vars.embedding_model }}"
          MANIFEST: "{{ (outputs.manifest.value ?? '') | toJson }}"
          FULL_RESYNC: "{{ inputs.full_resync }}"
          SMOKE_QUESTION: "{{ inputs.smoke_question }}"
          HF_HUB_DISABLE_PROGRESS_BARS: "1"
        script: |
          import glob
          import hashlib
          import json
          import os
          import re
          import time
          import warnings

          warnings.filterwarnings("ignore", message="Cannot enable progress bars")
          MAX_CHARS = 1200
          repo_url = os.environ["REPOSITORY_URL"].removesuffix(".git")
          ref = os.environ["GIT_REF"]
          model_name = os.environ["MODEL"]

          def as_obj(value, default):
              while isinstance(value, str):
                  if value == "":
                      return default
                  try:
                      value = json.loads(value)
                  except ValueError:
                      return default
              return default if value is None else value

          def chunks_of(text):
              # Split on Markdown headings, keep the heading trail as context,
              # then cut long sections on paragraph boundaries.
              trail, section, out = [], [], []
              def flush():
                  body = "\n".join(section).strip()
                  section.clear()
                  if not body:
                      return
                  title = " > ".join(t for _, t in trail)
                  piece = ""
                  for para in re.split(r"\n\s*\n", body):
                      if piece and len(piece) + len(para) > MAX_CHARS:
                          out.append((title, piece.strip()))
                          piece = ""
                      piece += para + "\n\n"
                  if piece.strip():
                      out.append((title, piece.strip()))
              in_code = False
              for line in text.splitlines():
                  if line.lstrip().startswith("```"):
                      in_code = not in_code
                  heading = None if in_code else re.match(r"^(#{1,4})\s+(.*)", line)
                  if heading:
                      flush()
                      level = len(heading.group(1))
                      trail[:] = [h for h in trail if h[0] < level] + [(level, heading.group(2).strip())]
                  section.append(line)
              flush()
              return out

          chunks = {}
          files = sorted(glob.glob(os.path.join("repo", os.environ["DOCS_GLOB"]), recursive=True))
          for path in files:
              rel = os.path.relpath(path, "repo").replace(os.sep, "/")
              with open(path, encoding="utf-8", errors="replace") as handle:
                  pieces = chunks_of(handle.read())
              for index, (section, body) in enumerate(pieces):
                  # Same text in the same file always gets the same id, wherever it moves.
                  chunk_id = hashlib.sha1("{}\n{}\n{}".format(rel, section, body).encode()).hexdigest()
                  chunks.setdefault(chunk_id, {
                      "source": rel,
                      "section": section,
                      "chunk_index": index,
                      "url": "{}/blob/{}/{}".format(repo_url, ref, rel),
                      "document": body,
                  })

          previous = as_obj(os.environ["MANIFEST"], {})
          model_changed = previous.get("model") not in (None, model_name)
          old_ids = set() if os.environ["FULL_RESYNC"].lower() == "true" or model_changed else set(previous.get("ids", []))
          new_ids = set(chunks)
          to_add = sorted(new_ids - old_ids)
          to_delete = sorted(set(previous.get("ids", [])) - new_ids)

          from fastembed import TextEmbedding
          started = time.time()
          model = TextEmbedding(model_name=model_name)
          texts = ["{}\n{}".format(chunks[i]["section"], chunks[i]["document"]) for i in to_add]
          vectors = list(model.embed(texts, batch_size=64)) if texts else []
          question_vector = next(iter(model.query_embed([os.environ["SMOKE_QUESTION"]]))).tolist()
          embed_seconds = round(time.time() - started, 1)

          with open("records.ion", "w", encoding="utf-8") as out:
              for chunk_id, vector in zip(to_add, vectors):
                  c = chunks[chunk_id]
                  out.write(json.dumps({
                      "id": chunk_id,
                      "embedding": [round(float(x), 6) for x in vector],
                      "document": c["document"],
                      "metadata": {k: c[k] for k in ("source", "section", "chunk_index", "url")} | {"ref": ref},
                  }) + "\n")

          changed_sources = sorted({chunks[i]["source"] for i in to_add})
          removed_sources = sorted({s for s in previous.get("sources", {}) if s not in {c["source"] for c in chunks.values()}})
          shown = changed_sources[:8]
          changed_summary = ", ".join(shown) + (" and {} more".format(len(changed_sources) - len(shown)) if len(changed_sources) > len(shown) else "")
          print("files={} chunks={} add={} delete={} unchanged={} model_changed={} embed_seconds={}".format(
              len(files), len(new_ids), len(to_add), len(to_delete), len(new_ids & old_ids), model_changed, embed_seconds))
          for source in changed_sources:
              print("  changed: " + source)
          for source in removed_sources:
              print("  removed file: " + source)

          print("::" + json.dumps({"outputs": {
              "files": len(files),
              "total_chunks": len(new_ids),
              "add_count": len(to_add),
              "delete_ids": to_delete,
              "unchanged": len(new_ids & old_ids),
              "changed_sources": changed_sources,
              "changed_summary": changed_summary,
              "removed_sources": removed_sources,
              "embed_seconds": embed_seconds,
              "question_vector": question_vector,
              "manifest": {
                  "model": model_name,
                  "ref": ref,
                  "ids": sorted(new_ids),
                  "sources": sorted({c["source"] for c in chunks.values()}),
              },
          }}) + "::")

  - id: collection
    type: io.kestra.plugin.chroma.CreateCollection
    description: Idempotent; creates the collection with cosine distance on the first run.
    baseUrl: "{{ secret('CHROMA_URL') }}"
    collectionName: "{{ inputs.collection }}"
    distance: COSINE
    getOrCreate: true
    metadata:
      embedding_model: "{{ vars.embedding_model }}"
      source: "{{ inputs.repository_url }}"

  - id: upsert_changes
    type: io.kestra.plugin.core.flow.If
    condition: "{{ outputs.plan.vars.add_count > 0 }}"
    then:
      - id: upsert
        type: io.kestra.plugin.chroma.Upsert
        baseUrl: "{{ secret('CHROMA_URL') }}"
        collectionName: "{{ inputs.collection }}"
        from: "{{ outputs.plan.outputFiles['records.ion'] }}"
        batchSize: 200
        retry:
          type: exponential
          interval: PT5S
          maxInterval: PT1M
          maxAttempts: 4

  - id: delete_stale
    type: io.kestra.plugin.core.flow.If
    description: Remove chunks whose text no longer exists, so answers never cite
      deleted or outdated sections.
    condition: "{{ outputs.plan.vars.delete_ids | length > 0 }}"
    then:
      - id: delete
        type: io.kestra.plugin.chroma.Delete
        baseUrl: "{{ secret('CHROMA_URL') }}"
        collectionName: "{{ inputs.collection }}"
        ids: "{{ outputs.plan.vars.delete_ids }}"

  - id: smoke_query
    type: io.kestra.plugin.chroma.Query
    description: Retrieve the nearest chunks for the canary question from the
      freshly synced collection.
    baseUrl: "{{ secret('CHROMA_URL') }}"
    collectionName: "{{ inputs.collection }}"
    vector: "{{ outputs.plan.vars.question_vector }}"
    maxResults: 5
    include:
      - metadatas
      - distances
    fetchType: FETCH

  - id: smoke_check
    type: io.kestra.plugin.core.flow.If
    description: Fail before the manifest is saved if retrieval regressed, so the
      next run retries the same diff.
    condition: "{{ outputs.smoke_query.rows | jq('[.[].metadata.source |
      contains(\"' ~ inputs.smoke_expected_source ~ '\")] | any') | first }}"
    then:
      - id: smoke_passed
        type: io.kestra.plugin.core.log.Log
        message: "Smoke test passed: top hit {{
          outputs.smoke_query.rows[0].metadata.source }} (distance {{
          outputs.smoke_query.rows[0].distance | numberFormat('0.000') }})."
    else:
      - id: smoke_failed
        type: io.kestra.plugin.core.execution.Fail
        errorMessage: "Retrieval smoke test failed: '{{ inputs.smoke_question }}' did
          not return a chunk from '{{ inputs.smoke_expected_source }}'. Top
          sources: {{ outputs.smoke_query.rows | jq('[.[].metadata.source]') |
          first | join(', ') }}"

  - id: save_manifest
    type: io.kestra.plugin.core.kv.Set
    description: Record what the collection now contains. Saved last, so a failed
      run is retried in full.
    key: "{{ render(vars.manifest_key) }}"
    value: "{{ outputs.plan.vars.manifest | toJson }}"
    kvType: JSON
    overwrite: true

  - id: report
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "text": {{ (':books: *Docs RAG sync* `' ~ inputs.collection ~ '` @ `' ~ inputs.git_ref ~ '`' ~ vars.nl ~ outputs.plan.vars.add_count ~ ' chunks embedded, ' ~ (outputs.plan.vars.delete_ids | length) ~ ' removed, ' ~ outputs.plan.vars.unchanged ~ ' unchanged (' ~ outputs.plan.vars.total_chunks ~ ' in the collection, from ' ~ outputs.plan.vars.files ~ ' files)' ~ vars.nl ~ ((outputs.plan.vars.changed_sources | length) > 0 ? 'Changed: ' ~ outputs.plan.vars.changed_summary ~ vars.nl : '') ~ ((outputs.plan.vars.removed_sources | length) > 0 ? 'Files gone: ' ~ (outputs.plan.vars.removed_sources | join(', ')) ~ vars.nl : '') ~ 'Smoke test: ' ~ inputs.smoke_question ~ ' -> ' ~ outputs.smoke_query.rows[0].metadata.source ~ ' (distance ' ~ (outputs.smoke_query.rows[0].distance | numberFormat('0.000')) ~ ')' ~ vars.nl ~ 'Execution: ' ~ execution.id) | toJson }}
      }

outputs:
  - id: embedded
    type: INT
    description: Chunks embedded and upserted by this run.
    value: "{{ outputs.plan.vars.add_count }}"

  - id: removed
    type: INT
    description: Stale chunks deleted by this run.
    value: "{{ outputs.plan.vars.delete_ids | length }}"

  - id: changed_sources
    type: JSON
    description: Files that contributed new or edited chunks.
    value: "{{ outputs.plan.vars.changed_sources }}"

  - id: smoke_top_hit
    type: STRING
    description: Source of the best match for the smoke test question.
    value: "{{ outputs.smoke_query.rows[0].metadata.source }}"

errors:
  - id: alert_on_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "text": {{ (':rotating_light: *Docs RAG sync FAILED* for `' ~ inputs.collection ~ '` @ `' ~ inputs.git_ref ~ '`. The collection may be partially updated; the manifest was not saved, so the next run retries the same diff. Execution: ' ~ execution.id) | toJson }}
      }