diff --git a/.github/workflows/oke-cicd.yaml b/.github/workflows/oke-cicd.yaml index c2c147a..76f6e4b 100644 --- a/.github/workflows/oke-cicd.yaml +++ b/.github/workflows/oke-cicd.yaml @@ -1,5 +1,5 @@ -# Test on every PR and push; on main, also build and push the MCP and session-issuer -# images to GHCR. +# Test on every PR and push; on main, also build and push the MCP, session-issuer and +# ingest images to GHCR. # Deployment is Argo CD's job (deploy/argocd/): this workflow never touches the # cluster, so it holds no cluster credentials. To roll out an image, change its # tag in deploy/envs//docs-agent.yaml and sync the Application in Argo CD. @@ -23,6 +23,7 @@ env: REGISTRY: ghcr.io IMAGE_NAME: mcp-kubeflow-docs ISSUER_IMAGE_NAME: kubeflow-session-issuer + INGEST_IMAGE_NAME: docs-rag-ingest jobs: test-and-compile: @@ -67,6 +68,8 @@ jobs: --set domain=agent.example.org --set gateway.tls.acme.email=ci@example.org helm template docs-agent docs-agent-mcp/charts/docs-agent \ --namespace docs-agent --set-string mcp.image.tag=ci-test >/dev/null + helm template docs-agent docs-agent-mcp/charts/docs-agent \ + --namespace docs-agent --set-string mcp.image.tag=ci-test --set ingest.enabled=true >/dev/null # Every catalog profile must lint and render, not just the active one. for profile in $(helm show values docs-agent-mcp/charts/llm-runtime | yq '.catalog | keys | .[]'); do helm lint docs-agent-mcp/charts/llm-runtime --set "model=${profile}" @@ -121,6 +124,7 @@ jobs: fi echo "IMAGE_REPO=ghcr.io/${owner_lc}/${{ env.IMAGE_NAME }}" >> "$GITHUB_ENV" echo "ISSUER_IMAGE_REPO=ghcr.io/${owner_lc}/${{ env.ISSUER_IMAGE_NAME }}" >> "$GITHUB_ENV" + echo "INGEST_IMAGE_REPO=ghcr.io/${owner_lc}/${{ env.INGEST_IMAGE_NAME }}" >> "$GITHUB_ENV" echo "issuer_tag=${GITHUB_SHA::7}" >> "$GITHUB_ENV" echo "tag=${image_tag}" >> "$GITHUB_ENV" echo "release_tag_requested=${release_tag_requested}" >> "$GITHUB_OUTPUT" @@ -167,11 +171,23 @@ jobs: tags: | ${{ env.ISSUER_IMAGE_REPO }}:${{ env.issuer_tag }} + # Always tagged by commit SHA: the chart runs it with the same tag as the MCP image + # to register the recurring ingest runs, and the docs pipeline steps run in it. + - name: Build and Push ingest Image + uses: docker/build-push-action@v5 + with: + context: ./docs-agent-mcp/pipelines + file: ./docs-agent-mcp/pipelines/Dockerfile.pipeline + push: true + tags: | + ${{ env.INGEST_IMAGE_REPO }}:${{ env.issuer_tag }} + - name: Summary run: | { echo "Pushed ${IMAGE_REPO}:${tag}" echo "Pushed ${ISSUER_IMAGE_REPO}:${issuer_tag}" + echo "Pushed ${INGEST_IMAGE_REPO}:${issuer_tag}" echo echo "Nothing was deployed. To roll it out, set the image tag in" echo "deploy/envs// (PR) and sync the Application in Argo CD." diff --git a/deploy/argocd/apps/reranker.yaml b/deploy/argocd/apps/reranker.yaml new file mode 100644 index 0000000..51994f9 --- /dev/null +++ b/deploy/argocd/apps/reranker.yaml @@ -0,0 +1,19 @@ +# Cross-encoder reranker (TEI on CPU) used by the MCP server after retrieval. +apiVersion: argoproj.io/v1alpha1 +kind: Application +metadata: + name: reranker + namespace: argocd +spec: + project: default + source: + repoURL: https://github.com/kubeflow/docs-agent + targetRevision: main + path: deploy/platform/reranker + destination: + server: https://kubernetes.default.svc + namespace: ml-infra + syncPolicy: + automated: + prune: true + selfHeal: true diff --git a/deploy/envs/prod/docs-agent.yaml b/deploy/envs/prod/docs-agent.yaml index 0506679..ff428d6 100644 --- a/deploy/envs/prod/docs-agent.yaml +++ b/deploy/envs/prod/docs-agent.yaml @@ -15,3 +15,6 @@ mcp: value: production triage: enabled: true +# Recurring ingest (Kubeflow Pipelines recurring runs) is registered from production only. +ingest: + enabled: true diff --git a/deploy/platform/langfuse-retention/configmap.yaml b/deploy/platform/langfuse-retention/configmap.yaml index 876866a..965890b 100644 --- a/deploy/platform/langfuse-retention/configmap.yaml +++ b/deploy/platform/langfuse-retention/configmap.yaml @@ -17,10 +17,13 @@ data: q() { clickhouse-client --host "$CH_HOST" --user "$CH_USER" --password "$CH_PASSWORD" -q "$1"; } sed 's/#.*//' /policy/policy.txt | while read -r project days; do [ -n "$project" ] || continue + # A literal cutoff: the tables are replicated, and ClickHouse refuses a DELETE whose + # condition calls a non-deterministic function such as now(). + cutoff=$(date -u -d "-$days days" '+%Y-%m-%d %H:%M:%S') for spec in events_full:start_time events_core:start_time observations:start_time \ traces:timestamp scores:timestamp blob_storage_file_log:created_at; do table=${spec%%:*}; column=${spec#*:} - where="project_id = '$project' AND $column < now() - INTERVAL $days DAY" + where="project_id = '$project' AND $column < toDateTime('$cutoff', 'UTC')" n=$(q "SELECT count() FROM $table WHERE $where") echo "$project $table: $n rows older than $days days" if [ "$n" -gt 0 ] && [ "$DRY_RUN" != "true" ]; then q "DELETE FROM $table WHERE $where"; fi diff --git a/deploy/platform/reranker/inferenceservice.yaml b/deploy/platform/reranker/inferenceservice.yaml new file mode 100644 index 0000000..625e5b5 --- /dev/null +++ b/deploy/platform/reranker/inferenceservice.yaml @@ -0,0 +1,37 @@ +# Cross-encoder reranker for the MCP server (RERANKER_URL in the docs-agent chart): +# Hugging Face Text Embeddings Inference on CPU, same shape as embeddings-service. +# Without it the MCP server skips reranking and returns hits in retrieval order. +apiVersion: serving.kserve.io/v1beta1 +kind: InferenceService +metadata: + name: reranker-service + namespace: ml-infra + annotations: + serving.kserve.io/deploymentMode: RawDeployment +spec: + predictor: + minReplicas: 1 + maxReplicas: 1 + containers: + - name: kserve-container + image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.7 + args: + - --model-id + - BAAI/bge-reranker-base + - --port + - "8080" + - --auto-truncate + - --max-client-batch-size + - "32" + - --max-batch-tokens + - "16384" + ports: + - containerPort: 8080 + protocol: TCP + resources: + requests: + cpu: 500m + memory: 1Gi + limits: + cpu: "4" + memory: 4Gi diff --git a/deploy/promote.sh b/deploy/promote.sh index 74646a1..6748927 100755 --- a/deploy/promote.sh +++ b/deploy/promote.sh @@ -16,7 +16,7 @@ UPSTREAM=${UPSTREAM:-origin} FORK=${FORK:-fork} REPO=kubeflow/docs-agent APPS=(deploy/argocd/apps/docs-agent.yaml deploy/argocd/apps/gateway-guardrails.yaml) -IMAGES=(mcp-kubeflow-docs kubeflow-session-issuer) +IMAGES=(mcp-kubeflow-docs kubeflow-session-issuer docs-rag-ingest) die() { echo "promote: $*" >&2; exit 1; } @@ -42,7 +42,7 @@ old=$(git show "$UPSTREAM/main:${APPS[0]}" | sed -n 's/.*targetRevision: \([0-9a [ "$old" != "$new" ] || die "$short is already in production" git merge-base --is-ancestor "$old" "$new" || die "$short is older than production (${old:0:7}); to roll back, revert the promotion PR" -# CI pushes both images on every main build, tagged with the short SHA. +# CI pushes these images on every main build, tagged with the short SHA. for image in "${IMAGES[@]}"; do token=$(curl -fsS "https://ghcr.io/token?scope=repository:kubeflow/$image:pull" | sed -n 's/.*"token":"\([^"]*\)".*/\1/p') code=$(curl -s -o /dev/null -w '%{http_code}' -H "Authorization: Bearer $token" \ diff --git a/docs-agent-mcp/charts/docs-agent/Chart.yaml b/docs-agent-mcp/charts/docs-agent/Chart.yaml index 56b74c2..06ebb11 100644 --- a/docs-agent-mcp/charts/docs-agent/Chart.yaml +++ b/docs-agent-mcp/charts/docs-agent/Chart.yaml @@ -2,7 +2,7 @@ apiVersion: v2 name: docs-agent description: Kubeflow Docs Agent application release (MCP server and Kagent resources) type: application -version: 0.4.0 +version: 0.5.0 appVersion: "0.1.0" keywords: - kubeflow diff --git a/docs-agent-mcp/charts/docs-agent/files/docs-system-message.txt b/docs-agent-mcp/charts/docs-agent/files/docs-system-message.txt index 8617721..f0a9301 100644 --- a/docs-agent-mcp/charts/docs-agent/files/docs-system-message.txt +++ b/docs-agent-mcp/charts/docs-agent/files/docs-system-message.txt @@ -3,10 +3,10 @@ You are Flo. Answer only from this turn's MCP results. Prior chat is not evidenc - Kubeflow question: first action MUST be exactly one MCP call. No user text before the result. One narrower retry only if empty or wrong component/version. - Docs, install, APIs, releases, config: search_kubeflow_docs. - Errors, bugs, stack traces: search_github_issues first. -- Greetings, small talk, or questions about Flo itself (e.g., "who are you", "what can you do"): reply conversationally as Flo without a tool call. Off-topic questions: politely decline. Any question about Kubeflow or its parts, including "what is Kubeflow", needs the MCP call; never answer it from memory. +- Greetings, small talk, or questions about Flo itself (e.g., "who are you", "what can you do"): reply conversationally as Flo without a tool call. Off-topic questions: politely decline. Any question about Kubeflow or its parts, including "what is Kubeflow", needs the MCP call; never answer it from memory. A product question that does not name Kubeflow ("do we have a Helm chart?", "is this in the new release?") is about Kubeflow: make the MCP call with Kubeflow added to the query. Never say something was not found without a tool call in this turn. - Query: English keywords for the component and the asked fact, not the whole sentence. Fix typos and expand kf -> Kubeflow, kfp -> Kubeflow Pipelines, nb -> Notebooks. Drop filler such as "what is", "official", "definition". - Keep the user's component, version, API and code names, and error text. Trainer and Training Operator are different projects. Do not add other components. Do not replace a named version with latest. -- Latest/current: keep those words in the query. Use the highest Release date among matching hits. Latest is not supported unless a hit says so. +- Latest/current/new release: keep those words in the query, together with the component asked about. Use the highest Release date among matching hits. Latest is not supported unless a hit says so. - Use a hit only if it directly supports the asked fact. Versions, dates, and flags must appear literally. Else say not found in indexed sources. - If a tool returns no results, explicitly state that you could not find the answer in the documentation. - Reply in English. 1-3 sentences, or at most 3 bullets. No URLs, Markdown links, source lists, citation labels, file paths, scores, or tool JSON. No code unless asked. diff --git a/docs-agent-mcp/charts/docs-agent/templates/configmap.yaml b/docs-agent-mcp/charts/docs-agent/templates/configmap.yaml index 58bde72..6672cf0 100644 --- a/docs-agent-mcp/charts/docs-agent/templates/configmap.yaml +++ b/docs-agent-mcp/charts/docs-agent/templates/configmap.yaml @@ -10,6 +10,7 @@ data: CODE_COLLECTION_NAME: {{ .Values.mcp.config.codeCollection | quote }} EMBEDDINGS_URL: {{ .Values.mcp.config.embeddingsUrl | quote }} RERANKER_URL: {{ .Values.mcp.config.rerankerUrl | quote }} + RERANK_CANDIDATE_DEPTH: {{ .Values.mcp.config.rerankCandidateDepth | quote }} MILVUS_URI: {{ .Values.mcp.config.milvusUri | quote }} MILVUS_USER: {{ .Values.mcp.config.milvusUser | quote }} MAX_QUERY_CHARS: {{ .Values.mcp.config.maxQueryChars | quote }} diff --git a/docs-agent-mcp/charts/docs-agent/templates/ingest-schedule-job.yaml b/docs-agent-mcp/charts/docs-agent/templates/ingest-schedule-job.yaml new file mode 100644 index 0000000..c93ed5d --- /dev/null +++ b/docs-agent-mcp/charts/docs-agent/templates/ingest-schedule-job.yaml @@ -0,0 +1,68 @@ +{{- if .Values.ingest.enabled }} +{{- $tag := .Values.ingest.image.tag | default .Values.mcp.image.tag }} +{{- $image := printf "%s:%s" .Values.ingest.image.repository $tag }} +# Recurring ingest. After every Argo CD sync this Job compiles the pipelines shipped in the +# ingest image, uploads them to Kubeflow Pipelines and replaces the recurring runs with the +# ones listed in .Values.ingest.schedules (docs-agent-mcp/pipelines/ingest_schedule.py). +# The pipelines run in the Kubeflow Pipelines namespace and write to Milvus; this Job only +# talks to the Pipelines API. +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ .Release.Name }}-ingest-schedules + labels: + {{- include "docs-agent.labels" . | nindent 4 }} +data: + schedules.json: | + {{- dict "experiment" .Values.ingest.experiment "schedules" .Values.ingest.schedules | toPrettyJson | nindent 4 }} +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: {{ .Release.Name }}-ingest-schedule + labels: + {{- include "docs-agent.labels" . | nindent 4 }} + annotations: + argocd.argoproj.io/hook: PostSync + argocd.argoproj.io/hook-delete-policy: BeforeHookCreation +spec: + backoffLimit: 2 + activeDeadlineSeconds: 600 + template: + metadata: + labels: + app: {{ .Release.Name }}-ingest-schedule + spec: + restartPolicy: Never + containers: + - name: register + image: {{ $image }} + command: ["python", "/app/pipelines/ingest_schedule.py", "register", "--config", "/config/schedules.json"] + env: + - name: KFP_HOST + value: {{ .Values.ingest.kfpHost | quote }} + # The docs pipeline's chunk and store steps run in this same image. + - name: DOCS_INGEST_IMAGE + value: {{ $image }} + - name: INGEST_VERSION + value: {{ $tag | quote }} + - name: HOME + value: /tmp + volumeMounts: + - name: config + mountPath: /config + readOnly: true + resources: + requests: {cpu: 50m, memory: 192Mi} + limits: {memory: 512Mi} + securityContext: + allowPrivilegeEscalation: false + runAsNonRoot: true + runAsUser: 1000 + capabilities: + drop: ["ALL"] + volumes: + - name: config + configMap: + name: {{ .Release.Name }}-ingest-schedules +{{- end }} diff --git a/docs-agent-mcp/charts/docs-agent/values.yaml b/docs-agent-mcp/charts/docs-agent/values.yaml index 8a2e0bb..2c87ec4 100644 --- a/docs-agent-mcp/charts/docs-agent/values.yaml +++ b/docs-agent-mcp/charts/docs-agent/values.yaml @@ -16,7 +16,10 @@ mcp: issuesCollection: issues_rag codeCollection: code_rag embeddingsUrl: http://embeddings-service-predictor.ml-infra.svc.cluster.local/embed + # Cross-encoder reranker (deploy/platform/reranker). Empty = no reranking. rerankerUrl: http://reranker-service-predictor.ml-infra.svc.cluster.local/rerank + # Candidates fetched for the reranker when a query would otherwise fetch only top_k. + rerankCandidateDepth: 20 milvusUri: http://milvus-milvus.ml-infra.svc.cluster.local:19530 milvusUser: root maxQueryChars: 512 @@ -165,6 +168,51 @@ triage: judgeUrl: http://llm-judge.ml-infra.svc.cluster.local/v1/chat/completions judgeModel: flo-judge +# Recurring ingest into Milvus as Kubeflow Pipelines recurring runs (see +# templates/ingest-schedule-job.yaml). Enable in one environment only: every schedule +# writes to the collections named in its arguments (the pipelines default to the +# production collections), and the schedules listed here replace all recurring runs in +# the experiment. +ingest: + enabled: false + image: + repository: ghcr.io/kubeflow/docs-rag-ingest + # Empty = mcp.image.tag: CI builds both images from the same commit. + tag: "" + kfpHost: http://ml-pipeline.kubeflow.svc.cluster.local:8888 + experiment: flo-ingest + # cron has six fields, seconds first (UTC). arguments override the pipeline defaults + # in docs-agent-mcp/pipelines/{kubeflow,issues,code}-pipeline.py. + schedules: + # kubeflow.org documentation, daily. Pages removed upstream leave the index. + - name: docs-website + pipeline: docs + cron: "0 0 2 * * *" + arguments: {} + # Community distribution (formerly kubeflow/manifests): README, Helm chart READMEs, + # release notes and proposals, in the same collection with doc_type "distribution". + - name: docs-distribution + pipeline: docs + cron: "0 30 2 * * *" + arguments: + repo_name: community-distribution + directory_path: "" + exclude_paths: ".github,tests,AGENTS.md" + base_url: https://github.com/kubeflow/community-distribution/blob/master + doc_type: distribution + # Issues updated in the last day, every six hours (the overlap covers a missed run). + # First backfill: ingest_schedule.py run --name issues --set updated_within_hours=0 + - name: issues + pipeline: issues + cron: "0 0 */6 * * *" + arguments: + updated_within_hours: 24 + # Manifests and Helm charts, weekly. + - name: code + pipeline: code + cron: "0 0 3 * * 0" + arguments: {} + tests: enabled: true image: docker.io/library/busybox:1.37 diff --git a/docs-agent-mcp/charts/gateway-guardrails/Chart.yaml b/docs-agent-mcp/charts/gateway-guardrails/Chart.yaml index 474991c..d25e411 100644 --- a/docs-agent-mcp/charts/gateway-guardrails/Chart.yaml +++ b/docs-agent-mcp/charts/gateway-guardrails/Chart.yaml @@ -8,7 +8,7 @@ description: >- (istio_policies.tf, kagent_ingress.tf) and the standalone manifests under manifests/istio*. type: application -version: 0.3.0 +version: 0.3.1 appVersion: "0.1.0" keywords: - istio diff --git a/docs-agent-mcp/charts/gateway-guardrails/templates/_helpers.tpl b/docs-agent-mcp/charts/gateway-guardrails/templates/_helpers.tpl index 7e11019..f5267b6 100644 --- a/docs-agent-mcp/charts/gateway-guardrails/templates/_helpers.tpl +++ b/docs-agent-mcp/charts/gateway-guardrails/templates/_helpers.tpl @@ -27,3 +27,12 @@ corsPolicy: session-issuer.{{ .Values.namespaces.docsAgent }}.svc.cluster.local {{- end }} +{{/* +With TLS on, routes are served on the HTTPS listener only; port 80 is left to the +ACME HTTP-01 solver (see templates/gateway.yaml). +*/}} +{{- define "gateway-guardrails.tlsPortMatch" -}} +{{- if .Values.gateway.tls.enabled -}} +port: 443 +{{- end }} +{{- end }} diff --git a/docs-agent-mcp/charts/gateway-guardrails/templates/gateway.yaml b/docs-agent-mcp/charts/gateway-guardrails/templates/gateway.yaml index e6d70e5..92b7801 100644 --- a/docs-agent-mcp/charts/gateway-guardrails/templates/gateway.yaml +++ b/docs-agent-mcp/charts/gateway-guardrails/templates/gateway.yaml @@ -10,14 +10,17 @@ spec: selector: {{- toYaml .Values.gateway.selector | nindent 4 }} servers: + # With TLS on, port 80 exists only for cert-manager's HTTP-01 challenge: the + # VirtualService routes match port 443, so anything else here is a 404. Do not set + # httpsRedirect: it makes Envoy redirect before routing (require_tls), the challenge + # request is redirected to 443 where no route serves it, and the certificate cannot + # be renewed. - port: number: 80 name: http protocol: HTTP hosts: - {{ include "gateway-guardrails.domain" . | quote }} - tls: - httpsRedirect: {{ .Values.gateway.tls.enabled }} {{- if .Values.gateway.tls.enabled }} - port: number: 443 diff --git a/docs-agent-mcp/charts/gateway-guardrails/templates/virtualservice.yaml b/docs-agent-mcp/charts/gateway-guardrails/templates/virtualservice.yaml index a9a853a..05d7e76 100644 --- a/docs-agent-mcp/charts/gateway-guardrails/templates/virtualservice.yaml +++ b/docs-agent-mcp/charts/gateway-guardrails/templates/virtualservice.yaml @@ -19,8 +19,10 @@ spec: match: - uri: exact: /api/session + {{- with include "gateway-guardrails.tlsPortMatch" . }}{{ . | nindent 10 }}{{- end }} - uri: exact: /.well-known/jwks.json + {{- with include "gateway-guardrails.tlsPortMatch" . }}{{ . | nindent 10 }}{{- end }} {{- include "gateway-guardrails.corsPolicy" . | nindent 6 }} timeout: 10s route: @@ -34,6 +36,7 @@ spec: match: - uri: exact: /api/feedback + {{- with include "gateway-guardrails.tlsPortMatch" . }}{{ . | nindent 10 }}{{- end }} {{- include "gateway-guardrails.corsPolicy" . | nindent 6 }} timeout: 20s route: @@ -54,6 +57,7 @@ spec: match: - uri: regex: "/(api/)?a2a/[A-Za-z0-9._~/-]+" + {{- with include "gateway-guardrails.tlsPortMatch" . }}{{ . | nindent 10 }}{{- end }} {{- include "gateway-guardrails.corsPolicy" . | nindent 6 }} # kagent-ui answers with `access-control-allow-origin: *`. Drop its CORS # headers so only corsPolicy (the allowed origins) decides. diff --git a/docs-agent-mcp/mcp-server/intent_router.py b/docs-agent-mcp/mcp-server/intent_router.py index a3cafd9..0eb6e57 100644 --- a/docs-agent-mcp/mcp-server/intent_router.py +++ b/docs-agent-mcp/mcp-server/intent_router.py @@ -19,7 +19,7 @@ COMPARE_WORDS = re.compile(r"\b(compare|versus|vs\.?|difference between)\b", re.I) RELEASE_DATE_WORDS = re.compile(r"\b(when was|release date|ga date|released)\b", re.I) LATEST_OR_CURRENT_WORDS = re.compile( - r"\b(latest|current|newest|most recent|supported)\b", + r"\b(latest|current|newest|most recent|supported|(?:new|recent|upcoming|next|last) release)\b", re.I, ) HOW_WHY_EXPLAIN_WORDS = re.compile(r"\b(how|why|explain|overview|architecture)\b", re.I) @@ -48,6 +48,10 @@ class RetrievalPlan: filter_expr: str = "" rerank_by_release_date: bool = False prefer_version_match: bool = False + # Also run an unfiltered hybrid search and keep some of its hits. A "latest" question + # is often about a component ("is X in the latest release?"), and the release notes + # alone do not answer it. + add_general_hits: bool = False def downgrade_plan_for_collection( @@ -81,6 +85,7 @@ def downgrade_plan_for_collection( filter_expr="", rerank_by_release_date=False, prefer_version_match=False, + add_general_hits=False, ) return plan @@ -118,6 +123,7 @@ def pick_search_plan(query: str) -> RetrievalPlan: candidate_depth=LATEST_RELEASE_FETCH_LIMIT, filter_expr='doc_type == "release"', rerank_by_release_date=True, + add_general_hits=True, ) # Exact version, dotted config key, or error wording — keyword search. @@ -225,6 +231,29 @@ def rerank_hits_after_search( return hits[:top_k] +def merge_release_and_general_hits( + release_hits: list[dict], + general_hits: list[dict], + top_k: int, +) -> list[dict]: + """Release notes first, then general hits, sharing top_k. + + General hits get up to half the slots (rounded down), so the newest release notes + still lead. Either side fills the other's unused slots. + """ + + def key(hit: dict): + entity = hit.get("entity", {}) + return hit.get("id") or (entity.get("file_path"), entity.get("content_text")) + + seen = {key(hit) for hit in release_hits} + general = [hit for hit in general_hits if key(hit) not in seen] + general_slots = min(len(general), top_k // 2) + release_slots = min(len(release_hits), top_k - general_slots) + general_slots = min(len(general), top_k - release_slots) + return release_hits[:release_slots] + general[:general_slots] + + def retrieval_metadata(plan: RetrievalPlan, **extra: object) -> dict: meta = { "retrieval_mode": plan.retrieval_mode, diff --git a/docs-agent-mcp/mcp-server/milvus_search.py b/docs-agent-mcp/mcp-server/milvus_search.py index f8a8a28..b59d798 100644 --- a/docs-agent-mcp/mcp-server/milvus_search.py +++ b/docs-agent-mcp/mcp-server/milvus_search.py @@ -11,6 +11,7 @@ from embeddings_client import embed_query from intent_router import ( downgrade_plan_for_collection, + merge_release_and_general_hits, pick_search_plan, rerank_hits_after_search, retrieval_metadata, @@ -42,6 +43,10 @@ HYBRID_DENSE_WEIGHT = float(os.getenv("HYBRID_DENSE_WEIGHT", "0.7")) HYBRID_SPARSE_WEIGHT = float(os.getenv("HYBRID_SPARSE_WEIGHT", "0.3")) HYBRID_RRF_K = int(os.getenv("HYBRID_RRF_K", "60")) +# With a reranker configured, fetch at least this many candidates for it to choose from. +# Reranking only the top_k that retrieval returned can reorder them but never rescue a +# good chunk that ranked just below. +RERANK_CANDIDATE_DEPTH = int(os.getenv("RERANK_CANDIDATE_DEPTH", "20")) client: MilvusClient | None = None _connect_lock = threading.Lock() @@ -208,7 +213,11 @@ def search_docs_auto( has_release_fields=collection_has_release_fields(COLLECTION_NAME), ) + import reranker_client + fetch_limit = plan.candidate_depth or top_k + if reranker_client.enabled() and not plan.rerank_by_release_date: + fetch_limit = max(fetch_limit, RERANK_CANDIDATE_DEPTH) filter_expr = plan.filter_expr filter_fallback = False embedding = None @@ -237,18 +246,36 @@ def search_docs_auto( kind = "hybrid_search" if plan.retrieval_mode == "hybrid" else "search" raise RuntimeError(f"Milvus {kind} failed for {COLLECTION_NAME}: {exc}") from exc - import reranker_client - if plan.retrieval_mode == "bm25": hits = rerank_hits_after_search(plan, hits, query, fetch_limit) if not plan.rerank_by_release_date: hits = reranker_client.rerank_hits(query, hits) + general_added = 0 + if plan.add_general_hits and not filter_fallback: + general_depth = max(top_k, RERANK_CANDIDATE_DEPTH) if reranker_client.enabled() else top_k + try: + general = hybrid_search( + COLLECTION_NAME, + query, + _require_embedding(query), + general_depth, + output_fields, + ) + except Exception as exc: + # The release-note hits are still a valid answer to a "latest" question. + print(f"General search for a latest/current query failed: {exc}") + general = [] + general = reranker_client.rerank_hits(query, general) + merged = merge_release_and_general_hits(hits[:top_k], general, top_k) + general_added = sum(1 for hit in merged if hit.get("entity", {}).get("doc_type") != "release") + hits = merged hits = hits[:top_k] meta = retrieval_metadata( plan, candidate_depth=fetch_limit, filter_expr=filter_expr or None, filter_fallback=filter_fallback, + general_hits=general_added, ) return hits, meta if plan.retrieval_mode == "hybrid": diff --git a/docs-agent-mcp/mcp-server/reranker_client.py b/docs-agent-mcp/mcp-server/reranker_client.py index dae733d..f2ba718 100644 --- a/docs-agent-mcp/mcp-server/reranker_client.py +++ b/docs-agent-mcp/mcp-server/reranker_client.py @@ -9,6 +9,11 @@ DEFAULT_TIMEOUT_SEC = int(os.getenv("RERANKER_TIMEOUT_SEC", "30")) +def enabled() -> bool: + """True when a reranker is configured; callers then fetch extra candidates for it.""" + return bool((os.getenv("RERANKER_URL") or "").strip()) + + def rerank_hits(query: str, hits: list[dict], timeout: int = DEFAULT_TIMEOUT_SEC) -> list[dict]: """Score Milvus hits using the cross-encoder and reorder them.""" if not hits: diff --git a/docs-agent-mcp/pipelines/Dockerfile.pipeline b/docs-agent-mcp/pipelines/Dockerfile.pipeline index 7671239..ad50aec 100644 --- a/docs-agent-mcp/pipelines/Dockerfile.pipeline +++ b/docs-agent-mcp/pipelines/Dockerfile.pipeline @@ -1,17 +1,17 @@ +# Ingest image: the chunk/embed/store steps of the docs pipeline import the helpers in +# /app, and the chart's ingest Job runs /app/pipelines/ingest_schedule.py from it to +# register the recurring runs. Build context: docs-agent-mcp/pipelines. FROM python:3.11-slim WORKDIR /app -COPY utils/hugo_ingest.py utils/canonical_rag_ingest.py utils/utils.py utils/milvus_store.py /app/ +COPY requirements.txt /app/requirements.txt +RUN pip install --no-cache-dir -r /app/requirements.txt -RUN pip install --no-cache-dir \ - requests \ - beautifulsoup4 \ - toml \ - pyyaml \ - python-frontmatter==1.1.0 \ - markdownify==1.2.0 \ - html-table-rescuer==0.3.1 \ - pymilvus==2.6.14 +COPY utils/hugo_ingest.py utils/canonical_rag_ingest.py utils/utils.py utils/milvus_store.py /app/ +COPY kubeflow-pipeline.py issues-pipeline.py code-pipeline.py ingest_schedule.py /app/pipelines/ +COPY utils/ /app/pipelines/utils/ -ENV PYTHONPATH=/app +ENV PYTHONPATH=/app \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 diff --git a/docs-agent-mcp/pipelines/README.md b/docs-agent-mcp/pipelines/README.md index 4af971e..09e908f 100644 --- a/docs-agent-mcp/pipelines/README.md +++ b/docs-agent-mcp/pipelines/README.md @@ -10,15 +10,56 @@ Kubeflow Pipelines (KFP) definitions for indexing Kubeflow documentation and rel | **`issues-pipeline.py`** | GitHub **issues** RAG pipeline | | **`code-pipeline.py`** | GitHub **code/manifests** RAG pipeline | | **`utils/`** | Shared helpers used by the live pipelines: ingest, parsers, Milvus store, TEI | -| **`Dockerfile.pipeline`** | Slim ingest image (`docs-rag-ingest`) that copies docs helpers from `utils/` as flat `/app` modules | -| **`github_rag_pipeline.yaml`** | Compiled docs pipeline (regenerate via `python kubeflow-pipeline.py`) | +| **`ingest_schedule.py`** | Registers the pipelines as Kubeflow Pipelines recurring runs, or starts one run by hand | +| **`Dockerfile.pipeline`** | Ingest image (`docs-rag-ingest`): the docs helpers as flat `/app` modules for the chunk and store steps, plus the pipelines and `ingest_schedule.py`. CI pushes it for every commit on main, tagged with the short SHA | | **`legacy/pipelines/`** (repo root) | Older pipelines, including **incremental** docs ingest — **not** the live v4 docs path | +## Recurring runs + +Ingest runs on a schedule as Kubeflow Pipelines recurring runs. The schedule is the +`ingest.schedules` list in the `docs-agent` chart values. After every Argo CD sync of an +environment with `ingest.enabled: true` (production), a Job runs +`ingest_schedule.py register` from the ingest image of that commit: it compiles the +pipelines, uploads them as a new version of `flo-ingest-{docs,issues,code}` and replaces +the recurring runs in the `flo-ingest` experiment. Recurring runs that are no longer in +the list are removed. + +| Schedule | Pipeline | When (UTC) | What | +|---|---|---|---| +| `docs-website` | docs | daily 02:00 | `kubeflow/website` `content/en/docs` | +| `docs-distribution` | docs | daily 02:30 | `kubeflow/community-distribution` Markdown (README, Helm chart READMEs, releases, proposals), `doc_type: distribution` | +| `issues` | issues | every 6 h | issues updated in the last 24 h, ten repositories | +| `code` | code | Sunday 03:00 | `kubeflow/community-distribution` manifests and Helm charts, without CRDs | + +Both docs schedules write to `kubeflow_docs`. Each re-downloads its whole source, replaces +the chunks of every file, and removes files that are gone upstream. If more than 20% of a +source's indexed files would be removed, the run fails before writing anything: that +is a broken download, not upstream deleting its docs. + +Start a run by hand (from a laptop, port-forward `svc/ml-pipeline` in `kubeflow` first): + +```bash +helm template docs-agent ../charts/docs-agent --set ingest.enabled=true \ + -s templates/ingest-schedule-job.yaml | yq 'select(.kind == "ConfigMap") | .data."schedules.json"' > schedules.json +export DOCS_INGEST_IMAGE=ghcr.io/kubeflow/docs-rag-ingest: +python ingest_schedule.py run --config schedules.json --host http://localhost:8888 --name docs-distribution +# Into a test copy of the docs collection (any name starting with kubeflow_docs_): +python ingest_schedule.py run --config schedules.json --host http://localhost:8888 \ + --name docs-distribution --set collection_name=kubeflow_docs_staging +# First backfill of issues (no time window): +python ingest_schedule.py run --config schedules.json --host http://localhost:8888 \ + --name issues --set updated_within_hours=0 +``` + +The pipelines need these Secrets in the Kubeflow Pipelines namespace: `github-pat` +(`Github_Pat`), `milvus-auth` (`MILVUS_USER`, `MILVUS_PASSWORD`) and +`milvus-maintenance-lock`. + ## Docs pipeline (core) -**Purpose:** Full rebuild of the documentation corpus into the v4 hybrid Milvus collection (`kubeflow_docs`). +**Purpose:** Sync one documentation source (a directory of Markdown in a GitHub repository) into the v4 hybrid Milvus collection (`kubeflow_docs`). -1. **Download GitHub directory** — `.md` / `.html` under a repo path (stock Python image) +1. **Download GitHub directory** — every `.md` / `.html` under a repo path at the default branch head; any unreadable file fails the run (stock Python image) 2. **Chunk and embed** — thin step that imports `hugo_ingest` / `canonical_rag_ingest` from the ingest image, then POSTs to TEI 3. **Store in Milvus** — thin step that imports `milvus_store` helpers from the same image @@ -26,8 +67,8 @@ Compile writes the YAML recipe (no Docker required). A **cluster run** needs the ```bash cd docs-agent-mcp/pipelines -# Optional: pin a built tag (default is ghcr.io/kubeflow/docs-rag-ingest:v0.1.0) -# set DOCS_INGEST_IMAGE=ghcr.io/kubeflow/docs-rag-ingest: +# The chunk and store steps run in this image; CI tags it with the short commit SHA. +export DOCS_INGEST_IMAGE=ghcr.io/kubeflow/docs-rag-ingest: python kubeflow-pipeline.py ``` diff --git a/docs-agent-mcp/pipelines/code-pipeline.py b/docs-agent-mcp/pipelines/code-pipeline.py index c6c6e05..674110f 100644 --- a/docs-agent-mcp/pipelines/code-pipeline.py +++ b/docs-agent-mcp/pipelines/code-pipeline.py @@ -157,12 +157,17 @@ def chunk_and_embed_code( code_data: dsl.Input[dsl.Dataset], chunk_size: int, chunk_overlap: int, + exclude_resource_kinds: str, embeddings_service_url: str, embedding_batch_size: int, embedded_data: dsl.Output[dsl.Dataset], ): """Chunk code files using YAML-aware and Python AST-aware parsing, then embed. + exclude_resource_kinds: comma-separated Kubernetes kinds to leave out. Generated + CustomResourceDefinitions were more than half of the index and crowd out the + manifests people ask about. + Routes files to the appropriate parser: - .yaml/.yml: Split at --- boundaries, extract K8s metadata - .py: Split at function/class boundaries using ast @@ -339,6 +344,7 @@ def chunk_code_file(content, file_path, c_size, c_overlap): # --- End inline chunking logic --- + excluded_kinds = {kind.strip() for kind in exclude_resource_kinds.split(",") if kind.strip()} records = [] with open(code_data.path, "r", encoding="utf-8") as f: @@ -362,6 +368,7 @@ def chunk_code_file(content, file_path, c_size, c_overlap): citation_url = file_data.get("citation_url") or f"https://github.com/{repo}/blob/main/{file_path}" chunks = chunk_code_file(content, file_path, chunk_size, chunk_overlap) + chunks = [chunk for chunk in chunks if chunk["resource_kind"] not in excluded_kinds] print(f"File: {file_path} -> {len(chunks)} chunks") if not chunks: @@ -559,12 +566,15 @@ def store_code_milvus( description="RAG pipeline for ingesting Kubeflow code and YAML manifests" ) def code_rag_pipeline( - repos: str = "kubeflow/manifests", - directory_paths: str = "apps/pipeline/upstream,apps/katib,common/istio,apps/jupyter", + # kubeflow/manifests was renamed; chunks indexed under the old repo name need a + # one-time delete (repo_name == "kubeflow/manifests") after the first run. + repos: str = "kubeflow/community-distribution", + directory_paths: str = "applications,common,experimental/helm", file_extensions: str = "yaml,yml,py,json", github_token: str = "", chunk_size: int = 1000, chunk_overlap: int = 100, + exclude_resource_kinds: str = "CustomResourceDefinition", embeddings_service_url: str = ( "http://embeddings-service-predictor.ml-infra.svc.cluster.local/embed" ), @@ -591,6 +601,7 @@ def code_rag_pipeline( code_data=download_task.outputs["code_data"], chunk_size=chunk_size, chunk_overlap=chunk_overlap, + exclude_resource_kinds=exclude_resource_kinds, embeddings_service_url=embeddings_service_url, embedding_batch_size=embedding_batch_size, ) diff --git a/docs-agent-mcp/pipelines/ingest_schedule.py b/docs-agent-mcp/pipelines/ingest_schedule.py new file mode 100644 index 0000000..72bb1ca --- /dev/null +++ b/docs-agent-mcp/pipelines/ingest_schedule.py @@ -0,0 +1,192 @@ +"""Register the ingest pipelines in Kubeflow Pipelines as recurring runs, or run one now. + + python ingest_schedule.py register --config schedules.json + python ingest_schedule.py run --config schedules.json --name docs-website + python ingest_schedule.py run --config schedules.json --name docs-distribution \ + --set collection_name=kubeflow_docs_staging + +`register` is what the chart's ingest Job runs after each sync (templates/ingest-schedule-job.yaml): +it compiles the pipelines in this directory, uploads them as a new pipeline version, and +replaces the recurring runs named in the config. Recurring runs that are in the experiment +but no longer in the config are deleted, so the config is the whole schedule. + +Config (JSON): + + { + "experiment": "flo-ingest", + "schedules": [ + {"name": "docs-website", "pipeline": "docs", "cron": "0 0 2 * * *", "arguments": {}} + ] + } + +`cron` has six fields (seconds first), as Kubeflow Pipelines expects. `arguments` override +the pipeline's defaults. From a laptop, port-forward first: +kubectl port-forward svc/ml-pipeline 8888:8888 -n kubeflow, then pass --host http://localhost:8888. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import sys +import tempfile +from datetime import datetime, timezone +from pathlib import Path + +PIPELINES_DIR = Path(__file__).resolve().parent +DEFAULT_HOST = "http://ml-pipeline.kubeflow.svc.cluster.local:8888" + +# key in the config -> (source file, pipeline function) +PIPELINES = { + "docs": ("kubeflow-pipeline.py", "github_rag_pipeline"), + "issues": ("issues-pipeline.py", "github_issues_rag_pipeline"), + "code": ("code-pipeline.py", "code_rag_pipeline"), +} + + +def load_config(path: str) -> dict: + config = json.loads(Path(path).read_text(encoding="utf-8")) + schedules = config.get("schedules") or [] + names = [schedule.get("name") for schedule in schedules] + if len(set(names)) != len(names) or not all(names): + raise ValueError("every schedule needs a unique name") + for schedule in schedules: + if schedule.get("pipeline") not in PIPELINES: + raise ValueError(f"{schedule['name']}: pipeline must be one of {sorted(PIPELINES)}") + if len(str(schedule.get("cron", "")).split()) != 6: + raise ValueError(f"{schedule['name']}: cron needs six fields (seconds first)") + if not isinstance(schedule.get("arguments", {}), dict): + raise ValueError(f"{schedule['name']}: arguments must be an object") + config.setdefault("experiment", "flo-ingest") + return config + + +def compile_pipeline(key: str, out_dir: str) -> str: + """Compile one pipeline to YAML and return the file path.""" + from kfp import compiler + + file_name, func_name = PIPELINES[key] + for extra in (PIPELINES_DIR, PIPELINES_DIR / "utils"): + if str(extra) not in sys.path: + sys.path.insert(0, str(extra)) + spec = importlib.util.spec_from_file_location(f"flo_ingest_{key}", PIPELINES_DIR / file_name) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + package_path = os.path.join(out_dir, f"{key}.yaml") + compiler.Compiler().compile(pipeline_func=getattr(module, func_name), package_path=package_path) + return package_path + + +def upload_version(client, key: str, package_path: str, version: str): + """Upload the compiled pipeline as a new version of flo-ingest-.""" + pipeline_name = f"flo-ingest-{key}" + pipeline_id = client.get_pipeline_id(pipeline_name) + if pipeline_id is None: + pipeline = client.upload_pipeline(package_path, pipeline_name=pipeline_name) + pipeline_id = pipeline.pipeline_id + uploaded = client.upload_pipeline_version(package_path, version, pipeline_id=pipeline_id) + return pipeline_id, uploaded.pipeline_version_id + + +def existing_recurring_runs(client, experiment_id: str) -> list: + runs, token = [], "" + while True: + page = client.list_recurring_runs(page_token=token, page_size=100, experiment_id=experiment_id) + runs.extend(page.recurring_runs or []) + token = page.next_page_token or "" + if not token: + return runs + + +def register(client, config: dict, version: str) -> None: + experiment = client.create_experiment(config["experiment"]) + experiment_id = experiment.experiment_id + + with tempfile.TemporaryDirectory() as out_dir: + versions = {} + for key in sorted({schedule["pipeline"] for schedule in config["schedules"]}): + versions[key] = upload_version(client, key, compile_pipeline(key, out_dir), version) + print(f"uploaded flo-ingest-{key} version {version}") + + # Create the new recurring runs first, then delete the old ones, so a failure half + # way leaves the previous schedule in place rather than none. + old = existing_recurring_runs(client, experiment_id) + for schedule in config["schedules"]: + pipeline_id, version_id = versions[schedule["pipeline"]] + client.create_recurring_run( + experiment_id=experiment_id, + job_name=schedule["name"], + description=f"flo ingest ({version})", + cron_expression=schedule["cron"], + max_concurrency=1, + no_catchup=True, + params=schedule.get("arguments", {}), + pipeline_id=pipeline_id, + version_id=version_id, + enabled=schedule.get("enabled", True), + enable_caching=False, + ) + print(f"scheduled {schedule['name']}: {schedule['cron']}") + for run in old: + client.delete_recurring_run(run.recurring_run_id) + print(f"removed previous recurring run {run.display_name}") + + +def run_once(client, config: dict, name: str, overrides: dict) -> str: + schedule = next((s for s in config["schedules"] if s["name"] == name), None) + if schedule is None: + raise SystemExit(f"no schedule named {name!r}; have {[s['name'] for s in config['schedules']]}") + arguments = {**schedule.get("arguments", {}), **overrides} + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + with tempfile.TemporaryDirectory() as out_dir: + result = client.create_run_from_pipeline_package( + pipeline_file=compile_pipeline(schedule["pipeline"], out_dir), + arguments=arguments, + run_name=f"{name}-manual-{stamp}", + experiment_name=config["experiment"], + enable_caching=False, + ) + print(f"started run {result.run_id} ({name}) with {arguments}") + return result.run_id + + +def parse_overrides(pairs: list[str]) -> dict: + """key=value pairs; values are parsed as JSON when possible (true, 300, "text").""" + overrides = {} + for pair in pairs: + key, sep, value = pair.partition("=") + if not sep or not key: + raise SystemExit(f"--set expects key=value, got {pair!r}") + try: + overrides[key] = json.loads(value) + except json.JSONDecodeError: + overrides[key] = value + return overrides + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("command", choices=["register", "run"]) + parser.add_argument("--config", required=True) + parser.add_argument("--host", default=os.environ.get("KFP_HOST", DEFAULT_HOST)) + parser.add_argument("--name", help="schedule to run (run only)") + parser.add_argument("--set", action="append", default=[], metavar="KEY=VALUE", help="argument override (run only)") + args = parser.parse_args(argv) + + import kfp + + config = load_config(args.config) + client = kfp.Client(host=args.host) + if args.command == "register": + stamp = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S") + register(client, config, f"{os.environ.get('INGEST_VERSION', 'dev')}-{stamp}") + else: + if not args.name: + raise SystemExit("run needs --name") + run_once(client, config, args.name, parse_overrides(args.set)) + + +if __name__ == "__main__": + main() diff --git a/docs-agent-mcp/pipelines/issues-pipeline.py b/docs-agent-mcp/pipelines/issues-pipeline.py index f91bcb4..c527e0f 100644 --- a/docs-agent-mcp/pipelines/issues-pipeline.py +++ b/docs-agent-mcp/pipelines/issues-pipeline.py @@ -37,6 +37,7 @@ def download_github_issues( labels: str, state: str, max_issues_per_repo: int, + updated_within_hours: int, github_token: str, issues_data: dsl.Output[dsl.Dataset] ): @@ -53,6 +54,8 @@ def download_github_issues( labels: Comma-separated labels to filter (e.g., "kind/bug,kind/question") state: Issue state - "open", "closed", or "all" max_issues_per_repo: Maximum issues to fetch per repository + updated_within_hours: Only issues updated in the last N hours (0 = no limit). + Scheduled runs use a window longer than their interval; a first backfill uses 0. github_token: GitHub personal access token for API authentication issues_data: Output dataset path (JSONL) """ @@ -148,6 +151,14 @@ def format_issue_markdown(title, repo_name, issue_number, url, labels_str, issue ) return content + since = "" + if int(updated_within_hours) > 0: + from datetime import datetime, timedelta, timezone + + cutoff = datetime.now(timezone.utc) - timedelta(hours=int(updated_within_hours)) + since = cutoff.strftime("%Y-%m-%dT%H:%M:%SZ") + print(f"Only issues updated since {since}") + for repo in repos.split(","): repo = repo.strip() if "/" not in repo: @@ -165,9 +176,14 @@ def format_issue_markdown(title, repo_name, issue_number, url, labels_str, issue params = { "state": state, "labels": labels, + # Most recently updated first, so the per-repo cap keeps active issues. + "sort": "updated", + "direction": "desc", "per_page": 100, "page": page } + if since: + params["since"] = since issues = api_request(url, params) if not issues: @@ -597,10 +613,16 @@ def store_issues_milvus( description="RAG pipeline for GitHub issues ingestion" ) def github_issues_rag_pipeline( - repos: str = "kubeflow/kubeflow,kubeflow/pipelines,kubeflow/manifests", + repos: str = ( + "kubeflow/pipelines,kubeflow/trainer,kubeflow/notebooks,kubeflow/katib,kubeflow/sdk," + "kubeflow/hub,kubeflow/spark-operator,kubeflow/community-distribution,kubeflow/dashboard," + "kubeflow/website" + ), labels: str = "", state: str = "all", max_issues_per_repo: int = 200, + # 0 = no limit (backfill). Scheduled runs pass a window longer than their interval. + updated_within_hours: int = 0, github_token: str = "", chunk_size: int = 1500, chunk_overlap: int = 150, @@ -617,6 +639,7 @@ def github_issues_rag_pipeline( labels=labels, state=state, max_issues_per_repo=max_issues_per_repo, + updated_within_hours=updated_within_hours, github_token=github_token, ) diff --git a/docs-agent-mcp/pipelines/kubeflow-pipeline.py b/docs-agent-mcp/pipelines/kubeflow-pipeline.py index 2f985bc..c7b86a0 100644 --- a/docs-agent-mcp/pipelines/kubeflow-pipeline.py +++ b/docs-agent-mcp/pipelines/kubeflow-pipeline.py @@ -29,100 +29,97 @@ ) @dsl.component( - base_image="docker.io/library/python:3.9", - packages_to_install=["requests", "beautifulsoup4"] + base_image="docker.io/library/python:3.11-slim", + packages_to_install=["requests==2.33.0", "beautifulsoup4==4.15.0"], ) def download_github_directory( repo_owner: str, repo_name: str, directory_path: str, + exclude_paths: str, github_token: str, - github_data: dsl.Output[dsl.Dataset] + github_data: dsl.Output[dsl.Dataset], ): - import os - import requests + """Download every .md/.html file under directory_path at the default branch head. + + One tree listing, then one raw download per file. Any file that cannot be read fails + the step: a partial download must not reach the store step, which removes documents + that are missing from the run. + """ import json - import base64 + import os import time - from bs4 import BeautifulSoup - - def resolve_github_token(token): - for candidate in (token, os.environ.get("Github_Pat"), os.environ.get("GITHUB_TOKEN")): - if candidate and str(candidate).strip(): - return str(candidate).strip() - return "" - - github_token = resolve_github_token(github_token) - if github_token: - print("Using authenticated GitHub API requests") - else: - print("WARNING: No github_token or Github_Pat env set; rate limits will be low (60 req/hr)") - headers = {"Authorization": f"token {github_token}"} if github_token else {} + import requests + from bs4 import BeautifulSoup - def api_request(url, params=None): - max_retries = 3 - for attempt in range(max_retries): + token = "" + for candidate in (github_token, os.environ.get("Github_Pat"), os.environ.get("GITHUB_TOKEN")): + if candidate and str(candidate).strip(): + token = str(candidate).strip() + break + if not token: + print("WARNING: no github_token or Github_Pat; GitHub allows 60 API requests per hour") + headers = {"Authorization": f"token {token}"} if token else {} + + def get(url, **kwargs): + last_error = "" + for attempt in range(4): try: - resp = requests.get(url, params=params, headers=headers) - - if resp.status_code == 403: - remaining = resp.headers.get("X-RateLimit-Remaining", "0") - if remaining == "0": - reset_time = int(resp.headers.get("X-RateLimit-Reset", 0)) - wait_time = max(reset_time - int(time.time()), 60) - print(f"Rate limited. Waiting {wait_time}s...") - time.sleep(min(wait_time, 300)) - continue - print(f"Forbidden (403) for {url}: {resp.text[:200]}") - - if resp.status_code == 200: - return resp.json() - - print(f"API error: HTTP {resp.status_code} for {url}") - return None - - except Exception as e: - print(f"Request failed (attempt {attempt + 1}): {e}") - time.sleep(2 ** attempt) - - return None - - def get_files_recursive(url): - files = [] - items = api_request(url) - if not items or not isinstance(items, list): - return files - - for item in items: - if item['type'] == 'file' and (item['name'].endswith('.md') or item['name'].endswith('.html')): - file_data = api_request(item['url']) - if not file_data or 'content' not in file_data: - print(f"Skipping unreadable file: {item['path']}") - continue - content = base64.b64decode(file_data['content']).decode('utf-8') - - if item['name'].endswith('.html'): - soup = BeautifulSoup(content, 'html.parser') - content = soup.get_text(separator=' ', strip=True) - - files.append({ - 'path': item['path'], - 'content': content, - 'file_name': item['name'] - }) - elif item['type'] == 'dir': - files.extend(get_files_recursive(item['url'])) - - return files - - api_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/contents/{directory_path}" - files = get_files_recursive(api_url) - print(f"Downloaded {len(files)} files") - - with open(github_data.path, 'w', encoding='utf-8') as f: - for file_data in files: - f.write(json.dumps(file_data, ensure_ascii=False) + '\n') + resp = requests.get(url, headers=headers, timeout=60, **kwargs) + except requests.RequestException as exc: + last_error = str(exc) + time.sleep(2**attempt) + continue + if resp.status_code == 200: + return resp + last_error = f"HTTP {resp.status_code}: {resp.text[:200]}" + if resp.status_code in (403, 429) and resp.headers.get("X-RateLimit-Remaining") == "0": + reset = int(resp.headers.get("X-RateLimit-Reset", "0") or 0) + wait = min(max(reset - int(time.time()), 30), 600) + print(f"GitHub rate limit reached; waiting {wait}s") + time.sleep(wait) + elif resp.status_code >= 500 or resp.status_code == 429: + time.sleep(2**attempt) + else: + break + raise RuntimeError(f"GET {url} failed: {last_error}") + + repo_api = f"https://api.github.com/repos/{repo_owner}/{repo_name}" + branch = get(repo_api).json()["default_branch"] + commit_sha = get(f"{repo_api}/commits/{branch}").json()["sha"] + tree = get(f"{repo_api}/git/trees/{commit_sha}", params={"recursive": "1"}).json() + if tree.get("truncated"): + raise RuntimeError(f"Tree listing of {repo_owner}/{repo_name} is truncated; narrow directory_path") + + prefix = directory_path.strip("/") + prefix = f"{prefix}/" if prefix else "" + excluded = tuple(part.strip().strip("/") for part in exclude_paths.split(",") if part.strip()) + + def wanted(path): + if not path.startswith(prefix) or not path.endswith((".md", ".html")): + return False + return not any(path == ex or path.startswith(f"{ex}/") for ex in excluded) + + paths = sorted(item["path"] for item in tree["tree"] if item["type"] == "blob" and wanted(item["path"])) + if not paths: + raise RuntimeError(f"No .md/.html files under '{prefix}' in {repo_owner}/{repo_name}@{commit_sha[:7]}") + print(f"{repo_owner}/{repo_name}@{commit_sha[:7]} ({branch}): {len(paths)} files under '{prefix}'") + + with open(github_data.path, "w", encoding="utf-8") as out: + for path in paths: + raw_url = f"https://raw.githubusercontent.com/{repo_owner}/{repo_name}/{commit_sha}/{path}" + content = get(raw_url).content.decode("utf-8", errors="replace") + if path.endswith(".html"): + content = BeautifulSoup(content, "html.parser").get_text(separator=" ", strip=True) + record = { + "path": path, + "content": content, + "file_name": path.rsplit("/", 1)[-1], + "commit_sha": commit_sha, + } + out.write(json.dumps(record, ensure_ascii=False) + "\n") + print(f"Downloaded {len(paths)} files") @dsl.component(base_image=INGEST_IMAGE) @@ -132,6 +129,7 @@ def chunk_and_embed( base_url: str, target_tokens: int, overlap_tokens: int, + doc_type: str, embeddings_service_url: str, embedding_batch_size: int, embedded_data: dsl.Output[dsl.Dataset], @@ -147,6 +145,7 @@ def chunk_and_embed( base_url=base_url, target_tokens=int(target_tokens), overlap_tokens=int(overlap_tokens), + doc_type=doc_type, ) records = embed_chunk_records( records, @@ -168,9 +167,16 @@ def store_milvus( clean_rebuild: bool, clean_rebuild_confirmation: str, maintenance_lock_token: str, + prune_removed: bool, + repo_name: str, + directory_path: str, ): from milvus_store import store_embedded_records + # document_id is ":"; exclude_paths under it are pruned too. + directory = directory_path.strip("/") + prune_prefix = f"{repo_name}:{directory}/" if directory else f"{repo_name}:" + store_embedded_records( embedded_data.path, milvus_host=milvus_host, @@ -179,6 +185,7 @@ def store_milvus( clean_rebuild=clean_rebuild, clean_rebuild_confirmation=clean_rebuild_confirmation, maintenance_lock_token=maintenance_lock_token, + prune_prefix=prune_prefix if prune_removed else "", ) @@ -190,8 +197,16 @@ def github_rag_pipeline( repo_owner: str = "kubeflow", repo_name: str = "website", directory_path: str = "content/en/docs", + # Comma-separated paths (files or directories) to skip under directory_path. + exclude_paths: str = "", github_token: str = "", base_url: str = "https://www.kubeflow.org/docs", + # Empty = infer from the Hugo path (documentation, release, nav, redirect). Set it for + # other sources, e.g. "distribution" for kubeflow/community-distribution. + doc_type: str = "", + # Remove indexed pages of this source that are gone upstream. Needs a full download + # of the source, which is what this pipeline does. + prune_removed: bool = True, target_tokens: int = DEFAULT_TARGET_TOKENS, overlap_tokens: int = DEFAULT_OVERLAP_TOKENS, embeddings_service_url: str = ( @@ -210,7 +225,8 @@ def github_rag_pipeline( repo_owner=repo_owner, repo_name=repo_name, directory_path=directory_path, - github_token=github_token + exclude_paths=exclude_paths, + github_token=github_token, ) if k8s is not None: @@ -227,6 +243,7 @@ def github_rag_pipeline( base_url=base_url, target_tokens=target_tokens, overlap_tokens=overlap_tokens, + doc_type=doc_type, embeddings_service_url=embeddings_service_url, embedding_batch_size=embedding_batch_size, ) @@ -240,6 +257,9 @@ def github_rag_pipeline( clean_rebuild=clean_rebuild, clean_rebuild_confirmation=clean_rebuild_confirmation, maintenance_lock_token=maintenance_lock_token, + prune_removed=prune_removed, + repo_name=repo_name, + directory_path=directory_path, ) if k8s is not None: diff --git a/docs-agent-mcp/pipelines/submit_run.py b/docs-agent-mcp/pipelines/submit_run.py deleted file mode 100644 index 18b321a..0000000 --- a/docs-agent-mcp/pipelines/submit_run.py +++ /dev/null @@ -1,50 +0,0 @@ -""" -Submit and run the github_rag_pipeline on the local KFP API server. -Port-forward must be active: kubectl port-forward svc/ml-pipeline 8888:8888 -n kubeflow -""" -import sys -from pathlib import Path - -import kfp - -_UTILS_DIR = Path(__file__).resolve().parent / "utils" -if str(_UTILS_DIR) not in sys.path: - sys.path.insert(0, str(_UTILS_DIR)) - -from utils import ( - DEFAULT_DOCS_CHUNK_OVERLAP, - DEFAULT_DOCS_CHUNK_SIZE, - DEFAULT_DOCS_MAX_TEI_CHARS, - DEFAULT_EMBEDDING_DIM, - DOCS_COLLECTION, -) - -KFP_HOST = "http://localhost:8888" -PIPELINE_YAML = "github_rag_pipeline.yaml" - -client = kfp.Client(host=KFP_HOST) - -run = client.create_run_from_pipeline_package( - pipeline_file=PIPELINE_YAML, - arguments={ - "repo_owner": "kubeflow", - "repo_name": "website", - "directory_path": "content/en/docs", - "github_token": "", # pass a token if you hit rate limits - "base_url": "https://www.kubeflow.org/docs", - "chunk_size": DEFAULT_DOCS_CHUNK_SIZE, - "chunk_overlap": DEFAULT_DOCS_CHUNK_OVERLAP, - "max_tei_chars": DEFAULT_DOCS_MAX_TEI_CHARS, - "embedding_dim": DEFAULT_EMBEDDING_DIM, - "milvus_host": "milvus-milvus.ml-infra.svc.cluster.local", - "milvus_port": "19530", - "collection_name": DOCS_COLLECTION, - }, - run_name="kubeflow-docs-rag-run-1", - experiment_name="kubeflow-docs-rag", - enable_caching=False, -) - -print("Run submitted!") -print(f"Run ID : {run.run_id}") -print(f"View at : {KFP_HOST}/#/runs/details/{run.run_id}") diff --git a/docs-agent-mcp/pipelines/utils/canonical_rag_ingest.py b/docs-agent-mcp/pipelines/utils/canonical_rag_ingest.py index 5a28be3..20f6e98 100644 --- a/docs-agent-mcp/pipelines/utils/canonical_rag_ingest.py +++ b/docs-agent-mcp/pipelines/utils/canonical_rag_ingest.py @@ -672,6 +672,17 @@ def extract_release_date( return None +# "[Figure: FOSSA Status](https://...)" (a linked badge) or a bare "Figure: alt text". +BADGE_OR_FIGURE_RE = re.compile(r"\[Figure: [^\]]*\]\([^)]*\)|Figure: [^\n\]]*") +MIN_CHUNK_WORD_CHARS = 20 + + +def is_badge_only_chunk(text: str) -> bool: + """True for a chunk that is only badges or figure captions (no sentence to retrieve).""" + remaining = BADGE_OR_FIGURE_RE.sub("", text or "") + return sum(ch.isalnum() for ch in remaining) < MIN_CHUNK_WORD_CHARS + + def build_milvus_records( file_data: dict[str, Any], *, @@ -679,15 +690,20 @@ def build_milvus_records( base_url: str, target_tokens: int = DEFAULT_TARGET_TOKENS, overlap_tokens: int = DEFAULT_OVERLAP_TOKENS, + doc_type: str = "", ) -> list[dict[str, Any]]: - """Parse and chunk one downloaded file record into Milvus-ready dicts.""" + """Parse and chunk one downloaded file record into Milvus-ready dicts. + + doc_type: set for sources that are not the Hugo website (for example "distribution" + for kubeflow/community-distribution). Empty = infer from the path and frontmatter. + """ content = file_data.get("content", "") or "" file_path = file_data.get("path", "") or "" file_name = file_data.get("file_name") or os.path.basename(file_path) parsed = parse_canonical_document(content) frontmatter = parsed.frontmatter - doc_type = infer_doc_type(file_path, frontmatter) + doc_type = (doc_type or "").strip() or infer_doc_type(file_path, frontmatter) release_date = extract_release_date( doc_type=doc_type, source_text=content, @@ -730,11 +746,18 @@ def build_milvus_records( target_tokens=target_tokens, overlap_tokens=overlap_tokens, ) + chunks = [chunk for chunk in chunks if not is_badge_only_chunk(chunk["content_text"])] citation_url = build_citation_url(file_path, base_url) if doc_type == "redirect" and frontmatter.get("manualLink"): citation_url = str(frontmatter["manualLink"])[:1024] + # Plain READMEs have no frontmatter title: use the first top-level heading. + title = parsed.title + if not title: + first_heading = re.search(r"^#\s+(.+?)\s*$", content, flags=re.MULTILINE) + title = first_heading.group(1) if first_heading else file_path + file_unique_id = f"{repo_name}:{file_path}" doc_status = infer_doc_status(file_path) records: list[dict[str, Any]] = [] @@ -749,7 +772,7 @@ def build_milvus_records( "citation_url": citation_url, "chunk_index": chunk_idx, "content_text": chunk["content_text"][:MAX_CONTENT_TEXT_CHARS], - "title": parsed.title[:256], + "title": title[:256], "weight": parsed.weight, "doc_type": doc_type[:32], "version": str(frontmatter.get("version") or "")[:32], @@ -808,6 +831,7 @@ def dumps_jsonl(records: Iterable[dict[str, Any]]) -> str: "CanonicalSection", "build_citation_url", "build_milvus_records", + "is_badge_only_chunk", "chunk_canonical_document", "extract_release_date", "dumps_jsonl", diff --git a/docs-agent-mcp/pipelines/utils/milvus_store.py b/docs-agent-mcp/pipelines/utils/milvus_store.py index 61ad6bb..434364b 100644 --- a/docs-agent-mcp/pipelines/utils/milvus_store.py +++ b/docs-agent-mcp/pipelines/utils/milvus_store.py @@ -25,8 +25,13 @@ ) APPROVED_DOCS_COLLECTION = DOCS_COLLECTION +# Test copies of the docs collection (kubeflow_docs_staging, kubeflow_docs_preview, ...). +APPROVED_DOCS_COLLECTION_PREFIX = f"{DOCS_COLLECTION}_" CLEAN_REBUILD_CONFIRMATION = f"DROP {APPROVED_DOCS_COLLECTION}" MAINTENANCE_LOCK_ENV = "MILVUS_MAINTENANCE_LOCK" +# A sync that would remove more than this share of a source's indexed documents is +# treated as a failed download, not as upstream deleting its docs. +MAX_PRUNE_FRACTION = 0.2 DENSE_FIELD = "vector" SPARSE_FIELD = "sparse_vector" @@ -53,24 +58,32 @@ def truncate_utf8(value: Any, max_bytes: int) -> str: def validate_production_collection_name(collection_name: str) -> None: name = (collection_name or "").strip() - if name != APPROVED_DOCS_COLLECTION: + suffix = name[len(APPROVED_DOCS_COLLECTION_PREFIX) :] + is_test_copy = name.startswith(APPROVED_DOCS_COLLECTION_PREFIX) and suffix.replace("_", "").isalnum() + if name != APPROVED_DOCS_COLLECTION and not is_test_copy: raise ValueError( f"Refusing unexpected collection name '{collection_name}'. " - f"Production docs indexing must target exactly '{APPROVED_DOCS_COLLECTION}'." + f"Docs indexing must target '{APPROVED_DOCS_COLLECTION}' or a test copy named " + f"'{APPROVED_DOCS_COLLECTION_PREFIX}'." ) +def clean_rebuild_confirmation_for(collection_name: str) -> str: + return f"DROP {(collection_name or '').strip()}" + + def validate_clean_rebuild_gates( *, clean_rebuild: bool, clean_rebuild_confirmation: str, maintenance_lock_token: str, maintenance_lock_env: str | None = None, + collection_name: str = APPROVED_DOCS_COLLECTION, ) -> None: if not clean_rebuild: return - expected_confirmation = CLEAN_REBUILD_CONFIRMATION + expected_confirmation = clean_rebuild_confirmation_for(collection_name) provided = (clean_rebuild_confirmation or "").strip() if provided != expected_confirmation: raise ValueError( @@ -179,6 +192,7 @@ def chunk_github_jsonl( base_url: str, target_tokens: int = DEFAULT_TARGET_TOKENS, overlap_tokens: int = DEFAULT_OVERLAP_TOKENS, + doc_type: str = "", ingest_module: Any | None = None, ) -> list[dict[str, Any]]: """Parse/chunk downloaded GitHub JSONL using canonical ingestion logic.""" @@ -191,12 +205,14 @@ def chunk_github_jsonl( if not stripped: continue file_data = json.loads(stripped) + kwargs = {"doc_type": doc_type} if doc_type else {} file_records = build_records( file_data, repo_name=repo_name, base_url=base_url, target_tokens=target_tokens, overlap_tokens=overlap_tokens, + **kwargs, ) records.extend(file_records) @@ -290,6 +306,56 @@ def compute_validation_metrics( } +def _quoted(values: Sequence[str]) -> str: + return ", ".join(json.dumps(value) for value in values) + + +def find_stale_document_ids( + milvus: MilvusClient, + collection_name: str, + *, + prune_prefix: str, + current_ids: set[str], +) -> list[str]: + """Indexed documents under prune_prefix that the source no longer has.""" + indexed: set[str] = set() + iterator = milvus.query_iterator( + collection_name=collection_name, + filter=f'{DOCUMENT_ID_FIELD} like "{prune_prefix}%"', + output_fields=[DOCUMENT_ID_FIELD], + batch_size=1000, + ) + try: + while True: + batch = iterator.next() + if not batch: + break + indexed.update(row[DOCUMENT_ID_FIELD] for row in batch) + finally: + iterator.close() + + stale = sorted(indexed - current_ids) + if indexed and len(stale) > MAX_PRUNE_FRACTION * len(indexed): + raise RuntimeError( + f"Refusing to prune {len(stale)} of {len(indexed)} indexed documents under " + f"'{prune_prefix}' (limit {MAX_PRUNE_FRACTION:.0%}). The download is probably " + "incomplete; nothing was written." + ) + return stale + + +def delete_documents(milvus: MilvusClient, collection_name: str, document_ids: Sequence[str]) -> int: + deleted = 0 + for start in range(0, len(document_ids), DELETE_BATCH_SIZE): + batch_ids = document_ids[start : start + DELETE_BATCH_SIZE] + result = milvus.delete( + collection_name=collection_name, + filter=f"{DOCUMENT_ID_FIELD} in [{_quoted(batch_ids)}]", + ) + deleted += int(result.get("delete_count", 0) or 0) + return deleted + + def store_embedded_records( embedded_data_path: str, *, @@ -299,14 +365,21 @@ def store_embedded_records( clean_rebuild: bool, clean_rebuild_confirmation: str, maintenance_lock_token: str, + prune_prefix: str = "", client: MilvusClient | None = None, ) -> dict[str, Any]: - """Read embedded JSONL and write the docs collection.""" + """Read embedded JSONL and write the docs collection. + + prune_prefix: when set (for example "website:content/en/docs/"), documents indexed + under that document_id prefix that are not in this run are deleted, so pages removed + upstream leave the index. Only pass it when the run downloaded that whole source. + """ validate_production_collection_name(collection_name) validate_clean_rebuild_gates( clean_rebuild=clean_rebuild, clean_rebuild_confirmation=clean_rebuild_confirmation, maintenance_lock_token=maintenance_lock_token, + collection_name=collection_name, ) milvus_user = os.environ.get("MILVUS_USER", "root") @@ -339,6 +412,8 @@ def store_embedded_records( if clean_rebuild and milvus.has_collection(collection_name): milvus.drop_collection(collection_name) + collection_created = False + if milvus.has_collection(collection_name): info = milvus.describe_collection(collection_name) description = info.get("description") or "" @@ -349,6 +424,7 @@ def store_embedded_records( "Use clean_rebuild=true with typed confirmation to recreate." ) else: + collection_created = True milvus.create_collection( collection_name=collection_name, schema=build_lean_v4_schema(), @@ -358,19 +434,26 @@ def store_embedded_records( rows = [prepare_v4_insert_row(record) for record in records] + pruned_documents = 0 if not clean_rebuild: unique_ids = sorted({row[DOCUMENT_ID_FIELD] for row in rows}) - deleted = 0 - for start in range(0, len(unique_ids), DELETE_BATCH_SIZE): - batch_ids = unique_ids[start : start + DELETE_BATCH_SIZE] - quoted = ", ".join(f'"{doc_id}"' for doc_id in batch_ids) - result = milvus.delete( - collection_name=collection_name, - filter=f"{DOCUMENT_ID_FIELD} in [{quoted}]", + stale_ids: list[str] = [] + if prune_prefix and not collection_created: + # Checked before any write, so a refused prune leaves the index untouched. + milvus.load_collection(collection_name) + stale_ids = find_stale_document_ids( + milvus, + collection_name, + prune_prefix=prune_prefix, + current_ids=set(unique_ids), ) - deleted += int(result.get("delete_count", 0) or 0) + deleted = delete_documents(milvus, collection_name, unique_ids) if deleted: print(f"Deleted {deleted} existing chunks for {len(unique_ids)} documents") + if stale_ids: + delete_documents(milvus, collection_name, stale_ids) + pruned_documents = len(stale_ids) + print(f"Pruned {pruned_documents} documents no longer in the source: {stale_ids[:20]}") inserted_count = 0 batch: list[dict[str, Any]] = [] @@ -407,6 +490,7 @@ def store_embedded_records( sparse_ready=SPARSE_FIELD in index_names, clean_rebuild=clean_rebuild, ) + metrics["pruned_documents"] = pruned_documents print("VALIDATION_METRICS=" + json.dumps(metrics, ensure_ascii=False, separators=(",", ":"))) return metrics diff --git a/tests/eval/README.md b/tests/eval/README.md index bea1ffe..b872034 100644 --- a/tests/eval/README.md +++ b/tests/eval/README.md @@ -40,7 +40,7 @@ cd docs-agent-mcp/pipelines Docs (only if gate 1 missed the Katib page): ```python -# same pattern as submit_run.py +# docs-agent-mcp/pipelines/ingest_schedule.py run --name docs-website --set ... # directory_path="content/en/docs/components/katib" # collection_name="kubeflow_docs" ``` diff --git a/tests/eval/flo_eval_dataset.json b/tests/eval/flo_eval_dataset.json index 78ba85a..df19d89 100644 --- a/tests/eval/flo_eval_dataset.json +++ b/tests/eval/flo_eval_dataset.json @@ -135,6 +135,8 @@ {"id": "docs-katib-trial-limits", "agent": "docs", "group": "docs_routing", "q": "Where do I set resource limits for a Katib trial template?", "expect": {"type": "call", "tools": ["search_kubeflow_docs"], "keep": ["trial"]}}, {"id": "docs-kfp-caching", "agent": "docs", "group": "docs_routing", "q": "How does KFP caching work and how do I disable it for one task?", "expect": {"type": "call", "tools": ["search_kubeflow_docs"], "keep": ["cach"]}}, {"id": "docs-notebook-default-size", "agent": "docs", "group": "docs_routing", "q": "What is the default workspace volume size for Kubeflow notebooks?", "expect": {"type": "call", "tools": ["search_kubeflow_docs"], "keep": ["notebook"]}}, + {"id": "docs-helm-chart-short", "agent": "docs", "group": "docs_routing", "q": "do we have a helm chart?", "expect": {"type": "call", "tools": ["search_kubeflow_docs"], "keep": ["helm"]}}, + {"id": "docs-workspaces-new-release", "agent": "docs", "group": "docs_routing", "q": "Are the new workspaces in the new release?", "expect": {"type": "call", "tools": ["search_kubeflow_docs"], "keep": ["workspace", "release"]}}, {"id": "fid-latest-release", "agent": "docs", "group": "query_fidelity", "q": "What is the latest Kubeflow release?", "expect": {"type": "call", "tools": ["search_kubeflow_docs"], "keep": ["latest"]}}, {"id": "fid-current-kfp", "agent": "docs", "group": "query_fidelity", "q": "What is the current version of Kubeflow Pipelines?", "expect": {"type": "call", "tools": ["search_kubeflow_docs"], "keep_any": ["current", "latest"]}}, diff --git a/tests/test_rag_quality_ingest.py b/tests/test_rag_quality_ingest.py new file mode 100644 index 0000000..93e0a24 --- /dev/null +++ b/tests/test_rag_quality_ingest.py @@ -0,0 +1,341 @@ +"""Tests for latest-query retrieval, distribution docs ingest, pruning and ingest schedules.""" + +import json +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +sys.modules.setdefault("pymilvus", MagicMock()) +ROOT = Path(__file__).resolve().parents[1] / "docs-agent-mcp" +for path in (ROOT / "mcp-server", ROOT / "pipelines", ROOT / "pipelines" / "utils"): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +import canonical_rag_ingest # noqa: E402 +import ingest_schedule # noqa: E402 +import intent_router # noqa: E402 +import milvus_search # noqa: E402 +import milvus_store # noqa: E402 +import reranker_client # noqa: E402 + + +def hit(doc_id, doc_type="documentation", release_date=None, text=""): + return { + "id": doc_id, + "distance": 1.0, + "entity": { + "doc_type": doc_type, + "release_date": release_date, + "content_text": text or f"chunk {doc_id}", + "file_path": f"docs/{doc_id}.md", + }, + } + + +# --- latest/current queries --------------------------------------------------------- + + +@pytest.mark.parametrize( + "query", + [ + "Are the new workspaces in the new release?", + "latest Kubeflow release", + "what changed in the last release", + ], +) +def test_release_questions_also_search_general_docs(query): + plan = intent_router.pick_search_plan(query) + assert plan.intent == "temporal" + assert plan.add_general_hits is True + + +def test_new_alone_is_not_a_release_question(): + assert intent_router.pick_search_plan("create a new notebook server").intent == "general" + + +def test_merge_keeps_release_notes_first_and_shares_slots(): + release = [hit(f"r{i}", "release", 100 - i) for i in range(5)] + general = [hit(f"g{i}") for i in range(5)] + merged = intent_router.merge_release_and_general_hits(release, general, 5) + assert [h["id"] for h in merged] == ["r0", "r1", "r2", "g0", "g1"] + + +def test_merge_fills_unused_slots_and_drops_duplicates(): + release = [hit("r0", "release", 1)] + general = [hit("r0", "release", 1), hit("g0"), hit("g1"), hit("g2"), hit("g3"), hit("g4")] + merged = intent_router.merge_release_and_general_hits(release, general, 5) + assert [h["id"] for h in merged] == ["r0", "g0", "g1", "g2", "g3"] + assert intent_router.merge_release_and_general_hits([], general[1:], 3) == general[1:4] + + +@pytest.fixture +def search(monkeypatch): + """milvus_search with Milvus, embeddings and the reranker stubbed out.""" + calls = SimpleNamespace(bm25=[], hybrid=[]) + release = [hit(f"r{i}", "release", 100 - i) for i in range(6)] + general = [hit("workspaces-guide"), hit("katib")] + + def bm25(collection, query, limit, fields, filter_expr=""): + calls.bm25.append(filter_expr) + return release + + def hybrid(collection, query, embedding, limit, fields, filter_expr="", candidate_depth=None): + calls.hybrid.append(filter_expr) + return general + + monkeypatch.setattr(milvus_search, "_load", lambda name: None) + monkeypatch.setattr(milvus_search, "collection_has_bm25", lambda name: True) + monkeypatch.setattr(milvus_search, "collection_has_release_fields", lambda name: True) + monkeypatch.setattr(milvus_search, "_require_embedding", lambda query: [0.0]) + monkeypatch.setattr(milvus_search, "bm25_search", bm25) + monkeypatch.setattr(milvus_search, "hybrid_search", hybrid) + monkeypatch.setattr(reranker_client, "rerank_hits", lambda query, hits: hits) + return calls + + +def test_latest_query_returns_release_notes_and_component_docs(search): + hits, meta = milvus_search.search_docs_auto("Are the new workspaces in the new release?", 5, []) + assert [h["id"] for h in hits] == ["r0", "r1", "r2", "workspaces-guide", "katib"] + assert search.bm25 == ['doc_type == "release"'] + assert search.hybrid == [""] + assert meta["intent"] == "temporal" + assert meta["general_hits"] == 2 + + +def test_latest_query_survives_a_failed_general_search(search, monkeypatch): + def broken(*args, **kwargs): + raise RuntimeError("milvus down") + + monkeypatch.setattr(milvus_search, "hybrid_search", broken) + hits, meta = milvus_search.search_docs_auto("latest Kubeflow release", 5, []) + assert [h["id"] for h in hits] == ["r0", "r1", "r2", "r3", "r4"] + assert meta["general_hits"] == 0 + + +def test_reranker_gets_more_candidates_than_top_k(search, monkeypatch): + limits = [] + + def hybrid(collection, query, embedding, limit, fields, filter_expr="", candidate_depth=None): + limits.append(limit) + return [hit(f"g{i}") for i in range(limit)] + + monkeypatch.setattr(milvus_search, "hybrid_search", hybrid) + monkeypatch.setattr(reranker_client, "rerank_hits", lambda query, hits: list(reversed(hits))) + + monkeypatch.setenv("RERANKER_URL", "") + hits, _ = milvus_search.search_docs_auto("Kubeflow Workspaces overview", 5, []) + assert limits == [5] and len(hits) == 5 + + monkeypatch.setenv("RERANKER_URL", "http://reranker/rerank") + hits, _ = milvus_search.search_docs_auto("Kubeflow Workspaces overview", 5, []) + assert limits == [5, milvus_search.RERANK_CANDIDATE_DEPTH] + assert [h["id"] for h in hits] == ["g19", "g18", "g17", "g16", "g15"] + + +# --- distribution docs --------------------------------------------------------------- + +README = """[![FOSSA Status](https://app.fossa.com/badge.svg)](https://app.fossa.com/projects/x) + +# Kubeflow Community Distribution + +## Overview + +All components are deployable with `kustomize`. In addition, experimental Helm charts for +selected components are available under `experimental/helm/`. +""" + + +def test_readme_is_indexed_as_distribution_with_a_title_and_github_citation(): + records = canonical_rag_ingest.build_milvus_records( + {"path": "README.md", "content": README}, + repo_name="community-distribution", + base_url="https://github.com/kubeflow/community-distribution/blob/master", + doc_type="distribution", + ) + assert records + assert {r["doc_type"] for r in records} == {"distribution"} + assert {r["title"] for r in records} == {"Kubeflow Community Distribution"} + assert records[0]["citation_url"] == "https://github.com/kubeflow/community-distribution/blob/master/README.md" + assert records[0]["file_unique_id"] == "community-distribution:README.md" + assert any("Helm charts" in r["content_text"] for r in records) + assert not any(canonical_rag_ingest.is_badge_only_chunk(r["content_text"]) for r in records) + + +def test_website_pages_still_infer_their_doc_type(): + records = canonical_rag_ingest.build_milvus_records( + { + "path": "content/en/docs/kubeflow-distribution/releases/kubeflow-26.03.md", + "content": "---\ntitle: Kubeflow 26.03\n---\n\nRelease notes for this version of Kubeflow.\n", + }, + repo_name="website", + base_url="https://www.kubeflow.org/docs", + ) + assert {r["doc_type"] for r in records} == {"release"} + assert {r["title"] for r in records} == {"Kubeflow 26.03"} + + +@pytest.mark.parametrize( + ("text", "badge_only"), + [ + ("[Figure: FOSSA Status](https://app.fossa.com/projects/x?ref=badge_shield)", True), + ("Figure: architecture diagram", True), + ("Kubeflow Workspaces provides on-demand development environments.", False), + ("[Figure: logo](https://x/y.png)\n\nKubeflow Hub stores models and their metadata.", False), + ], +) +def test_badge_only_chunks_are_detected(text, badge_only): + assert canonical_rag_ingest.is_badge_only_chunk(text) is badge_only + + +# --- store: collection names and pruning -------------------------------------------- + + +@pytest.mark.parametrize("name", ["kubeflow_docs", "kubeflow_docs_staging", "kubeflow_docs_pr12"]) +def test_docs_collection_and_test_copies_are_allowed(name): + milvus_store.validate_production_collection_name(name) + + +@pytest.mark.parametrize("name", ["issues_rag", "kubeflow_docs_", "kubeflow_docs_a;drop", "", "other_kubeflow_docs"]) +def test_other_collections_are_refused(name): + with pytest.raises(ValueError): + milvus_store.validate_production_collection_name(name) + + +def test_clean_rebuild_confirmation_names_the_target_collection(monkeypatch): + monkeypatch.setenv(milvus_store.MAINTENANCE_LOCK_ENV, "lock") + with pytest.raises(ValueError): + milvus_store.validate_clean_rebuild_gates( + clean_rebuild=True, + clean_rebuild_confirmation="DROP kubeflow_docs", + maintenance_lock_token="lock", + collection_name="kubeflow_docs_staging", + ) + milvus_store.validate_clean_rebuild_gates( + clean_rebuild=True, + clean_rebuild_confirmation="DROP kubeflow_docs_staging", + maintenance_lock_token="lock", + collection_name="kubeflow_docs_staging", + ) + + +class FakeIterator: + def __init__(self, rows): + self.pages = [rows, []] + + def next(self): + return self.pages.pop(0) + + def close(self): + pass + + +def fake_milvus(indexed_ids): + client = MagicMock() + client.query_iterator.return_value = FakeIterator([{"document_id": doc_id} for doc_id in indexed_ids]) + client.delete.return_value = {"delete_count": 1} + return client + + +def test_stale_documents_are_found_under_the_source_prefix(): + indexed = [f"website:content/en/docs/p{i}.md" for i in range(10)] + client = fake_milvus(indexed) + stale = milvus_store.find_stale_document_ids( + client, + "kubeflow_docs", + prune_prefix="website:content/en/docs/", + current_ids=set(indexed[:9]), + ) + assert stale == ["website:content/en/docs/p9.md"] + assert client.query_iterator.call_args.kwargs["filter"] == 'document_id like "website:content/en/docs/%"' + + +def test_prune_refuses_when_too_much_of_the_source_is_missing(): + indexed = [f"website:content/en/docs/p{i}.md" for i in range(10)] + with pytest.raises(RuntimeError, match="Refusing to prune 5 of 10"): + milvus_store.find_stale_document_ids( + fake_milvus(indexed), + "kubeflow_docs", + prune_prefix="website:content/en/docs/", + current_ids=set(indexed[:5]), + ) + + +def test_first_ingest_of_a_source_has_nothing_to_prune(): + stale = milvus_store.find_stale_document_ids( + fake_milvus([]), + "kubeflow_docs", + prune_prefix="community-distribution:", + current_ids={"community-distribution:README.md"}, + ) + assert stale == [] + + +# --- recurring runs ------------------------------------------------------------------- + + +def write_config(tmp_path, schedules): + path = tmp_path / "schedules.json" + path.write_text(json.dumps({"schedules": schedules})) + return str(path) + + +def test_schedule_config_is_validated(tmp_path): + good = {"name": "docs-website", "pipeline": "docs", "cron": "0 0 2 * * *", "arguments": {}} + assert ingest_schedule.load_config(write_config(tmp_path, [good]))["experiment"] == "flo-ingest" + for bad in ( + {**good, "cron": "0 2 * * *"}, # five fields: KFP wants seconds first + {**good, "pipeline": "unknown"}, + {**good, "arguments": []}, + ): + with pytest.raises(ValueError): + ingest_schedule.load_config(write_config(tmp_path, [bad])) + with pytest.raises(ValueError): + ingest_schedule.load_config(write_config(tmp_path, [good, good])) + + +def test_register_creates_new_recurring_runs_before_removing_old_ones(tmp_path, monkeypatch): + config = ingest_schedule.load_config( + write_config( + tmp_path, + [ + {"name": "docs-website", "pipeline": "docs", "cron": "0 0 2 * * *"}, + { + "name": "docs-distribution", + "pipeline": "docs", + "cron": "0 30 2 * * *", + "arguments": {"doc_type": "distribution"}, + }, + ], + ) + ) + monkeypatch.setattr(ingest_schedule, "compile_pipeline", lambda key, out_dir: f"{out_dir}/{key}.yaml") + events = [] + client = MagicMock() + client.create_experiment.return_value = SimpleNamespace(experiment_id="exp") + client.get_pipeline_id.return_value = "pipe" + client.upload_pipeline_version.return_value = SimpleNamespace(pipeline_version_id="ver") + client.list_recurring_runs.return_value = SimpleNamespace( + recurring_runs=[SimpleNamespace(recurring_run_id="old", display_name="docs-website")], + next_page_token="", + ) + client.create_recurring_run.side_effect = lambda **kw: events.append(("create", kw["job_name"])) + client.delete_recurring_run.side_effect = lambda run_id: events.append(("delete", run_id)) + + ingest_schedule.register(client, config, "abc1234-1") + + assert events == [("create", "docs-website"), ("create", "docs-distribution"), ("delete", "old")] + client.upload_pipeline_version.assert_called_once() # both schedules share the docs pipeline + created = client.create_recurring_run.call_args_list[1].kwargs + assert created["params"] == {"doc_type": "distribution"} + assert created["max_concurrency"] == 1 and created["no_catchup"] is True + assert (created["pipeline_id"], created["version_id"]) == ("pipe", "ver") + + +def test_set_overrides_parse_json_values(): + assert ingest_schedule.parse_overrides(["collection_name=kubeflow_docs_staging", "prune_removed=false"]) == { + "collection_name": "kubeflow_docs_staging", + "prune_removed": False, + }