Chroma Vector Store Dangling Document and Quarantine Eviction Guard

id: chroma-vector-store-dangling-document-and-quarantine-eviction-guard
namespace: company.ai

inputs:
  - id: base_url
    type: STRING
    displayName: Chroma Base URL
    description: HTTP endpoint URL of the Chroma vector store.
    defaults: "http://localhost:8000"

  - id: tenant
    type: STRING
    displayName: Chroma Tenant
    description: Chroma tenant identifier.
    defaults: "default_tenant"

  - id: database
    type: STRING
    displayName: Chroma Database
    description: Target database name within the Chroma tenant.
    defaults: "default_database"

  - id: collection_name
    type: STRING
    displayName: Collection Name
    description: Name of the Chroma vector collection to audit and evict dangling
      records from.
    defaults: "customer-knowledge-base"

  - id: batch_size
    type: INT
    displayName: Inspection Batch Size
    description: Maximum number of candidate vectors to retrieve and evaluate in a
      single inspection pass.
    defaults: 200

  - id: stale_retention_days
    type: INT
    displayName: Stale Retention Grace Period (Days)
    description: Maximum age in days before soft-deleted or tombstoned embeddings
      are eligible for permanent eviction.
    defaults: 7

  - id: enable_quarantine_archive
    type: BOOL
    displayName: Archive Quarantined Embeddings to KV
    description: When true, archives the full payload of quarantined vectors into
      Kestra KV before permanent deletion.
    defaults: true

  - id: dry_run
    type: BOOL
    displayName: Dry Run Mode
    description: When true, generates the eviction plan and logs diagnostics without
      deleting vectors from Chroma.
    defaults: true

triggers:
  - id: periodic_eviction_sweep
    type: io.kestra.plugin.core.trigger.Schedule
    description: Run automated Chroma vector store dangling document audit and
      eviction every 6 hours.
    cron: "0 */6 * * *"
    disabled: true

  - id: on_demand_eviction_webhook
    type: io.kestra.plugin.core.trigger.Webhook
    description: Authenticated HTTP webhook endpoint invoked by RAG document
      pipelines or upstream tombstone events.
    key: "{{ secret('CHROMA_EVICTION_GUARD_WEBHOOK_KEY') }}"

tasks:
  - id: validate_inputs
    type: io.kestra.plugin.core.execution.Assert
    description: Assert operational inputs satisfy threshold and sanity constraints.
    conditions:
      - "{{ inputs.collection_name | length > 0 }}"
      - "{{ inputs.base_url | length > 0 }}"
      - "{{ inputs.batch_size >= 1 }}"
      - "{{ inputs.stale_retention_days >= 0 }}"
    errorMessage: "Invalid configuration inputs. Provide a valid Chroma collection
      name, accessible base URL, batch_size >= 1, and non-negative
      stale_retention_days."

  - id: resolve_collection
    type: io.kestra.plugin.chroma.CreateCollection
    description: Resolve target Chroma collection and retrieve its internal collection UUID.
    baseUrl: "{{ inputs.base_url }}"
    tenant: "{{ inputs.tenant }}"
    database: "{{ inputs.database }}"
    authToken: "{{ secret('CHROMA_AUTH_TOKEN') ?? '' }}"
    collectionName: "{{ inputs.collection_name }}"
    getOrCreate: true

  - id: get_pre_eviction_count
    type: io.kestra.plugin.core.http.Request
    description: Query Chroma collection count API endpoint before initiating eviction pass.
    uri: "{{ inputs.base_url }}/api/v2/tenants/{{ inputs.tenant }}/databases/{{
      inputs.database }}/collections/{{ outputs.resolve_collection.collectionId
      }}/count"
    method: GET
    headers:
      Accept: application/json
    options:
      allowFailed: true

  - id: inspect_and_quarantine_embeddings
    type: io.kestra.plugin.scripts.python.Script
    description: Retrieve candidate embeddings from Chroma, classify dangling and
      tombstoned vectors, and evaluate cross-run accumulation trends.
    containerImage: python:3.11-slim
    beforeCommands:
      - pip install -q requests
    env:
      BASE_URL: "{{ inputs.base_url }}"
      TENANT: "{{ inputs.tenant }}"
      DATABASE: "{{ inputs.database }}"
      COLLECTION_NAME: "{{ inputs.collection_name }}"
      COLLECTION_ID: "{{ outputs.resolve_collection.collectionId }}"
      BATCH_SIZE: "{{ inputs.batch_size }}"
      STALE_RETENTION_DAYS: "{{ inputs.stale_retention_days }}"
      DRY_RUN: "{{ inputs.dry_run }}"
      AUTH_TOKEN: "{{ secret('CHROMA_AUTH_TOKEN') ?? '' }}"
      PREVIOUS_STATE_JSON: "{{ kv('chroma_tracker_' ~ (inputs.collection_name |
        slugify), errorOnMissing=false) ?? '{}' }}"
      EXECUTION_ID: "{{ execution.id }}"
    script: |
      import os
      import sys
      import json
      import requests
      from datetime import datetime, timezone, timedelta

      now_utc = datetime.now(timezone.utc)
      now_iso = now_utc.isoformat()

      base_url = os.environ.get("BASE_URL", "http://localhost:8000").rstrip("/")
      tenant = os.environ.get("TENANT", "default_tenant")
      database = os.environ.get("DATABASE", "default_database")
      collection_name = os.environ.get("COLLECTION_NAME", "customer-knowledge-base")
      collection_id = os.environ.get("COLLECTION_ID", "")
      batch_size = int(os.environ.get("BATCH_SIZE", "200"))
      stale_retention_days = int(os.environ.get("STALE_RETENTION_DAYS", "7"))
      dry_run = os.environ.get("DRY_RUN", "true").lower() == "true"
      execution_id = os.environ.get("EXECUTION_ID", "manual")
      auth_token = os.environ.get("AUTH_TOKEN", "").strip()

      prev_state_raw = os.environ.get("PREVIOUS_STATE_JSON", "{}")
      try:
          prev_state = json.loads(prev_state_raw) if prev_state_raw.strip() else {}
      except Exception:
          prev_state = {}

      headers = {
          "Accept": "application/json",
          "Content-Type": "application/json"
      }
      if auth_token:
          headers["X-Chroma-Token"] = auth_token

      get_url = f"{base_url}/api/v2/tenants/{tenant}/databases/{database}/collections/{collection_id}/get"
      payload = {
          "limit": batch_size,
          "include": ["metadatas", "documents", "embeddings"]
      }

      total_inspected = 0
      dangling_vector_ids = []
      quarantine_records = []
      # Chroma's Query task requires a query vector of the exact dimensionality the collection
      # was created with -- a mismatched dimension is rejected by the server. Derive a real
      # zero-vector of the right length from an actual fetched embedding instead of guessing one,
      # so the post-eviction verification query below never fails on dimension mismatch.
      verification_vector = [0.0]
      retention_cutoff = now_utc - timedelta(days=stale_retention_days)

      try:
          resp = requests.post(get_url, headers=headers, json=payload, timeout=15)
          if resp.status_code == 200:
              data = resp.json()
              ids = data.get("ids", [])
              documents = data.get("documents", [])
              metadatas = data.get("metadatas", [])
              embeddings = data.get("embeddings", [])
              total_inspected = len(ids)

              for emb in embeddings:
                  if emb:
                      verification_vector = [0.0] * len(emb)
                      break

              for i, vid in enumerate(ids):
                  meta = metadatas[i] if i < len(metadatas) and isinstance(metadatas[i], dict) else {}
                  doc_text = documents[i] if i < len(documents) and isinstance(documents[i], str) else ""

                  is_dangling = False
                  eviction_reason = ""

                  status_val = str(meta.get("status", "")).upper()
                  if status_val in ["DELETED", "ARCHIVED", "TOMBSTONE", "INACTIVE"]:
                      is_dangling = True
                      eviction_reason = f"Upstream lifecycle status is '{status_val}'"
                  elif str(meta.get("is_tombstone", "")).lower() in ["true", "1", "yes"]:
                      is_dangling = True
                      eviction_reason = "Tombstone flag 'is_tombstone=true' set by CDC sync"
                  elif str(meta.get("quarantine", "")).lower() in ["true", "1", "yes"]:
                      is_dangling = True
                      eviction_reason = "Flagged for security or quality quarantine"
                  elif meta.get("expires_at"):
                      try:
                          exp_str = str(meta.get("expires_at"))
                          exp_dt = datetime.fromisoformat(exp_str.replace("Z", "+00:00"))
                          if exp_dt < now_utc:
                              is_dangling = True
                              eviction_reason = f"Embedding expired at {exp_str}"
                      except Exception:
                          pass

                  if is_dangling and meta.get("deleted_at"):
                      try:
                          del_str = str(meta.get("deleted_at"))
                          del_dt = datetime.fromisoformat(del_str.replace("Z", "+00:00"))
                          if del_dt > retention_cutoff:
                              continue
                      except Exception:
                          pass

                  if is_dangling:
                      dangling_vector_ids.append(vid)
                      quarantine_records.append({
                          "id": vid,
                          "collection_name": collection_name,
                          "collection_id": collection_id,
                          "eviction_reason": eviction_reason,
                          "document_preview": doc_text[:120] if doc_text else "",
                          "metadata": meta,
                          "detected_at": now_iso,
                          "execution_id": execution_id
                      })
          else:
              print(f"Chroma GET records API returned status code {resp.status_code}: {resp.text}")
      except Exception as e:
          print(f"Error querying Chroma records: {e}")

      dangling_count = len(dangling_vector_ids)
      has_dangling = dangling_count > 0

      prev_dangling_count = prev_state.get("dangling_count", 0)
      if has_dangling:
          if prev_dangling_count > 0 and dangling_count >= prev_dangling_count:
              trend_status = "ACCUMULATING"
          else:
              trend_status = "DETECTED"
      else:
          trend_status = "CLEARED" if prev_dangling_count > 0 else "HEALTHY"

      summary_text = (
          f"Chroma Vector Eviction Guard: {trend_status} | Collection '{collection_name}' (ID: {collection_id}). "
          f"Inspected: {total_inspected} vector(s). Dangling / tombstone embeddings detected: {dangling_count} "
          f"(Previous: {prev_dangling_count}). Grace retention window: {stale_retention_days} days. Dry run: {dry_run}."
      )
      print(summary_text)

      new_state_snapshot = {
          "collection_name": collection_name,
          "collection_id": collection_id,
          "total_inspected": total_inspected,
          "dangling_count": dangling_count,
          "trend_status": trend_status,
          "dry_run": dry_run,
          "timestamp": now_iso,
          "execution_id": execution_id
      }

      quarantine_manifest = {
          "collection_name": collection_name,
          "collection_id": collection_id,
          "quarantine_count": dangling_count,
          "records": quarantine_records,
          "exported_at": now_iso,
          "execution_id": execution_id
      }

      output_payload = {
          "total_inspected": total_inspected,
          "dangling_count": dangling_count,
          "has_dangling": has_dangling,
          "dangling_vector_ids": dangling_vector_ids,
          "trend_status": trend_status,
          "summary": summary_text,
          "verification_vector": verification_vector,
          "new_state_snapshot": new_state_snapshot,
          "quarantine_manifest": quarantine_manifest
      }

      print("::" + json.dumps({"outputs": output_payload}) + "::")

  - id: save_state_tracker_kv
    type: io.kestra.plugin.core.kv.Set
    description: Persist updated collection audit snapshot in Kestra KV store for
      sequential trend tracking.
    key: "chroma_tracker_{{ inputs.collection_name | slugify }}"
    value: "{{ outputs.inspect_and_quarantine_embeddings.vars.new_state_snapshot |
      toJson }}"
    kvType: JSON
    overwrite: true

  - id: evaluate_eviction_findings
    type: io.kestra.plugin.core.flow.If
    description: Branch execution based on whether dangling or tombstoned embeddings
      were identified for eviction.
    condition: "{{ outputs.inspect_and_quarantine_embeddings.vars.has_dangling == true }}"
    then:
      - id: archive_to_quarantine_kv
        type: io.kestra.plugin.core.flow.If
        description: Archive full metadata and document payloads of quarantined vectors
          into Kestra KV before deletion.
        condition: "{{ inputs.enable_quarantine_archive }}"
        then:
          - id: save_quarantine_manifest_kv
            type: io.kestra.plugin.core.kv.Set
            description: Store durable quarantine manifest in Kestra KV for audit compliance
              and disaster recovery.
            key: "chroma_quarantine_{{ inputs.collection_name | slugify }}_{{ execution.id
              }}"
            value: "{{ outputs.inspect_and_quarantine_embeddings.vars.quarantine_manifest |
              toJson }}"
            kvType: JSON
            overwrite: true

      - id: remediation_gate
        type: io.kestra.plugin.core.flow.If
        description: Execute verified deletion of dangling vector IDs when dry_run is
          disabled.
        condition: "{{ not inputs.dry_run and
          (outputs.inspect_and_quarantine_embeddings.vars.dangling_count > 0)
          }}"
        then:
          - id: pause_for_eviction_approval
            type: io.kestra.plugin.core.flow.Pause
            description: Operator approval gate. Review identified dangling vectors before
              modifying live Chroma collection. Resume to apply or abort/kill to
              cancel.

          - id: evict_dangling_vectors
            type: io.kestra.plugin.chroma.Delete
            description: Permanently purge identified dangling and tombstoned vector IDs
              from the Chroma collection.
            baseUrl: "{{ inputs.base_url }}"
            tenant: "{{ inputs.tenant }}"
            database: "{{ inputs.database }}"
            authToken: "{{ secret('CHROMA_AUTH_TOKEN') ?? '' }}"
            collectionName: "{{ inputs.collection_name }}"
            ids: "{{ outputs.inspect_and_quarantine_embeddings.vars.dangling_vector_ids }}"

          - id: verify_post_eviction_count
            type: io.kestra.plugin.core.http.Request
            description: Post-remediation verification querying collection count endpoint to
              confirm vector count reduction.
            uri: "{{ inputs.base_url }}/api/v2/tenants/{{ inputs.tenant }}/databases/{{
              inputs.database }}/collections/{{
              outputs.resolve_collection.collectionId }}/count"
            method: GET
            headers:
              Accept: application/json
            options:
              allowFailed: true

          - id: verify_sample_tombstone_absence
            type: io.kestra.plugin.chroma.Query
            description: Execute verification similarity query filtered on DELETED status to
              assert complete absence of tombstone records.
            baseUrl: "{{ inputs.base_url }}"
            tenant: "{{ inputs.tenant }}"
            database: "{{ inputs.database }}"
            authToken: "{{ secret('CHROMA_AUTH_TOKEN') ?? '' }}"
            collectionName: "{{ inputs.collection_name }}"
            vector: "{{ outputs.inspect_and_quarantine_embeddings.vars.verification_vector
              }}"
            maxResults: 5
            where:
              status: "DELETED"
            fetchType: FETCH

          - id: log_verified_eviction
            type: io.kestra.plugin.core.log.Log
            description: Log confirmation of verified vector eviction and updated collection
              metrics.
            message: "Verified eviction: successfully removed {{
              outputs.evict_dangling_vectors.requestedIdCount }} vector(s) from
              collection '{{ inputs.collection_name }}'. Initial count: {{
              outputs.get_pre_eviction_count.body }}, post-eviction count: {{
              outputs.verify_post_eviction_count.body }}. Tombstone verification
              query returned {{ outputs.verify_sample_tombstone_absence.rows |
              length }} rows."

        else:
          - id: log_dry_run_plan
            type: io.kestra.plugin.core.log.Log
            description: Log planned vector evictions in dry-run mode without modifying
              vector store state.
            message: "Dry run mode active. No vectors deleted from '{{
              inputs.collection_name }}'. Planned evictions: {{
              outputs.inspect_and_quarantine_embeddings.vars.dangling_count }}
              vector(s). Sample IDs: {{
              outputs.inspect_and_quarantine_embeddings.vars.dangling_vector_ids
              | slice(0, 5) | toJson }}."

      - id: notify_slack_eviction_alert
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Deliver actionable incident notification to Slack with affected
          vector counts, trend telemetry, and remediation status.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "attachments": [
              {
                "color": "{{ outputs.inspect_and_quarantine_embeddings.vars.trend_status == 'ACCUMULATING' ? 'danger' : 'warning' }}",
                "title": ":wastebasket: Chroma Vector Store Dangling Document & Tombstone Eviction Alert",
                "text": "*Collection:* `{{ inputs.collection_name }}` (ID: `{{ outputs.resolve_collection.collectionId }}`)\n*Dangling Vectors Detected:* {{ outputs.inspect_and_quarantine_embeddings.vars.dangling_count }}\n*Trend Status:* `{{ outputs.inspect_and_quarantine_embeddings.vars.trend_status }}`\n*Dry Run:* `{{ inputs.dry_run }}`\n*Quarantine Archive Saved:* `{{ inputs.enable_quarantine_archive }}`\n*Execution ID:* `{{ execution.id }}`\n\n*Summary:* {{ outputs.inspect_and_quarantine_embeddings.vars.summary }}"
              }
            ]
          }

    else:
      - id: log_healthy_collection
        type: io.kestra.plugin.core.log.Log
        description: Log healthy collection status when no dangling or tombstoned
          embeddings are detected.
        message: "Chroma collection '{{ inputs.collection_name }}' is healthy. No
          dangling, tombstoned, or quarantined vectors detected across {{
          outputs.get_pre_eviction_count.body }} total vector(s)."

errors:
  - id: alert_on_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: Dispatch emergency alert if Chroma vector store eviction guard
      encounters an unhandled execution error.
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "attachments": [
          {
            "color": "danger",
            "title": ":rotating_light: Chroma Vector Store Eviction Guard Execution Failed",
            "text": "Failed execution `{{ execution.id }}` while evaluating Chroma collection `{{ inputs.collection_name }}` on `{{ inputs.base_url }}`.\nVerify Chroma vector database availability, network reachability, and worker logs."
          }
        ]
      }