diff --git a/.github/workflows/build-images.yml b/.github/workflows/build-images.yml index 22fb405e6..66a5b2554 100644 --- a/.github/workflows/build-images.yml +++ b/.github/workflows/build-images.yml @@ -17,22 +17,12 @@ on: - all - api - worker - deployment_target: - description: Staging deployment target - required: false - type: choice - default: ecs-staging - options: - - ecs-staging - - eks-staging-rollback env: ACR_REGISTRY: ${{ secrets.ALIYUN_ACR_REGISTRY }} ACR_NAMESPACE: ${{ secrets.ALIYUN_ACR_NAMESPACE }} ECR_REGISTRY: 107424103509.dkr.ecr.us-east-1.amazonaws.com ECR_REPOSITORY: knowhere - AWS_EKS_PROD_CLUSTER_NAME: ${{ secrets.AWS_EKS_PROD_CLUSTER_NAME }} - AWS_EKS_PROD_REGION: ${{ secrets.AWS_EKS_PROD_REGION }} AWS_ECS_STAGING_CLUSTER_NAME: knowhere-fargate AWS_ECS_STAGING_REGION: us-east-1 AWS_ECS_STAGING_EXECUTION_ROLE_ARN: arn:aws:iam::107424103509:role/knowhere-fargate-staging-execution-role @@ -41,6 +31,14 @@ env: AWS_ECS_STAGING_SECRETS_ARN: ${{ secrets.AWS_ECS_STAGING_SECRETS_ARN }} AWS_ECS_STAGING_API_SERVICE_NAME: knowhere-api-staging AWS_ECS_STAGING_WORKER_SERVICE_NAME: knowhere-worker-staging + AWS_ECS_PROD_CLUSTER_NAME: knowhere-fargate + AWS_ECS_PROD_REGION: us-east-1 + AWS_ECS_PROD_EXECUTION_ROLE_ARN: ${{ secrets.AWS_ECS_PROD_EXECUTION_ROLE_ARN }} + AWS_ECS_PROD_API_TASK_ROLE_ARN: ${{ secrets.AWS_ECS_PROD_API_TASK_ROLE_ARN }} + AWS_ECS_PROD_WORKER_TASK_ROLE_ARN: ${{ secrets.AWS_ECS_PROD_WORKER_TASK_ROLE_ARN }} + AWS_ECS_PROD_SECRETS_ARN: ${{ secrets.AWS_ECS_PROD_SECRETS_ARN }} + AWS_ECS_PROD_API_SERVICE_NAME: knowhere-api-prod + AWS_ECS_PROD_WORKER_SERVICE_NAME: knowhere-worker-prod jobs: build-and-publish: @@ -79,6 +77,22 @@ jobs: with: persist-credentials: false + # Production releases are cut from main so the ECS path uses a reviewed + # source line. + - name: Verify production release source + if: ${{ github.event_name == 'release' }} + env: + RELEASE_TAG: ${{ github.event.release.tag_name }} + shell: bash + run: | + set -euo pipefail + git fetch origin main --no-tags + release_commit="$(git rev-parse "${RELEASE_TAG}^{commit}")" + if ! git merge-base --is-ancestor "${release_commit}" origin/main; then + echo "::error::Release ${RELEASE_TAG} does not point to a commit on main" + exit 1 + fi + - name: Decide build context id: context shell: bash @@ -95,7 +109,7 @@ jobs: should_push="true" - if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ "${{ github.event.inputs.deployment_target }}" != "ecs-staging" ] && [ -n "${{ github.event.inputs.service }}" ] && [ "${{ github.event.inputs.service }}" != "all" ] && [ "${{ github.event.inputs.service }}" != "${{ matrix.service }}" ]; then + if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ -n "${{ github.event.inputs.service }}" ] && [ "${{ github.event.inputs.service }}" != "all" ] && [ "${{ github.event.inputs.service }}" != "${{ matrix.service }}" ]; then should_build="false" else should_build="true" @@ -242,8 +256,6 @@ jobs: ${{ (github.ref == 'refs/heads/staging' || github.event_name == 'workflow_dispatch') && (github.event_name != 'workflow_dispatch' || - github.event.inputs.deployment_target == 'ecs-staging') && - (github.event_name != 'workflow_dispatch' || github.event.inputs.service == '' || github.event.inputs.service == 'all' || github.event.inputs.service == 'api') }} @@ -299,38 +311,25 @@ jobs: "$image_uri" \ -m alembic upgrade heads - migrate-eks-staging: - name: Run EKS rollback database migration + migrate-ecs-production: + name: Run production Aurora database migration runs-on: ubuntu-latest needs: build-and-publish - if: >- - ${{ github.event_name == 'workflow_dispatch' && - github.event.inputs.deployment_target == 'eks-staging-rollback' && - (github.event.inputs.service == '' || - github.event.inputs.service == 'all' || - github.event.inputs.service == 'api') }} + if: ${{ github.event_name == 'release' && github.event.action == 'published' }} permissions: contents: read steps: - - name: Decide migration context - id: context - shell: bash - run: | - set -euo pipefail - short_sha="${GITHUB_SHA::8}" - echo "namespace=knowhere-staging" >> "$GITHUB_OUTPUT" - echo "image_uri=${{ env.ECR_REGISTRY }}/${{ env.ECR_REPOSITORY }}/knowhere-backend:staging-${short_sha}" >> "$GITHUB_OUTPUT" - - - name: Validate deployment configuration + - name: Validate production migration configuration shell: bash run: | set -euo pipefail - if [ -z "${{ secrets.AWS_ACCESS_KEY_ID }}" ] || \ - [ -z "${{ secrets.AWS_SECRET_ACCESS_KEY }}" ] || \ - [ -z "${{ env.AWS_EKS_PROD_CLUSTER_NAME }}" ] || \ - [ -z "${{ env.AWS_EKS_PROD_REGION }}" ]; then - echo "::error::EKS rollback migration requires AWS deployment credentials and cluster configuration." + if [ -z "${{ secrets.AWS_ACCESS_KEY_ID }}" ] || [ -z "${{ secrets.AWS_SECRET_ACCESS_KEY }}" ]; then + echo "::error::Production migration requires AWS ECR credentials." + exit 1 + fi + if [ -z "${{ secrets.PRODUCTION_MIGRATION_DATABASE_URL }}" ]; then + echo "::error::PRODUCTION_MIGRATION_DATABASE_URL is not configured." exit 1 fi @@ -339,61 +338,38 @@ jobs: with: aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - aws-region: ${{ env.AWS_EKS_PROD_REGION }} + aws-region: ${{ env.AWS_ECS_PROD_REGION }} - - name: Setup kubectl - uses: azure/setup-kubectl@v3 - - - name: Update kubeconfig + - name: Login to ECR shell: bash run: | - aws eks update-kubeconfig \ - --name "${AWS_EKS_PROD_CLUSTER_NAME}" \ - --region "${AWS_EKS_PROD_REGION}" + set -euo pipefail + aws ecr get-login-password --region "${AWS_ECS_PROD_REGION}" \ + | docker login --username AWS --password-stdin "${ECR_REGISTRY}" - - name: Run migration job - shell: bash + - name: Run migration against production Aurora env: - IMAGE_URI: ${{ steps.context.outputs.image_uri }} - NAMESPACE: ${{ steps.context.outputs.namespace }} + DATABASE_URL: ${{ secrets.PRODUCTION_MIGRATION_DATABASE_URL }} + # The image build metadata uses the short release label `prod`, but + # AppConfig validates the runtime environment as `production`. + ENVIRONMENT: production + S3_BUCKET_NAME: knowhere-storage-prod + S3_TEMP_PATH: /tmp + TMP_PATH: /tmp/aismart_bid + shell: bash run: | set -euo pipefail - - job_name="knowhere-api-migrate-${GITHUB_RUN_ID}" - migration_manifest="$(kubectl get deployment/knowhere-api --namespace "$NAMESPACE" -o json \ - | jq --arg job_name "$job_name" --arg image_uri "$IMAGE_URI" ' - if ((.spec.template.spec.containers // []) | length) == 0 then - error("knowhere-api deployment has no containers") - else - { - apiVersion: "batch/v1", - kind: "Job", - metadata: {name: $job_name, namespace: .metadata.namespace}, - spec: { - backoffLimit: 0, - ttlSecondsAfterFinished: 300, - template: { - metadata: {labels: {"app": "knowhere-api-migrate"}}, - spec: ( - .spec.template.spec - | .restartPolicy = "Never" - | .containers[0].image = $image_uri - | .containers[0].command = ["python", "-m", "alembic", "upgrade", "heads"] - | del(.containers[0].args) - ) - } - } - } - end - ')" - - printf '%s\n' "$migration_manifest" | kubectl apply -f - - if ! kubectl wait --for=condition=complete "job/$job_name" --namespace "$NAMESPACE" --timeout=900s; then - kubectl describe job "$job_name" --namespace "$NAMESPACE" || true - kubectl logs "job/$job_name" --namespace "$NAMESPACE" --all-containers=true || true - exit 1 - fi - kubectl delete job "$job_name" --namespace "$NAMESPACE" --ignore-not-found + image_uri="${ECR_REGISTRY}/${ECR_REPOSITORY}/knowhere-backend:${{ github.event.release.tag_name }}-prod" + docker run --rm \ + --env DATABASE_URL \ + --env S3_BUCKET_NAME \ + --env TMP_PATH \ + --env S3_TEMP_PATH \ + --env DB_SSL_MODE=disable \ + --env ENVIRONMENT \ + --entrypoint python \ + "${image_uri}" \ + -m alembic upgrade heads deploy-ecs-staging: name: Deploy staging services to ECS @@ -403,8 +379,7 @@ jobs: ${{ always() && needs.build-and-publish.result == 'success' && (needs['migrate-ecs-staging'].result == 'success' || needs['migrate-ecs-staging'].result == 'skipped') && ((github.ref == 'refs/heads/staging' && github.event_name == 'push') || - (github.event_name == 'workflow_dispatch' && - github.event.inputs.deployment_target == 'ecs-staging')) }} + github.event_name == 'workflow_dispatch') }} permissions: contents: read @@ -556,11 +531,29 @@ jobs: EXECUTION_ROLE_ARN: ${{ env.AWS_ECS_STAGING_EXECUTION_ROLE_ARN }} API_TASK_ROLE_ARN: ${{ env.AWS_ECS_STAGING_API_TASK_ROLE_ARN }} WORKER_TASK_ROLE_ARN: ${{ env.AWS_ECS_STAGING_WORKER_TASK_ROLE_ARN }} - STAGING_SECRETS_ARN: ${{ env.AWS_ECS_STAGING_SECRETS_ARN }} + SECRETS_ARN: ${{ env.AWS_ECS_STAGING_SECRETS_ARN }} + DEPLOYMENT_ENVIRONMENT: staging + RUNTIME_ENVIRONMENT: staging + APP_ENV: staging + DB_SSL_MODE: require + API_DB_POOL_SIZE: "5" + API_DB_MAX_OVERFLOW: "5" + WORKER_DB_SYNC_POOL_SIZE: "2" + WORKER_DB_SYNC_MAX_OVERFLOW: "2" + S3_BUCKET_NAME: knowhere-storage-staging + INTERNAL_DASHBOARD_ENDPOINT: https://staging.knowhereto.ai + FRONTEND_URL: https://staging.knowhereto.ai + API_WEBHOOK_ENDPOINT: https://api-staging.knowhereto.ai/v1/internal/s3-events + SNS_TOPIC_ARN: arn:aws:sns:us-east-1:107424103509:knowhere-staging-s3-events + QSTASH_CALLBACK_BASE_URL: https://api-staging.knowhereto.ai/api/v1 + WORKER_CPU: "2048" + WORKER_MEMORY: "4096" shell: bash run: | set -euo pipefail - python3 deploy/ecs/render_task_definitions.py --output-dir "$RUNNER_TEMP/ecs-task-definitions" + python3 deploy/ecs/render_task_definitions.py \ + --environment staging \ + --output-dir "$RUNNER_TEMP/ecs-task-definitions" - name: Register ECS task definitions id: task-definitions @@ -628,143 +621,230 @@ jobs: echo "API image: ${{ steps.images.outputs.api_image }}" echo "Worker image: ${{ steps.images.outputs.worker_image }}" - deploy-eks: + deploy-ecs-production: + name: Deploy production services to ECS runs-on: ubuntu-latest - needs: [build-and-publish, migrate-eks-staging] + needs: [build-and-publish, migrate-ecs-production] if: >- - ${{ always() && github.event_name != 'pull_request' && + ${{ always() && github.event_name == 'release' && needs.build-and-publish.result == 'success' && - (needs['migrate-eks-staging'].result == 'success' || needs['migrate-eks-staging'].result == 'skipped') && - (github.event_name == 'release' || - (github.event_name == 'workflow_dispatch' && - github.event.inputs.deployment_target == 'eks-staging-rollback')) }} + needs['migrate-ecs-production'].result == 'success' }} permissions: contents: read - strategy: - matrix: - service: [api, worker] - include: - - service: api - ecr_repo_name: knowhere-backend - - service: worker - ecr_repo_name: knowhere-worker - steps: - - name: Decide deployment context - id: context + - name: Checkout released source + uses: actions/checkout@v4 + with: + persist-credentials: false + ref: ${{ github.event.release.tag_name }} + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ env.AWS_ECS_PROD_REGION }} + + - name: Validate ECS production prerequisites shell: bash run: | - if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then - environment="staging" - elif [ "${{ github.event_name }}" = "release" ]; then - environment="prod" - elif [[ "${{ github.ref }}" == refs/tags/* ]]; then - environment="prod" - elif [ "${{ github.ref }}" = "refs/heads/staging" ]; then - environment="staging" - else - environment="staging" - should_deploy="false" + set -euo pipefail + + required_values=( + AWS_ECS_PROD_EXECUTION_ROLE_ARN + AWS_ECS_PROD_API_TASK_ROLE_ARN + AWS_ECS_PROD_WORKER_TASK_ROLE_ARN + AWS_ECS_PROD_SECRETS_ARN + ) + missing=() + for variable_name in "${required_values[@]}"; do + if [ -z "${!variable_name}" ]; then + missing+=("${variable_name}") + fi + done + if [ "${#missing[@]}" -gt 0 ]; then + printf 'Missing ECS production configuration: %s\n' "${missing[*]}" >&2 + exit 1 fi - if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ -n "${{ github.event.inputs.service }}" ] && [ "${{ github.event.inputs.service }}" != "all" ] && [ "${{ github.event.inputs.service }}" != "${{ matrix.service }}" ]; then - should_deploy="false" - elif [ -z "${should_deploy:-}" ]; then - should_deploy="true" + cluster_status="$(aws ecs describe-clusters \ + --cluster "${AWS_ECS_PROD_CLUSTER_NAME}" \ + --query 'clusters[0].status' --output text)" + if [ "${cluster_status}" != "ACTIVE" ]; then + echo "::error::ECS cluster ${AWS_ECS_PROD_CLUSTER_NAME} is not ACTIVE (status: ${cluster_status})." + exit 1 fi - if [ -z "${{ secrets.AWS_ACCESS_KEY_ID }}" ] || \ - [ -z "${{ secrets.AWS_SECRET_ACCESS_KEY }}" ] || \ - [ -z "${{ env.AWS_EKS_PROD_CLUSTER_NAME }}" ] || \ - [ -z "${{ env.AWS_EKS_PROD_REGION }}" ]; then - if [ "$environment" = "prod" ] && [ "$should_deploy" = "true" ]; then - echo "::error::Production deployment credentials are not configured." + for role_arn in \ + "${AWS_ECS_PROD_EXECUTION_ROLE_ARN}" \ + "${AWS_ECS_PROD_API_TASK_ROLE_ARN}" \ + "${AWS_ECS_PROD_WORKER_TASK_ROLE_ARN}"; do + aws iam get-role --role-name "${role_arn##*/}" --query 'Role.Arn' --output text >/dev/null + done + + aws secretsmanager describe-secret \ + --secret-id "${AWS_ECS_PROD_SECRETS_ARN}" \ + --query 'ARN' --output text >/dev/null + + for log_group in /ecs/knowhere-api-prod /ecs/knowhere-worker-prod; do + found="$(aws logs describe-log-groups \ + --log-group-name-prefix "${log_group}" \ + --query 'logGroups[0].logGroupName' --output text)" + if [ "${found}" != "${log_group}" ]; then + echo "::error::Required CloudWatch log group ${log_group} does not exist." exit 1 fi - should_deploy="false" - fi + done - if [ "$environment" = "prod" ]; then - namespace="knowhere-prod" - else - namespace="knowhere-staging" - fi + for service in \ + "${AWS_ECS_PROD_API_SERVICE_NAME}" \ + "${AWS_ECS_PROD_WORKER_SERVICE_NAME}"; do + service_json="$(aws ecs describe-services \ + --cluster "${AWS_ECS_PROD_CLUSTER_NAME}" \ + --services "${service}" --output json)" + status="$(jq -r '.services[0].status // "MISSING"' <<<"${service_json}")" + if [ "${status}" != "ACTIVE" ]; then + echo "::error::Required ECS service ${service} is not ACTIVE (status: ${status})." + exit 1 + fi + network_count="$(jq '[.services[0].networkConfiguration.awsvpcConfiguration.subnets // [] | length] | add' <<<"${service_json}")" + security_group_count="$(jq '[.services[0].networkConfiguration.awsvpcConfiguration.securityGroups // [] | length] | add' <<<"${service_json}")" + if [ "${network_count}" -lt 1 ] || [ "${security_group_count}" -lt 1 ]; then + echo "::error::ECS service ${service} has no usable awsvpc subnet/security-group configuration." + exit 1 + fi + done - short_sha="${GITHUB_SHA::8}" - if [ "${{ github.event_name }}" = "release" ]; then - git_tag="${{ github.event.release.tag_name }}" - image_tag="${git_tag}-${environment}" - elif [[ "${{ github.ref }}" == refs/tags/* ]]; then - git_tag="${GITHUB_REF#refs/tags/}" - image_tag="${git_tag}-${environment}" - else - image_tag="${environment}-${short_sha}" + api_load_balancer_count="$(aws ecs describe-services \ + --cluster "${AWS_ECS_PROD_CLUSTER_NAME}" \ + --services "${AWS_ECS_PROD_API_SERVICE_NAME}" \ + --query 'length(services[0].loadBalancers)' --output text)" + if [ "${api_load_balancer_count}" -lt 1 ]; then + echo "::error::Production API ECS service has no load balancer target configured." + exit 1 fi - echo "environment=$environment" >> "$GITHUB_OUTPUT" - echo "namespace=$namespace" >> "$GITHUB_OUTPUT" - echo "should_deploy=$should_deploy" >> "$GITHUB_OUTPUT" - echo "image_uri=${{ env.ECR_REGISTRY }}/${{ env.ECR_REPOSITORY }}/${{ matrix.ecr_repo_name }}:${image_tag}" >> "$GITHUB_OUTPUT" + - name: Resolve immutable ECR image digests + id: images + shell: bash + run: | + set -euo pipefail + image_tag="${{ github.event.release.tag_name }}-prod" - - name: Configure AWS credentials - if: steps.context.outputs.should_deploy == 'true' - uses: aws-actions/configure-aws-credentials@v4 - with: - aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - aws-region: ${{ env.AWS_EKS_PROD_REGION }} + resolve_image() { + local repository="$1" + local output_name="$2" + local digest + digest="$(aws ecr describe-images \ + --repository-name "${ECR_REPOSITORY}/${repository}" \ + --image-ids "imageTag=${image_tag}" \ + --query 'imageDetails[0].imageDigest' --output text)" + if [ -z "${digest}" ] || [ "${digest}" = "None" ]; then + echo "::error::No ECR image found for ${repository}:${image_tag}." + exit 1 + fi + echo "${output_name}=${ECR_REGISTRY}/${ECR_REPOSITORY}/${repository}@${digest}" >> "${GITHUB_OUTPUT}" + } - - name: Update kubeconfig - if: steps.context.outputs.should_deploy == 'true' + resolve_image knowhere-backend api_image + resolve_image knowhere-worker worker_image + + - name: Render ECS task definitions + env: + API_IMAGE: ${{ steps.images.outputs.api_image }} + WORKER_IMAGE: ${{ steps.images.outputs.worker_image }} + EXECUTION_ROLE_ARN: ${{ env.AWS_ECS_PROD_EXECUTION_ROLE_ARN }} + API_TASK_ROLE_ARN: ${{ env.AWS_ECS_PROD_API_TASK_ROLE_ARN }} + WORKER_TASK_ROLE_ARN: ${{ env.AWS_ECS_PROD_WORKER_TASK_ROLE_ARN }} + SECRETS_ARN: ${{ env.AWS_ECS_PROD_SECRETS_ARN }} + DEPLOYMENT_ENVIRONMENT: prod + RUNTIME_ENVIRONMENT: production + APP_ENV: production + DB_SSL_MODE: disable + API_DB_POOL_SIZE: "50" + API_DB_MAX_OVERFLOW: "50" + WORKER_DB_SYNC_POOL_SIZE: "5" + WORKER_DB_SYNC_MAX_OVERFLOW: "5" + S3_BUCKET_NAME: knowhere-storage-prod + INTERNAL_DASHBOARD_ENDPOINT: https://knowhereto.ai + FRONTEND_URL: https://knowhereto.ai + API_WEBHOOK_ENDPOINT: https://api.knowhereto.ai/v1/internal/s3-events + SNS_TOPIC_ARN: arn:aws:sns:us-east-1:107424103509:knowhere-prod-s3-events + QSTASH_CALLBACK_BASE_URL: https://api.knowhereto.ai/api/v1 + WORKER_CPU: "2048" + WORKER_MEMORY: "4096" shell: bash run: | - aws eks update-kubeconfig \ - --name "${{ env.AWS_EKS_PROD_CLUSTER_NAME }}" \ - --region "${{ env.AWS_EKS_PROD_REGION }}" + set -euo pipefail + python3 deploy/ecs/render_task_definitions.py \ + --environment production \ + --output-dir "${RUNNER_TEMP}/ecs-task-definitions" - - name: Roll out service on AWS EKS - if: steps.context.outputs.should_deploy == 'true' + - name: Register ECS task definitions + id: task-definitions shell: bash run: | - deployment_name="knowhere-${{ matrix.service }}" - container_name="${{ matrix.service }}" - namespace="${{ steps.context.outputs.namespace }}" - image_uri="${{ steps.context.outputs.image_uri }}" + set -euo pipefail + api_task_definition_arn="$(aws ecs register-task-definition \ + --cli-input-json "file://${RUNNER_TEMP}/ecs-task-definitions/knowhere-api-prod.json" \ + --query 'taskDefinition.taskDefinitionArn' --output text)" + worker_task_definition_arn="$(aws ecs register-task-definition \ + --cli-input-json "file://${RUNNER_TEMP}/ecs-task-definitions/knowhere-worker-prod.json" \ + --query 'taskDefinition.taskDefinitionArn' --output text)" + echo "api_task_definition_arn=${api_task_definition_arn}" >> "${GITHUB_OUTPUT}" + echo "worker_task_definition_arn=${worker_task_definition_arn}" >> "${GITHUB_OUTPUT}" - echo "Deploying ${deployment_name} to ${namespace} with ${image_uri}" + - name: Update API ECS service on on-demand Fargate + shell: bash + run: | + set -euo pipefail + aws ecs update-service \ + --cluster "${AWS_ECS_PROD_CLUSTER_NAME}" \ + --service "${AWS_ECS_PROD_API_SERVICE_NAME}" \ + --task-definition "${{ steps.task-definitions.outputs.api_task_definition_arn }}" \ + --capacity-provider-strategy capacityProvider=FARGATE,weight=1 - kubectl set image "deployment/${deployment_name}" \ - "${container_name}=${image_uri}" \ - --namespace "${namespace}" + - name: Update worker ECS service on Fargate Spot + shell: bash + run: | + set -euo pipefail + aws ecs update-service \ + --cluster "${AWS_ECS_PROD_CLUSTER_NAME}" \ + --service "${AWS_ECS_PROD_WORKER_SERVICE_NAME}" \ + --task-definition "${{ steps.task-definitions.outputs.worker_task_definition_arn }}" \ + --capacity-provider-strategy capacityProvider=FARGATE_SPOT,weight=1 - if [ "${{ matrix.service }}" = "worker" ]; then - kubectl patch "deployment/${deployment_name}" \ - --namespace "${namespace}" \ - --type='strategic' \ - -p='{"spec":{"template":{"spec":{"containers":[{"name":"worker","readinessProbe":{"exec":{"command":["python","-c","from shared.services.worker_health import assert_worker_healthy; assert_worker_healthy()"]}},"livenessProbe":{"exec":{"command":["python","-c","from shared.services.worker_health import assert_worker_healthy; assert_worker_healthy()"]}}}]}}}}' - fi + - name: Wait for API ECS service to stabilize + shell: bash + run: | + aws ecs wait services-stable \ + --cluster "${AWS_ECS_PROD_CLUSTER_NAME}" \ + --services "${AWS_ECS_PROD_API_SERVICE_NAME}" - kubectl rollout status "deployment/${deployment_name}" \ - --namespace "${namespace}" \ - --timeout=300s + - name: Wait for worker ECS service to stabilize + shell: bash + run: | + aws ecs wait services-stable \ + --cluster "${AWS_ECS_PROD_CLUSTER_NAME}" \ + --services "${AWS_ECS_PROD_WORKER_SERVICE_NAME}" - - name: Summarize deployment - if: steps.context.outputs.should_deploy == 'true' + - name: Summarize ECS production deployment shell: bash run: | - echo "Service: ${{ matrix.service }}" - echo "Environment: ${{ steps.context.outputs.environment }}" - echo "Namespace: ${{ steps.context.outputs.namespace }}" - echo "Image: ${{ steps.context.outputs.image_uri }}" + echo "Environment: production" + echo "Cluster: ${AWS_ECS_PROD_CLUSTER_NAME}" + echo "API image: ${{ steps.images.outputs.api_image }}" + echo "Worker image: ${{ steps.images.outputs.worker_image }}" release: name: Attach deployment release assets runs-on: ubuntu-latest - needs: deploy-eks + needs: deploy-ecs-production if: >- ${{ github.event_name == 'release' && - github.event.action == 'published' }} + github.event.action == 'published' && + needs.deploy-ecs-production.result == 'success' }} permissions: contents: write steps: diff --git a/.github/workflows/manage-staging.yml b/.github/workflows/manage-staging.yml index 085ff8dd3..31432856b 100644 --- a/.github/workflows/manage-staging.yml +++ b/.github/workflows/manage-staging.yml @@ -112,6 +112,68 @@ jobs: --services "$service_name" } + wait_for_healthy_workers() { + local worker_health_deadline="$((SECONDS + 300))" + local healthy_worker_count="0" + local -a worker_tasks=() + + # The AWS CLI services-stable waiter checks deployment count and + # runningCount only. A cold task can therefore be RUNNING while + # container health remains UNKNOWN during its health-check + # startPeriod. Poll the task health explicitly before opening API + # admission, with a bounded timeout so a bad rollout still fails. + while (( SECONDS < worker_health_deadline )); do + healthy_worker_count="0" + worker_tasks=() + mapfile -t worker_tasks < <( + aws --profile knowhere ecs list-tasks \ + --cluster knowhere-fargate \ + --service-name knowhere-worker-staging \ + --desired-status RUNNING \ + --query 'taskArns[]' \ + --output text | tr '\t' '\n' | sed '/^None$/d;/^$/d' + ) + + if [ "${#worker_tasks[@]}" -eq 2 ]; then + healthy_worker_count="$(aws --profile knowhere ecs describe-tasks \ + --cluster knowhere-fargate \ + --tasks "${worker_tasks[@]}" \ + --query 'length(tasks[?lastStatus==`RUNNING` && healthStatus==`HEALTHY`])' \ + --output text)" + if [ "$healthy_worker_count" -eq 2 ]; then + return 0 + fi + fi + + echo "Waiting for two healthy worker tasks; running=${#worker_tasks[@]}, healthy=$healthy_worker_count" + sleep 15 + done + + echo "Timed out waiting for two healthy worker tasks" >&2 + return 1 + } + + wait_for_public_api_health() { + local api_health_deadline="$((SECONDS + 300))" + + # services-stable can return as soon as the API task is RUNNING, + # before the load balancer has registered a healthy target. Treat + # transient 502/503 responses as cold-start progress, but keep a + # hard timeout so a broken target never reports a successful start. + while (( SECONDS < api_health_deadline )); do + if curl --fail --silent --show-error --max-time 10 \ + https://api-staging.knowhereto.ai/health >/dev/null; then + return 0 + fi + + echo "Waiting for the public staging API health endpoint" + sleep 10 + done + + echo "Timed out waiting for the public staging API health endpoint" >&2 + return 1 + } + read_services() { aws --profile knowhere ecs describe-services \ --cluster knowhere-fargate \ @@ -129,32 +191,11 @@ jobs: # traffic; ECS stability alone does not prove container health. update_service knowhere-worker-staging 2 wait_for_service knowhere-worker-staging - mapfile -t worker_tasks < <( - aws --profile knowhere ecs list-tasks \ - --cluster knowhere-fargate \ - --service-name knowhere-worker-staging \ - --desired-status RUNNING \ - --query 'taskArns[]' \ - --output text | tr '\t' '\n' - ) - if [ "${#worker_tasks[@]}" -ne 2 ]; then - echo "Expected two running worker tasks, found ${#worker_tasks[@]}" >&2 - exit 1 - fi - healthy_workers="$(aws --profile knowhere ecs describe-tasks \ - --cluster knowhere-fargate \ - --tasks "${worker_tasks[@]}" \ - --query 'length(tasks[?lastStatus==`RUNNING` && healthStatus==`HEALTHY`])' \ - --output text)" - if [ "$healthy_workers" -ne 2 ]; then - echo "Expected two healthy workers, found $healthy_workers" >&2 - exit 1 - fi + wait_for_healthy_workers update_service knowhere-api-staging 1 wait_for_service knowhere-api-staging - curl --fail --silent --show-error --max-time 30 \ - https://api-staging.knowhereto.ai/health >/dev/null + wait_for_public_api_health startup_seconds="$(($(date +%s) - start_started_epoch))" ;; stop) diff --git a/.gitignore b/.gitignore index 859ba378a..4dbbc7299 100644 --- a/.gitignore +++ b/.gitignore @@ -68,8 +68,7 @@ test_*.csv *.csv !requirements.csv -# Local debugging scripts -apps/worker/scripts/ +# Local debugging scripts (apps/worker/scripts/ is tracked) apps/worker/experiments/ apps/worker/start_celery_worker.py apps/worker/start_celery_debug.sh diff --git a/apps/api/app/services/document_ingestion/handoff_service.py b/apps/api/app/services/document_ingestion/handoff_service.py index c8de180bd..fe6af5029 100644 --- a/apps/api/app/services/document_ingestion/handoff_service.py +++ b/apps/api/app/services/document_ingestion/handoff_service.py @@ -29,6 +29,7 @@ class _UploadedFileJob(Protocol): job_id: str job_type: str + status: str class DocumentIngestionHandoffService: @@ -66,6 +67,17 @@ async def start_uploaded_file_workflow( ], ) + # Upload completion can arrive through both the S3 notification and the + # confirm-upload endpoint. Once either path has moved the job out of + # waiting-file, the other path must be a no-op instead of dispatching a + # second worker task for the same logical job. + if job.status != JobStatus.WAITING_FILE.value: + logger.info( + "Upload handoff already completed: " + f"job_id={job.job_id}, status={job.status}" + ) + return + outcome = await self._state_machine.transition_outcome( db, job.job_id, @@ -75,6 +87,18 @@ async def start_uploaded_file_workflow( "system", ) if not outcome.succeeded: + # A concurrent handoff may have won the CAS transition after this + # caller loaded the waiting-file snapshot. The state machine + # reports the winner's state as ``from_state``; treat that result + # as an idempotent no-op and do not enqueue another task. + if outcome.reason == "invalid_transition" and outcome.from_state != ( + JobStatus.WAITING_FILE.value + ): + logger.info( + "Upload handoff won by another trigger: " + f"job_id={job.job_id}, status={outcome.from_state}" + ) + return logger.warning( "Upload handoff transition rejected: " f"job_id={job.job_id}, reason={outcome.reason}" diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index d91c37951..6507016b1 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -119,6 +119,93 @@ async def start_uploaded_file_parse( ] +@pytest.mark.asyncio +async def test_should_not_dispatch_a_second_task_for_a_replayed_upload_event( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + monkeypatch: MonkeyPatch, +) -> None: + workflow_calls: list[dict[str, str]] = [] + + class FakeDocumentIngestionWorkerDispatcher: + async def start_uploaded_file_parse( + self, + *, + job_id: str, + user_id: str, + ) -> str: + workflow_calls.append({"job_id": job_id, "user_id": user_id}) + return "contract-task-id" + + async with api_client_factory() as api_client: + user_id, job_id = await _insert_waiting_file_job() + handoff_service = importlib.import_module( + "app.services.document_ingestion.handoff_service" + ) + monkeypatch.setattr( + handoff_service, + "DocumentIngestionWorkerDispatcher", + FakeDocumentIngestionWorkerDispatcher, + ) + + first_response = await api_client.post( + "/api/v1/internal/s3-events", + json=_build_s3_event_payload(job_id), + ) + replay_response = await api_client.post( + "/api/v1/internal/s3-events", + json=_build_s3_event_payload(job_id), + ) + + assert first_response.status_code == 200 + assert replay_response.status_code == 200 + assert workflow_calls == [{"job_id": job_id, "user_id": user_id}] + + +@pytest.mark.asyncio +async def test_should_treat_a_concurrent_upload_handoff_cas_winner_as_a_no_op() -> None: + from app.services.document_ingestion.handoff_service import ( + DocumentIngestionHandoffService, + ) + from shared.core.state_machine.transition_outcome import JobTransitionOutcome + + class FakeStateMachine: + async def transition_outcome(self, *args: object, **kwargs: object) -> object: + del args, kwargs + return JobTransitionOutcome.rejected( + job_id="job-race", + to_state="pending", + reason="invalid_transition", + attempts=1, + from_state="pending", + ) + + class FakeDispatcher: + async def start_uploaded_file_parse( + self, + *, + job_id: str, + user_id: str, + ) -> str: + del job_id, user_id + raise AssertionError("CAS loser must not dispatch a duplicate task") + + service = DocumentIngestionHandoffService( + state_machine=FakeStateMachine(), + worker_dispatcher=FakeDispatcher(), + ) + + await service.start_uploaded_file_workflow( + db=cast(object, None), + job=SimpleNamespace( + job_id="job-race", + job_type="document_ingestion", + status="waiting-file", + ), + user_id="contract-user", + trigger="s3_upload_completed", + ) + + @pytest.mark.asyncio async def test_should_accept_a_pre_rename_waiting_file_job_type_during_upload_handoff( api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], diff --git a/apps/worker/app/services/document_agent/agents/calibration/SKILL.md b/apps/worker/app/services/document_agent/agents/calibration/SKILL.md index 9910ac271..cc3c3f932 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/SKILL.md +++ b/apps/worker/app/services/document_agent/agents/calibration/SKILL.md @@ -23,8 +23,6 @@ into **page-numbering regimes** (distinct numbering systems / label shapes: decimal digits, roman numerals, prefixed folio labels, etc.). - Do not mix samples across regimes when computing an offset. -- Include `entry_indices` (0-based indices into `toc_region.entries`) for each - regime you submit. - Run the same initial-calibration procedure independently for each regime that has usable entries. @@ -53,32 +51,13 @@ offset: treat that sample / regime as **not found**, submit whatever regimes you already confirmed (or `status=failed`), and let production fallback handle the rest. Do not guess pages. -## Phase 2 — Completion (deterministic after submit; production path) - -For each TOC region, every regime with a candidate offset is completed -independently, then merged by **physical page**: - -1. Build TitleNodes via production `extract_toc_nodes` (regime-aware parse: - decimal / roman / prefixed labels → `printed_page` + `page_kind`). -2. For **each** regime with an offset: - - Project leaves belonging to that regime - - Run production Phase-2: prune → tail verify → binary-search → - small-step recalibrate (single-leaf regimes apply offset directly) -3. Merge all regime `match_overrides` (physical pages), then null-page parent - locate once on the combined tree. -4. Emit production `SkeletonAnchor` (`offset` = primary decimal summary, - `match_overrides` = union of all regimes, `null_page_report`, `bulk_count`, - `pruned_count`, `locate_agent`). -5. On recalibrate/budget failure inside one regime: keep that regime's complete - **prefix**; **drop** unresolved **suffix** leaves from the TOC tree (no TOC), - then run null-page parent locate on what remains. Never fall back to a fixed - post-TOC window. - -## Usability bar - -- Coarse structure may use the result when `SkeletonAnchor.offset_status=ok` - and `bulk_count > 0` (at least one complete production segment). -- Otherwise downstream treats the document as no-TOC / Root fallback. +## Phase 2 — Completion (deterministic after submit) + +Not your job and not yours to describe. After submit, production completes each +regime independently (prune → tail verify → binary search → small-step +recalibrate), merges the regimes by physical page, and emits the +`SkeletonAnchor`. It recomputes segment coverage, per-regime status and the +no-TOC entry set itself, so do not submit those. ## Tools @@ -86,16 +65,22 @@ independently, then merged by **physical page**: your question. Prefer the progressive 1→3→5 schedule above. Per-call page count is capped; overall spend is limited by the calibration visual token budget and `max_rounds`. -- `calibration.submit`: finish Phase 1. Pass the full result under +- `calibration.submit`: finish Phase 1. Pass the result under `tool_args.result` (or result fields directly in `tool_args`). ## Output rules -- Submit `status`, `regimes`, top-level `offset` / `offset_status` for the - primary decimal-digit regime when identifiable, `tool_calls`, `notes`. -- Each regime must include `kind`, candidate `offset`, `offset_status`, - `entry_indices`, `samples` (with `title`, `printed_label`, `physical` when - known), and `posterior` if you already inspected a late check. +Submit exactly the fields in the `calibration.submit` schema — `status`, +`regimes`, `notes` — and nothing else: + +- Per regime: `kind` and the candidate `offset`. Add `entry_indices` only when + the regime is not simply the entries whose printed-label shape matches `kind`, + and `samples` (`title` + `physical`) only for anchors you actually confirmed. +- `notes`: one short sentence saying why. When you found no offset, submit + `status=failed` and say why in that one sentence. - Keep `kind` values consistent within one run (`decimal`, `roman`, `prefixed`, or `other`). +- Anything else — per-regime status, segment coverage, no-TOC entries, tool call + counts, region index — is recomputed after submit; emitting it only risks the + submit being cut off by the output limit, which ends the run with no result. - Stay within the token / round budgets announced in the payload. diff --git a/apps/worker/app/services/document_agent/agents/calibration/loop.py b/apps/worker/app/services/document_agent/agents/calibration/loop.py index 0ebb26039..d43c12431 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/loop.py +++ b/apps/worker/app/services/document_agent/agents/calibration/loop.py @@ -19,6 +19,13 @@ finalize_calibration_result, ) from app.services.document_agent.agents.calibration.types import ( + FAILURE_BUDGET_EXHAUSTED, + FAILURE_INVALID_JSON, + FAILURE_LLM_ERROR, + FAILURE_MAX_ROUNDS, + FAILURE_MODEL_MISSING, + FAILURE_NO_OFFSET, + FAILURE_TOC_EMPTY, CalibrationResult, calibration_result_from_dict, ) @@ -33,6 +40,8 @@ _SKILL_PATH = Path(__file__).resolve().parent / "SKILL.md" +_DECISION_MAX_TOKENS = 2500 + _DECISION_INSTRUCTIONS = """ You are the calibration SubAgent. Follow the Skill strictly. Each turn return a JSON object with keys: @@ -45,6 +54,8 @@ Phase 2 (tail verify, binary search, small-step recalibrate) runs automatically after submit. Do not use a fixed post-TOC page window. Hard limits are token budgets and max_rounds — not a total page-count ledger. +Emit exactly the fields in the calibration.submit schema and nothing else; an +oversized submit is cut off by the output limit and ends the run. Include the word json in your response. """.strip() @@ -125,7 +136,11 @@ def run_calibration_phase1( if no_links: hierarchies = strip_toc_links(hierarchies) if not hierarchies: - return CalibrationResult(status="failed", notes="toc_hierarchies empty") + return CalibrationResult( + status="failed", + notes="toc_hierarchies empty", + failure_kind=FAILURE_TOC_EMPTY, + ) region_payload = _toc_region_payload(hierarchies, region_index) resolved_page_count = int( page_count or ctx.blackboard.page_count or 0 @@ -139,9 +154,6 @@ def run_calibration_phase1( blackboard = ctx.blackboard blackboard.global_signals["calibration_region_index"] = region_index blackboard.global_signals["calibration_tool_calls"] = 0 - blackboard.global_signals["calibration_inspect_pages_used"] = int( - blackboard.global_signals.get("calibration_inspect_pages_used") or 0 - ) blackboard.global_signals["calibration_done"] = False blackboard.global_signals.pop("calibration_result", None) @@ -174,9 +186,6 @@ def run_calibration_phase1( "calibration_visual": calib_stage, "plan": snap.get("plan") if isinstance(snap, dict) else None, "visual": snap.get("visual") if isinstance(snap, dict) else None, - "inspect_pages_used_diagnostic": blackboard.global_signals.get( - "calibration_inspect_pages_used" - ), }, "toc_region": region_payload, "history_tail": history[-8:], @@ -191,6 +200,7 @@ def run_calibration_phase1( CalibrationResult( status="failed", notes="planner model missing", + failure_kind=FAILURE_MODEL_MISSING, region_index=region_index, ), history, @@ -202,6 +212,7 @@ def run_calibration_phase1( CalibrationResult( status="failed", notes="planner budget exhausted", + failure_kind=FAILURE_BUDGET_EXHAUSTED, region_index=region_index, tool_calls=int( blackboard.global_signals.get("calibration_tool_calls") or 0 @@ -218,19 +229,51 @@ def run_calibration_phase1( messages=[{"role": "user", "content": prompt}], model=model, temperature=0.0, - max_tokens=2500, + max_tokens=_DECISION_MAX_TOKENS, response_format={"type": "json_object"}, usage_task="calibration.react_loop", ) - ctx.budget.commit("plan", actual=usage.get("total_tokens", est), est=est) - decision = _parse_decision(raw) except Exception as exc: ctx.budget.refund("plan", est=est) - logger.warning("[calibration] decision failed round={}: {}", round_index, exc) + logger.warning("[calibration] llm call failed round={}: {}", round_index, exc) + return _attach_history( + CalibrationResult( + status="failed", + notes=f"llm call failed: {exc}", + failure_kind=FAILURE_LLM_ERROR, + region_index=region_index, + tool_calls=int( + blackboard.global_signals.get("calibration_tool_calls") or 0 + ), + ), + history, + ) + + ctx.budget.commit("plan", actual=usage.get("total_tokens", est), est=est) + try: + decision = _parse_decision(raw) + except json.JSONDecodeError as exc: + history.append( + { + "round": round_index, + "error": f"decision output not parseable: {exc}", + "completion_tokens": usage.get("completion_tokens"), + "max_tokens": _DECISION_MAX_TOKENS, + } + ) + logger.warning( + "[calibration] decision output not parseable round={} " + "completion_tokens={} max_tokens={}: {}", + round_index, + usage.get("completion_tokens"), + _DECISION_MAX_TOKENS, + exc, + ) return _attach_history( CalibrationResult( status="failed", - notes=f"decision failed: {exc}", + notes=f"decision output not parseable: {exc}", + failure_kind=FAILURE_INVALID_JSON, region_index=region_index, tool_calls=int( blackboard.global_signals.get("calibration_tool_calls") or 0 @@ -283,6 +326,7 @@ def run_calibration_phase1( CalibrationResult( status="failed", notes=f"budget exhausted: {tool_result.error}", + failure_kind=FAILURE_BUDGET_EXHAUSTED, region_index=region_index, tool_calls=int( blackboard.global_signals.get("calibration_tool_calls") or 0 @@ -305,6 +349,7 @@ def run_calibration_phase1( CalibrationResult( status="failed", notes="max rounds reached without calibration.submit", + failure_kind=FAILURE_MAX_ROUNDS, region_index=region_index, tool_calls=int(blackboard.global_signals.get("calibration_tool_calls") or 0), ), @@ -411,6 +456,7 @@ def run_calibration_for_all_regions( "regimes": [], "regions": [], "notes": "toc_hierarchies empty", + "failure_kind": FAILURE_TOC_EMPTY, "tool_calls": 0, "no_links": no_links, } @@ -463,7 +509,9 @@ def run_calibration_for_all_regions( } if primary_result is None: primary_result = CalibrationResult( - status="failed", notes="no region produced offset" + status="failed", + notes="no region produced offset", + failure_kind=FAILURE_NO_OFFSET, ) anchor = deserialize_skeleton_anchor(primary_anchor) @@ -481,6 +529,7 @@ def run_calibration_for_all_regions( offset_status=anchor.offset_status, tool_calls=tool_calls, notes=primary_result.notes, + failure_kind=primary_result.failure_kind, ), no_links=no_links, region_payloads=region_results, diff --git a/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py b/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py index 88736bf5d..bb9b80c4c 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py +++ b/apps/worker/app/services/document_agent/agents/calibration/orchestrator.py @@ -34,7 +34,7 @@ def anchor_hierarchy( page_texts=page_texts, page_count=page_count, ) - if phase1.status == "failed" and not phase1.regimes and phase1.offset is None: + if phase1.status == "failed" and not phase1.regimes: return anchor_hierarchy_from_offset( nodes=nodes, offset_hint=None, diff --git a/apps/worker/app/services/document_agent/agents/calibration/procedure.py b/apps/worker/app/services/document_agent/agents/calibration/procedure.py index a3ed401d4..2fac004ea 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/procedure.py +++ b/apps/worker/app/services/document_agent/agents/calibration/procedure.py @@ -77,8 +77,6 @@ def pick_primary_offset(result: CalibrationResult) -> int | None: for regime in result.regimes: if regime.offset is not None: return int(regime.offset) - if result.offset is not None: - return int(result.offset) return None @@ -115,8 +113,7 @@ def seed_overrides_from_samples( candidates=[int(sample.physical)], evidence={ "calibration": True, - "printed_label": sample.printed_label, - "method": sample.method or "agent_phase1", + "method": "agent_phase1", "regime_kind": regime.kind, }, ) @@ -284,14 +281,6 @@ def anchor_hierarchy_from_regimes( for regime in result.regimes if regime.offset is not None ] - if not usable_regimes and result.offset is not None: - usable_regimes = [ - CalibrationRegime( - kind="decimal", - offset=int(result.offset), - offset_status="ok", - ) - ] for regime in usable_regimes: kind = normalize_kind(regime.kind) @@ -493,7 +482,6 @@ def _annotate_regimes_from_anchor( offset_status="ok" if segments else "failed", entry_indices=indices, samples=list(regime.samples), - posterior=list(regime.posterior), segments=segments, no_toc_entry_indices=no_toc, notes=( @@ -564,6 +552,7 @@ def finalize_calibration_result( offset_status=anchor.offset_status, tool_calls=result.tool_calls, notes="; ".join(p for p in notes_parts if p), + failure_kind=result.failure_kind, region_index=result.region_index, history_tail=list(result.history_tail), ) @@ -590,6 +579,7 @@ def build_calibration_payload( "regions": list(region_payloads or []), "tool_calls": int(tool_calls if tool_calls is not None else result.tool_calls), "notes": result.notes, + "failure_kind": result.failure_kind, "no_links": no_links, } ) diff --git a/apps/worker/app/services/document_agent/agents/calibration/service.py b/apps/worker/app/services/document_agent/agents/calibration/service.py index d1f743767..288531000 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/service.py +++ b/apps/worker/app/services/document_agent/agents/calibration/service.py @@ -54,9 +54,9 @@ def calibrate_offset( return CalibrationResult(status="failed", notes=str(exc)) logger.info( - "[calibration] Phase-1 status={} regimes={} primary_offset={}", + "[calibration] Phase-1 status={} failure_kind={} regime_offsets={}", phase1.status, - len(phase1.regimes), - phase1.offset, + phase1.failure_kind, + [regime.offset for regime in phase1.regimes], ) return phase1 diff --git a/apps/worker/app/services/document_agent/agents/calibration/tools.py b/apps/worker/app/services/document_agent/agents/calibration/tools.py index 37cf26c6a..952a6e86c 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/tools.py +++ b/apps/worker/app/services/document_agent/agents/calibration/tools.py @@ -6,6 +6,7 @@ from typing import Any from app.services.document_agent.agents.calibration.types import ( + FAILURE_NO_OFFSET, calibration_result_from_dict, ) from app.services.document_agent.manifest import ToolContext, ToolResult @@ -44,16 +45,72 @@ def build_calibration_registry() -> ToolRegistry: registry.register( ToolSpec( name="calibration.submit", - description="Submit the final CalibrationResult and finish.", + description=( + "Submit the Phase-1 result and finish. Only the fields below are " + "read; Phase-2 recomputes everything else." + ), parameters={ "type": "object", "properties": { "result": { "type": "object", - "description": ( - "Phase-1 CalibrationResult with candidate regime " - "offsets and entry_indices; Phase-2 completion runs after submit" - ), + "properties": { + "status": { + "type": "string", + "enum": ["ok", "failed"], + }, + "regimes": { + "type": "array", + "items": { + "type": "object", + "properties": { + "kind": { + "type": "string", + "description": ( + "Page-numbering system of this " + "regime's printed labels" + ), + }, + "offset": { + "type": "integer", + "description": "physical - printed", + }, + "entry_indices": { + "type": "array", + "items": {"type": "integer"}, + "description": ( + "0-based indices into " + "toc_region.entries; omit when the " + "regime is exactly the entries " + "whose label shape matches kind" + ), + }, + "samples": { + "type": "array", + "items": { + "type": "object", + "properties": { + "title": {"type": "string"}, + "physical": {"type": "integer"}, + }, + "required": ["title", "physical"], + }, + "description": ( + "Anchors you confirmed with " + "inspect.pages, so Phase-2 does not " + "re-verify them" + ), + }, + }, + "required": ["kind", "offset"], + }, + }, + "notes": { + "type": "string", + "description": "One short sentence: why this result", + }, + }, + "required": ["status", "regimes"], }, }, "required": ["result"], @@ -78,18 +135,10 @@ def calibration_submit(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: start = time.monotonic() raw = args.get("result") if not isinstance(raw, dict): - if any(key in args for key in ("status", "regimes", "offset")): + if any(key in args for key in ("status", "regimes")): raw = { key: args.get(key) - for key in ( - "status", - "regimes", - "offset", - "offset_status", - "tool_calls", - "notes", - "region_index", - ) + for key in ("status", "regimes", "notes") if key in args } else: @@ -105,6 +154,8 @@ def calibration_submit(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: error="ok result must include regimes", latency_ms=int((time.monotonic() - start) * 1000), ) + if not result.regimes: + result.failure_kind = FAILURE_NO_OFFSET tool_calls = int(ctx.blackboard.global_signals.get("calibration_tool_calls") or 0) result.tool_calls = tool_calls region_index = ctx.blackboard.global_signals.get("calibration_region_index") diff --git a/apps/worker/app/services/document_agent/agents/calibration/types.py b/apps/worker/app/services/document_agent/agents/calibration/types.py index 3b22cd8b5..5970efd99 100644 --- a/apps/worker/app/services/document_agent/agents/calibration/types.py +++ b/apps/worker/app/services/document_agent/agents/calibration/types.py @@ -1,25 +1,37 @@ -"""Calibration SubAgent result types.""" +"""Calibration SubAgent result types. + +The ``calibration.submit`` payload carries only what Phase-2 cannot recompute: +``status``, per-regime numbering ``kind`` + candidate ``offset`` (plus the +anchor ``samples`` already confirmed by vision), and one short ``notes`` reason. +``segments`` / ``no_toc_entry_indices`` / ``offset_status`` / per-regime +``notes`` / ``tool_calls`` / ``region_index`` are Phase-2 or harness outputs and +are never read back from a submit payload. +""" from __future__ import annotations from dataclasses import asdict, dataclass, field from typing import Any +# Failure classes recorded on ``CalibrationResult.failure_kind``, one per +# failure exit of the ReAct loop. ``INVALID_JSON`` means the decision payload +# did not parse (history carries completion_tokens vs max_tokens so truncation +# can be diagnosed offline); the rest mean the episode ran without an offset. +FAILURE_INVALID_JSON = "invalid_json" +FAILURE_LLM_ERROR = "llm_error" +FAILURE_MODEL_MISSING = "model_missing" +FAILURE_BUDGET_EXHAUSTED = "budget_exhausted" +FAILURE_MAX_ROUNDS = "max_rounds" +FAILURE_NO_OFFSET = "no_offset" +FAILURE_TOC_EMPTY = "toc_empty" + @dataclass class CalibrationSample: - title: str - printed_label: str | int | None = None - physical: int | None = None - method: str | None = None + """A printed→physical anchor the agent confirmed with ``inspect.pages``.""" - -@dataclass -class CalibrationPosterior: title: str - expected_physical: int | None = None - confirmed: bool | None = None - method: str | None = None + physical: int | None = None @dataclass @@ -37,10 +49,10 @@ class CalibrationSegment: class CalibrationRegime: kind: str offset: int | None = None - offset_status: str = "failed" entry_indices: list[int] = field(default_factory=list) samples: list[CalibrationSample] = field(default_factory=list) - posterior: list[CalibrationPosterior] = field(default_factory=list) + # Phase-2 outputs below; never parsed from the agent submit payload. + offset_status: str = "failed" segments: list[CalibrationSegment] = field(default_factory=list) no_toc_entry_indices: list[int] = field(default_factory=list) notes: str = "" @@ -54,6 +66,9 @@ class CalibrationResult: offset_status: str = "failed" tool_calls: int = 0 notes: str = "" + # Empty on success; otherwise one of the FAILURE_* constants, so a submit + # that never parsed is never read as "this document has no offset". + failure_kind: str = "" region_index: int | None = None # Debug-only trail from the ReAct loop (not part of submit schema). history_tail: list[dict[str, Any]] = field(default_factory=list) @@ -92,6 +107,7 @@ def _as_int_list(value: Any) -> list[int]: def calibration_result_from_dict(data: dict[str, Any]) -> CalibrationResult: + """Parse the minimal submit payload; ignore anything Phase-2 recomputes.""" regimes: list[CalibrationRegime] = [] for raw in data.get("regimes") or []: if not isinstance(raw, dict): @@ -99,53 +115,22 @@ def calibration_result_from_dict(data: dict[str, Any]) -> CalibrationResult: samples = [ CalibrationSample( title=str(s.get("title") or ""), - printed_label=s.get("printed_label"), physical=_as_optional_int(s.get("physical")), - method=s.get("method") if isinstance(s.get("method"), str) else None, ) for s in (raw.get("samples") or []) if isinstance(s, dict) ] - posterior = [ - CalibrationPosterior( - title=str(p.get("title") or ""), - expected_physical=_as_optional_int(p.get("expected_physical")), - confirmed=p.get("confirmed") if isinstance(p.get("confirmed"), bool) else None, - method=p.get("method") if isinstance(p.get("method"), str) else None, - ) - for p in (raw.get("posterior") or []) - if isinstance(p, dict) - ] - segments = [ - CalibrationSegment( - offset=_as_optional_int(seg.get("offset")) or 0, - leaf_start=_as_optional_int(seg.get("leaf_start")) or 0, - leaf_end=_as_optional_int(seg.get("leaf_end")) or 0, - entry_indices=_as_int_list(seg.get("entry_indices")), - status=str(seg.get("status") or "ok"), - ) - for seg in (raw.get("segments") or []) - if isinstance(seg, dict) and _as_optional_int(seg.get("offset")) is not None - ] regimes.append( CalibrationRegime( kind=str(raw.get("kind") or "other"), offset=_as_optional_int(raw.get("offset")), - offset_status=str(raw.get("offset_status") or "failed"), entry_indices=_as_int_list(raw.get("entry_indices")), samples=samples, - posterior=posterior, - segments=segments, - no_toc_entry_indices=_as_int_list(raw.get("no_toc_entry_indices")), - notes=str(raw.get("notes") or ""), ) ) return CalibrationResult( status=str(data.get("status") or "failed"), regimes=regimes, - offset=_as_optional_int(data.get("offset")), - offset_status=str(data.get("offset_status") or "failed"), - tool_calls=_as_optional_int(data.get("tool_calls")) or 0, notes=str(data.get("notes") or ""), - region_index=_as_optional_int(data.get("region_index")), + failure_kind=str(data.get("failure_kind") or ""), ) diff --git a/apps/worker/app/services/document_agent/executor/react_loop.py b/apps/worker/app/services/document_agent/executor/react_loop.py index 1bdfeff45..994f6390c 100644 --- a/apps/worker/app/services/document_agent/executor/react_loop.py +++ b/apps/worker/app/services/document_agent/executor/react_loop.py @@ -38,11 +38,6 @@ def _compact_blackboard(ctx: ToolContext) -> dict[str, Any]: ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else [] ), "toc_hierarchies_count": len(ctx.blackboard.toc_hierarchies or []), - "h1_count": ( - len(ctx.blackboard.h1_result.h1_candidates) - if ctx.blackboard.h1_result - else 0 - ), "shard_plan": ctx.blackboard.shard_plan.to_dict() if ctx.blackboard.shard_plan else None, diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index a8c136024..b3517b50b 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -55,7 +55,6 @@ class DocumentProfile: is_scanned: bool category: str routing_category: str = "generic" - category_rationale: str = "" language: str = "unknown" rationale: str = "" # Content-band margins as fractions of page height (top origin, y down). diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py index 6cb266fda..48983f92f 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -168,7 +168,6 @@ def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDec is_scanned=is_scanned, category=category or "unknown document", routing_category=routing_category, - category_rationale=str(data.get("category_rationale") or ""), language=str(data.get("language") or "unknown"), rationale=str(data.get("rationale") or ""), header_y=header_y, @@ -280,14 +279,6 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: [], ), "sampled_page_features": feature_summary, - "h1_pages": [ - {"title": item.title, "page": item.page} - for item in ( - self.ctx.blackboard.h1_result.h1_candidates - if self.ctx.blackboard.h1_result - else [] - ) - ], "available_actions": [ "grep.text", "propose.shard_plan", diff --git a/apps/worker/app/services/document_agent/planner/prompts.py b/apps/worker/app/services/document_agent/planner/prompts.py index 0cd83f03c..6772e59f5 100644 --- a/apps/worker/app/services/document_agent/planner/prompts.py +++ b/apps/worker/app/services/document_agent/planner/prompts.py @@ -4,8 +4,7 @@ "You are a document profile agent. Use page-feature statistics " "and the provided page screenshots to classify the PDF. " "Return strict JSON only with keys: is_scanned, category, routing_category, " - "category_rationale, language, rationale, header_y, footer_y, next_action, " - "grep_query. " + "language, rationale, header_y, footer_y, next_action, grep_query. " "category is a concise semantic document type in at most 5 English words. " "routing_category must be one of atlas, scanned, slides, generic. " "Set routing_category=atlas only when pages are primarily drawing/detail " diff --git a/apps/worker/app/services/document_agent/tools/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py index 7e1b2c7b3..974a3b609 100644 --- a/apps/worker/app/services/document_agent/tools/inspect_pages.py +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -96,9 +96,7 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: prompt = ( "Answer the question about the provided PDF page image(s). " - "Return strict JSON object with keys: " - '{"answer": string, "page_notes": [{"page": number, "note": string}], ' - '"confidence": number}. ' + 'Return strict JSON object with keys: {"answer": string}. ' "Include the word json in your reasoning.\n\n" f"Pages: {pages}\nQuestion: {question}\n" ) @@ -146,24 +144,12 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - # Diagnostic counter only (not a hard budget). - used = int(ctx.blackboard.global_signals.get("inspect_pages_used") or 0) - used = max( - used, - int(ctx.blackboard.global_signals.get("calibration_inspect_pages_used") or 0), - ) - next_used = used + len(pages) - ctx.blackboard.global_signals["inspect_pages_used"] = next_used - ctx.blackboard.global_signals["calibration_inspect_pages_used"] = next_used return ToolResult( status="ok", payload={ "pages": pages, "question": question, "answer": payload.get("answer"), - "page_notes": payload.get("page_notes") or [], - "confidence": payload.get("confidence"), - "raw": payload, }, latency_ms=int((time.monotonic() - start) * 1000), tokens_used=tokens_used, diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index 4eacc259d..cd5bcfce0 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -32,7 +32,10 @@ cleanup_stage_tracker, init_stage_tracker, ) -from shared.core.exceptions.domain_exceptions import ValidationException +from shared.core.exceptions.domain_exceptions import ( + UnavailableException, + ValidationException, +) from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.ai.llm_overrides import cleanup_llm_overrides, init_llm_overrides from shared.services.ai.token_tracking import cleanup_token_tracker, init_token_tracker @@ -63,21 +66,39 @@ def execute(self, job_id: str, user_id: str | None) -> dict[str, object]: "reason": "job_already_terminal", } - with RedisJobLock(job_context.redis_service, job_id): - task_workspace = TemporaryParseWorkspace.create(job_id) - try: - result = _run_parse_job( - job_id=job_id, - job_context=job_context, - lifecycle_service=lifecycle_service, - task_workspace=task_workspace, - ) - finally: - task_workspace.cleanup() + try: + with RedisJobLock(job_context.redis_service, job_id): + task_workspace = TemporaryParseWorkspace.create(job_id) + try: + result = _run_parse_job( + job_id=job_id, + job_context=job_context, + lifecycle_service=lifecycle_service, + task_workspace=task_workspace, + ) + finally: + task_workspace.cleanup() + except UnavailableException as exc: + if not _is_processing_lock_contention(exc): + raise + logger.info( + "Skipping duplicate parse delivery while another worker owns " + f"the processing lock: job_id={job_id}" + ) + return { + "status": "skipped", + "job_id": job_id, + "reason": "job_already_processing", + } return result +def _is_processing_lock_contention(error: UnavailableException) -> bool: + """Identify lock contention without swallowing unrelated 503 failures.""" + return error.internal_message.startswith("Could not acquire processing lock") + + def _run_parse_job( *, job_id: str, diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index 0f3a2e31e..d041b0d5b 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -153,7 +153,6 @@ def _profile_pdf_with_db( page_count=coordinator.blackboard.page_count, language=agent_profile.language, reasoning=agent_profile.rationale, - category_rationale=agent_profile.category_rationale, metrics={ "doc_stats": coordinator.blackboard.doc_stats, "doc_shape": coordinator.blackboard.global_signals.get("doc_shape", {}), diff --git a/apps/worker/app/services/document_parser/profiling/profile_model.py b/apps/worker/app/services/document_parser/profiling/profile_model.py index 2c557f56a..1d3fdde5a 100644 --- a/apps/worker/app/services/document_parser/profiling/profile_model.py +++ b/apps/worker/app/services/document_parser/profiling/profile_model.py @@ -40,7 +40,6 @@ class ParserDocumentProfile: page_count: int = 0 language: str = "unknown" reasoning: str = "" - category_rationale: str = "" toc: ParserTocProfile = field(default_factory=ParserTocProfile) granularity: str = "page" anatomy: Any | None = None diff --git a/apps/worker/scripts/_debug_publish.py b/apps/worker/scripts/_debug_publish.py new file mode 100644 index 000000000..61f0c1970 --- /dev/null +++ b/apps/worker/scripts/_debug_publish.py @@ -0,0 +1,341 @@ +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from loguru import logger + +DEFAULT_DEBUG_USER_ID = "debug_local_user" +DEFAULT_NAMESPACE = "default" + + +@dataclass(frozen=True) +class DebugPublishResult: + job_id: str + document_id: str | None + user_id: str + namespace: str + source_file_name: str + chunk_count: int + referenced_asset_count: int + uploaded_asset_count: int + + def to_dict(self) -> dict[str, Any]: + return { + "job_id": self.job_id, + "document_id": self.document_id, + "user_id": self.user_id, + "namespace": self.namespace, + "source_file_name": self.source_file_name, + "chunk_count": self.chunk_count, + "referenced_asset_count": self.referenced_asset_count, + "uploaded_asset_count": self.uploaded_asset_count, + } + + +def load_chunks_from_result_dir(result_dir: str | os.PathLike[str]) -> list[dict[str, Any]]: + chunks_path = Path(result_dir).expanduser().resolve() / "chunks.json" + with chunks_path.open(encoding="utf-8") as f: + raw = json.load(f) + chunks = raw.get("chunks", raw) if isinstance(raw, dict) else raw + if not isinstance(chunks, list): + raise ValueError(f"Invalid chunks.json shape: {chunks_path}") + return chunks + + +def publish_debug_result_dir( + *, + result_dir: str | os.PathLike[str], + source_file_name: str, + chunks: list[dict[str, Any]] | None = None, + job_id: str | None = None, + user_id: str = DEFAULT_DEBUG_USER_ID, + namespace: str = DEFAULT_NAMESPACE, + parse_track: str | None = None, + upload_assets: bool = True, + upload_only: bool = False, +) -> DebugPublishResult: + """Publish an already materialized parse result directory to local debug DB.""" + add_dir = str(Path(result_dir).expanduser().resolve()) + resolved_chunks = chunks if chunks is not None else load_chunks_from_result_dir(add_dir) + resolved_job_id = job_id or f"debug_{uuid4().hex[:8]}" + resolved_parse_track = parse_track or _infer_parse_track(resolved_chunks) + + from app.services.document_ingestion.artifact_refs import ( + collect_referenced_artifact_refs, + ) + + refs = collect_referenced_artifact_refs(resolved_chunks) + logger.info( + "debug publish: chunks={} parse_track={} referenced_assets={}", + len(resolved_chunks), + resolved_parse_track, + len(refs), + ) + + if upload_assets: + _ensure_buckets() + + uploaded_count = 0 + if upload_only: + if not job_id: + raise ValueError("upload_only requires an explicit job_id") + if upload_assets: + uploaded_count = _upload_assets( + job_id=resolved_job_id, + result_dir=add_dir, + refs=refs, + ) + return DebugPublishResult( + job_id=resolved_job_id, + document_id=None, + user_id=user_id, + namespace=namespace, + source_file_name=source_file_name, + chunk_count=len(resolved_chunks), + referenced_asset_count=len(refs), + uploaded_asset_count=uploaded_count, + ) + + document_id = _publish_chunks_to_db( + chunks=resolved_chunks, + add_dir=add_dir, + source_file_name=source_file_name, + job_id=resolved_job_id, + user_id=user_id, + namespace=namespace, + parse_track=resolved_parse_track, + ) + + if upload_assets: + uploaded_count = _upload_assets( + job_id=resolved_job_id, + result_dir=add_dir, + refs=refs, + ) + + return DebugPublishResult( + job_id=resolved_job_id, + document_id=document_id, + user_id=user_id, + namespace=namespace, + source_file_name=source_file_name, + chunk_count=len(resolved_chunks), + referenced_asset_count=len(refs), + uploaded_asset_count=uploaded_count, + ) + + +def _infer_parse_track(chunks: list[dict[str, Any]]) -> str: + return ( + "page_memory" + if any(str(c.get("type") or c.get("chunk_type") or "").lower() == "page" for c in chunks) + else "chunk" + ) + + +def _ensure_buckets() -> None: + from shared.core.config import settings + + client = settings.get_s3_client() + for bucket in { + settings.S3_BUCKET_NAME, + getattr(settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME), + }: + try: + client.head_bucket(Bucket=bucket) + logger.info(" bucket exists: {}", bucket) + except Exception: + client.create_bucket(Bucket=bucket) + logger.info(" bucket created: {}", bucket) + + +def _upload_assets(*, job_id: str, result_dir: str, refs: set[str]) -> int: + """Upload only referenced client artifacts to results/{job_id}/.""" + from concurrent.futures import ThreadPoolExecutor + + from shared.services.storage.result_storage import get_result_storage + + storage = get_result_storage() + result_path = Path(result_dir) + bucket = storage.results_bucket + + tasks: list[tuple[str, str]] = [] + missing = 0 + for relative in sorted(refs): + local = result_path / relative + if not local.is_file(): + missing += 1 + logger.warning(" referenced asset missing on disk: {}", relative) + continue + raw_key = storage.build_raw_key(job_id=job_id, relative_path=relative) + tasks.append((str(local), raw_key)) + + def _put(item: tuple[str, str]) -> None: + local_path, raw_key = item + storage._job_file_storage.upload_local_file( # noqa: SLF001 + local_path, + raw_key, + bucket=bucket, + ) + + with ThreadPoolExecutor(max_workers=16) as pool: + list(pool.map(_put, tasks)) + + logger.info( + " uploaded {} referenced asset files under results/{} (missing={})", + len(tasks), + job_id, + missing, + ) + return len(tasks) + + +def _publish_chunks_to_db( + *, + chunks: list[dict[str, Any]], + add_dir: str, + source_file_name: str, + job_id: str, + user_id: str, + namespace: str, + parse_track: str, +) -> str: + from sqlalchemy import select, text as sql_text + + from app.services.connect_builder.summary_builder import ( + build_section_summary_lookup, + enrich_doc_nav_summaries, + ) + from shared.core.database_sync import get_sync_db_context + from shared.models.database.document import DocumentSection, GraphNode + from shared.models.database.job import Job + from shared.models.database.job_result import JobResult + from shared.services.retrieval.publication_service import RetrievalPublicationService + + file_dir_name = os.path.basename(add_dir) + logger.info("Step 1: enrich_doc_nav_summaries") + enrich_doc_nav_summaries( + os.path.dirname(add_dir), + source_file=file_dir_name, + use_llm=False, + ) + + logger.info("Step 2: build_section_summary_lookup") + section_summaries = build_section_summary_lookup(add_dir) + logger.info(" section_summaries entries: {}", len(section_summaries)) + + logger.info("Step 3: inject document_top_summary into chunk metadata") + _inject_navigation_metadata( + chunks=chunks, + add_dir=add_dir, + source_file_name=source_file_name, + ) + + logger.info("Step 4-6: DB publication (job_id={})", job_id) + with get_sync_db_context() as db: + db.execute( + sql_text( + 'INSERT INTO "user" (id, name, email) VALUES (:uid, :name, :email) ' + "ON CONFLICT DO NOTHING" + ), + {"uid": user_id, "name": "Debug User", "email": "debug@local.test"}, + ) + db.flush() + + job = Job( + job_id=job_id, + user_id=user_id, + status="PROCESSING", + job_type="parse", + source_type="direct_upload", + file_path=source_file_name, + job_metadata={ + "namespace": namespace, + "source_file_name": source_file_name, + "parse_track": parse_track, + }, + ) + db.add(job) + db.flush() + + job_result = JobResult( + job_id=job_id, + delivery_mode="url", + document_metadata={}, + ) + db.add(job_result) + db.flush() + job_result_id = job_result.id + + pub_service = RetrievalPublicationService() + published = pub_service.publish_document_state( + db, + job_id=job_id, + job_result_id=job_result_id, + chunks=chunks, + section_summaries=section_summaries, + ) + document_id = published.document_id if published else None + if not document_id: + db.rollback() + raise RuntimeError("publish_document_state returned None") + logger.info(" published document_id: {}", document_id) + + pub_service.publish_document_graph( + db, + job_id=job_id, + job_result_id=job_result_id, + ) + + sections = list( + db.execute( + select( + DocumentSection.section_level, + DocumentSection.section_title, + DocumentSection.summary, + ) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + .order_by(DocumentSection.sort_order) + ).all() + ) + with_summary = sum(1 for _, _, summary in sections if summary) + logger.info( + " DocumentSection rows: {}, with_summary: {}", + len(sections), + with_summary, + ) + + graph_node = db.execute( + select(GraphNode) + .where(GraphNode.owner_document_id == document_id) + .where(GraphNode.node_kind == "document") + ).scalar_one_or_none() + logger.info(" GraphNode published: {}", bool(graph_node)) + + db.commit() + logger.info(" DB transaction committed") + return str(document_id) + + +def _inject_navigation_metadata( + *, + chunks: list[dict[str, Any]], + add_dir: str, + source_file_name: str, +) -> None: + from app.services.connect_builder.summary_builder import load_nav_top_summary + + document_top_summary = load_nav_top_summary(add_dir, source_file_name) + for chunk in chunks: + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + chunk["metadata"] = metadata + if document_top_summary: + metadata["document_top_summary"] = document_top_summary diff --git a/apps/worker/scripts/debug_parse.py b/apps/worker/scripts/debug_parse.py new file mode 100644 index 000000000..6ad95f2de --- /dev/null +++ b/apps/worker/scripts/debug_parse.py @@ -0,0 +1,704 @@ +#!/usr/bin/env python3 +"""Unified production-style document parsing debug script. + +Supports all chunk-track Knowhere formats (PDF, DOCX, XLSX, PPTX, MD, Image, +Fragment) through the same checkerboard parser entry used by the worker. Use +``debug_page_memory.py`` for page-track step debugging. + +Pipeline stages: + 1. checkerboard_parse_output → DataFrame + 2. dataframe_to_chunks → list[ChunkPayload] + 3. ZipResultService → chunks.json / manifest.json / doc_nav.json / *.zip + 4. enrich_doc_nav → summary enrichment + top_summary + 5. DB publication → DocumentSection + DocumentChunk (optional, --run-db) + +Output directory: + default → ~/.knowhere/chengke_kb// + +Usage: + cd apps/worker + + # All formats (full pipeline) + python scripts/debug_parse.py --file /path/to/any.pdf + python scripts/debug_parse.py --file /path/to/doc.docx + python scripts/debug_parse.py --file /path/to/sheet.xlsx + python scripts/debug_parse.py --fragment "粘贴的文本..." + + # Options + python scripts/debug_parse.py --spacex --run-db # Enable DB publication +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import sys +import time +import zipfile +from pathlib import Path +from typing import Any + +# ── Bootstrap: path + env ────────────────────────────────────────────────────── +ROOT = Path(__file__).resolve().parents[3] +WORKER_ROOT = ROOT / "apps" / "worker" +sys.path.insert(0, str(WORKER_ROOT)) +sys.path.insert(0, str(ROOT / "packages" / "shared-python")) + +from dotenv import load_dotenv # noqa: E402 + +load_dotenv(WORKER_ROOT / ".env") +os.environ.setdefault("LOCAL_DEBUG", "1") +os.environ.setdefault("OVERSIZED_PDF_SHARD_ENABLED", "true") + +from loguru import logger # noqa: E402 + +from shared.core.config import settings # noqa: E402 + +# ── Constants ────────────────────────────────────────────────────────────────── +DEFAULT_SPACEX_PDF = Path("/Users/wuchengke/Desktop/temp/test_docs/spacex-s1.pdf") +DEFAULT_SJSYJ_PDF = Path( + "/Users/wuchengke/Desktop/temp/test_docs/" + "SJSYJ-SC-2024 企业制度汇编(上册).pdf" +) + +PRODUCTION_OUTPUT_ROOT = Path("~/.knowhere/chengke_kb").expanduser() +AGENT_TRANSIENT_DIRS = ( + "_doc_agent", + "planner_pages", + "toc_pages", + "inspect_pages", + "verify_pages", + "agent_visuals", +) + +# ══════════════════════════════════════════════════════════════════════════════ +# Section A: DB Publication (Stage 10) — preserved from original debug_parse.py +# ══════════════════════════════════════════════════════════════════════════════ + +def _run_db_publication( + chunks: list, + add_dir: str, + source_file_name: str, +): + """Stage 10: publish an already finalized debug result to local DB/S3.""" + from scripts._debug_publish import publish_debug_result_dir + + result = publish_debug_result_dir( + result_dir=add_dir, + source_file_name=source_file_name, + chunks=chunks, + upload_assets=True, + ) + logger.info(" ✅ DB transaction committed (job_id={})", result.job_id) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Section C: Common Post-Parse Pipeline (Stage 7-10) +# ══════════════════════════════════════════════════════════════════════════════ + +def _finalize_output( + parsed_df, + add_dir: str, + source_file_name: str, + *, + run_db: bool = False, + job_metadata: dict[str, Any] | None = None, +) -> list[dict[str, Any]]: + """Stage 7-10: chunks → ZIP → enrich → optional DB. + + Mirrors production flow: + - parse_result_package.py L49 → dataframe_to_chunks + - success_finalization.py L196 → ZipResultService + - success_finalization.py L126 → enrich_doc_nav_summaries + - debug_parse.py _run_db_publication → DB write + + Returns the chunks list. + """ + from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks + from shared.services.storage.zip_result_service import ZipResultService + from app.services.connect_builder.summary_builder import ( + build_section_summary_lookup, + enrich_doc_nav_summaries, + ensure_doc_nav_json, + load_nav_top_summary, + ) + + timings: dict[str, float] = {} + + # ── Stage 7: DataFrame → chunks (mirrors parse_result_package.py L49) ── + logger.info("=" * 60) + logger.info("📦 Stage 7: dataframe_to_chunks") + logger.info("=" * 60) + + t0 = time.time() + chunks = dataframe_to_chunks(parsed_df) + timings["Stage 7: chunks"] = time.time() - t0 + + text_count = sum(1 for c in chunks if c.get("type") == "text") + image_count = sum(1 for c in chunks if c.get("type") == "image") + table_count = sum(1 for c in chunks if c.get("type") == "table") + page_count = sum(1 for c in chunks if c.get("type") == "page") + table_ref_count = sum( + 1 + for c in chunks + if c.get("type") == "table" + and str(c.get("content") or "").strip().startswith("tables/") + ) + table_inline_html_count = sum( + 1 + for c in chunks + if c.get("type") == "table" + and " 0: + logger.info("") + logger.info("═" * 58) + logger.info(" 📊 POST-PARSE TIMELINE") + logger.info("═" * 58) + for phase, elapsed in timings.items(): + pct = elapsed / t_total * 100 + logger.info(f" {phase:<35s} │ {elapsed:>7.2f}s ({pct:>5.1f}%)") + logger.info(" " + "─" * 55) + logger.info(f" {'TOTAL':<35s} │ {t_total:>7.2f}s (100.0%)") + logger.info("═" * 58) + + return chunks + + +def _cleanup_agent_transient_dirs(add_dir: str) -> None: + """Remove VLM render caches before packaging debug output.""" + removed: list[str] = [] + for dirname in AGENT_TRANSIENT_DIRS: + path = os.path.join(add_dir, dirname) + if os.path.isdir(path): + shutil.rmtree(path) + removed.append(dirname) + nested_doc_agent = os.path.join(add_dir, "_doc_agent") + if os.path.isdir(nested_doc_agent): + for dirname in AGENT_TRANSIENT_DIRS: + path = os.path.join(nested_doc_agent, dirname) + if os.path.isdir(path): + shutil.rmtree(path) + removed.append(f"_doc_agent/{dirname}") + if not os.listdir(nested_doc_agent): + os.rmdir(nested_doc_agent) + removed.append("_doc_agent") + if removed: + logger.info(f" Cleaned transient agent dirs: {', '.join(removed)}") + + +# ══════════════════════════════════════════════════════════════════════════════ +# Section D: Pipeline Entry Points +# ══════════════════════════════════════════════════════════════════════════════ + +def _run_standard_pipeline( + file_path: str, + source_file_name: str, + output_root: str, + *, + run_db: bool = False, + fragment_content: str = "", +) -> dict[str, Any]: + """Standard pipeline for all formats: checkerboard_parse_output → finalize. + + Uses the production black-box entry point. Handles all formats including + oversized PDFs (which are routed internally by parse_pdfs). + + Token/time tracking mirrors production parse_execution.py exactly: + init trackers → parse → collect stats → cleanup. + """ + from app.services.document_parser.parse_service import checkerboard_parse_output + from app.services.document_parser.support.stage_profiler import ( + init_stage_tracker, + cleanup_stage_tracker, + get_current_stage_tracker, + ) + from shared.services.ai.token_tracking import ( + init_token_tracker, + cleanup_token_tracker, + get_current_token_tracker, + ) + + filename = source_file_name + is_fragment = ".fragment" in file_path.lower() + + logger.info("=" * 60) + logger.info(f"📄 Standard pipeline: {filename}") + logger.info(f" Output root: {output_root}") + logger.info("=" * 60) + + # ── Init trackers (same as parse_execution.py; reuse run_pipeline tracker) ── + token_usage_dict = get_current_token_tracker() + owns_token_tracker = token_usage_dict is None + if token_usage_dict is None: + token_usage_dict = init_token_tracker() + + stage_timing_dict = get_current_stage_tracker() + owns_stage_tracker = stage_timing_dict is None + if stage_timing_dict is None: + stage_timing_dict = init_stage_tracker() + + try: + t0 = time.time() + result = checkerboard_parse_output( + file_full_path=file_path, + filename=filename, + output_dir=output_root, + internal_output_filename=filename, + smart_title_parse=True, + summary_image=True, + summary_table=True, + summary_txt=True, + doc_type="auto", + fragment_content=fragment_content if is_fragment else "", + ) + parse_elapsed = time.time() - t0 + + # ── Snapshot stats before cleanup ── + stages_snapshot = { + "timing_ms": dict(stage_timing_dict), + "token_usage": dict(token_usage_dict), + } + finally: + if owns_token_tracker: + cleanup_token_tracker() + if owns_stage_tracker: + cleanup_stage_tracker() + + add_dir = result.output_dir + parsed_df = result.parsed_df + + logger.info("=" * 60) + logger.info(f"✅ Parse complete in {parse_elapsed:.1f}s") + logger.info(f" Output path: {add_dir}") + if parsed_df is not None: + logger.info(f" DataFrame rows: {len(parsed_df)}") + logger.info("=" * 60) + + # ── Print consumption stats ── + logger.info("") + logger.info("═" * 58) + logger.info(" 📊 CONSUMPTION STATS (mirrors manifest.processing.stages)") + logger.info("═" * 58) + token_usage = stages_snapshot["token_usage"] + logger.info( + f" Token usage: prompt={token_usage['prompt_tokens']}, " + f"completion={token_usage['completion_tokens']}, " + f"total={token_usage['total_tokens']}" + ) + timing_ms = stages_snapshot["timing_ms"] + if timing_ms: + logger.info(" Stage timings:") + for stage, ms in sorted(timing_ms.items()): + logger.info(f" {stage:<45s} │ {ms:>8,}ms") + else: + logger.info(" Stage timings: (none recorded)") + logger.info("═" * 58) + + if not add_dir or not os.path.exists(add_dir) or parsed_df is None or parsed_df.empty: + logger.error("❌ Parse returned empty result, cannot proceed") + return {"status": "error", "parse_elapsed": parse_elapsed} + + # Build job_metadata matching production parse_execution.py + success_finalization.py + from datetime import datetime, timezone + + processing_completed_at = datetime.now(timezone.utc) + debug_job_metadata: dict[str, Any] = { + "stages": stages_snapshot, + "processing_started_at": processing_completed_at.isoformat(), + "processing_completed_at": processing_completed_at.isoformat(), + "processing_duration_ms": int(parse_elapsed * 1000), + } + + try: + # Stage 7-10 + chunks = _finalize_output( + parsed_df, add_dir, source_file_name, + run_db=run_db, + job_metadata=debug_job_metadata, + ) + finally: + debug_job_metadata["stages"] = { + "timing_ms": dict(stage_timing_dict), + "token_usage": dict(token_usage_dict), + } + + return { + "status": "success", + "parse_elapsed": round(parse_elapsed, 1), + "output_dir": add_dir, + "chunks_count": len(chunks) if chunks else 0, + "stages": stages_snapshot, + } + + + +def run_pipeline( + file_path: str, + source_file_name: str, + *, + run_db: bool = False, + fragment_content: str = "", + output_root_override: str | None = None, +) -> dict[str, Any]: + """Unified production-style E2E parser entry point.""" + from app.services.document_parser.support.stage_profiler import ( + cleanup_stage_tracker, + get_current_stage_tracker, + init_stage_tracker, + ) + from shared.services.ai.token_tracking import ( + cleanup_token_tracker, + get_current_token_tracker, + init_token_tracker, + ) + + owns_token_tracker = get_current_token_tracker() is None + if owns_token_tracker: + init_token_tracker() + + owns_stage_tracker = get_current_stage_tracker() is None + if owns_stage_tracker: + init_stage_tracker() + + try: + if output_root_override: + output_root = output_root_override + else: + output_root = str(PRODUCTION_OUTPUT_ROOT) + + return _run_standard_pipeline( + file_path, + source_file_name, + output_root, + run_db=run_db, + fragment_content=fragment_content, + ) + finally: + if owns_token_tracker: + cleanup_token_tracker() + if owns_stage_tracker: + cleanup_stage_tracker() + + +# ══════════════════════════════════════════════════════════════════════════════ +# Section E: CLI +# ══════════════════════════════════════════════════════════════════════════════ + +def test_config(): + """Print current configuration.""" + logger.info("=== Current Config ===") + logger.info(f"ENVIRONMENT: {getattr(settings, 'ENVIRONMENT', 'N/A')}") + logger.info( + f"DATABASE_URL: {getattr(settings, 'DATABASE_URL', 'N/A')[:50]}..." + ) + logger.info(f"REDIS_HOST: {getattr(settings, 'REDIS_HOST', 'N/A')}") + logger.info( + f"DS_KEY: {'set' if getattr(settings, 'DS_KEY', None) else 'unset'}" + ) + logger.info( + f"ALI_API_KEYS: {'set' if getattr(settings, 'ALI_API_KEYS', None) else 'unset'}" + ) + logger.info(f"IMAGE_MODEL: {getattr(settings, 'IMAGE_MODEL', 'N/A')}") + logger.info(f"NORMOL_MODEL: {getattr(settings, 'NORMOL_MODEL', 'N/A')}") + logger.info( + f"HIERARCHY_LLM_MODEL: {getattr(settings, 'HIERARCHY_LLM_MODEL', 'N/A')}" + ) + logger.info( + f"MAX_PDF_PAGE_LIMIT: {getattr(settings, 'MAX_PDF_PAGE_LIMIT', 'N/A')}" + ) + + +def _parse_cases(args: argparse.Namespace) -> list[tuple[str, str, str]]: + """Parse CLI args into [(name, file_path, fragment_content), ...]. + + Returns a list of tuples: (job_name, file_path, fragment_content). + For non-fragment cases, fragment_content is "". + """ + cases: list[tuple[str, str, str]] = [] + + for raw_case in args.case or []: + if "=" not in raw_case: + raise ValueError("--case must use name=/path/to/file") + name, path = raw_case.split("=", 1) + resolved = str(Path(path).expanduser().resolve()) + cases.append((name.strip(), resolved, "")) + + if args.file: + file_path = str(Path(args.file).expanduser().resolve()) + job_id = args.job_id or Path(args.file).stem + cases.append((job_id, file_path, "")) + + if args.fragment: + cases.append(("fragment", ".fragment", args.fragment)) + + if args.spacex: + cases.append( + ("spacex-s1", str(DEFAULT_SPACEX_PDF.expanduser().resolve()), "") + ) + if args.sjsyj: + cases.append( + ("sjsyj", str(DEFAULT_SJSYJ_PDF.expanduser().resolve()), "") + ) + + return cases + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Unified document parsing debug script — supports all Knowhere formats.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog="""examples: + # Full pipeline (any format) + python scripts/debug_parse.py --file /path/to/doc.pdf + python scripts/debug_parse.py --file /path/to/doc.docx + python scripts/debug_parse.py --fragment "粘贴的文本..." + + # Optional DB publication + python scripts/debug_parse.py --spacex --run-db +""", + ) + + # Input sources + input_group = parser.add_argument_group("input") + input_group.add_argument("--file", help="Path to file to parse (any format)") + input_group.add_argument("--job-id", help="Job id override for --file") + input_group.add_argument( + "--fragment", help="Text content for fragment mode parsing" + ) + input_group.add_argument( + "--spacex", + action="store_true", + help=f"SpaceX S-1 fixture: {DEFAULT_SPACEX_PDF}", + ) + input_group.add_argument( + "--sjsyj", + action="store_true", + help=f"企业制度汇编 fixture: {DEFAULT_SJSYJ_PDF}", + ) + input_group.add_argument( + "--case", + action="append", + help="Named fixture: name=/path/to/file", + ) + + # Post-processing + post_group = parser.add_argument_group("post-processing") + post_group.add_argument( + "--run-db", + action="store_true", + help="Enable Stage 10: DB publication (requires running database)", + ) + + # Output control + output_group = parser.add_argument_group("output") + output_group.add_argument( + "--output-root", + default=None, + help=( + "Override output root directory. Default: " + f"{PRODUCTION_OUTPUT_ROOT}" + ), + ) + output_group.add_argument( + "--clean", + action="store_true", + help="Delete existing output before running", + ) + output_group.add_argument( + "--test-config", + action="store_true", + help="Print current configuration and exit", + ) + + args = parser.parse_args() + + if args.test_config: + test_config() + return 0 + + cases = _parse_cases(args) + if not cases: + parser.error( + "provide --file, --fragment, --spacex, --sjsyj, or at least one --case" + ) + + summaries = [] + + for name, file_path, fragment_content in cases: + if file_path != ".fragment" and not os.path.exists(file_path): + raise FileNotFoundError(f"File not found: {file_path}") + + source_file_name = ( + os.path.basename(file_path) if file_path != ".fragment" else "" + ) + + # Clean if requested + if args.clean: + from app.services.document_parser.orchestration.path_segment import ( + build_parser_path_segment, + ) + dir_name = build_parser_path_segment(source_file_name) + clean_root = Path(args.output_root) if args.output_root else PRODUCTION_OUTPUT_ROOT + clean_dir = clean_root / dir_name + if clean_dir.exists(): + logger.info(f"🗑️ Cleaning {clean_dir}") + shutil.rmtree(clean_dir) + + logger.info("") + logger.info("█" * 60) + logger.info(f" CASE: {name}") + logger.info(f" FILE: {file_path}") + logger.info("█" * 60) + + result = run_pipeline( + file_path, + source_file_name, + run_db=args.run_db, + fragment_content=fragment_content, + output_root_override=args.output_root, + ) + result["job_id"] = name + summaries.append(result) + + # Final JSON summary + print(json.dumps({"cases": summaries}, ensure_ascii=False, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/debug_retrieval.py b/apps/worker/scripts/debug_retrieval.py new file mode 100644 index 000000000..28bf75551 --- /dev/null +++ b/apps/worker/scripts/debug_retrieval.py @@ -0,0 +1,1406 @@ +""" +Agentic retrieval E2E debug runner for the evidence-only flow. + +This script uses real DB data and a real LLM. It prints every LLM prompt and +response in full so the routing context is inspectable. + +Usage: + cd apps/worker + python debug_agentic_e2e.py +""" +from __future__ import annotations + +import asyncio +import json +import os +import re +import sys +import time +from contextlib import contextmanager +from typing import Any +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../../packages/shared-python')) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from dotenv import load_dotenv +from loguru import logger +from shared.utils.token_estimate import estimate_tokens + +load_dotenv(os.path.join(os.path.dirname(__file__), '../.env')) +os.environ.setdefault('LOCAL_DEBUG', '0') +os.environ.setdefault('LLM_MOCK_ENABLED', 'false') + +USER_ID = 'debug_local_user' +NAMESPACE = 'default' +TOP_K = 10 +CHUNK_SCOPE_DATA_TYPE = { + 'all': 1, + 'text': 2, + 'image': 3, + 'table': 4, + 'text-image': 5, + 'text-table': 6, + 'page': 7, + 'chunk': 8, +} + +captured_interactions: list[dict[str, Any]] = [] + + +@contextmanager +def temporary_env(overrides: dict[str, str] | None): + """Apply per-test env overrides and restore them afterwards.""" + overrides = overrides or {} + old_values = {key: os.environ.get(key) for key in overrides} + os.environ.update(overrides) + try: + yield + finally: + for key, value in old_values.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def _print_box(title: str, body: Any) -> None: + line = '=' * 100 + print(f'\n{line}\n{title}\n{line}') + print(str(body)) + print(line) + + +def _summarize_tool_payload(payload: dict[str, Any] | None) -> str: + """Keep tool logs readable; full LLM prompts/responses are printed separately.""" + if not isinstance(payload, dict): + return str(payload) + + lines: list[str] = [] + if 'top_doc_ids' in payload: + lines.append(f'top_doc_ids={payload.get("top_doc_ids")}') + if 'channel_counts' in payload: + lines.append(f'channel_counts={payload.get("channel_counts")}') + if 'fused_rows' in payload: + rows = payload.get('fused_rows') or [] + lines.append(f'fused_rows={len(rows)}') + for row in rows[:5]: + lines.append( + f' - score={row.get("score", 0):.4f} ' + f'doc={str(row.get("document_id", ""))[:12]} ' + f'path="{row.get("section_path") or row.get("source_chunk_path")}" ' + f'chunk={row.get("chunk_id")}' + ) + if 'candidate_docs' in payload: + docs = payload.get('candidate_docs') or [] + lines.append(f'candidate_docs={len(docs)}') + for doc in docs: + lines.append( + f' - doc={doc.get("document_id")} ' + f'name="{doc.get("source_file_name", "")}" ' + f'confidence={doc.get("confidence")}' + ) + if 'document_id' in payload: + lines.append(f'document_id={payload.get("document_id")}') + if 'has_outline' in payload: + lines.append(f'has_outline={payload.get("has_outline")} ' + f'leaf_count={payload.get("leaf_count", 0)} ' + f'children_count={payload.get("children_count", 0)}') + for key in ('reason', 'document_id', 'raw_response', 'overflowed'): + if key in payload: + lines.append(f'{key}={payload.get(key)}') + return '\n'.join(lines) or str(payload) + + +def _decision_stage(kind: str) -> tuple[str, str]: + return { + 'kg_document_select': ( + 'Phase 1B: KG Document Select', + 'LLM reads the KB inventory and chooses candidate documents. Bottom discovery is merged separately so missed high-recall docs are still protected.', + ), + 'navigate': ( + 'Phase 2A: Navigate (unified)', + 'LLM decides action (NAVIGATE/STOP), optional asset tools (SEARCH_IMAGES/SEARCH_TABLES), and section selections in a single call.', + ), + 'asset_filter': ( + 'Phase 2A-1: Asset Filter', + 'LLM filters asset candidates (images/tables) by semantic relevance to the search query. Only matching assets are added to pending evidence.', + ), + 'workflow_planner': ( + 'Phase 0: Workflow Planner', + 'Thinking-model planner decides whether to decompose the query into multiple retrieval sub-steps or pass through as a single step.', + ), + }.get(kind, ('Unknown Phase', 'Unclassified LLM call.')) + + +def _fence(text: Any, lang: str = 'text') -> str: + body = str(text).replace('```', '``\\`') + return f'```{lang}\n{body}\n```\n' + + +def _extract_resource_status(prompt: str) -> dict[str, str]: + """Extract budget lines from LLM prompts for compact trace summaries.""" + status: dict[str, str] = {} + for raw_line in str(prompt).splitlines(): + line = raw_line.strip() + if line.startswith('Planning Budget:'): + status['planning'] = line.removeprefix('Planning Budget:').strip() + elif line.startswith('Context Budget:'): + status['context'] = line.removeprefix('Context Budget:').strip() + elif line.startswith('Planning budget:'): + status['planning'] = line.removeprefix('Planning budget:').strip() + elif line.startswith('Context budget:'): + status['context'] = line.removeprefix('Context budget:').strip() + return status + + +def _extract_context_projection(prompt: str) -> dict[str, int | str]: + """Estimate context budget as the debugger sees this exact prompt.""" + resource_status = _extract_resource_status(prompt) + projection: dict[str, int | str] = { + 'prompt_tokens_estimate': estimate_tokens(prompt), + } + context = resource_status.get('context', '') + match = re.search(r'(\d+)\s*/\s*(\d+)\s+remaining', context) + if match: + remaining_in_prompt = int(match.group(1)) + capacity = int(match.group(2)) + remaining_pct = 0 if capacity <= 0 else max( + 0, + min(100, round(remaining_in_prompt * 100 / capacity)), + ) + projection.update({ + 'context_remaining_in_prompt': remaining_in_prompt, + 'context_capacity': capacity, + 'context_used_pct_in_prompt': 100 - remaining_pct, + 'context_remaining_pct_in_prompt': remaining_pct, + }) + return projection + + +def _charge_pool_for_kind(kind: str) -> str: + return { + 'kg_document_select': 'bootstrap', + 'navigate': 'planning', + 'asset_filter': 'planning', + 'workflow_planner': 'bootstrap', + }.get(kind, 'unknown') + + +def make_verbose_llm(real_llm_fn): + counter = {'n': 0} + + async def verbose_llm(prompt) -> str: + counter['n'] += 1 + n = counter['n'] + if '=== Document Corpus Overview ===' in prompt: + kind = 'kg_document_select' + elif '=== Rules ===' in prompt and '=== Actionable Observation ===' in prompt: + kind = 'navigate' + elif 'retrieval workflow planner' in prompt.lower(): + kind = 'workflow_planner' + elif 'You are an asset relevance filter' in prompt: + kind = 'asset_filter' + else: + kind = 'unknown' + _print_box(f'LLM PROMPT #{n} [{kind}] chars={len(prompt)}', prompt) + context_projection = _extract_context_projection(str(prompt)) + charge_pool = _charge_pool_for_kind(kind) + logger.info( + 'LLM PROMPT #{} [{}] pool={} token_estimate={}{}', + n, + kind, + charge_pool, + context_projection.get('prompt_tokens_estimate'), + ( + ' context_in_prompt=' + f"{context_projection.get('context_remaining_in_prompt')}/" + f"{context_projection.get('context_capacity')}" + ) + if 'context_remaining_in_prompt' in context_projection else '', + ) + t0 = time.monotonic() + response = await real_llm_fn(prompt) + elapsed_ms = int((time.monotonic() - t0) * 1000) + _print_box( + f'LLM RESPONSE #{n} [{kind}] elapsed={elapsed_ms}ms chars={len(response)}', + response, + ) + captured_interactions.append({ + 'call_index': n, + 'kind': kind, + 'charge_pool': charge_pool, + 'prompt_chars': len(prompt), + 'prompt': prompt, + 'resource_status': _extract_resource_status(str(prompt)), + 'context_projection': context_projection, + 'response_chars': len(response), + 'response': response, + 'latency_ms': elapsed_ms, + }) + return response + + return verbose_llm + + +async def phase1_contract() -> dict[str, dict[str, Any]]: + from sqlalchemy import text + from shared.core.database import get_db_context + + logger.info('\n' + '█' * 80) + logger.info('Phase 1: DB contract check for current design') + logger.info('█' * 80) + + docs: dict[str, dict[str, Any]] = {} + async with get_db_context() as db: + rows = (await db.execute(text( + """ + SELECT d.document_id, d.source_file_name, d.current_job_result_id, + count(DISTINCT dc.id) AS chunk_count, + count(DISTINCT ds.section_id) AS section_count + FROM documents d + LEFT JOIN document_chunks dc + ON dc.document_id = d.document_id + AND dc.job_result_id = d.current_job_result_id + LEFT JOIN document_sections ds + ON ds.document_id = d.document_id + AND ds.job_result_id = d.current_job_result_id + WHERE d.user_id=:u AND d.namespace=:n AND d.status='active' + GROUP BY d.document_id, d.source_file_name, d.current_job_result_id + ORDER BY section_count DESC, chunk_count DESC + """ + ), {'u': USER_ID, 'n': NAMESPACE})).all() + + for doc_id, fname, job_result_id, chunk_count, section_count in rows: + props = (await db.execute(text( + """ + SELECT properties + FROM graph_nodes + WHERE owner_document_id=:d AND node_kind='document' + LIMIT 1 + """ + ), {'d': doc_id})).scalar() + has_nav_sections = isinstance(props, dict) and 'nav_sections' in props + logger.info( + 'doc={} name={} chunks={} sections={} graph_has_nav_sections={}', + doc_id, + fname, + chunk_count, + section_count, + has_nav_sections, + ) + docs[doc_id] = { + 'fname': fname, + 'job_result_id': job_result_id, + 'chunk_count': int(chunk_count or 0), + 'section_count': int(section_count or 0), + 'graph_has_nav_sections': has_nav_sections, + } + + return docs + + +async def phase2_scope_candidates(docs: dict[str, dict[str, Any]]) -> str | None: + from shared.core.database import get_db_context + from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections + + logger.info('\n' + '█' * 80) + logger.info('Phase 2: _load_child_sections root L1/L2 candidates') + logger.info('█' * 80) + + target_doc_id = next( + ( + doc_id for doc_id, info in docs.items() + if info['section_count'] > 0 and info['job_result_id'] + ), + None, + ) + if not target_doc_id: + logger.warning('No document with sections found') + return None + + info = docs[target_doc_id] + async with get_db_context() as db: + items = await load_child_sections(db, target_doc_id, info['job_result_id'], None) + logger.info('selected_doc={} name={} items={}', target_doc_id, info['fname'], len(items)) + for item in items[:30]: + logger.info( + ' L{} text={} image={} table={} path="{}"', + item.get('level'), + item.get('chunk_count'), + item.get('image_count'), + item.get('table_count'), + item.get('path'), + ) + if len(items) > 30: + logger.info(' ... and {} more', len(items) - 30) + return target_doc_id + + +async def run_test( + query: str, + label: str, + env_overrides: dict[str, str] | None = None, + expected_decision: str = '', +) -> dict[str, Any]: + """Unified test runner — all queries go through WorkflowOrchestrator. + + Mirrors the production agentic path (use_agentic=True) which routes + through WorkflowOrchestrator (planner → per-step RetrievalAgent). + """ + from shared.core.database import get_db_context + from shared.services.retrieval.llm_adapter import create_retrieval_llm_fn + from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator + from shared.services.retrieval.agentic.navigation.section_tree import ( + load_child_sections, + ) + + logger.info('\n' + '█' * 80) + logger.info(f'TEST [{label}] | QUERY: {query}') + logger.info('█' * 80) + + captured_interactions.clear() + + real_llm = create_retrieval_llm_fn() + if real_llm is None: + logger.error('LLM is not configured') + return {'error': 'no_llm'} + llm_fn = make_verbose_llm(real_llm) + + import shared.services.retrieval.agentic.navigation.section_tree as nav_mod + import shared.services.retrieval.agentic.navigation.tools as nav_tools_mod + orig_load_child_sections = nav_mod.load_child_sections + orig_tools_load_child_sections = nav_tools_mod.load_child_sections + + async def verbose_load_child_sections( + db, + document_id, + job_result_id, + scope_path=None, + exclude_paths=None, + limit_depth=True, + section_rows=None, + ): + items = await load_child_sections( + db, + document_id, + job_result_id, + scope_path, + exclude_paths=exclude_paths, + limit_depth=limit_depth, + section_rows=section_rows, + ) + _print_box( + f'SCOPE CANDIDATES doc={document_id} scope={scope_path or "root"} ' + f'count={len(items)} exclude_paths={exclude_paths}', + '\n'.join( + f'- L{item.get("level")} text={item.get("chunk_count")} ' + f'image={item.get("image_count")} table={item.get("table_count")} ' + f'is_leaf={item.get("is_leaf")} ' + f'path="{item.get("path")}"\n' + f' summary={(item.get("summary") or "")[:240]}' + for item in items + ), + ) + return items + + # Patch on both modules since tools.py caches the direct import + nav_mod.load_child_sections = verbose_load_child_sections + nav_tools_mod.load_child_sections = verbose_load_child_sections + + try: + t0 = time.monotonic() + with temporary_env(env_overrides): + async with get_db_context() as db: + result = await WorkflowOrchestrator().run( + db, + user_id=USER_ID, + namespace=NAMESPACE, + query=query, + top_k=TOP_K, + exclude_document_ids=[], + exclude_sections=[], + llm_fn=llm_fn, + ) + elapsed_ms = int((time.monotonic() - t0) * 1000) + finally: + # Restore originals for next test + nav_mod.load_child_sections = orig_load_child_sections + nav_tools_mod.load_child_sections = orig_tools_load_child_sections + + # Extract action list from captured interactions + actions = [c['kind'] for c in captured_interactions] + workflow_steps = [step.to_api_dict() for step in result.steps] + budget_accounting = _build_budget_accounting( + interactions=captured_interactions, + workflow_steps=workflow_steps, + ) + key_decisions = _build_key_decisions( + result=result, + workflow_steps=workflow_steps, + expected_decision=expected_decision, + env_overrides=env_overrides or {}, + budget_accounting=budget_accounting, + ) + + return { + 'label': label, + 'query': query, + 'expected_decision': expected_decision, + 'env_overrides': env_overrides or {}, + 'key_decisions': key_decisions, + 'budget_accounting': budget_accounting, + 'router_used': result.router_used, + 'stop_reason': '', + 'total_ms': elapsed_ms, + 'actions': actions, + 'evidence_text_chars': sum(len(step.evidence_text or '') for step in result.steps), + 'evidence_text': '\n\n'.join(step.evidence_text or '' for step in result.steps), + 'answer_text_chars': len(result.answer_text), + 'answer_text': result.answer_text, + 'referenced_chunks_count': len(result.referenced_chunks), + 'referenced_chunks': result.referenced_chunks, + 'budget_snapshot': result.wallet_snapshot, + 'workflow_plan': result.plan.to_dict() if result.plan else None, + 'workflow_steps': workflow_steps, + 'wallet_snapshot': result.wallet_snapshot, + 'planner_snapshot': result.planner_snapshot, + 'llm_interactions': len(captured_interactions), + 'llm_interaction_details': [ + { + 'call_index': c['call_index'], + 'kind': c['kind'], + 'charge_pool': c.get('charge_pool', 'unknown'), + 'latency_ms': c['latency_ms'], + 'prompt_chars': c['prompt_chars'], + 'response_chars': c['response_chars'], + 'prompt': c['prompt'], + 'resource_status': c.get('resource_status', {}), + 'context_projection': c.get('context_projection', {}), + 'response': c['response'], + } + for c in captured_interactions + ], + } + + +async def run_single_query( + *, + query: str, + label: str, + top_k: int, + data_type: int, +) -> dict[str, Any]: + """Run one query through the public production retrieval entry.""" + from shared.core.database import get_db_context + from shared.services.retrieval.app_service import run_retrieval_query + + _DATA_TYPE_MAP: dict[int, set[str] | None] = { + 1: None, 2: {"text"}, 3: {"image"}, 4: {"table"}, + 5: {"text", "image"}, 6: {"text", "table"}, + 7: {"page"}, 8: {"text", "image", "table"}, + } + chunk_types = _DATA_TYPE_MAP.get(data_type) + + t0 = time.monotonic() + async with get_db_context() as db: + result = await run_retrieval_query( + db=db, + user_id=USER_ID, + namespace=NAMESPACE, + query=query, + top_k=top_k, + exclude_document_ids=[], + exclude_sections=[], + chunk_types=chunk_types, + # --query mode should exercise the production agentic path; omitting + # this falls through to classic_topk and skips WorkflowOrchestrator. + use_agentic=True, + ) + elapsed_ms = int((time.monotonic() - t0) * 1000) + return { + 'label': label, + 'query': query, + 'top_k': top_k, + 'data_type': data_type, + 'namespace': NAMESPACE, + 'total_ms': elapsed_ms, + 'result': result, + } + + +def _asset_url_values(item: dict[str, Any]) -> list[str]: + if item.get('asset_url'): + return [str(item['asset_url'])] + return [] + + +def _render_single_query_report(report: dict[str, Any]) -> str: + result = report.get('result') or {} + refs = result.get('referenced_chunks') or [] + rows = result.get('results') or [] + evidence = result.get('evidence_text') or '' + md = [ + '# Retrieval Debug Trace\n', + f"Query: `{report.get('query', '')}`\n", + f"Label: `{report.get('label', '')}`\n", + f"Namespace: `{report.get('namespace', '')}`\n", + f"data_type: `{report.get('data_type')}`\n", + ( + f"Run summary: router=`{result.get('router_used')}`, " + f"stop_reason=`{result.get('stop_reason', '')}`, " + f"elapsed={report.get('total_ms')}ms, " + f"evidence={len(evidence)} chars, refs={len(refs)}, results={len(rows)}.\n" + ), + '\n## Referenced Chunks\n', + ] + if refs: + for i, ref in enumerate(refs, 1): + urls = _asset_url_values(ref) + md.append( + f"{i}. type=`{ref.get('chunk_type') or ref.get('type') or ''}`, " + f"section=`{ref.get('section_path', '')}`, " + f"file_path=`{ref.get('file_path', '')}`, asset_url_count={len(urls)}\n" + ) + for url in urls: + md.append(f" - {url}\n") + else: + md.append('No referenced chunks.\n') + + md.append('\n## Results\n') + if rows: + for i, row in enumerate(rows, 1): + urls = _asset_url_values(row) + md.append( + f"{i}. type=`{row.get('chunk_type') or row.get('type') or ''}`, " + f"section=`{row.get('section_path', '')}`, " + f"score=`{row.get('score', '')}`, asset_url_count={len(urls)}\n" + ) + for url in urls: + md.append(f" - {url}\n") + else: + md.append('No result rows.\n') + + md.append(f"\n## Evidence Text\n\n```text\n{evidence}\n```\n") + return ''.join(md) + + +def _budget_pool_line(snapshot: dict[str, Any] | None, pool_name: str) -> str: + if not isinstance(snapshot, dict): + return 'n/a' + pool = snapshot.get(pool_name) + if not isinstance(pool, dict): + return 'n/a' + capacity = pool.get('capacity', '?') + remaining = pool.get('remaining', '?') + try: + capacity_int = int(capacity) + remaining_int = int(remaining) + remaining_pct: int | str = 0 if capacity_int <= 0 else max( + 0, + min(100, round(remaining_int * 100 / capacity_int)), + ) + except (TypeError, ValueError): + remaining_pct = '?' + return ( + f"{pool.get('status', '?')} " + f"used={pool.get('used_pct', '?')}% " + f"remaining={remaining_pct}% ({remaining}/{capacity})" + ) + + +def _format_trimmed_path(item: Any, index: int) -> str: + if not isinstance(item, dict): + return f"{index}. path=`{item}`" + return ( + f"{index}. doc=`{item.get('document_name') or item.get('document_id', '')}`, " + f"path=`{item.get('path', '')}`, " + f"confidence={item.get('confidence_score', '?')}, " + f"discovery={item.get('discovery_score', '?')}, " + f"importance={item.get('importance_score', '?')}, " + f"tokens~={item.get('token_estimate', '?')}" + ) + + +def _compact_json(value: Any) -> str: + if not value: + return "" + return json.dumps(value, ensure_ascii=False, separators=(',', ': ')) + + +def _budget_line(snapshot: dict[str, Any] | None) -> str: + if not isinstance(snapshot, dict): + return "" + parts: list[str] = [] + for pool_name in ('bootstrap', 'planning', 'context'): + pool = snapshot.get(pool_name) + if isinstance(pool, dict): + parts.append( + f"{pool_name}={pool.get('status', '?')} " + f"{pool.get('used_pct', 0)}% " + f"({pool.get('remaining', 0)}/{pool.get('capacity', 0)})" + ) + return '; '.join(parts) + + +def _trace_observation_summary(observation: dict[str, Any]) -> str: + if not isinstance(observation, dict) or not observation: + return "" + summary: dict[str, Any] = {} + for key in ( + 'candidate_count', + 'visible_count', + 'total_steps', + 'collected_count', + 'guard_triggered', + 'exit_reason', + 'query', + ): + if key in observation: + summary[key] = observation[key] + if 'exclude_set' in observation: + exclude_set = observation.get('exclude_set') or [] + summary['exclude_set_count'] = len(exclude_set) + if not summary: + return "" + return _compact_json(summary) + + +def _append_trace_collected( + md: list[str], + collected: Any, + *, + label: str = 'Collected', + limit: int = 20, +) -> None: + if not isinstance(collected, list) or not collected: + return + md.append(f" - {label}: {len(collected)}\n") + for item in collected[:limit]: + if isinstance(item, dict): + path = item.get('path', '') + confidence = item.get('confidence') + suffix = f" confidence={confidence}" if confidence is not None else "" + hydrate_mode = item.get('hydrate_mode') + if hydrate_mode: + suffix += f" mode={hydrate_mode}" + md.append(f" - `{path}`{suffix}\n") + else: + md.append(f" - `{item}`\n") + if len(collected) > limit: + md.append(f" - (+{len(collected) - limit} more)\n") + + +def _pool_remaining(snapshot: dict[str, Any] | None, pool_name: str) -> int | None: + if not isinstance(snapshot, dict): + return None + pool = snapshot.get(pool_name) + if not isinstance(pool, dict): + return None + remaining = pool.get('remaining') + if remaining is None: + return None + try: + return int(remaining) + except (TypeError, ValueError): + return None + + +def _build_budget_accounting( + *, + interactions: list[dict[str, Any]], + workflow_steps: list[dict[str, Any]], +) -> dict[str, Any]: + pool_estimates = {'bootstrap': 0, 'planning': 0, 'context': 0, 'unknown': 0} + rows = [] + for interaction in interactions: + pool = str(interaction.get('charge_pool') or 'unknown') + if pool not in pool_estimates: + pool = 'unknown' + estimate = int( + (interaction.get('context_projection') or {}).get('prompt_tokens_estimate') + or 0 + ) + before = pool_estimates[pool] + pool_estimates[pool] = before + estimate + rows.append({ + 'call_index': interaction.get('call_index'), + 'kind': interaction.get('kind'), + 'pool': pool, + 'prompt_tokens_estimate': estimate, + 'pool_estimate_before': before, + 'pool_estimate_after': pool_estimates[pool], + }) + + final_context_remaining = next( + ( + _pool_remaining(step.get('budget_snapshot'), 'context') + for step in workflow_steps + if step.get('budget_snapshot') + ), + None, + ) + evidence_tokens = max( + estimate_tokens(step.get('evidence_text') or '') + for step in workflow_steps + ) if workflow_steps else 0 + + return { + 'rows': rows, + 'pool_estimates': pool_estimates, + 'evidence_tokens_estimate': evidence_tokens, + 'final_context_remaining': final_context_remaining, + } + + +def _build_key_decisions( + *, + result, + workflow_steps: list[dict[str, Any]], + expected_decision: str, + env_overrides: dict[str, str], + budget_accounting: dict[str, Any], +) -> list[str]: + planner_snapshot = result.planner_snapshot or {} + wallet_snapshot = result.wallet_snapshot or {} + step_statuses = [str(step.get('status') or '') for step in workflow_steps] + step_stop_reasons = [str(step.get('stop_reason') or '') for step in workflow_steps] + step_failure_reasons = [ + str(step.get('failure_reason') or '') + for step in workflow_steps + if step.get('failure_reason') + ] + + decisions = [ + ( + 'Planner inventory: ' + f"{planner_snapshot.get('total_docs', 0)} docs / " + f"{planner_snapshot.get('total_chunks', 0)} chunks" + ), + ( + 'Workflow wallet: ' + f"total={wallet_snapshot.get('total', 'n/a')} " + f"remaining={wallet_snapshot.get('remaining', 'n/a')} " + f"allocated={wallet_snapshot.get('allocated', 'n/a')}" + ), + f"Step statuses: {', '.join(step_statuses) or 'none'}", + f"Step stop reasons: {', '.join(step_stop_reasons) or 'none'}", + ] + if env_overrides: + decisions.append( + 'Env overrides: ' + + ', '.join(f'{key}={value}' for key, value in sorted(env_overrides.items())) + ) + for step in workflow_steps: + snap = step.get('budget_snapshot') or {} + trimmed_paths = snap.get('trimmed_paths') or [] + decisions.append( + f"Step {step.get('step_id')} budget: " + f"planning={_budget_pool_line(snap, 'planning')}; " + f"context={_budget_pool_line(snap, 'context')}; " + f"inventory={snap.get('total_docs', 0)} docs/{snap.get('total_chunks', 0)} chunks" + ) + if trimmed_paths: + preview = '; '.join( + str(item.get('path') or item)[:160] + for item in trimmed_paths[:3] + ) + decisions.append( + f"Trimmed paths: {len(trimmed_paths)} section(s) removed before answer. " + f"Preview: {preview}" + ) + if step_failure_reasons: + decisions.append('Failure reason propagated: ' + ';'.join(step_failure_reasons)) + decisions.append( + 'Evidence-only contract: final answer is not generated inside KNOWHERE.' + ) + + if expected_decision == 'budget_stop': + decisions.append( + 'Decision check: PASS' if any(status == 'budget_stop' for status in step_statuses) + else 'Decision check: FAIL (expected budget_stop)' + ) + elif expected_decision == 'not_found': + decisions.append( + 'Decision check: PASS' if any(status == 'not_found' for status in step_statuses) + else 'Decision check: FAIL (expected not_found)' + ) + return decisions + + + +def _render_md_report(all_reports: list[dict[str, Any]]) -> str: + """Render a readable algorithm trace report for the latest agentic flow.""" + md = [ + '# Agentic Retrieval E2E Trace Report\n', + 'This report is generated from real DB data and real LLM calls. It follows the current evidence-only algorithm flow: bottom discovery, KG document selection, per-document navigation, discovery merge, and evidence rendering.\n', + 'The debug runner writes only this Markdown report; no JSON artifact is emitted.\n', + ] + + for r in all_reports: + md.append(f"\n## {r['label']}\n") + md.append(f"Query: `{r['query']}`\n") + md.append( + f"Run summary: router=`{r['router_used']}`, stop_reason=`{r.get('stop_reason', '')}`, " + f"elapsed={r['total_ms']}ms, " + f"LLM calls={r.get('llm_interactions', 0)}, evidence={r.get('evidence_text_chars', 0)} chars, " + f"answer={r.get('answer_text_chars', 0)} chars, referenced_chunks={r.get('referenced_chunks_count', 0)}.\n" + ) + md.append(f"Action sequence: `{' -> '.join(r['actions'])}`\n\n") + + key_decisions = r.get('key_decisions') or [] + if key_decisions: + md.append('### Key Decision Checks\n') + for item in key_decisions: + md.append(f"- {item}\n") + md.append('\n') + + budget_accounting = r.get('budget_accounting') or {} + accounting_rows = budget_accounting.get('rows') or [] + if accounting_rows: + md.append('### Budget Accounting\n') + for row in accounting_rows: + md.append( + f"- Call {row.get('call_index')} `{row.get('kind')}` " + f"pool=`{row.get('pool')}` prompt_est={row.get('prompt_tokens_estimate')} " + f"pool_estimate={row.get('pool_estimate_before')}→{row.get('pool_estimate_after')}\n" + ) + md.append( + f"- Evidence tokens estimate: {budget_accounting.get('evidence_tokens_estimate')} tokens; " + f"final context remaining: {budget_accounting.get('final_context_remaining')} tokens; " + "no answer prompt is built inside KNOWHERE.\n\n" + ) + + if r.get('workflow_plan'): + md.append('### Workflow Plan\n') + md.append(_fence(r['workflow_plan'], 'json')) + md.append('### Workflow Steps\n') + for step in r.get('workflow_steps', []): + md.append( + f"- `{step.get('step_id')}` kind=`{step.get('step_kind')}` " + f"status=`{step.get('status')}` role=`{step.get('output_role')}` " + f"depends_on=`{step.get('depends_on')}` refs={len(step.get('referenced_chunks') or [])} " + f"stop_reason=`{step.get('stop_reason')}`\n" + ) + if step.get('failure_reason'): + md.append(f" failure_reason: `{step.get('failure_reason')}`\n") + if step.get('answer_text'): + md.append(_fence(step.get('answer_text'))) + if r.get('wallet_snapshot'): + md.append('### Wallet Snapshot\n') + md.append(_fence(r['wallet_snapshot'], 'json')) + if r.get('planner_snapshot'): + md.append('### Planner Snapshot\n') + md.append(_fence(r['planner_snapshot'], 'json')) + + # Decision Route Summary + md.append('### Decision Route\n') + interactions = r.get('llm_interaction_details', []) + for interaction in interactions: + kind = interaction['kind'] + idx = interaction['call_index'] + latency = interaction['latency_ms'] + resp = interaction.get('response', '') + resource_status = interaction.get('resource_status') or {} + context_projection = interaction.get('context_projection') or {} + stage, _purpose = _decision_stage(kind) + resp_preview = resp[:200].replace('\n', ' ').strip() + budget_bits = [] + prompt_estimate = context_projection.get('prompt_tokens_estimate') + charge_pool = interaction.get('charge_pool') or 'unknown' + if prompt_estimate is not None: + budget_bits.append(f"{charge_pool}_prompt_est={prompt_estimate}") + if resource_status.get('planning'): + budget_bits.append(f"planning={resource_status.get('planning')}") + if resource_status.get('context'): + budget_bits.append(f"context={resource_status.get('context')}") + if 'context_remaining_in_prompt' in context_projection: + budget_bits.append( + 'context_prompt=' + f"{context_projection.get('context_remaining_in_prompt')}/" + f"{context_projection.get('context_capacity')} " + f"used={context_projection.get('context_used_pct_in_prompt')}% " + f"remaining={context_projection.get('context_remaining_pct_in_prompt')}%" + ) + budget_suffix = f" budget: {'; '.join(budget_bits)}" if budget_bits else '' + md.append( + f"{idx}. **{stage}** ({latency}ms){budget_suffix} — `{resp_preview}`\n" + ) + md.append('\n') + + md.append('### Algorithm Trace\n') + md.append('Phase 1A Bottom Discovery always runs before these LLM calls. It performs high-recall lexical discovery and contributes document/path hints to later phases.\n\n') + + for interaction in interactions: + kind = interaction['kind'] + idx = interaction['call_index'] + latency = interaction['latency_ms'] + stage, purpose = _decision_stage(kind) + + md.append(f"#### Step {idx}: {stage}\n") + md.append(f"Kind: `{kind}`. Latency: {latency}ms. Prompt chars: {interaction['prompt_chars']}. Response chars: {interaction['response_chars']}.\n") + resource_status = interaction.get('resource_status') or {} + if resource_status: + md.append( + 'Budget: ' + f"planning=`{resource_status.get('planning', '?')}`, " + f"context=`{resource_status.get('context', '?')}`.\n" + ) + context_projection = interaction.get('context_projection') or {} + if context_projection: + md.append( + 'Debug budget: ' + f"prompt_tokens_estimate=`{context_projection.get('prompt_tokens_estimate', '?')}`" + ) + if 'context_remaining_in_prompt' in context_projection: + md.append( + ', ' + f"context_remaining_in_prompt=`{context_projection.get('context_remaining_in_prompt')}/" + f"{context_projection.get('context_capacity')}` " + f"({context_projection.get('context_used_pct_in_prompt')}% used, " + f"{context_projection.get('context_remaining_pct_in_prompt')}% remaining)" + ) + md.append('.\n') + md.append(f"Purpose: {purpose}\n") + md.append('
Full Prompt\n\n') + md.append(_fence(interaction['prompt'])) + md.append('\n
\n\n') + md.append('LLM response:\n') + md.append(_fence(interaction['response'], 'json')) + md.append('\n') + + md.append('### Answer Contract\n') + md.append('`answer_text` is intentionally empty. Downstream agents synthesize answers from `evidence_text`.\n\n') + + evidence_text = r.get('evidence_text', '') + md.append('### Rendered Evidence\n') + if evidence_text: + md.append(f'
Full evidence_text ({len(evidence_text)} chars)\n\n') + md.append(_fence(evidence_text)) + md.append('\n
\n\n') + else: + md.append('No rendered evidence collected.\n\n') + + refs = r.get('referenced_chunks', []) + md.append('### Referenced Chunks\n') + if refs: + for i, ref in enumerate(refs): + md.append( + f"{i + 1}. type=`{ref.get('chunk_type', '')}`, " + f"section=`{ref.get('section_path', '')}`, " + f"file_path=`{ref.get('file_path', '')}`\n" + ) + md.append("\n") + else: + md.append('No referenced chunks.\n\n') + + # Decision Trace — navigation decisions per step + all_decision_trace: list[dict] = [] + for step in r.get('workflow_steps', []): + dt = step.get('decision_trace') or [] + for entry in dt: + entry['_step_id'] = step.get('step_id', '?') + all_decision_trace.extend(dt) + if all_decision_trace: + md.append('### Decision Trace\n') + md.append('Navigation decisions made during agentic retrieval. ' + 'Each row follows observe -> decide -> result.\n\n') + for entry in all_decision_trace: + phase = entry.get('phase', '?') + agent = entry.get('agent', '?') + doc = entry.get('document') or '' + step_id = entry.get('_step_id', '?') + trace_index = entry.get('step_index', '?') + parent_index = entry.get('parent_step_index') + scope = entry.get('scope') or 'root' + observation = entry.get('observation') or {} + decision = entry.get('decision') or {} + result = entry.get('result') or {} + action = decision.get('action', '?') + args = decision.get('args') or {} + reason = decision.get('reason') or '' + status = result.get('status', '?') + + md.append( + f"- **[{step_id}] #{trace_index} {agent}.{phase}** " + f"doc=`{doc}` scope=`{scope}` action=`{action}` " + f"status=`{status}`" + ) + if parent_index is not None: + md.append(f" parent=`#{parent_index}`") + if reason: + md.append(f" — {reason}") + md.append('\n') + + if args: + md.append(f" - Args: `{_compact_json(args)}`\n") + + observation_summary = _trace_observation_summary(observation) + if observation_summary: + md.append(f" - Observation: `{observation_summary}`\n") + + _append_trace_collected(md, result.get('collected')) + + excluded = result.get('excluded_hints') or [] + if excluded: + md.append(f" - Excluded hints: {len(excluded)}\n") + for hint in excluded[:5]: + md.append( + f" - `{hint.get('path', '')}` " + f"(covered by `{hint.get('covered_by', '')})`\n" + ) + if len(excluded) > 5: + md.append(f" - (+{len(excluded) - 5} more)\n") + + for key in ( + 'hydrated_count', + 'matched_assets', + 'tool_status', + 'sub_agent_assessment', + 'note', + 'new_scope', + 'error', + ): + value = result.get(key) + if value not in (None, '', [], {}): + md.append(f" - {key}: `{value}`\n") + + budget_summary = _budget_line(entry.get('budget')) + if budget_summary: + md.append(f" - Budget: {budget_summary}\n") + + md.append('\n') + + md.append('\n') + + # Final Budget Snapshot + budget = r.get('budget_snapshot') + if budget: + md.append('### Final Budget Snapshot\n') + for pool_name in ('bootstrap', 'planning', 'context'): + pool = budget.get(pool_name) + if isinstance(pool, dict): + md.append( + f"- **{pool_name}**: {pool.get('status', '?')} " + f"({pool.get('used_pct', 0)}% used, " + f"remaining={pool.get('remaining', 0)}/{pool.get('capacity', 0)})\n" + ) + md.append( + f"- **Coverage**: {budget.get('explored_chunks', 0)}/{budget.get('total_chunks', 0)} chunks, " + f"{budget.get('explored_docs', 0)}/{budget.get('total_docs', 0)} docs\n" + ) + trimmed = budget.get('trimmed_paths', []) + if trimmed: + md.append(f"- **Trimmed paths**: {len(trimmed)} sections removed for budget\n") + for i, item in enumerate(trimmed, 1): + md.append(f" - {_format_trimmed_path(item, i)}\n") + md.append('\n') + + md.append("\n---\n") + + return '\n'.join(md) + + +async def main() -> None: + from datetime import datetime + + # Enable verbose logging to see full LLM prompts and responses + os.environ['RETRIEVAL_AGENTIC_VERBOSE'] = 'true' + os.environ['RETRIEVAL_AGENTIC_TRACE_ENABLED'] = 'false' + os.environ['RETRIEVAL_DECOMPOSITION_MAX_STEPS'] = '5' + + tests = [ + # ── Case 1: Original failing case (NAVIGATE → STOP → empty [DrillDown]) + # LLM navigates into "四、市场分析" then STOPs with no_relevant_child. + # Before fix: produces empty `▸ 四、 市场分析 [DrillDown]` + # After fix: child removed, evidence shows full document outline. + { + 'query': '安全大模型市场规模预测 2024 2025', + 'label': 'T1_Market_Size_NAVIGATE_STOP', + 'expected_decision': '', + 'env_overrides': { + 'RETRIEVAL_WALLET_TOTAL_BUDGET': '200000', + 'RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET': '40000', + 'RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET': '2000', + 'RETRIEVAL_AGENTIC_PLANNING_RATIO': '0.5', + 'RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS': '30000', + }, + }, + # ── Case 2: Deep hierarchy drill — should NAVIGATE → NAVIGATE → leaf hydration + # Targets a specific vendor inside "四、市场分析 / (二) 国内..." → should drill + # to L3 leaf and hydrate actual content. + { + 'query': '深信服安全大模型的技术方案和部署形态', + 'label': 'T2_Deep_Drill_Vendor', + 'expected_decision': '', + 'env_overrides': { + 'RETRIEVAL_WALLET_TOTAL_BUDGET': '200000', + 'RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET': '40000', + 'RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET': '2000', + 'RETRIEVAL_AGENTIC_PLANNING_RATIO': '0.5', + 'RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS': '30000', + }, + }, + # ── Case 3: Leaf-level match — query matches a specific L1 leaf section + # "法律声明" is a [Leaf] L1 node. Should SELECT directly, no drill-down needed. + { + 'query': '这份安全报告的法律声明和版权信息', + 'label': 'T3_Leaf_Direct_Select', + 'expected_decision': '', + 'env_overrides': { + 'RETRIEVAL_WALLET_TOTAL_BUDGET': '200000', + 'RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET': '40000', + 'RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET': '2000', + 'RETRIEVAL_AGENTIC_PLANNING_RATIO': '0.5', + 'RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS': '30000', + }, + }, + # ── Case 4: No relevant document — query about a topic not in any document + # Should result in no docs selected or empty evidence, testing the empty + # evidence path and downstream notification. + { + 'query': '量子计算对密码学的影响和后量子加密标准', + 'label': 'T4_No_Relevant_Doc', + 'expected_decision': '', + 'env_overrides': { + 'RETRIEVAL_WALLET_TOTAL_BUDGET': '200000', + 'RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET': '40000', + 'RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET': '2000', + 'RETRIEVAL_AGENTIC_PLANNING_RATIO': '0.5', + 'RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS': '30000', + }, + }, + # ── Case 5: Broad overview — should STOP at root with sufficient_outline + # Asks for an overview of the entire document, which should be answerable + # from the root outline alone. + { + 'query': '安全大模型技术与市场研究报告有哪些主要章节', + 'label': 'T5_Broad_Overview_STOP', + 'expected_decision': '', + 'env_overrides': { + 'RETRIEVAL_WALLET_TOTAL_BUDGET': '200000', + 'RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET': '40000', + 'RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET': '2000', + 'RETRIEVAL_AGENTIC_PLANNING_RATIO': '0.5', + 'RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS': '30000', + }, + }, + # ── Case 6: Cross-doc — query touching construction doc (not security) + # Tests that KG select picks the right document from a diverse corpus. + { + 'query': '土方开挖施工安全保证措施有哪些', + 'label': 'T6_Cross_Doc_Construction', + 'expected_decision': '', + 'env_overrides': { + 'RETRIEVAL_WALLET_TOTAL_BUDGET': '200000', + 'RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET': '40000', + 'RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET': '2000', + 'RETRIEVAL_AGENTIC_PLANNING_RATIO': '0.5', + 'RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS': '30000', + }, + }, + # ── Case 7: Asset search — SEARCH_IMAGES tool for chart/image queries + # Tests the new SEARCH_IMAGES tool: Navigator should use SEARCH_IMAGES + # with a semantic query, LLM filters from all images, only matching + # charts are added to evidence via reconcile_deferred_assets. + { + 'query': '帮我找出所有金融股票相关的图和折线图', + 'label': 'T7_Chart_Search_Images', + 'expected_decision': '', + 'env_overrides': { + 'RETRIEVAL_WALLET_TOTAL_BUDGET': '200000', + 'RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET': '40000', + 'RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET': '2000', + 'RETRIEVAL_AGENTIC_PLANNING_RATIO': '0.5', + 'RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS': '30000', + }, + }, + # ── Case 8: 冯荣州身份证图片 — image asset search + { + 'query': '冯荣州 的身份证图片发我', + 'label': 'T8_FRZ_ID_Image', + 'expected_decision': '', + 'env_overrides': { + 'RETRIEVAL_WALLET_TOTAL_BUDGET': '200000', + 'RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET': '40000', + 'RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET': '2000', + 'RETRIEVAL_AGENTIC_PLANNING_RATIO': '0.5', + 'RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS': '30000', + }, + }, + ] + + # ── CLI parsing ────────────────────────────────────────────────────── + import argparse + parser = argparse.ArgumentParser(description='Agentic retrieval E2E debug runner') + parser.add_argument( + '--test', '-t', action='append', dest='test_filters', default=[], + help='Filter tests by label substring (repeatable, e.g. --test T2 --test T6)', + ) + parser.add_argument( + '--output-dir', '-o', dest='output_dir', default=None, + help='Base output directory for trace runs. Default: ~/Desktop/agentic_traces', + ) + parser.add_argument( + '--namespace', '-n', dest='namespace', default=None, + help='Override the NAMESPACE used for DB queries.', + ) + parser.add_argument( + '--query', '-q', dest='query', default=None, + help='Run one ad-hoc query through shared.services.retrieval.app_service.', + ) + parser.add_argument( + '--label', dest='label', default='adhoc', + help='Trace label for --query output.', + ) + parser.add_argument( + '--top-k', dest='top_k', type=int, default=TOP_K, + help=f'Top K for --query mode. Default: {TOP_K}.', + ) + parser.add_argument( + '--chunk-scope', + choices=sorted(CHUNK_SCOPE_DATA_TYPE), + default='all', + help=( + 'Chunk type scope for --query mode: all=mixed, page=PAGE only, ' + 'chunk=text/image/table only.' + ), + ) + parser.add_argument( + '--data-type', + dest='data_type', + type=int, + default=None, + help='Override retrieval data_type for --query mode.', + ) + parser.add_argument( + '--print-evidence', + action='store_true', + help='Print full evidence_text to stdout in --query mode.', + ) + # Also accept a positional arg for backward compat: `python debug_retrieval.py T6` + parser.add_argument('positional_filter', nargs='?', default=None) + args = parser.parse_args() + + # Merge positional filter into test_filters for backward compat + if args.positional_filter and args.positional_filter not in args.test_filters: + args.test_filters.append(args.positional_filter) + + if args.namespace: + global NAMESPACE + NAMESPACE = args.namespace + + if args.query: + data_type = ( + args.data_type + if args.data_type is not None + else CHUNK_SCOPE_DATA_TYPE[args.chunk_scope] + ) + timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') + output_base_dir = args.output_dir or os.path.expanduser('~/Desktop/agentic_traces') + safe_label = re.sub(r'[^\w\-]', '_', args.label) + output_dir = os.path.join( + os.path.expanduser(output_base_dir), + f'{timestamp}_{safe_label}', + ) + os.makedirs(output_dir, exist_ok=True) + + report = await run_single_query( + query=args.query, + label=args.label, + top_k=args.top_k, + data_type=data_type, + ) + result_path = os.path.join(output_dir, 'result.json') + trace_path = os.path.join(output_dir, 'trace.md') + with open(result_path, 'w', encoding='utf-8') as f: + json.dump(report, f, ensure_ascii=False, indent=2, default=str) + with open(trace_path, 'w', encoding='utf-8') as f: + f.write(_render_single_query_report(report)) + + result = report.get('result') or {} + evidence = result.get('evidence_text') or '' + refs = result.get('referenced_chunks') or [] + rows = result.get('results') or [] + print('=' * 90) + print(f"QUERY: {args.query}") + print( + f"router={result.get('router_used')} " + f"stop={result.get('stop_reason', '')} " + f"data_type={data_type} scope={args.chunk_scope}" + ) + print(f"evidence_text chars: {len(evidence)}") + print(f"referenced_chunks: {len(refs)} results: {len(rows)}") + print(f"TRACE: {output_dir}") + if args.print_evidence: + print("\n----- EVIDENCE TEXT -----\n") + print(evidence) + return + + docs = await phase1_contract() + await phase2_scope_candidates(docs) + + # ── Filter tests ───────────────────────────────────────────────────── + if args.test_filters: + tests = [ + t for t in tests + if any(f in t['label'] for f in args.test_filters) + ] + logger.info(f' Filtered to {len(tests)} test(s) matching {args.test_filters}') + if not tests: + logger.error('No tests matched the filter(s). Available labels:') + for t in tests: + logger.error(f' - {t["label"]}') + return + + # ── Timestamped output folder for the entire run ───────────────────── + timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') + output_base_dir = args.output_dir or os.path.expanduser('~/Desktop/agentic_traces') + output_dir = os.path.join(os.path.expanduser(output_base_dir), timestamp) + os.makedirs(output_dir, exist_ok=True) + + all_reports = [] + + for test in tests: + report = await run_test( + test['query'], + test['label'], + env_overrides=test.get('env_overrides'), + expected_decision=test.get('expected_decision', ''), + ) + all_reports.append(report) + + # Write per-query trace file into the SAME folder + safe_name = re.sub(r'[^\w\-]', '_', report['label']) + file_path = os.path.join(output_dir, f'{safe_name}.md') + try: + with open(file_path, 'w', encoding='utf-8') as f: + f.write(_render_md_report([report])) + logger.info(f' Trace saved: {file_path}') + except Exception as e: + logger.error(f' Failed to save trace {file_path}: {e}') + + # Write combined summary index + index_path = os.path.join(output_dir, '_index.md') + with open(index_path, 'w', encoding='utf-8') as f: + f.write('# Agentic E2E Trace Index\n\n') + f.write(f'Generated: {datetime.now().isoformat()}\n\n') + f.write('| # | Label | Query | Router | LLM Calls | Refs | Elapsed |\n') + f.write('|:--|:------|:------|:-------|:----------|:-----|:--------|\n') + for i, r in enumerate(all_reports, 1): + safe_name = re.sub(r'[^\w\-]', '_', r['label']) + f.write( + f"| {i} | [{r['label']}]({safe_name}.md) " + f"| {r['query'][:40]}… " + f"| `{r['router_used']}` " + f"| {r.get('llm_interactions', 0)} " + f"| {r.get('referenced_chunks_count', 0)} " + f"| {r['total_ms']}ms |\n" + ) + + logger.info('\n' + '=' * 80) + logger.info(f'TRACE FOLDER saved to: {output_dir}') + logger.info(f' Index: {index_path}') + logger.info(f' Total tests: {len(all_reports)} traces') + logger.info('=' * 80) + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/apps/worker/scripts/debug_text_track.py b/apps/worker/scripts/debug_text_track.py new file mode 100644 index 000000000..eaae97bae --- /dev/null +++ b/apps/worker/scripts/debug_text_track.py @@ -0,0 +1,858 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Staged text-track document parsing debug script. + +Supports PDF (shard-aware), DOCX, and MD formats with four breakpoints: + 1. profile — production-aligned DOC_AGENT profile: + run_coarse → lightweight (≤MAX) / structural (>MAX) + 2. mineru — shard splitting + MinerU extraction + 3. hierarchy — heading prediction → merged hierarchy tree + 4. full — complete extraction → chunks/doc_nav/manifest + +Output directory: + ~/.knowhere/_debug_parse//text_track/ + +Usage: + cd apps/worker + uv run python scripts/debug_text_track.py --file /path/to/doc.pdf + uv run python scripts/debug_text_track.py --file /path/to/doc.pdf --stop-at profile + uv run python scripts/debug_text_track.py --sjsyj --stop-at mineru --reuse-profile + uv run python scripts/debug_text_track.py --sjsyj --stop-at hierarchy --reuse-mineru + uv run python scripts/debug_text_track.py --file /path/to/doc.docx --stop-at hierarchy +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import sys +import time +from pathlib import Path +from typing import Any + +# ── Bootstrap ─────────────────────────────────────────────────────────────── +ROOT = Path(__file__).resolve().parents[3] +WORKER_ROOT = ROOT / "apps" / "worker" +sys.path.insert(0, str(WORKER_ROOT)) +sys.path.insert(0, str(ROOT / "packages" / "shared-python")) + +from dotenv import load_dotenv + +load_dotenv(WORKER_ROOT / ".env") +os.environ.setdefault("LOCAL_DEBUG", "1") +os.environ.setdefault("OVERSIZED_PDF_SHARD_ENABLED", "true") + +from loguru import logger + +from shared.services.ai.token_tracking import ( + init_token_tracker, + get_current_token_tracker, +) + +# ── Constants ─────────────────────────────────────────────────────────────── +DEFAULT_SJSYJ_PDF = Path( + "/Users/wuchengke/Desktop/temp/test_docs/" + "SJSYJ-SC-2024 企业制度汇编(上册).pdf" +) +DEFAULT_SPACEX_PDF = Path("/Users/wuchengke/Desktop/temp/test_docs/spacex-s1.pdf") +OUTPUT_ROOT = Path("~/.knowhere/_debug_parse").expanduser() + + +def _write_json(path: Path, data: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") + logger.info(" → {}", path) + + +# ── Stage 1: Profile + Shard Plan (PDF only) ─────────────────────────────── + +def _stage_profile(pdf_path: str, filename: str, out_dir: Path, model: str | None): + """Production-aligned profile: coarse → lightweight / structural.""" + from app.services.document_agent.coordinator import ProfileCoordinator + from app.services.document_parser.profiling.taxonomy import PdfRoutingCategory + from shared.core.config import settings + + logger.info("=" * 70) + logger.info("🧬 Stage 1: DOC_AGENT profile (production-aligned)") + logger.info("=" * 70) + + doc_agent_dir = out_dir / "_doc_agent" + doc_agent_dir.mkdir(parents=True, exist_ok=True) + + vlm_model = model or settings.IMAGE_MODEL + coordinator = ProfileCoordinator( + pdf_path=pdf_path, + job_id=filename, + output_dir=str(doc_agent_dir), + model=vlm_model, + settings={ + "planner_model": vlm_model, + "vlm_model": vlm_model, + "model": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, + "toc_profile_enabled": settings.PDF_PROFILE_TOC_ENABLED, + "toc_before_coarse": settings.PDF_PROFILE_TOC_ENABLED, + }, + ) + t0 = time.time() + agent_profile = coordinator.run_coarse() + page_count = int(coordinator.blackboard.page_count or 0) + routing = PdfRoutingCategory.normalize(agent_profile.routing_category) + if routing is PdfRoutingCategory.ATLAS: + raise RuntimeError( + "Atlas routing skips text-track anatomy (same as production). " + "Use page-memory / atlas debug paths for this document." + ) + if page_count > settings.MAX_PDF_PAGE_LIMIT: + logger.info( + " page_count={} > MAX_PDF_PAGE_LIMIT={} → run_structural", + page_count, + settings.MAX_PDF_PAGE_LIMIT, + ) + anatomy = coordinator.run_structural() + else: + logger.info( + " page_count={} ≤ MAX_PDF_PAGE_LIMIT={} → run_lightweight_anatomy", + page_count, + settings.MAX_PDF_PAGE_LIMIT, + ) + anatomy = coordinator.run_lightweight_anatomy() + elapsed = time.time() - t0 + + profile = coordinator.blackboard.document_profile or agent_profile + margins = coordinator.blackboard.global_signals.get("content_margins") or {} + header_y = getattr(profile, "header_y", None) + footer_y = getattr(profile, "footer_y", None) + if header_y is None: + header_y = margins.get("header_y") + if footer_y is None: + footer_y = margins.get("footer_y") + + logger.info(" profile done in {:.1f}s", elapsed) + logger.info( + " category={!r} routing={} header_y={} footer_y={}", + getattr(profile, "category", None), + getattr(profile, "routing_category", None), + header_y, + footer_y, + ) + logger.info(" page_count={}", anatomy.page_count) + logger.info(" toc_pages={}", anatomy.toc_result.toc_pages) + logger.info(" shard_plan.enabled={}", anatomy.shard_plan.enabled) + logger.info(" shard_count={}", len(anatomy.shard_plan.shards)) + for shard in anatomy.shard_plan.shards: + logger.info( + " shard_{}: p{}-{} ({} pages, anchor={})", + shard.shard_index, shard.page_start, shard.page_end, + shard.page_end - shard.page_start + 1, shard.anchor_type, + ) + + return anatomy, elapsed, { + "profile": profile.to_dict() if profile is not None else None, + "path": ( + "structural" + if page_count > settings.MAX_PDF_PAGE_LIMIT + else "lightweight" + ), + "asset_pages": sum( + 1 for feature in coordinator.blackboard.page_features if feature.has_asset + ), + } + + +def _load_anatomy_cache(out_dir: Path, pdf_path: str, filename: str): + from app.services.document_agent.manifest import ( + PageAnatomyMap, PageFeature, PageLabel, + Shard, ShardPlan, TocResult, ValidationReport, + ) + + cache_path = out_dir / "_doc_agent" / "anatomy_map.json" + if not cache_path.exists(): + # Fallback: check sibling page_memory dir (shared profile output) + sibling = out_dir.parent / "page_memory" / "_doc_agent" / "anatomy_map.json" + if sibling.exists(): + cache_path = sibling + else: + raise FileNotFoundError(f"No cached anatomy: {cache_path}") + + logger.info("⏩ Reusing cached anatomy: {}", cache_path) + data = json.loads(cache_path.read_text(encoding="utf-8")) + toc = data.get("toc_result") or {} + sp = data.get("shard_plan") or {} + + page_features = [ + PageFeature( + page=int(pf.get("page", 0)), + raw_text_length=int(pf.get("raw_text_length", 0)), + text_density=float(pf.get("text_density", 0)), + image_coverage=float(pf.get("image_coverage", 0)), + image_count=int(pf.get("image_count", 0)), + table_count=int(pf.get("table_count", 0)), + drawings_count=int(pf.get("drawings_count", 0)), + orientation=pf.get("orientation", "portrait"), + width=float(pf.get("width", 0)), + height=float(pf.get("height", 0)), + has_asset=bool(pf.get("has_asset", False)), + is_blank_like=bool(pf.get("is_blank_like", False)), + ) + for pf in data.get("page_features", []) + ] + page_labels = [ + PageLabel( + page=int(pl.get("page", 0)), + kind=pl.get("kind", "normal"), + confidence=float(pl.get("confidence", 0)), + evidence=pl.get("evidence", {}), + ) + for pl in data.get("page_labels", []) + ] + + return PageAnatomyMap( + job_id=data.get("job_id", filename), + file_path=data.get("file_path", pdf_path), + page_count=int(data.get("page_count", 0)), + page_features=page_features, + page_labels=page_labels, + toc_result=TocResult( + toc_pages=list(toc.get("toc_pages", [])), + method=toc.get("method", "none"), + ), + shard_plan=ShardPlan( + enabled=bool(sp.get("enabled", False)), + reason=sp.get("reason", "not_needed"), + shards=[ + Shard( + shard_index=int(s.get("shard_index", i)), + page_start=int(s.get("page_start", 1)), + page_end=int(s.get("page_end", 1)), + page_offset=int(s.get("page_offset", 0)), + anchor_type=s.get("anchor_type", "forced_max_size"), + anchor_evidence=s.get("anchor_evidence", ""), + confidence=float(s.get("confidence", 0) or 0), + ) + for i, s in enumerate(sp.get("shards", [])) + ], + validation=ValidationReport(valid=True), + ), + toc_hierarchies=data.get("toc_hierarchies"), + toc_page_offset=data.get("toc_page_offset"), + ) + + +# ── Stage 2: MinerU Extraction (PDF only) ──────────────────────────────────── + +def _stage_mineru_pdf( + pdf_path: str, + filename: str, + out_dir: Path, + anatomy, +) -> tuple[list[str], float]: + """Split PDF into shards and run MinerU extraction only (no heading prediction).""" + from app.services.document_parser.formats.pdf.shard_splitter import ( + bin_pack_shards, + split_pdf, + ) + from app.services.document_parser.providers.mineru.pdf_service import parse_via_full + from shared.core.config import settings + + logger.info("=" * 70) + logger.info("🔄 Stage 2: Shard splitting + MinerU extraction") + logger.info("=" * 70) + + t0 = time.time() + agent_shards = anatomy.shard_plan.shards + + toc_pages: set[int] = set() + if anatomy.toc_result and anatomy.toc_result.toc_pages: + toc_pages = set(anatomy.toc_result.toc_pages) + + max_pages = int(os.environ.get("MAX_PDF_PAGE_LIMIT", getattr(settings, "MAX_PDF_PAGE_LIMIT", 200))) + merged_shards = bin_pack_shards(agent_shards, max_pages=max_pages) + logger.info(" {} agent shards → {} MinerU shards", len(agent_shards), len(merged_shards)) + + work_dir = str(out_dir / "_shards") + os.makedirs(work_dir, exist_ok=True) + + fast_path = len(merged_shards) == 1 and not toc_pages + + if fast_path: + logger.info(" single shard, no TOC exclusion → fast path (no split)") + shard_out = os.path.join(work_dir, "shard_0") + os.makedirs(shard_out, exist_ok=True) + parse_via_full(pdf_path, filename, shard_out) + shard_output_dirs = [shard_out] + else: + shard_pdf_paths, _page_remap = split_pdf( + pdf_path, merged_shards, work_dir, + exclude_pages=toc_pages if toc_pages else None, + ) + logger.info(" split into {} shard PDFs (excluded {} TOC pages)", + len(shard_pdf_paths), len(toc_pages)) + + shard_output_dirs: list[str] = [None] * len(shard_pdf_paths) # type: ignore[list-item] + for i, shard_pdf in enumerate(shard_pdf_paths): + shard_out = os.path.join(work_dir, f"shard_{i}") + os.makedirs(shard_out, exist_ok=True) + shard_filename = f"{os.path.splitext(filename)[0]}_shard{i}.pdf" + logger.info(" 🔄 MinerU shard_{}: parsing...", i) + parse_via_full(shard_pdf, shard_filename, shard_out) + shard_output_dirs[i] = shard_out + + elapsed = time.time() - t0 + logger.info(" ✅ MinerU extraction done in {:.1f}s → {} shard dirs", elapsed, len(shard_output_dirs)) + return shard_output_dirs, elapsed + + +# ── Stage 3: Per-Shard Heading Prediction + Hierarchy Tree ──────────────────── + +def _stage_hierarchy_pdf( + out_dir: Path, + anatomy, + model: str | None, +) -> tuple[list[str], float]: + """Run heading prediction on each shard's full.md, merge, and output hierarchy tree.""" + from app.services.document_parser.formats.markdown.parser import ( + eval_md_headings, + merge_html_tables, + ) + from app.services.document_parser.formats.pdf.shard_merger import ( + merge_images, + merge_shard_lines, + ) + from app.services.document_parser.formats.pdf.shard_splitter import bin_pack_shards + from app.services.document_agent.tools.propose_shard_plan import split_toc_for_shard + from shared.core.config import settings + + logger.info("=" * 70) + logger.info("🔬 Stage 3: Per-shard heading prediction → merged hierarchy") + logger.info("=" * 70) + + t0 = time.time() + agent_shards = anatomy.shard_plan.shards + toc_hierarchies = anatomy.toc_hierarchies + + max_pages = int(os.environ.get("MAX_PDF_PAGE_LIMIT", getattr(settings, "MAX_PDF_PAGE_LIMIT", 200))) + merged_shards = bin_pack_shards(agent_shards, max_pages=max_pages) + + work_dir = out_dir / "_shards" + + # Discover shard output dirs + shard_output_dirs: list[str] = [] + for i in range(len(merged_shards)): + shard_dir = str(work_dir / f"shard_{i}") + if not os.path.isdir(shard_dir): + raise FileNotFoundError(f"shard_{i} dir not found: {shard_dir} (run --stop-at mineru first)") + shard_output_dirs.append(shard_dir) + + hierarchy_model = model or os.environ.get("NORMOL_MODEL") + + def _predict_shard(shard_idx: int, shard_out_dir: str) -> list[str]: + md_path = os.path.join(shard_out_dir, "full.md") + if not os.path.exists(md_path): + raise FileNotFoundError(f"shard_{shard_idx}: full.md not found at {md_path}") + + with open(md_path, "r", encoding="utf-8") as f: + md_lines = f.readlines() + md_lines = [line.strip() for line in md_lines if line.strip()] + md_lines = merge_html_tables(md_lines) + + is_first = shard_idx == 0 + shard = merged_shards[shard_idx] + shard_toc = ( + toc_hierarchies if is_first + else split_toc_for_shard( + toc_hierarchies, shard.page_start, shard.page_end, + offset_override=getattr(anatomy, "toc_page_offset", None), + ) + ) + + lines_with_heading = eval_md_headings( + md_lines, + source_type="md", + toc_hierarchies=shard_toc, + smart_parse=True, + model_name=hierarchy_model, + output_dir=shard_out_dir, + layout_json_path=( + os.path.join(shard_out_dir, "layout.json") + if os.path.exists(os.path.join(shard_out_dir, "layout.json")) + else None + ), + is_first_shard=is_first, + ) + + heading_count = sum(1 for line in lines_with_heading if line.startswith("#")) + logger.info(" ✅ shard_{}: {} headings from {} lines", + shard_idx, heading_count, len(lines_with_heading)) + + _write_json( + Path(shard_out_dir) / "lines_with_heading.json", + lines_with_heading, + ) + return lines_with_heading + + all_shard_lines: list[list[str]] = [] + for i, shard_dir in enumerate(shard_output_dirs): + lines = _predict_shard(i, shard_dir) + all_shard_lines.append(lines) + + # Merge (boundary-heading dedup; signature matches production parser) + merged_lines = merge_shard_lines(all_shard_lines) + + # Always copy shard images into the package root. Unlike production's MinerU + # fast path (which writes directly to output_dir), this debug script always + # stages MinerU under ``_shards/shard_*`` — even for the 1-shard/no-TOC case. + # Skipping merge_images on fast_path left ``text_track/images/`` empty and + # dropped all image chunks in Phase B. + merge_images(shard_output_dirs, str(out_dir)) + + total_headings = sum(1 for line in merged_lines if line.startswith("#")) + logger.info(" 📎 Merged: {} lines, {} headings", len(merged_lines), total_headings) + + _write_json(work_dir / "merged_lines.json", merged_lines) + + # Output the full hierarchy tree as a structured JSON + hierarchy_tree = _build_hierarchy_tree(merged_lines) + _write_json(out_dir / "hierarchy.json", hierarchy_tree) + logger.info(" 🌲 Hierarchy tree: {} top-level nodes, {} total nodes", + len(hierarchy_tree), _count_tree_nodes(hierarchy_tree)) + + elapsed = time.time() - t0 + return merged_lines, elapsed + + +def _build_hierarchy_tree(merged_lines: list[str]) -> dict[str, Any]: + """Build a nested hierarchy dict from merged lines (same format as manifest HIERARCHY). + + Sibling titles that collide get the same ``_2`` / ``_3`` … suffix used by + markdown ``ParseState.enter_heading`` / chunk paths, so duplicate names + (e.g. repeated ``临床实践要点:``) remain visible in the tree. + """ + from shared.services.chunks.path_segments import ( + DOCUMENT_PATH_SEP, + escape_path_segment, + ) + + headings: list[tuple[int, str]] = [] + for line in merged_lines: + if line.startswith("#"): + level = 0 + for ch in line: + if ch == "#": + level += 1 + else: + break + title = line[level:].strip() + headings.append((level, title)) + + if not headings: + return {} + + root: dict[str, Any] = {} + # (level, node_dict, escaped_disambiguated_title) + stack: list[tuple[int, dict[str, Any], str]] = [] + path_counter: dict[str, int] = {} + + for level, title in headings: + node: dict[str, Any] = {} + while stack and stack[-1][0] >= level: + stack.pop() + + current_heading = escape_path_segment(title) + parent_names = [item_title for _, _, item_title in stack] + tentative_path = DOCUMENT_PATH_SEP.join([*parent_names, current_heading]) + if tentative_path in path_counter: + path_counter[tentative_path] += 1 + current_heading = f"{current_heading}_{path_counter[tentative_path]}" + else: + path_counter[tentative_path] = 1 + + parent = stack[-1][1] if stack else root + parent[current_heading] = node + stack.append((level, node, current_heading)) + + return root + + +def _count_tree_nodes(tree: dict[str, Any]) -> int: + count = 0 + for _key, children in tree.items(): + count += 1 + if children: + count += _count_tree_nodes(children) + return count + + +def _stage_hierarchy_docx( + file_path: str, + filename: str, + out_dir: Path, + model: str | None, +) -> tuple[Any, float]: + from app.services.document_parser.formats.docx.parser import parse_docx + + logger.info("=" * 70) + logger.info("🔬 Stage 2: DOCX parsing (parse_docx)") + logger.info("=" * 70) + + t0 = time.time() + base_llm_paras = { + "smart_title_parse": True, + "summary_image": True, + "summary_table": True, + "summary_txt": True, + "stopwords": [], + "model_name": model or os.environ.get("NORMOL_MODEL"), + } + # relative_root must be the document file name (production contract), + # never the absolute debug out_dir — otherwise chunk.path / section trees + # leak ~/.knowhere/_debug_parse/.../text_track into retrieval. + parsed_df = parse_docx( + file_path, + base_llm_paras, + str(out_dir), + filename, + file_url="", + relative_root=filename, + ) + elapsed = time.time() - t0 + logger.info(" DOCX parsed in {:.1f}s → {} rows", elapsed, len(parsed_df) if parsed_df is not None else 0) + return parsed_df, elapsed + + +def _stage_hierarchy_md( + file_path: str, + filename: str, + out_dir: Path, + model: str | None, +) -> tuple[Any, float]: + from app.services.document_parser.formats.markdown.parser import parse_md + + logger.info("=" * 70) + logger.info("🔬 Stage 2: Markdown parsing (parse_md)") + logger.info("=" * 70) + + t0 = time.time() + base_llm_paras = { + "smart_title_parse": True, + "summary_image": True, + "summary_table": True, + "summary_txt": True, + "stopwords": [], + "model_name": model or os.environ.get("NORMOL_MODEL"), + } + parsed_df = parse_md( + str(out_dir), + source_type="md", + file_path=file_path, + base_llm_paras=base_llm_paras, + relative_root=filename, + ) + elapsed = time.time() - t0 + logger.info(" MD parsed in {:.1f}s → {} rows", elapsed, len(parsed_df) if parsed_df is not None else 0) + return parsed_df, elapsed + + +# ── Stage 4: Full Extraction → Chunks ────────────────────────────────────── + +def _stage_full_pdf( + out_dir: Path, + filename: str, + merged_lines: list[str], + model: str | None, +): + from app.services.document_parser.formats.markdown.parser import parse_md + from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess + + logger.info("=" * 70) + logger.info("📦 Stage 4: parse_md Phase B → DataFrame → chunks") + logger.info("=" * 70) + + t0 = time.time() + base_llm_paras = { + "smart_title_parse": True, + "summary_image": True, + "summary_table": True, + "summary_txt": True, + "stopwords": [], + "model_name": model or os.environ.get("NORMOL_MODEL"), + } + parsed_df = parse_md( + str(out_dir), + source_type="md", + base_llm_paras=base_llm_paras, + relative_root=filename, + lines_with_heading=merged_lines, + ) + parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) + elapsed_parse = time.time() - t0 + logger.info(" parse_md Phase B done in {:.1f}s → {} rows", elapsed_parse, len(parsed_df) if parsed_df is not None else 0) + + return _finalize_df(out_dir, filename, parsed_df) + + +def _finalize_df(out_dir: Path, filename: str, parsed_df): + from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks + from shared.services.storage.zip_doc_navigation import ZipDocNavigationBuilder + from datetime import datetime, timezone + + t0 = time.time() + chunks = dataframe_to_chunks(parsed_df) + logger.info(" {} chunks generated", len(chunks)) + + _write_json(out_dir / "chunks.json", {"chunks": chunks}) + + doc_nav = ZipDocNavigationBuilder().build_doc_nav(chunks, filename) + _write_json(out_dir / "doc_nav.json", doc_nav) + + manifest = { + "version": "2.0", + "job_id": filename, + "source_file_name": filename, + "processing_date": datetime.now(timezone.utc).isoformat(), + "processing": { + "token_usage": get_current_token_tracker(), + }, + "statistics": doc_nav.get("stats", {}), + } + _write_json(out_dir / "manifest.json", manifest) + + try: + from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries + enrich_doc_nav_summaries(str(out_dir.parent), source_file=filename, use_llm=False) + except Exception as exc: + logger.warning(" enrich failed (non-fatal): {}", exc) + + elapsed = time.time() - t0 + logger.info(" finalize done in {:.1f}s", elapsed) + return chunks, elapsed + + +# ── Main ──────────────────────────────────────────────────────────────────── + +def main() -> int: + parser = argparse.ArgumentParser( + description=( + "Staged text-track debug: profile → mineru → hierarchy → full" + ), + ) + parser.add_argument("--file", default=None, help="Input file path (PDF/DOCX/MD)") + parser.add_argument("--model", default=None, help="Override LLM model") + parser.add_argument( + "--stop-at", + choices=["profile", "mineru", "hierarchy", "full"], + default="full", + help="Pipeline stopping point (default: full)", + ) + parser.add_argument("--reuse-profile", action="store_true", + help="Reuse cached _doc_agent/anatomy_map.json") + parser.add_argument("--reuse-mineru", action="store_true", + help="Reuse cached shard dirs (skip MinerU extraction)") + parser.add_argument("--reuse-hierarchy", action="store_true", + help="Reuse cached merged_lines.json (skip heading prediction)") + parser.add_argument("--sjsyj", action="store_true", + help=f"Use fixture: {DEFAULT_SJSYJ_PDF}") + parser.add_argument("--spacex", action="store_true", + help=f"Use fixture: {DEFAULT_SPACEX_PDF}") + parser.add_argument("--run-db", action="store_true", + help="Publish to DB after full extraction") + parser.add_argument("--clean", action="store_true", + help="Delete existing output before running") + args = parser.parse_args() + + # Resolve input file + file_path: str | None = args.file + if args.sjsyj: + file_path = str(DEFAULT_SJSYJ_PDF) + elif args.spacex: + file_path = str(DEFAULT_SPACEX_PDF) + + if not file_path: + parser.error("Provide --file, --sjsyj, or --spacex") + + file_path = str(Path(file_path).expanduser().resolve()) + if not os.path.exists(file_path): + raise FileNotFoundError(file_path) + + filename = os.path.basename(file_path) + ext = Path(file_path).suffix.lower() + + from app.services.document_parser.orchestration.path_segment import build_parser_path_segment + dir_name = build_parser_path_segment(filename) + out_dir = OUTPUT_ROOT / dir_name / "text_track" + + if args.clean and out_dir.exists(): + logger.info("🗑️ Cleaning {}", out_dir) + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + init_token_tracker() + + logger.info("█" * 70) + logger.info(" TEXT-TRACK DEBUG: {}", filename) + logger.info(" OUTPUT: {}", out_dir) + logger.info(" STOP-AT: {}", args.stop_at) + logger.info("█" * 70) + + trace: dict[str, Any] = { + "file": filename, + "format": ext, + "stop_at": args.stop_at, + "stages": {}, + } + t_start = time.time() + + # ── Format dispatch ────────────────────────────────────────────────────── + if ext == ".pdf": + # Stage 1: Profile + if args.reuse_profile: + anatomy = _load_anatomy_cache(out_dir, file_path, filename) + else: + anatomy, profile_elapsed, profile_meta = _stage_profile( + file_path, filename, out_dir, args.model + ) + trace["stages"]["profile"] = { + "elapsed_s": round(profile_elapsed, 1), + **profile_meta, + } + + if args.stop_at == "profile": + trace["stages"].setdefault("profile", {})["shard_count"] = len( + anatomy.shard_plan.shards + ) + _write_json(out_dir / "trace.json", trace) + logger.info("⏸️ Stopped at profile → {}", out_dir) + return 0 + + # Stage 2: MinerU extraction + if not args.reuse_mineru: + _shard_dirs, mineru_elapsed = _stage_mineru_pdf( + file_path, filename, out_dir, anatomy, + ) + trace["stages"]["mineru"] = { + "elapsed_s": round(mineru_elapsed, 1), + "shard_count": len(_shard_dirs), + } + else: + logger.info("⏩ Reusing cached MinerU shard dirs") + + if args.stop_at == "mineru": + _write_json(out_dir / "trace.json", trace) + logger.info("⏸️ Stopped at mineru → {}", out_dir) + return 0 + + # Stage 3: Heading prediction + hierarchy tree + if args.reuse_hierarchy: + merged_path = out_dir / "_shards" / "merged_lines.json" + if not merged_path.exists(): + raise FileNotFoundError(f"No cached hierarchy: {merged_path}") + logger.info("⏩ Reusing cached merged_lines: {}", merged_path) + merged_lines = json.loads(merged_path.read_text(encoding="utf-8")) + else: + merged_lines, hier_elapsed = _stage_hierarchy_pdf( + out_dir, anatomy, args.model, + ) + trace["stages"]["hierarchy"] = { + "elapsed_s": round(hier_elapsed, 1), + "merged_lines_count": len(merged_lines), + "heading_count": sum(1 for ln in merged_lines if ln.startswith("#")), + } + + if args.stop_at == "hierarchy": + _write_json(out_dir / "trace.json", trace) + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 + + # Stage 4: Full extraction + chunks, full_elapsed = _stage_full_pdf(out_dir, filename, merged_lines, args.model) + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), + } + + elif ext in (".docx", ".doc"): + if args.stop_at in ("profile", "mineru"): + logger.info("ℹ️ No profiling/MinerU for DOCX format. Nothing to do.") + return 0 + + # Stage 2: parse_docx + parsed_df, hier_elapsed = _stage_hierarchy_docx(file_path, filename, out_dir, args.model) + trace["stages"]["hierarchy"] = { + "elapsed_s": round(hier_elapsed, 1), + "row_count": len(parsed_df) if parsed_df is not None else 0, + } + + if args.stop_at == "hierarchy": + _write_json(out_dir / "trace.json", trace) + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 + + # Stage 4: DataFrame → chunks + from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess + parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) + chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), + } + + elif ext in (".md", ".markdown"): + if args.stop_at in ("profile", "mineru"): + logger.info("ℹ️ No profiling/MinerU for Markdown format. Nothing to do.") + return 0 + + # Stage 2: parse_md + parsed_df, hier_elapsed = _stage_hierarchy_md(file_path, filename, out_dir, args.model) + trace["stages"]["hierarchy"] = { + "elapsed_s": round(hier_elapsed, 1), + "row_count": len(parsed_df) if parsed_df is not None else 0, + } + + if args.stop_at == "hierarchy": + _write_json(out_dir / "trace.json", trace) + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 + + # Stage 4: DataFrame → chunks + from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess + parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) + chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), + } + + else: + logger.error("Unsupported format: {}", ext) + return 1 + + # ── Optional DB publication ────────────────────────────────────────────── + if args.run_db: + from scripts._debug_publish import publish_debug_result_dir + publish_result = publish_debug_result_dir( + result_dir=out_dir, + source_file_name=filename, + chunks=chunks, + parse_track="text_track", + upload_assets=True, + ) + trace["stages"]["db_publish"] = { + "job_id": publish_result.job_id, + "document_id": publish_result.document_id, + } + + # ── Final trace ────────────────────────────────────────────────────────── + trace["total_elapsed_s"] = round(time.time() - t_start, 1) + trace["token_usage"] = get_current_token_tracker() + _write_json(out_dir / "trace.json", trace) + + logger.info("") + logger.info("═" * 70) + logger.info(" ✅ DONE in {:.1f}s → {}", time.time() - t_start, out_dir) + logger.info("═" * 70) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/_debug_pm_shared.py b/apps/worker/scripts/page_memory/_debug_pm_shared.py new file mode 100644 index 000000000..173a18ffa --- /dev/null +++ b/apps/worker/scripts/page_memory/_debug_pm_shared.py @@ -0,0 +1,1957 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402, F401 +"""Shared utilities for the staged page-memory debug scripts. + +All stage scripts (debug_pm_stage0..6) import from here instead of +duplicating bootstrap, artifact I/O, and argparse helpers. +""" + +import gevent.monkey + +gevent.monkey.patch_all() + +import argparse +import json +import os +import sys +import time +from copy import deepcopy +from dataclasses import asdict, dataclass, field, is_dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +# ── Bootstrap ─────────────────────────────────────────────────────────────── +ROOT = Path(__file__).resolve().parents[4] +WORKER_ROOT = ROOT / "apps" / "worker" +sys.path.insert(0, str(WORKER_ROOT)) +sys.path.insert(0, str(ROOT / "packages" / "shared-python")) + +from dotenv import load_dotenv + +load_dotenv(WORKER_ROOT / ".env") +os.environ.setdefault("LOCAL_DEBUG", "1") +os.environ.setdefault("OVERSIZED_PDF_SHARD_ENABLED", "true") + +from loguru import logger + +from app.services.page_memory._utils import ( + page_scope_info, + scope_id_for_pages, + sort_skeletons, +) +from app.services.page_memory._serialization import ( + build_hierarchy_tree as _build_hierarchy_from_skeletons, + derive_hierarchy_page_scope as _derive_hierarchy_page_scope, + scope_manifest as _scope_manifest, + serialize_assets as _serialize_assets, + serialize_hierarchy_artifact as _serialize_hierarchy_artifact, + serialize_page_tags as _serialize_page_tags, + serialize_scope_skeletons as _serialize_scope_skeletons, +) +from shared.services.ai.token_tracking import ( + init_token_tracker, + get_current_token_tracker, +) +from shared.services.ai.token_costing import build_token_cost_estimate + +# Re-exported for staged debug scripts. Listing them here marks the imports as +# intentional so CodeQL does not treat them as unused. +__all__ = [ + "page_scope_info", + "scope_id_for_pages", + "_build_hierarchy_from_skeletons", + "_derive_hierarchy_page_scope", + "_scope_manifest", + "_serialize_scope_skeletons", +] + +DEFAULT_PDF = Path( + "/Users/wuchengke/Desktop/temp/test_docs/" + "SJSYJ-SC-2024 企业制度汇编(上册).pdf" +) +OUTPUT_ROOT = Path("~/.knowhere/_debug_parse").expanduser() + + +# ── Token cost tracker ──────────────────────────────────────────────────────── + + +def _usage_delta(prev: dict[str, Any], cur: dict[str, Any]) -> dict[str, Any]: + _NUM = ("prompt_tokens", "completion_tokens", "total_tokens", "calls") + + def _sub(a: dict, b: dict) -> dict: + return {f: int(b.get(f, 0)) - int(a.get(f, 0)) for f in _NUM} + + def _bucket(pa: dict, pb: dict) -> dict: + r: dict[str, Any] = {} + for k in set(pa) | set(pb): + pk, ck = pa.get(k, {}), pb.get(k, {}) + if not isinstance(pk, dict) or not isinstance(ck, dict): + continue + e = {f: v for f, v in _sub(pk, ck).items() if v} + pm, cm = pk.get("models", {}), ck.get("models", {}) + if pm or cm: + md = _bucket(pm, cm) + if md: + e["models"] = md + if e: + r[k] = e + return r + + d = _sub(prev, cur) + for bk in ("by_model", "by_task"): + bd = _bucket(prev.get(bk, {}), cur.get(bk, {})) + if bd: + d[bk] = bd + return d + + +class TokenCostTracker: + """Incremental token usage & cost tracker for debug pipeline stages.""" + + def __init__(self) -> None: + self._dict = init_token_tracker() + self._root_gid = self._gid() + self._prev: dict[str, Any] = deepcopy(self._dict) + self._stages: list[dict[str, Any]] = [] + + @staticmethod + def _gid() -> int: + from shared.services.ai.token_tracking import _current_greenlet_id + + return _current_greenlet_id() + + def register_child_thread(self) -> None: + from shared.services.ai.token_tracking import _root_ids, _lock + + gid = self._gid() + if gid != self._root_gid: + with _lock: + _root_ids[gid] = self._root_gid + + def snapshot_stage(self, stage: str) -> None: + cur = deepcopy(get_current_token_tracker() or {}) + delta = _usage_delta(self._prev, cur) + self._stages.append({ + "stage": stage, + "prompt_tokens": delta.get("prompt_tokens", 0), + "completion_tokens": delta.get("completion_tokens", 0), + "total_tokens": delta.get("total_tokens", 0), + "calls": delta.get("calls", 0), + "cost": build_token_cost_estimate(delta), + }) + self._prev = cur + + def total_cost(self) -> dict[str, Any]: + return build_token_cost_estimate(get_current_token_tracker() or {}) + + def stage_summary(self) -> list[dict[str, Any]]: + return list(self._stages) + + +@dataclass +class TraceStageAdapter: + stages: list[dict[str, Any]] + + def record_stage( + self, + stage: str, + *, + page_info: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, + ) -> None: + record_stage( + self.stages, + stage, + page_info=page_info, + variables=variables, + ) + + +@dataclass +class ScopeResult: + """Result returned by run_scope_pipeline() for one coarse scope.""" + + scope_id: str + skeletons: list[Any] + tags: list[Any] + assets_by_page: dict[int, list[Any]] + rendered: list[Any] + final_pages: list[int] + scope_manifest: dict[str, Any] + trace_stages: list[dict[str, Any]] = field(default_factory=list) + + +# ── Argparse helpers ────────────────────────────────────────────────────────── + + +def base_argparser(description: str) -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=description) + parser.add_argument("--file", default=str(DEFAULT_PDF), help="PDF path") + parser.add_argument("--model", default=None, help="Override hierarchy/profiler model") + parser.add_argument("--vlm-model", default=None, help="VLM model override") + parser.add_argument("--no-vlm", action="store_true", help="Disable VLM calls") + parser.add_argument( + "--out-suffix", + default="", + help=( + "Append to the debug doc directory name so runs do not overwrite " + "the default page_memory tree " + "(e.g. --out-suffix boundary_clip → …/doc__boundary_clip/page_memory)." + ), + ) + return parser + + +def resolve_paths(args: argparse.Namespace) -> tuple[str, str, Path]: + """Returns (pdf_path, filename, out_dir).""" + from app.services.document_parser.orchestration.path_segment import build_parser_path_segment + + pdf_path = str(Path(args.file).expanduser().resolve()) + if not os.path.exists(pdf_path): + raise FileNotFoundError(pdf_path) + filename = os.path.basename(pdf_path) + dir_name = build_parser_path_segment(filename) + suffix = str(getattr(args, "out_suffix", "") or "").strip() + if suffix: + safe = "".join(ch if ch.isalnum() or ch in "-_" else "_" for ch in suffix) + dir_name = f"{dir_name}__{safe}" + out_dir = OUTPUT_ROOT / dir_name / "page_memory" + out_dir.mkdir(parents=True, exist_ok=True) + return pdf_path, filename, out_dir + + +# ── ToolContext builder ─────────────────────────────────────────────────────── + + +def build_ctx( + *, pdf_path: str, job_id: str, out_dir: Path, + page_count: int, page_texts: dict[int, str], vlm_model: str | None, + asset_extraction_enabled: bool = False, +): + from app.services.document_agent.manifest import ToolContext + from app.services.document_agent.state import AgentBlackboard + from app.services.document_agent.budget import BudgetTracker + + blackboard = AgentBlackboard() + blackboard.page_count = page_count + blackboard.page_full_text_cache = dict(page_texts) + + vmodel = vlm_model or os.environ.get("IMAGE_MODEL") + reason_model = os.environ.get("PAGE_LOCATE_REASON_MODEL") or os.environ.get("NORMOL_MODEL") + + budget = BudgetTracker(plan_budget=50000, visual_budget=200000) + return ToolContext( + pdf_path=pdf_path, + job_id=job_id, + blackboard=blackboard, + budget=budget, + trace=None, + output_dir=str(out_dir / "_doc_agent"), + settings={ + "vlm_model": vmodel, + "model": reason_model, + "agent_png_dpi": os.environ.get("AGENT_PNG_DPI", "144"), + }, + ) + + +# ── Anatomy / doc-profile cache ────────────────────────────────────────────── + + +def resolve_anatomy_cache_path(out_dir: Path) -> Path: + """Prefer package-root ``doc_profile.json``; fall back to legacy paths.""" + from app.services.document_agent.persist import DOC_PROFILE_FILENAME + + candidates = ( + out_dir / DOC_PROFILE_FILENAME, + out_dir / "_doc_agent" / DOC_PROFILE_FILENAME, + out_dir / "_doc_agent" / "anatomy_map.json", + ) + for path in candidates: + if path.is_file(): + return path + return candidates[0] + + +def load_anatomy_cache(cache_path: Path, pdf_path: str, job_id: str): + from app.services.document_agent.manifest import ( + H1BoundaryResult, + H1Candidate, + PageAnatomyMap, + PageFeature, + PageLabel, + Shard, + ShardPlan, + TocAnchorPage, + TocEvidence, + TocResult, + ValidationReport, + ) + + logger.info(f"⏩ Reusing cached anatomy: {cache_path}") + data = json.loads(cache_path.read_text(encoding="utf-8")) + toc = data.get("toc_result") or {} + h1 = data.get("h1_result") or {} + sp = data.get("shard_plan") or {} + + page_features = [] + for pf in data.get("page_features", []): + page_features.append(PageFeature( + page=int(pf.get("page", 0)), + raw_text_length=int(pf.get("raw_text_length", 0)), + text_density=float(pf.get("text_density", 0)), + image_coverage=float(pf.get("image_coverage", 0)), + image_count=int(pf.get("image_count", 0)), + table_count=int(pf.get("table_count", 0)), + drawings_count=int(pf.get("drawings_count", 0)), + orientation=pf.get("orientation", "portrait"), + width=float(pf.get("width", 0)), + height=float(pf.get("height", 0)), + has_asset=bool(pf.get("has_asset", False)), + is_blank_like=bool(pf.get("is_blank_like", False)), + )) + page_labels = [] + for pl in data.get("page_labels", []): + page_labels.append(PageLabel( + page=int(pl.get("page", 0)), + kind=pl.get("kind", "normal"), + confidence=float(pl.get("confidence", 0)), + evidence=pl.get("evidence", {}), + )) + + candidates = [] + for candidate in toc.get("candidates") or []: + if not isinstance(candidate, dict): + continue + candidates.append( + TocAnchorPage( + page=int(candidate.get("page", 0)), + png_path=str(candidate.get("png_path") or ""), + source=candidate.get("source", "text_scan"), + ) + ) + evidence = [] + for item in toc.get("evidence") or []: + if not isinstance(item, dict): + continue + evidence.append( + TocEvidence( + page_index=int(item.get("page_index", 0)), + source=str(item.get("source") or ""), + confidence=float(item.get("confidence", 0) or 0), + reason=str(item.get("reason") or ""), + ) + ) + + return PageAnatomyMap( + job_id=data.get("job_id", job_id), + file_path=data.get("file_path", pdf_path), + page_count=int(data.get("page_count", 0)), + page_features=page_features, + page_labels=page_labels, + toc_result=TocResult( + toc_pages=list(toc.get("toc_pages", [])), + candidates=candidates, + evidence=evidence, + method=toc.get("method", "none"), + notes=str(toc.get("notes") or ""), + failure_kind=toc.get("failure_kind", "none"), + ), + h1_result=H1BoundaryResult( + h1_candidates=[ + H1Candidate( + title=c.get("title", ""), + page=int(c.get("page", 0)), + confidence=float(c.get("confidence", 0) or 0), + matched_line=c.get("matched_line", ""), + source=c.get("source", "none"), + ) + for c in h1.get("h1_candidates", []) + ], + ), + shard_plan=ShardPlan( + enabled=bool(sp.get("enabled", False)), + reason=sp.get("reason", "not_needed"), + shards=[ + Shard( + shard_index=int(s.get("shard_index", i)), + page_start=int(s.get("page_start", 1)), + page_end=int(s.get("page_end", 1)), + page_offset=int(s.get("page_offset", 0)), + anchor_type=s.get("anchor_type", "forced_max_size"), + anchor_evidence=s.get("anchor_evidence", ""), + confidence=float(s.get("confidence", 0) or 0), + ) + for i, s in enumerate(sp.get("shards", [])) + ], + validation=ValidationReport(valid=True), + ), + toc_hierarchies=data.get("toc_hierarchies"), + ) + + +# ── Profile ─────────────────────────────────────────────────────────────────── + + +def run_profile( + pdf_path: str, + job_id: str, + out_dir: Path, + model: str | None, + *, + skip_toc_anchoring: bool = False, +): + """Run page-memory profile exactly like production ``memory_service.run``. + + Uses ``profile_document(..., skip_shard_plan=True, oversized_policy="page_memory")`` + so coarse → anatomy matches the live track (no ReAct shard planning). + + ``skip_toc_anchoring=True`` stops after TOC extract + link attach (legacy + monolithic helper). Prefer staged debug: Stage-0 bootstrap then Stage-1 TOC. + """ + from app.services.document_parser.profiling.doc_profiler import profile_document + from shared.core.config import settings + + logger.info("=" * 70) + logger.info(f"🧬 DOC_PROFILE (page_memory, monolithic) — {job_id}") + logger.info("=" * 70) + if skip_toc_anchoring: + logger.info(" skip_toc_anchoring=True (TOC + links only; no calibration)") + + previous_image_model = settings.IMAGE_MODEL + if model: + settings.IMAGE_MODEL = model + logger.info(f" IMAGE_MODEL override → {model}") + + t0 = time.time() + try: + profile = profile_document( + pdf_path, + job_id, + job_id=job_id, + output_dir=str(out_dir), + skip_shard_plan=True, + oversized_policy="page_memory", + skip_toc_anchoring=skip_toc_anchoring, + ) + finally: + if model: + settings.IMAGE_MODEL = previous_image_model + + logger.info(f" profile done in {time.time() - t0:.1f}s") + logger.info( + " category={} routing={} page_count={} is_atlas={}", + profile.category, + getattr(profile.routing_category, "value", profile.routing_category), + profile.page_count, + profile.is_atlas, + ) + + anatomy = profile.anatomy + if anatomy is None: + raise RuntimeError( + "page_memory profile returned no anatomy " + f"(routing={profile.routing_category}). Atlas / no-anatomy path " + "cannot continue Stage 1." + ) + + from app.services.document_agent.persist import DOC_PROFILE_FILENAME + + profile_path = out_dir / DOC_PROFILE_FILENAME + if not profile_path.exists(): + write_debug_json(profile_path, anatomy.to_dict()) + + asset_pages = sum(1 for f in anatomy.page_features if getattr(f, "has_asset", False)) + logger.info(f" page_count={anatomy.page_count}") + logger.info(f" toc_pages={anatomy.toc_result.toc_pages}") + logger.info(f" has_asset_pages={asset_pages}/{anatomy.page_count}") + if anatomy.h1_result: + logger.info(f" h1_candidates={len(anatomy.h1_result.h1_candidates)}") + logger.info( + " shard_plan.enabled={} shards={}", + anatomy.shard_plan.enabled, + len(anatomy.shard_plan.shards), + ) + logger.info(f" doc_profile → {profile_path}") + return anatomy + + +def _build_debug_coordinator( + *, + pdf_path: str, + job_id: str, + out_dir: Path, + model: str | None, + settings_extra: dict[str, Any] | None = None, +): + """Build a ProfileCoordinator with the same models as production page_memory.""" + from app.services.document_agent.coordinator import ProfileCoordinator + from shared.core.config import settings + + if model: + settings.IMAGE_MODEL = model + agent_output_dir = out_dir / "_doc_agent" + agent_output_dir.mkdir(parents=True, exist_ok=True) + merged = { + "planner_model": settings.IMAGE_MODEL, + "vlm_model": settings.IMAGE_MODEL, + "toc_profile_enabled": True, + "model": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, + } + if settings_extra: + merged.update(settings_extra) + return ProfileCoordinator( + pdf_path=pdf_path, + job_id=job_id, + output_dir=str(agent_output_dir), + db=None, + model=settings.IMAGE_MODEL, + settings=merged, + ) + + +def _document_profile_from_dict(data: dict[str, Any] | None): + from app.services.document_agent.manifest import DocumentProfile + + raw = data or {} + return DocumentProfile( + is_scanned=bool(raw.get("is_scanned", False)), + category=str(raw.get("category") or "unknown"), + routing_category=str(raw.get("routing_category") or "generic"), + language=str(raw.get("language") or "unknown"), + rationale=str(raw.get("rationale") or ""), + header_y=raw.get("header_y"), + footer_y=raw.get("footer_y"), + ) + + +def _page_features_from_dicts(rows: list[Any]) -> list[Any]: + from app.services.document_agent.manifest import PageFeature + + out = [] + for pf in rows: + if not isinstance(pf, dict): + continue + out.append( + PageFeature( + page=int(pf.get("page", 0)), + raw_text_length=int(pf.get("raw_text_length", 0)), + text_density=float(pf.get("text_density", 0)), + image_coverage=float(pf.get("image_coverage", 0)), + image_count=int(pf.get("image_count", 0)), + table_count=int(pf.get("table_count", 0)), + drawings_count=int(pf.get("drawings_count", 0)), + orientation=pf.get("orientation", "portrait"), + width=float(pf.get("width", 0)), + height=float(pf.get("height", 0)), + has_asset=bool(pf.get("has_asset", False)), + is_blank_like=bool(pf.get("is_blank_like", False)), + invisible_text_length=int(pf.get("invisible_text_length", 0) or 0), + ) + ) + return out + + +def _page_labels_from_dicts(rows: list[Any]) -> list[Any]: + from app.services.document_agent.manifest import PageLabel + + out = [] + for pl in rows: + if not isinstance(pl, dict): + continue + out.append( + PageLabel( + page=int(pl.get("page", 0)), + kind=pl.get("kind", "normal"), + confidence=float(pl.get("confidence", 0)), + evidence=dict(pl.get("evidence") or {}), + ) + ) + return out + + +def persist_stage0_state(out_dir: Path, coordinator) -> Path: + """Persist Stage-0 blackboard so Stage-1 can resume at Find.""" + doc_agent_dir = out_dir / "_doc_agent" + doc_agent_dir.mkdir(parents=True, exist_ok=True) + + texts = { + str(page): text + for page, text in dict(coordinator.blackboard.page_full_text_cache or {}).items() + } + write_debug_json(page_text_cache_path(out_dir), texts) + + profile = coordinator.blackboard.document_profile + state = { + "version": "1.0", + "page_count": int(coordinator.blackboard.page_count or 0), + "document_profile": profile.to_dict() if profile is not None else None, + "page_features": [ + feature.to_dict() for feature in (coordinator.blackboard.page_features or []) + ], + "page_labels": [ + label.to_dict() for label in (coordinator.blackboard.page_labels or []) + ], + "doc_stats": dict(coordinator.blackboard.doc_stats or {}), + "global_signals": dict(coordinator.blackboard.global_signals or {}), + "page_full_text_cache_path": PAGE_TEXT_CACHE_NAME, + } + path = stage0_state_path(out_dir) + write_debug_json(path, state) + return path + + +def load_stage0_into_coordinator(coordinator, out_dir: Path) -> None: + """Restore Stage-0 outputs onto a coordinator blackboard.""" + state_path = stage0_state_path(out_dir) + require_file( + state_path, + hint=( + "Run Stage 0 first: uv run python " + "scripts/page_memory/debug_pm_stage0_bootstrap.py --file ..." + ), + ) + state = json.loads(state_path.read_text(encoding="utf-8")) + text_path = page_text_cache_path(out_dir) + require_file( + text_path, + hint="Stage-0 page_full_text_cache.json missing; re-run Stage 0", + ) + raw_texts = json.loads(text_path.read_text(encoding="utf-8")) + page_texts = {int(page): str(text) for page, text in dict(raw_texts or {}).items()} + + bb = coordinator.blackboard + bb.page_count = int(state.get("page_count") or 0) + bb.document_profile = _document_profile_from_dict(state.get("document_profile")) + bb.page_features = _page_features_from_dicts(list(state.get("page_features") or [])) + bb.page_labels = _page_labels_from_dicts(list(state.get("page_labels") or [])) + bb.doc_stats = dict(state.get("doc_stats") or {}) + bb.global_signals = dict(state.get("global_signals") or {}) + bb.page_full_text_cache = page_texts + # Stage-1 owns TOC + assets from here. + bb.toc_result = None + bb.toc_hierarchies = None + bb.toc_page_offset = None + bb.skeleton_anchor = None + bb.skeleton_nodes = None + bb.pending_skeleton_anchors = [] + bb.global_signals.pop("toc_profile_attempted", None) + # Keep assets_probed / has_asset from Stage-0; Stage-1 is TOC-only. + logger.info( + "⏩ Resumed Stage-0: pages={} text_pages={} labels={} assets_probed={}", + bb.page_count, + len(page_texts), + len(bb.page_labels), + bool(bb.global_signals.get("assets_probed")), + ) + + +def run_stage0_bootstrap( + pdf_path: str, + job_id: str, + out_dir: Path, + model: str | None, +): + """Production-aligned Stage-0: bootstrap → coarse VLM → text scan → asset probe.""" + from shared.core.config import settings + + logger.info("=" * 70) + logger.info( + f"🧬 Stage 0: BOOTSTRAP + COARSE VLM + TEXT SCAN + ASSET PROBE — {job_id}" + ) + logger.info("=" * 70) + + previous_image_model = settings.IMAGE_MODEL + t0 = time.time() + try: + coordinator = _build_debug_coordinator( + pdf_path=pdf_path, + job_id=job_id, + out_dir=out_dir, + model=model, + settings_extra={"stop_after_asset_probe": True}, + ) + profile = coordinator.run_coarse() + state_path = persist_stage0_state(out_dir, coordinator) + asset_pages = sum( + 1 + for feature in (coordinator.blackboard.page_features or []) + if getattr(feature, "has_asset", False) + ) + update_pipeline_state( + pipeline_state_path(out_dir), + stage=0, + document={ + "source_file_name": job_id, + "page_count": coordinator.blackboard.page_count, + "stage0_state": str(state_path), + }, + payload={ + "page_count": coordinator.blackboard.page_count, + "is_scanned": bool(getattr(profile, "is_scanned", False)), + "routing_category": getattr(profile, "routing_category", None), + "text_pages": len(coordinator.blackboard.page_full_text_cache or {}), + "assets_probed": bool( + coordinator.blackboard.global_signals.get("assets_probed") + ), + "has_asset_pages": asset_pages, + }, + ) + finally: + if model: + settings.IMAGE_MODEL = previous_image_model + + logger.info(f" stage0 done in {time.time() - t0:.1f}s → {state_path}") + logger.info( + " category={} routing={} page_count={} is_scanned={} has_asset_pages={}", + getattr(profile, "category", None), + getattr(profile, "routing_category", None), + coordinator.blackboard.page_count, + getattr(profile, "is_scanned", None), + asset_pages, + ) + return coordinator, profile, state_path + + +def run_stage1_toc( + pdf_path: str, + job_id: str, + out_dir: Path, + model: str | None, +): + """Production-aligned Stage-1: Find → extract → link attach (no calibration). + + Resumes Stage-0 blackboard (including asset probe). Skips + ``run_toc_anchoring`` (Stage-2). Does not re-run asset probe. + """ + from app.services.document_agent.persist import ( + DOC_PROFILE_FILENAME, + build_anatomy_map, + persist_anatomy_map, + ) + from app.services.document_agent.validators import single_shard_plan + from shared.core.config import settings + + logger.info("=" * 70) + logger.info(f"🧬 Stage 1: TOC FIND → EXTRACT → LINK — {job_id}") + logger.info("=" * 70) + + previous_image_model = settings.IMAGE_MODEL + t0 = time.time() + try: + coordinator = _build_debug_coordinator( + pdf_path=pdf_path, + job_id=job_id, + out_dir=out_dir, + model=model, + settings_extra={"skip_toc_anchoring": True}, + ) + load_stage0_into_coordinator(coordinator, out_dir) + coordinator._ensure_toc_profile(strict=False) + coordinator.blackboard.shard_plan = single_shard_plan( + coordinator.blackboard.page_count + ) + anatomy = build_anatomy_map(coordinator.ctx) + persist_anatomy_map(coordinator.ctx, {}) + profile_path = out_dir / DOC_PROFILE_FILENAME + write_debug_json(profile_path, anatomy.to_dict()) + update_pipeline_state( + pipeline_state_path(out_dir), + stage=1, + document={ + "source_file_name": job_id, + "page_count": anatomy.page_count, + "anatomy_path": str(profile_path), + }, + payload={ + "toc_pages": list(getattr(anatomy.toc_result, "toc_pages", []) or []), + "region_count": len(list(anatomy.toc_hierarchies or [])), + "skip_toc_anchoring": True, + }, + ) + finally: + if model: + settings.IMAGE_MODEL = previous_image_model + + logger.info(f" stage1 done in {time.time() - t0:.1f}s") + logger.info(f" toc_pages={anatomy.toc_result.toc_pages}") + logger.info(f" doc_profile → {profile_path}") + return anatomy + + +# ── JSON / artifact I/O ─────────────────────────────────────────────────────── + + +def jsonable(value: Any) -> Any: + if hasattr(value, "to_dict") and callable(value.to_dict): + return value.to_dict() + if is_dataclass(value) and not isinstance(value, type): + return jsonable(asdict(value)) + if isinstance(value, dict): + return {str(k): jsonable(v) for k, v in value.items()} + if isinstance(value, list | tuple): + return [jsonable(v) for v in value] + if hasattr(value, "value"): + return getattr(value, "value") + return value + + +def write_debug_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(jsonable(value), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + logger.info(f" debug json → {path}") + + +def toc_hierarchies_to_hierarchy_tree( + toc_hierarchies: list[dict[str, Any]] | None, +) -> dict[str, Any]: + """Build Stage-1 human TOC dump in final ``HIERARCHY`` shape. + + Uses original VLM headings (keeps numbering prefixes). Does **not** run + ``extract_toc_nodes`` / ``clean_toc_title`` — those are locate-time only. + Regions are concatenated in profile order so multi-TOC docs still form one + nested tree readable like ``hierarchy.json``. + """ + from app.services.document_agent.tools.vlm_toc_extractor import build_toc_tree + + flat_entries: list[dict[str, Any]] = [] + for region in toc_hierarchies or []: + if not isinstance(region, dict): + continue + rows = region.get("toc_with_level") or [] + if rows: + for entry in rows: + if not isinstance(entry, dict): + continue + heading = str(entry.get("heading") or entry.get("title") or "").strip() + if not heading: + continue + flat_entries.append( + { + "title": heading, + "level": entry.get("level", 1), + "page_number": entry.get("page_number"), + } + ) + continue + # Fallback: region only stored ``toc_tree`` (already nested, original keys). + tree = region.get("toc_tree") + if isinstance(tree, dict) and tree: + flat_entries.extend(_flatten_hierarchy_tree_entries(tree)) + return build_toc_tree(flat_entries) + + +def _flatten_hierarchy_tree_entries( + tree: dict[str, Any], + *, + level: int = 1, +) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] + for title, children in tree.items(): + heading = str(title or "").strip() + if not heading: + continue + entries.append({"title": heading, "level": level}) + if isinstance(children, dict) and children: + entries.extend( + _flatten_hierarchy_tree_entries(children, level=level + 1) + ) + return entries + + +def write_toc_hierarchy_artifact( + out_dir: Path, + *, + hierarchy_tree: dict[str, Any], + stats: dict[str, Any] | None = None, +) -> Path: + """Write Stage-1 TOC as human-readable ``HIERARCHY`` at package root. + + Debug-only for page_memory (not packaged into ZIP). Shape matches final + ``hierarchy.json`` / ``manifest.HIERARCHY``; titles keep TOC prefixes. + """ + path = out_dir / "toc_hierarchy.json" + # Legacy list dump from an earlier debug format. + for legacy_name in ("toc_hierarchies.json",): + (out_dir / legacy_name).unlink(missing_ok=True) + write_debug_json( + path, + { + "HIERARCHY": hierarchy_tree or {}, + "stats": dict(stats or {}), + }, + ) + return path + + +PIPELINE_STATE_VERSION = "1.0" +PIPELINE_STATE_NAME = "pipeline_state.json" +_PIPELINE_STAGES = tuple(f"stage{number}" for number in range(0, 7)) + + +def pipeline_state_path(out_dir: Path) -> Path: + return out_dir / "_doc_agent" / PIPELINE_STATE_NAME + + +STAGE0_STATE_NAME = "stage0_state.json" +PAGE_TEXT_CACHE_NAME = "page_full_text_cache.json" + + +def stage0_state_path(out_dir: Path) -> Path: + return out_dir / "_doc_agent" / STAGE0_STATE_NAME + + +def page_text_cache_path(out_dir: Path) -> Path: + return out_dir / "_doc_agent" / PAGE_TEXT_CACHE_NAME + + +def load_pipeline_state( + state_path: Path, + *, + legacy_locate_cache: Path | None = None, +) -> dict[str, Any]: + """Load the shared Stage 0-6 ledger, with locate-cache compatibility.""" + if state_path.exists(): + data = json.loads(state_path.read_text(encoding="utf-8")) + if not isinstance(data, dict): + raise ValueError(f"pipeline state must be an object: {state_path}") + data.setdefault("version", PIPELINE_STATE_VERSION) + data.setdefault("stages", {}) + return data + + if legacy_locate_cache is not None and legacy_locate_cache.exists(): + rows = json.loads(legacy_locate_cache.read_text(encoding="utf-8")) + if not isinstance(rows, list): + raise ValueError( + f"legacy locate cache must be a list: {legacy_locate_cache}" + ) + logger.warning( + "Legacy locate cache detected; it will be migrated on the next stage write: {}", + legacy_locate_cache, + ) + return { + "version": PIPELINE_STATE_VERSION, + "stages": { + "stage2": { + "status": "legacy", + "skeletons": rows, + } + }, + } + + raise FileNotFoundError(state_path) + + +def _pipeline_skeleton_rows(state: dict[str, Any]) -> list[dict[str, Any]]: + stages = state.get("stages") + stage2 = stages.get("stage2") if isinstance(stages, dict) else None + rows = stage2.get("skeletons") if isinstance(stage2, dict) else None + if rows is None: + # Compatibility with the short-lived ``stage2_state.json`` proposal. + rows = state.get("skeletons") + if not isinstance(rows, list): + raise ValueError("pipeline state is missing stages.stage2.skeletons[]") + return [row for row in rows if isinstance(row, dict)] + + +def load_pipeline_skeletons( + state_path: Path, + *, + legacy_locate_cache: Path | None = None, +) -> list[Any]: + from app.services.page_memory.skeleton_extractor import SectionSkeleton + + state = load_pipeline_state( + state_path, + legacy_locate_cache=legacy_locate_cache, + ) + return [ + SectionSkeleton( + section_path=str(row["section_path"]), + title=str(row["title"]), + level=int(row["level"]), + start_page=int(row["start_page"]), + end_page=int(row["end_page"]), + parent_path=row.get("parent_path"), + evidence=dict(row.get("evidence") or {}), + ) + for row in _pipeline_skeleton_rows(state) + ] + + +def update_pipeline_state( + state_path: Path, + *, + stage: int, + payload: dict[str, Any], + document: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Atomically update one stage and invalidate stale later-stage entries.""" + stage_key = f"stage{stage}" + if stage_key not in _PIPELINE_STAGES: + raise ValueError(f"unsupported pipeline stage: {stage}") + + if state_path.exists(): + state = load_pipeline_state(state_path) + else: + state = { + "version": PIPELINE_STATE_VERSION, + "stages": {}, + } + stages = state.setdefault("stages", {}) + if not isinstance(stages, dict): + raise ValueError(f"pipeline state stages must be an object: {state_path}") + + stage_index = _PIPELINE_STAGES.index(stage_key) + for stale_key in _PIPELINE_STAGES[stage_index + 1 :]: + stages.pop(stale_key, None) + + updated_at = datetime.now(timezone.utc).isoformat() + stages[stage_key] = { + "status": "complete", + "updated_at": updated_at, + **jsonable(payload), + } + if document: + state["document"] = { + **dict(state.get("document") or {}), + **jsonable(document), + } + state["updated_at"] = updated_at + + state_path.parent.mkdir(parents=True, exist_ok=True) + temp_path = state_path.with_suffix(f"{state_path.suffix}.tmp") + temp_path.write_text( + json.dumps(state, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + temp_path.replace(state_path) + logger.info(" pipeline state → {}", state_path) + return state + + +def remove_legacy_doc_agent_artifacts( + doc_agent_dir: Path, + *, + include_stage2: bool = False, +) -> None: + names = { + "parser_profile.json", + "toc_hierarchies.json", + } + if include_stage2: + names.update( + { + "calibration_result.json", + "null_page_parent_locate.json", + "locate_cache.json", + "stage2_state.json", + } + ) + for name in names: + (doc_agent_dir / name).unlink(missing_ok=True) + legacy_preview_dir = doc_agent_dir / "coarse_assets" + if legacy_preview_dir.is_dir(): + import shutil + + shutil.rmtree(legacy_preview_dir) + (doc_agent_dir / "coarse_assets.html").unlink(missing_ok=True) + + +def record_stage( + stages: list[dict[str, Any]], + stage: str, + *, + page_info: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, +) -> None: + stages.append( + { + "stage": stage, + "page_info": page_info or {}, + "variables": variables or {}, + "created_at": datetime.now(timezone.utc).isoformat(), + } + ) + + +STAGE_COSTS_VERSION = "1.0" +STAGE_COSTS_NAME = "stage_costs.json" +_COST_STAGE_KEYS = tuple(f"stage{number}" for number in range(0, 7)) + + +def stage_costs_path(out_dir: Path) -> Path: + return out_dir / "_doc_agent" / STAGE_COSTS_NAME + + +def load_stage_costs(out_dir: Path) -> dict[str, Any]: + path = stage_costs_path(out_dir) + if not path.exists(): + return {"version": STAGE_COSTS_VERSION, "stages": {}} + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, dict): + raise ValueError(f"stage costs must be an object: {path}") + data.setdefault("version", STAGE_COSTS_VERSION) + data.setdefault("stages", {}) + return data + + +def _stage_usage_snapshot(tracker: TokenCostTracker | None) -> dict[str, Any]: + usage = get_current_token_tracker() or {} + return { + "token_usage": deepcopy(usage), + "prompt_tokens": int(usage.get("prompt_tokens") or 0), + "completion_tokens": int(usage.get("completion_tokens") or 0), + "total_tokens": int(usage.get("total_tokens") or 0), + "calls": int(usage.get("calls") or 0), + "cost": ( + tracker.total_cost() + if tracker is not None + else build_token_cost_estimate(usage) + ), + "by_substage": tracker.stage_summary() if tracker is not None else [], + } + + +def _merge_token_usage( + destination: dict[str, Any], + source: dict[str, Any], +) -> None: + """Merge raw production token-tracker snapshots recursively.""" + for key, value in source.items(): + if isinstance(value, dict): + child = destination.setdefault(str(key), {}) + if isinstance(child, dict): + _merge_token_usage(child, value) + elif isinstance(value, int | float) and not isinstance(value, bool): + destination[str(key)] = destination.get(str(key), 0) + value + + +def record_stage_cost( + out_dir: Path, + *, + pipeline_stage: int, + elapsed_s: float, + token_cost_tracker: TokenCostTracker | None = None, + trace_stages: list[dict[str, Any]] | None = None, + stop_at: str | None = None, + extra: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Persist one pipeline stage's elapsed/cost ledger entry independently.""" + stage_key = f"stage{pipeline_stage}" + if stage_key not in _COST_STAGE_KEYS: + raise ValueError(f"unsupported cost pipeline stage: {pipeline_stage}") + + ledger = load_stage_costs(out_dir) + stages = ledger.setdefault("stages", {}) + if not isinstance(stages, dict): + raise ValueError(f"stage costs stages must be an object: {stage_costs_path(out_dir)}") + + stage_index = _COST_STAGE_KEYS.index(stage_key) + for stale_key in _COST_STAGE_KEYS[stage_index + 1 :]: + stages.pop(stale_key, None) + + usage = _stage_usage_snapshot(token_cost_tracker) + updated_at = datetime.now(timezone.utc).isoformat() + entry: dict[str, Any] = { + "status": "complete", + "elapsed_s": round(float(elapsed_s), 3), + "stop_at": stop_at, + "updated_at": updated_at, + **usage, + } + if trace_stages is not None: + entry["trace_stages"] = list(trace_stages) + if extra: + entry.update(jsonable(extra)) + stages[stage_key] = entry + ledger["updated_at"] = updated_at + + path = stage_costs_path(out_dir) + path.parent.mkdir(parents=True, exist_ok=True) + temp_path = path.with_suffix(f"{path.suffix}.tmp") + temp_path.write_text( + json.dumps(ledger, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + temp_path.replace(path) + logger.info( + " stage cost → {} ({} {:.1f}s / ${:.6f})", + path, + stage_key, + entry["elapsed_s"], + float((entry.get("cost") or {}).get("total_cost") or 0), + ) + return ledger + + +def aggregate_stage_costs(ledger: dict[str, Any]) -> dict[str, Any]: + """Roll up independently stored stage cost entries for TRACE.JSON.""" + stages = ledger.get("stages") if isinstance(ledger, dict) else {} + if not isinstance(stages, dict): + stages = {} + + by_pipeline_stage: dict[str, Any] = {} + by_substage: list[dict[str, Any]] = [] + merged_trace_stages: list[dict[str, Any]] = [] + usage = { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "calls": 0, + "by_model": {}, + "by_task": {}, + } + elapsed_s = 0.0 + + for stage_key in _COST_STAGE_KEYS: + row = stages.get(stage_key) + if not isinstance(row, dict): + continue + stage_elapsed = float(row.get("elapsed_s") or 0) + elapsed_s += stage_elapsed + for usage_key in ( + "prompt_tokens", + "completion_tokens", + "total_tokens", + "calls", + ): + usage[usage_key] += int(row.get(usage_key) or 0) + raw_usage = row.get("token_usage") + if isinstance(raw_usage, dict): + # Numeric top-level fields were already added above. + _merge_token_usage( + usage, + { + key: value + for key, value in raw_usage.items() + if key not in { + "prompt_tokens", + "completion_tokens", + "total_tokens", + "calls", + } + }, + ) + stage_cost = row.get("cost") if isinstance(row.get("cost"), dict) else {} + by_pipeline_stage[stage_key] = { + "elapsed_s": stage_elapsed, + "stop_at": row.get("stop_at"), + "prompt_tokens": int(row.get("prompt_tokens") or 0), + "completion_tokens": int(row.get("completion_tokens") or 0), + "total_tokens": int(row.get("total_tokens") or 0), + "calls": int(row.get("calls") or 0), + "cost": stage_cost, + "updated_at": row.get("updated_at"), + } + for item in row.get("by_substage") or []: + if isinstance(item, dict): + by_substage.append( + { + "pipeline_stage": stage_key, + **item, + } + ) + for item in row.get("trace_stages") or []: + if isinstance(item, dict): + merged_trace_stages.append( + { + **item, + "pipeline_stage": item.get("pipeline_stage") or stage_key, + } + ) + + total_estimate = build_token_cost_estimate(usage) + return { + "elapsed_s": round(elapsed_s, 3), + "completed_pipeline_stages": list(by_pipeline_stage), + "token_usage": usage, + "token_cost": { + "total": { + "prompt_tokens": usage["prompt_tokens"], + "completion_tokens": usage["completion_tokens"], + "total_tokens": usage["total_tokens"], + "calls": usage["calls"], + **total_estimate, + }, + "by_pipeline_stage": by_pipeline_stage, + "by_stage": by_substage, + }, + "stages": merged_trace_stages, + } + + +def build_production_job_metadata_from_stage_costs( + *, + page_count: int, + ledger: dict[str, Any], +) -> dict[str, Any]: + """Project split debug runs onto the production ZIP manifest contract.""" + aggregated = aggregate_stage_costs(ledger) + stages = ledger.get("stages") if isinstance(ledger, dict) else {} + if not isinstance(stages, dict): + stages = {} + + timing_ms: dict[str, int] = {} + completed_at: datetime | None = None + for stage_key in _COST_STAGE_KEYS: + row = stages.get(stage_key) + if not isinstance(row, dict): + continue + timing_ms[stage_key] = int( + round(float(row.get("elapsed_s") or 0) * 1000) + ) + raw_updated_at = row.get("updated_at") + if isinstance(raw_updated_at, str): + try: + timestamp = datetime.fromisoformat(raw_updated_at) + except ValueError: + continue + if completed_at is None or timestamp > completed_at: + completed_at = timestamp + + duration_ms = int(round(float(aggregated.get("elapsed_s") or 0) * 1000)) + completed_at = completed_at or datetime.now(timezone.utc) + started_at = completed_at - timedelta(milliseconds=duration_ms) + return { + "page_count": page_count, + "parse_track": "page_memory", + "billing_status": None, + "billing_amount_micro_dollars": None, + "billing_credits": None, + "processing_started_at": started_at.isoformat(), + "processing_completed_at": completed_at.isoformat(), + "processing_duration_ms": duration_ms, + "stages": { + "timing_ms": timing_ms, + "token_usage": aggregated.get("token_usage") or {}, + }, + } + + +def write_trace( + *, + out_dir: Path, + stages: list[dict[str, Any]], + final_status: str, + summary: dict[str, Any], +) -> None: + write_debug_json( + out_dir / "trace.json", + { + "final_status": final_status, + "summary": summary, + "stages": stages, + }, + ) + + +def stop_with_trace( + *, + out_dir: Path, + stages: list[dict[str, Any]], + stop_at: str, + page_count: int | None = None, + scope_id: str | None = None, + pipeline_stage: int | None = None, + elapsed_s: float | None = None, + token_cost_tracker: TokenCostTracker | None = None, + final_status: str | None = None, + extra_summary: dict[str, Any] | None = None, +) -> int: + """Write TRACE.JSON, optionally recording this pipeline stage's cost ledger.""" + if pipeline_stage is not None and elapsed_s is not None: + record_stage_cost( + out_dir, + pipeline_stage=pipeline_stage, + elapsed_s=elapsed_s, + token_cost_tracker=token_cost_tracker, + trace_stages=stages, + stop_at=stop_at, + ) + + aggregated = aggregate_stage_costs(load_stage_costs(out_dir)) + merged_stages = aggregated.get("stages") or stages + summary: dict[str, Any] = { + "page_count": page_count, + "scope_id": scope_id, + "rows_count": None, + "elapsed_s": aggregated.get("elapsed_s"), + "completed_pipeline_stages": aggregated.get("completed_pipeline_stages"), + "token_cost": aggregated.get("token_cost"), + } + if extra_summary: + summary.update(jsonable(extra_summary)) + + write_trace( + out_dir=out_dir, + stages=merged_stages, + final_status=final_status or f"stopped_at_{stop_at}", + summary=summary, + ) + remove_nested_doc_agent_trace(out_dir) + maybe_purge_debug_visuals(out_dir) + return 0 + + +def remove_nested_doc_agent_trace(out_dir: Path) -> None: + try: + (out_dir / "_doc_agent" / "trace.json").unlink(missing_ok=True) + except Exception as exc: + logger.debug(f"failed to remove nested doc-agent trace: {exc}") + + +def maybe_purge_debug_visuals(out_dir: Path) -> None: + from app.services.document_agent.visual import ( + purge_debug_visual_dirs, + visual_debug_enabled, + ) + + if not visual_debug_enabled(): + purge_debug_visual_dirs(str(out_dir)) + + +def write_scope_artifacts( + *, + out_dir: Path, + scope_id: str, + scope_manifest: dict[str, Any], + hierarchy: list[Any], + tags: list[Any] | None = None, + assets_by_page: dict[int, list[Any]] | None = None, +) -> None: + """Write per-scope viewing artifacts (no standalone scope.json). + + ``tags`` / ``assets_by_page`` of ``None`` leave existing files untouched. + """ + scope_dir = out_dir / "scopes" / scope_id + write_debug_json( + scope_dir / "fine_hierarchy.json", + _serialize_hierarchy_artifact(hierarchy, scope_manifest_data=scope_manifest), + ) + if tags is not None: + write_debug_json(scope_dir / "page_tags.json", _serialize_page_tags(tags)) + if assets_by_page is not None: + write_debug_json(scope_dir / "assets.json", _serialize_assets(assets_by_page)) + + +def write_top_level_artifacts( + *, + out_dir: Path, + hierarchy: list[Any], + tags: list[Any], + assets_by_page: dict[int, list[Any]] | None = None, +) -> None: + write_debug_json(out_dir / "hierarchy.json", _serialize_hierarchy_artifact(hierarchy)) + write_debug_json(out_dir / "page_tags.json", _serialize_page_tags(tags)) + if assets_by_page is not None: + write_debug_json(out_dir / "assets.json", _serialize_assets(assets_by_page)) + else: + (out_dir / "assets.json").unlink(missing_ok=True) + + +def cleanup_page_memory_artifacts(out_dir: Path) -> None: + stale_files = { + "assets.json", + "chunks.json", + "coarse_scopes.json", + "doc_nav.json", + "hierarchy.json", + "manifest.json", + "node_rows.csv", + "node_rows.json", + "page_plans.json", + "page_rendered.json", + "page_tags.json", + "report.md", + "trace.json", + } + for name in stale_files: + path = out_dir / name + try: + if path.is_file(): + path.unlink() + except Exception: + logger.debug(f"cleanup failed for {path}") + for name in ("asset_annotate", "debug", "images", "pages", "scopes", "tables"): + path = out_dir / name + try: + if path.is_dir(): + import shutil + + shutil.rmtree(path) + except Exception: + logger.debug(f"cleanup failed for {path}") + + +# ── Tree helpers ────────────────────────────────────────────────────────────── + + +def walk(nodes: list, depth: int = 0) -> list[tuple[int, Any]]: + rows: list[tuple[int, Any]] = [] + for node in nodes: + rows.append((depth, node)) + rows.extend(walk(node.children, depth + 1)) + return rows + + +def walk_node_count(nodes: list) -> int: + return len(walk(nodes)) + + +def hierarchy_metrics(nodes: list, *, source: str) -> dict[str, Any]: + rows = walk(nodes) + depths = [depth + 1 for depth, _node in rows] + return { + "hierarchy_source": source, + "title_node_count": len(rows), + "title_leaf_count": sum(1 for _depth, node in rows if not node.children), + "title_max_depth": max(depths) if depths else 0, + } + + +# ── Artifact loaders ────────────────────────────────────────────────────────── + + +def load_hierarchy_artifact(path: Path) -> tuple[dict[str, Any], list[Any]]: + """Load a fine hierarchy artifact and its canonical scope manifest.""" + from app.services.page_memory.skeleton_extractor import SectionSkeleton + + data = json.loads(path.read_text(encoding="utf-8")) + nodes = data.get("nodes") if isinstance(data, dict) else None + if not isinstance(nodes, list): + raise ValueError(f"fine hierarchy artifact missing nodes: {path}") + skeletons: list[SectionSkeleton] = [] + for node in nodes: + if not isinstance(node, dict): + continue + section_path = str(node.get("section_path") or "").strip() + title = str(node.get("title") or "").strip() + if not section_path or not title: + continue + skeletons.append( + SectionSkeleton( + section_path=section_path, + title=title, + level=int(node.get("level") or 1), + start_page=int(node.get("start_page") or 1), + end_page=int(node.get("end_page") or node.get("start_page") or 1), + parent_path=node.get("parent_path"), + evidence=dict(node.get("evidence") or {}), + ) + ) + scope = data.get("scope") if isinstance(data, dict) else None + return dict(scope) if isinstance(scope, dict) else {}, sort_skeletons(skeletons) + + +def load_skeletons_from_hierarchy_artifact(path: Path) -> list[Any]: + """Compatibility reader for callers that only need hierarchy nodes.""" + _scope, skeletons = load_hierarchy_artifact(path) + return skeletons + + +def load_page_tags_artifact(path: Path) -> list[Any]: + """Load ``page_tags.json`` written by ``serialize_page_tags``.""" + from app.services.page_memory.page_tagger import PageTagResult + + payload = json.loads(path.read_text(encoding="utf-8")) + if isinstance(payload, dict) and isinstance(payload.get("tags"), list): + rows = payload["tags"] + elif isinstance(payload, list): + rows = payload + else: + raise ValueError(f"page tags artifact has an unsupported schema: {path}") + tags: list[PageTagResult] = [] + for item in rows: + if not isinstance(item, dict): + continue + tags.append( + PageTagResult( + page_index=int(item.get("page_index") or 0), + summary=str(item.get("summary") or ""), + keywords=list(item.get("keywords") or []), + strategy_used=str(item.get("strategy_used") or ""), + entities=list(item.get("entities") or []), + observed_titles=list(item.get("observed_titles") or []), + ) + ) + return tags + + +def load_assets_artifact(path: Path) -> dict[int, list[Any]]: + from app.services.page_memory.page_assets import PageAsset + + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, list): + raise ValueError(f"assets artifact must be a list: {path}") + assets_by_page: dict[int, list[PageAsset]] = {} + for item in data: + if not isinstance(item, dict): + continue + page_index = int(item.get("page_index") or 0) + raw_source_pages = item.get("source_page_nums") + source_page_nums = ( + [int(page) for page in raw_source_pages] + if isinstance(raw_source_pages, list) + else [] + ) + asset = PageAsset( + asset_id=str(item.get("asset_id") or ""), + page_index=page_index, + asset_index=int(item.get("asset_index") or 0), + kind=str(item.get("kind") or "figure"), + bbox_px=[int(v) for v in (item.get("bbox_px") or [])], + width_px=int(item.get("width_px") or 0), + height_px=int(item.get("height_px") or 0), + width_pt=float(item.get("width_pt") or 0), + height_pt=float(item.get("height_pt") or 0), + confidence=float(item.get("confidence") or 0), + title=str(item.get("title") or ""), + summary=str(item.get("summary") or ""), + keywords=[ + str(keyword) + for keyword in (item.get("keywords") or []) + if str(keyword).strip() + ], + entities=[ + entity + for entity in (item.get("entities") or []) + if isinstance(entity, dict) + ], + image_uri=str(item.get("image_uri") or ""), + html_uri=str(item.get("html_uri") or ""), + image_path=str(item.get("image_path") or ""), + html_path=str(item.get("html_path") or ""), + extraction_status=str(item.get("extraction_status") or "loaded"), + source_page_nums=source_page_nums, + ) + if page_index > 0: + assets_by_page.setdefault(page_index, []).append(asset) + return assets_by_page + + +def load_scope_skeletons_artifact(path: Path) -> tuple[dict[str, Any], list[Any]]: + """Load Stage3 ``skeletons.json`` envelope → (meta, SectionSkeleton list).""" + from app.services.page_memory.skeleton_extractor import SectionSkeleton + + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, dict): + raise ValueError(f"scope skeletons artifact must be an object: {path}") + raw_rows = data.get("skeletons") + if not isinstance(raw_rows, list): + raise ValueError(f"scope skeletons artifact missing skeletons[]: {path}") + skeletons = [ + SectionSkeleton( + section_path=str(item.get("section_path") or ""), + title=str(item.get("title") or ""), + level=int(item.get("level") or 0), + start_page=int(item.get("start_page") or 0), + end_page=int(item.get("end_page") or 0), + parent_path=item.get("parent_path") + if isinstance(item.get("parent_path"), str) + else None, + evidence=dict(item.get("evidence") or {}) + if isinstance(item.get("evidence"), dict) + else {}, + ) + for item in raw_rows + if isinstance(item, dict) + ] + start_page = int(data.get("start_page") or 0) + end_page = int(data.get("end_page") or start_page) + meta = { + "scope_id": str(data.get("scope_id") or path.parent.name), + "start_page": start_page, + "end_page": end_page, + "page_count": int(data.get("page_count") or max(end_page - start_page + 1, 0)), + "strategy": str(data.get("strategy") or ""), + "skeleton_count": int(data.get("skeleton_count") or len(skeletons)), + "processing_pages": list(data.get("processing_pages") or []), + "excluded_toc_pages": list(data.get("excluded_toc_pages") or []), + } + return meta, skeletons + + +def _scope_meta_from_dir(scope_dir: Path) -> dict[str, Any]: + """Read coarse scope metadata from ``skeletons.json`` (or scope_id fallback).""" + skel_path = scope_dir / "skeletons.json" + if skel_path.exists(): + try: + meta, _ = load_scope_skeletons_artifact(skel_path) + return meta + except (ValueError, json.JSONDecodeError, OSError): + pass + start = end = 0 + name = scope_dir.name + if name.startswith("p") and "-" in name: + try: + left, right = name[1:].split("-", 1) + start, end = int(left), int(right) + except ValueError: + start = end = 0 + return { + "scope_id": name, + "start_page": start, + "end_page": end, + "page_count": max(end - start + 1, 0) if start and end else 0, + "strategy": "", + "skeleton_count": 0, + } + + +def load_locate_cache(locate_cache: Path) -> list[Any]: + from app.services.page_memory.skeleton_extractor import SectionSkeleton + + raw = json.loads(locate_cache.read_text(encoding="utf-8")) + return [ + SectionSkeleton( + section_path=r["section_path"], + title=r["title"], + level=r["level"], + start_page=r["start_page"], + end_page=r["end_page"], + parent_path=r.get("parent_path"), + evidence=r.get("evidence", {}), + ) + for r in raw + ] + + +def _serialize_skeletons(skeletons: list[Any]) -> list[dict[str, Any]]: + return [ + { + "section_path": skel.section_path, + "title": skel.title, + "level": skel.level, + "start_page": skel.start_page, + "end_page": skel.end_page, + "parent_path": skel.parent_path, + "evidence": skel.evidence, + } + for skel in skeletons + ] + + +# ── Coarse scope helpers ───────────────────────────────────────────────────── + + +def build_debug_coarse_scopes( + *, + skeletons: list[Any], + filename: str, + page_count: int, + anatomy: Any | None = None, +) -> list[dict[str, Any]]: + from app.services.page_memory._utils import build_hierarchy_scopes + from toc_page_policy import TocPagePolicy + + policy = TocPagePolicy.from_anatomy(anatomy) + scopes = build_hierarchy_scopes( + skeletons=skeletons, + filename=filename, + page_count=page_count, + ) + return [ + { + "scope_id": scope.scope_id, + "skeletons": scope.skeletons, + "start_page": scope.start_page, + "end_page": scope.end_page, + "strategy": scope.strategy, + "processing_pages": policy.filter_processing_pages( + list(range(scope.start_page, scope.end_page + 1)) + ), + "excluded_toc_pages": sorted( + page + for page in policy.pure_toc_pages + if scope.start_page <= page <= scope.end_page + ), + } + for scope in scopes + ] + + +def add_scope_selection_args(parser: argparse.ArgumentParser) -> None: + """Flags shared by stage 4/5/6 for picking one or more coarse scopes.""" + parser.add_argument( + "--scope-id", + default=None, + help="Process one scope only (e.g. p14-23). Repeat via comma: p14-23,p38-41", + ) + parser.add_argument( + "--all-scopes", + action="store_true", + help="Process every scope under scopes/ (default when no selector is set)", + ) + parser.add_argument( + "--page-range", + default=None, + help="Select scope(s) overlapping this page range (e.g. 14-23 or 225)", + ) + parser.add_argument( + "--fat-only", + action="store_true", + help="Select the single largest scope by page span", + ) + parser.add_argument( + "--list-scopes", + action="store_true", + help="Print available scopes and exit", + ) + + +def list_scope_dirs( + scopes_dir: Path, + *, + require_file: str = "skeletons.json", + nonempty_json: bool = False, +) -> list[dict[str, Any]]: + """Return scope metadata for directories that contain ``require_file``.""" + if not scopes_dir.exists(): + return [] + rows: list[dict[str, Any]] = [] + for path in sorted(scopes_dir.iterdir()): + if not path.is_dir(): + continue + required = path / require_file + if not required.exists(): + continue + if nonempty_json: + text = required.read_text(encoding="utf-8").strip() + if not text or text in {"[]", "{}", "null"}: + continue + info = _scope_meta_from_dir(path) + start = int(info.get("start_page") or 0) + end = int(info.get("end_page") or 0) + rows.append( + { + "scope_id": path.name, + "start_page": start, + "end_page": end, + "page_count": int(info.get("page_count") or max(end - start + 1, 0)), + "skeleton_count": int(info.get("skeleton_count") or 0), + "strategy": str(info.get("strategy") or ""), + "path": path, + } + ) + return rows + + +def resolve_debug_scope_ids( + *, + scopes_dir: Path, + scope_id: str | None = None, + page_range: str | None = None, + fat_only: bool = False, + all_scopes: bool = False, + list_scopes: bool = False, + require_file: str = "skeletons.json", + nonempty_json: bool = False, +) -> list[str]: + """Resolve which scope directories to process; exit on list/validate errors.""" + available = list_scope_dirs( + scopes_dir, + require_file=require_file, + nonempty_json=nonempty_json, + ) + if list_scopes: + if not available: + logger.error("❌ No scopes with {} under {}", require_file, scopes_dir) + raise SystemExit(1) + logger.info("Available scopes ({}):", len(available)) + for row in available: + logger.info( + " {} p{}-{} pages={} skeletons={} {}", + row["scope_id"], + row["start_page"], + row["end_page"], + row["page_count"], + row["skeleton_count"], + row["strategy"], + ) + raise SystemExit(0) + + if not available: + logger.error("❌ No scope directories with {} found under {}", require_file, scopes_dir) + logger.error( + " Run Stage 3 first: uv run python scripts/page_memory/" + "debug_pm_stage3_coarse_scope.py --file ..." + ) + raise SystemExit(1) + + by_id = {row["scope_id"]: row for row in available} + selected: list[str] = [] + + if scope_id: + requested = [part.strip() for part in str(scope_id).split(",") if part.strip()] + missing = [sid for sid in requested if sid not in by_id] + if missing: + logger.error("❌ Unknown scope-id(s): {}", ", ".join(missing)) + logger.error( + " Available: {}", + ", ".join(row["scope_id"] for row in available), + ) + raise SystemExit(1) + selected = requested + elif fat_only: + fattest = max( + available, + key=lambda row: int(row["end_page"]) - int(row["start_page"]), + ) + selected = [fattest["scope_id"]] + logger.info( + "🎯 --fat-only → {} p{}-{}", + fattest["scope_id"], + fattest["start_page"], + fattest["end_page"], + ) + elif page_range: + parts = str(page_range).split("-") + try: + pr_start = int(parts[0]) + pr_end = int(parts[1]) if len(parts) > 1 else pr_start + except ValueError: + logger.error("❌ Invalid --page-range {!r}; expected e.g. 14-23", page_range) + raise SystemExit(1) from None + selected = [ + row["scope_id"] + for row in available + if row["start_page"] <= pr_end and row["end_page"] >= pr_start + ] + if not selected: + logger.error( + "❌ No scope overlaps --page-range {}-{} under {}", + pr_start, + pr_end, + scopes_dir, + ) + raise SystemExit(1) + logger.info( + "📄 --page-range {}-{} → {} scope(s): {}", + pr_start, + pr_end, + len(selected), + ", ".join(selected), + ) + else: + # Default: all scopes (``--all-scopes`` is documented as the same). + selected = [row["scope_id"] for row in available] + if all_scopes: + logger.info(" --all-scopes: {} scopes", len(selected)) + + return selected + + +# ── Require cache helper ───────────────────────────────────────────────────── + + +def require_file(path: Path, *, hint: str) -> None: + """Abort with a clear message if a required cache file is missing.""" + if not path.exists(): + logger.error(f"❌ Required file not found: {path}") + logger.error(f" Hint: {hint}") + raise SystemExit(1) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage0_bootstrap.py b/apps/worker/scripts/page_memory/debug_pm_stage0_bootstrap.py new file mode 100644 index 000000000..6dc260a31 --- /dev/null +++ b/apps/worker/scripts/page_memory/debug_pm_stage0_bootstrap.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Stage 0: bootstrap + coarse VLM + text scan + asset probe (production-aligned). + +Matches ``ProfileCoordinator._run_coarse`` through ``stop_after_asset_probe``: + bootstrap → coarse VLM → text scan → asset probe + → persist stage0_state + page_full_text_cache + +TOC Find / extract / links belong to Stage 1 +(``debug_pm_stage1_hierarchy.py``). Calibration belongs to Stage 2. + +Usage: + cd apps/worker + uv run python scripts/page_memory/debug_pm_stage0_bootstrap.py --file /path/to/doc.pdf +""" + +import sys +from pathlib import Path as _Path + +sys.path.insert(0, str(_Path(__file__).resolve().parent)) + +import time + +from loguru import logger + +from _debug_pm_shared import ( + TokenCostTracker, + base_argparser, + page_text_cache_path, + record_stage, + resolve_paths, + run_stage0_bootstrap, + stage0_state_path, + stop_with_trace, +) + + +def main() -> int: + parser = base_argparser( + "Stage 0: bootstrap + coarse VLM + text scan + asset probe (no TOC)" + ) + args = parser.parse_args() + + pdf_path, filename, out_dir = resolve_paths(args) + + logger.info("█" * 70) + logger.info( + f" STAGE 0: BOOTSTRAP + COARSE VLM + TEXT SCAN + ASSET PROBE — {filename}" + ) + logger.info(f" OUTPUT: {out_dir}") + logger.info("█" * 70) + + t_start = time.time() + trace_stages: list[dict] = [] + token_cost_tracker = TokenCostTracker() + + coordinator, profile, state_path = run_stage0_bootstrap( + pdf_path, + filename, + out_dir, + args.model, + ) + page_count = int(coordinator.blackboard.page_count or 0) + text_pages = len(coordinator.blackboard.page_full_text_cache or {}) + asset_pages = sum( + 1 + for feature in (coordinator.blackboard.page_features or []) + if getattr(feature, "has_asset", False) + ) + record_stage( + trace_stages, + "bootstrap_coarse_scan_assets", + page_info={"page_count": page_count}, + variables={ + "source": "stop_after_asset_probe", + "category": getattr(profile, "category", None), + "routing_category": getattr(profile, "routing_category", None), + "is_scanned": bool(getattr(profile, "is_scanned", False)), + "text_pages": text_pages, + "has_asset_pages": asset_pages, + "assets_probed": bool( + coordinator.blackboard.global_signals.get("assets_probed") + ), + "stage0_state": str(state_path), + "page_full_text_cache": str(page_text_cache_path(out_dir)), + }, + ) + token_cost_tracker.snapshot_stage("bootstrap_coarse_scan_assets") + + elapsed = time.time() - t_start + logger.info(f"✅ Stage 0 done in {elapsed:.1f}s → {stage0_state_path(out_dir)}") + + return stop_with_trace( + out_dir=out_dir, + stages=trace_stages, + stop_at="bootstrap", + page_count=page_count, + pipeline_stage=0, + elapsed_s=elapsed, + token_cost_tracker=token_cost_tracker, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage1_hierarchy.py b/apps/worker/scripts/page_memory/debug_pm_stage1_hierarchy.py new file mode 100644 index 000000000..664a93616 --- /dev/null +++ b/apps/worker/scripts/page_memory/debug_pm_stage1_hierarchy.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Stage 1: TOC Find → extract → link attach (no calibration). + +Resumes Stage-0 blackboard (``stage0_state.json`` + ``page_full_text_cache.json``, +including asset-probe ``has_asset`` flags) and runs the production TOC segment: + find.toc_anchor_pages → extract.toc_with_boundaries → attach links + → persist doc_profile.json + +Requires Stage 0 first: + uv run python scripts/page_memory/debug_pm_stage0_bootstrap.py --file ... + +Calibration belongs to Stage 2 (``debug_pm_stage2_calibration.py``). + +Usage: + cd apps/worker + uv run python scripts/page_memory/debug_pm_stage1_hierarchy.py --file /path/to/doc.pdf + uv run python scripts/page_memory/debug_pm_stage1_hierarchy.py --reuse-anatomy +""" + +import sys +from pathlib import Path as _Path + +sys.path.insert(0, str(_Path(__file__).resolve().parent)) + +import time + +from loguru import logger + +from _debug_pm_shared import ( + TokenCostTracker, + base_argparser, + load_anatomy_cache, + resolve_anatomy_cache_path, + record_stage, + remove_legacy_doc_agent_artifacts, + require_file, + resolve_paths, + run_stage1_toc, + stage0_state_path, + stop_with_trace, + toc_hierarchies_to_hierarchy_tree, + write_toc_hierarchy_artifact, +) + + +def _count_hierarchy_keys(tree: dict) -> int: + total = 0 + for children in (tree or {}).values(): + total += 1 + if isinstance(children, dict): + total += _count_hierarchy_keys(children) + return total + + +def _count_linked_entries(toc_hierarchies: list | None) -> tuple[int, int]: + total = 0 + linked = 0 + for region in toc_hierarchies or []: + if not isinstance(region, dict): + continue + entries = region.get("toc_with_level") or [] + if not isinstance(entries, list): + continue + for entry in entries: + if not isinstance(entry, dict): + continue + total += 1 + link = entry.get("link") + if isinstance(link, dict) and link.get("physical_page") is not None: + linked += 1 + return linked, total + + +def main() -> int: + parser = base_argparser("Stage 1: TOC Find → extract → link (no calibration)") + parser.add_argument( + "--reuse-anatomy", + action="store_true", + help="Reuse cached Stage-1 doc_profile.json (skip Find/extract/link)", + ) + args = parser.parse_args() + + pdf_path, filename, out_dir = resolve_paths(args) + + logger.info("█" * 70) + logger.info(f" STAGE 1: TOC FIND → EXTRACT → LINK — {filename}") + logger.info(f" OUTPUT: {out_dir}") + logger.info("█" * 70) + + t_start = time.time() + trace_stages: list[dict] = [] + token_cost_tracker = TokenCostTracker() + + doc_agent_dir = out_dir / "_doc_agent" + anatomy_cache = resolve_anatomy_cache_path(out_dir) + + if args.reuse_anatomy and anatomy_cache.exists(): + anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) + profile_source = "reuse_anatomy" + else: + require_file( + stage0_state_path(out_dir), + hint=( + "Run Stage 0 first: uv run python " + "scripts/page_memory/debug_pm_stage0_bootstrap.py --file ..." + ), + ) + anatomy = run_stage1_toc(pdf_path, filename, out_dir, args.model) + profile_source = "stage0_resume_toc" + + page_count = anatomy.page_count + linked, total = _count_linked_entries(list(anatomy.toc_hierarchies or [])) + record_stage( + trace_stages, + "toc", + page_info={"page_count": page_count}, + variables={ + "source": profile_source, + "toc_pages": anatomy.toc_result.toc_pages, + "toc_entries": total, + "toc_entries_with_link": linked, + "skip_toc_anchoring": True, + "skeleton_anchor": getattr(anatomy, "skeleton_anchor", None), + }, + ) + token_cost_tracker.snapshot_stage("toc") + remove_legacy_doc_agent_artifacts(doc_agent_dir) + + logger.info("=" * 70) + logger.info("🧠 TOC hierarchy (Stage-1 debug dump)") + logger.info("=" * 70) + logger.info(" TOC entries with link: {}/{}", linked, total) + + hierarchy_tree = toc_hierarchies_to_hierarchy_tree(anatomy.toc_hierarchies) + toc_path = write_toc_hierarchy_artifact( + out_dir, + hierarchy_tree=hierarchy_tree, + stats={ + "source": "toc_hierarchies_raw", + "region_count": len(list(anatomy.toc_hierarchies or [])), + "hierarchy_key_count": _count_hierarchy_keys(hierarchy_tree), + "toc_entries_with_link": linked, + "toc_entries": total, + }, + ) + logger.info(f" toc_hierarchy → {toc_path}") + + record_stage( + trace_stages, + "C2.toc_hierarchy_dump", + variables={ + "toc_hierarchy_path": str(toc_path), + "region_count": len(list(anatomy.toc_hierarchies or [])), + "hierarchy_key_count": _count_hierarchy_keys(hierarchy_tree), + }, + ) + + elapsed = time.time() - t_start + logger.info(f"✅ Stage 1 done in {elapsed:.1f}s → {out_dir}") + + return stop_with_trace( + out_dir=out_dir, + stages=trace_stages, + stop_at="toc", + page_count=page_count, + pipeline_stage=1, + elapsed_s=elapsed, + token_cost_tracker=token_cost_tracker, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py b/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py new file mode 100644 index 000000000..a659e28c8 --- /dev/null +++ b/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Stage 2: Calibration over Stage-1 TOC ``doc_profile.json``. + +Reads Stage-1 output (TOC hierarchies + optional ``link.physical_page``), +runs calibration (Agent Phase-1 + Phase-2 completion), and writes +``skeleton_anchor`` / ``toc_page_offset`` back onto ``doc_profile.json``. + +Requires Stage 0 → Stage 1 first: + uv run python scripts/page_memory/debug_pm_stage0_bootstrap.py --file ... + uv run python scripts/page_memory/debug_pm_stage1_hierarchy.py --file ... + +Usage: + cd apps/worker + uv run python scripts/page_memory/debug_pm_stage2_calibration.py --file /path/to/doc.pdf + uv run python scripts/page_memory/debug_pm_stage2_calibration.py --file /path/to/doc.pdf --no-links +""" + +from __future__ import annotations + +import json +import os +import sys +import time +from pathlib import Path as _Path + +sys.path.insert(0, str(_Path(__file__).resolve().parent)) + +from loguru import logger + +from _debug_pm_shared import ( + TokenCostTracker, + base_argparser, + load_anatomy_cache, + resolve_anatomy_cache_path, + pipeline_state_path, + record_stage, + require_file, + resolve_paths, + stop_with_trace, + update_pipeline_state, + write_debug_json, +) + + +def main() -> int: + parser = base_argparser("Stage 2: Calibration SubAgent") + parser.add_argument( + "--no-links", + action="store_true", + help="Strip link.physical_page from TOC entries before the agent runs", + ) + parser.add_argument( + "--max-rounds", + type=int, + default=16, + help="Max ReAct rounds per TOC region", + ) + args = parser.parse_args() + + from app.services.document_agent.agents.calibration import ( + run_calibration_for_all_regions, + ) + from app.services.document_agent.pdf_text import read_page_texts + from app.services.document_agent.persist import DOC_PROFILE_FILENAME + + pdf_path, filename, out_dir = resolve_paths(args) + anatomy_cache = resolve_anatomy_cache_path(out_dir) + require_file( + anatomy_cache, + hint=( + "Run Stage 0 then Stage 1 first:\n" + " uv run python scripts/page_memory/debug_pm_stage0_bootstrap.py --file ...\n" + " uv run python scripts/page_memory/debug_pm_stage1_hierarchy.py --file ..." + ), + ) + + anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) + page_count = int(anatomy.page_count or 0) + hierarchies = list(getattr(anatomy, "toc_hierarchies", None) or []) + + logger.info("█" * 70) + logger.info(f" STAGE 2: CALIBRATION SUBAGENT — {filename}") + logger.info(f" OUTPUT: {out_dir}") + logger.info(f" no_links={bool(args.no_links)} regions={len(hierarchies)}") + logger.info("█" * 70) + + t_start = time.time() + trace_stages: list[dict] = [] + token_cost_tracker = TokenCostTracker() + + # Production null-page locate needs page texts (same as C4). + page_texts = read_page_texts(pdf_path, list(range(1, page_count + 1)), timeout=300) + body_pages = sorted(page_texts.keys()) + logger.info( + " read {} pages, {} non-empty", + len(page_texts), + sum(1 for text in page_texts.values() if str(text).strip()), + ) + + vlm_model = args.vlm_model or os.environ.get("IMAGE_MODEL") or "" + planner_model = ( + args.model + or os.environ.get("HIERARCHY_LLM_MODEL") + or os.environ.get("NORMOL_MODEL") + or vlm_model + ) + + doc_agent_dir = out_dir / "_doc_agent" + doc_agent_dir.mkdir(parents=True, exist_ok=True) + + calibration = run_calibration_for_all_regions( + pdf_path=pdf_path, + page_count=page_count, + toc_hierarchies=hierarchies, + output_dir=str(doc_agent_dir), + vlm_model=vlm_model, + planner_model=planner_model, + no_links=bool(args.no_links), + max_rounds=max(1, int(args.max_rounds)), + page_texts=page_texts, + body_pages=body_pages, + ) + + # Production-compatible core fields. + skeleton_anchor = { + "offset": calibration.get("offset"), + "offset_status": calibration.get("offset_status"), + "match_overrides": calibration.get("match_overrides") or {}, + "null_page_report": calibration.get("null_page_report") or [], + "bulk_count": calibration.get("bulk_count") or 0, + "pruned_count": calibration.get("pruned_count") or 0, + "locate_agent": calibration.get("locate_agent") or "offset_only", + } + + profile_path = out_dir / DOC_PROFILE_FILENAME + payload = json.loads(profile_path.read_text(encoding="utf-8")) + payload["skeleton_anchor"] = skeleton_anchor + payload["calibration"] = calibration + if calibration.get("offset") is not None: + payload["toc_page_offset"] = calibration.get("offset") + write_debug_json(profile_path, payload) + + state_path = pipeline_state_path(out_dir) + update_pipeline_state( + state_path, + stage=2, + document={ + "source_file_name": filename, + "page_count": page_count, + "anatomy_path": str(anatomy_cache), + }, + payload={"skeleton_anchor": skeleton_anchor, "calibration": calibration}, + ) + + record_stage( + trace_stages, + "calibration", + page_info={"page_count": page_count}, + variables={ + "status": calibration.get("status"), + "failure_kind": calibration.get("failure_kind"), + "offset": calibration.get("offset"), + "offset_status": calibration.get("offset_status"), + "bulk_count": calibration.get("bulk_count"), + "locate_agent": calibration.get("locate_agent"), + "regime_count": len(calibration.get("regimes") or []), + "tool_calls": calibration.get("tool_calls"), + "no_links": bool(args.no_links), + }, + ) + token_cost_tracker.snapshot_stage("calibration") + + elapsed = time.time() - t_start + logger.info( + "✅ Stage 2 done status={} offset={} bulk={} locate={} in {:.1f}s → {}", + calibration.get("status"), + calibration.get("offset"), + calibration.get("bulk_count"), + calibration.get("locate_agent"), + elapsed, + profile_path, + ) + + return stop_with_trace( + out_dir=out_dir, + stages=trace_stages, + stop_at="calibration", + page_count=page_count, + pipeline_stage=2, + elapsed_s=elapsed, + token_cost_tracker=token_cost_tracker, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage3_coarse_scope.py b/apps/worker/scripts/page_memory/debug_pm_stage3_coarse_scope.py new file mode 100644 index 000000000..e0a953ebe --- /dev/null +++ b/apps/worker/scripts/page_memory/debug_pm_stage3_coarse_scope.py @@ -0,0 +1,291 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Stage 3: Coarse scope generation + per-scope directory creation. + +Generates coarse hierarchy scopes from skeletons and creates per-scope +directories with ``skeletons.json`` (meta + coarse nodes) plus empty +``page_tags.json`` / ``assets.json`` placeholders for later stages. + +Requires Stage 2 output: _doc_agent/pipeline_state.json, doc_profile.json + +Usage: + cd apps/worker + uv run python scripts/page_memory/debug_pm_stage3_coarse_scope.py --file /path/to/doc.pdf + uv run python scripts/page_memory/debug_pm_stage3_coarse_scope.py --fat-only + uv run python scripts/page_memory/debug_pm_stage3_coarse_scope.py --page-range 225-302 +""" + +import sys +from pathlib import Path as _Path + +sys.path.insert(0, str(_Path(__file__).resolve().parent)) + +import time + +from loguru import logger + +from _debug_pm_shared import ( + TokenCostTracker, + base_argparser, + build_debug_coarse_scopes, + load_anatomy_cache, + resolve_anatomy_cache_path, + load_pipeline_skeletons, + pipeline_state_path, + record_stage, + require_file, + resolve_paths, + scope_id_for_pages, + stop_with_trace, + update_pipeline_state, + write_debug_json, + _serialize_skeletons, + _serialize_scope_skeletons, +) + + +def main() -> int: + parser = base_argparser("Stage 3: Coarse scope generation") + parser.add_argument( + "--fat-only", action="store_true", + help="Auto-select the largest coarse scope only", + ) + parser.add_argument( + "--page-range", default=None, + help="Only process page range, e.g. '225-302'", + ) + parser.add_argument( + "--all-scopes", action="store_true", + help="Process all coarse scopes (default behavior)", + ) + args = parser.parse_args() + + from app.services.page_memory.skeleton_extractor import SectionSkeleton + + pdf_path, filename, out_dir = resolve_paths(args) + doc_agent_dir = out_dir / "_doc_agent" + anatomy_cache = resolve_anatomy_cache_path(out_dir) + state_path = pipeline_state_path(out_dir) + legacy_locate_cache = doc_agent_dir / "locate_cache.json" + + if not state_path.exists() and not legacy_locate_cache.exists(): + require_file( + state_path, + hint="Run Stage 2 first: uv run python scripts/page_memory/debug_pm_stage2_calibration.py --file ...", + ) + require_file( + anatomy_cache, + hint="Run Stage 1 first: uv run python scripts/page_memory/debug_pm_stage1_hierarchy.py --file ...", + ) + + anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) + page_count = anatomy.page_count + skeletons = load_pipeline_skeletons( + state_path, + legacy_locate_cache=legacy_locate_cache, + ) + if not state_path.exists(): + update_pipeline_state( + state_path, + stage=2, + document={ + "source_file_name": filename, + "page_count": page_count, + "anatomy_path": str(anatomy_cache), + }, + payload={ + "calibration": {}, + "null_page_parent_locate": {}, + "skeletons": _serialize_skeletons(skeletons), + "migrated_from": str(legacy_locate_cache), + }, + ) + + logger.info("█" * 70) + logger.info(f" STAGE 3: COARSE SCOPE GENERATION — {filename}") + logger.info(f" OUTPUT: {out_dir}") + logger.info("█" * 70) + + t_start = time.time() + trace_stages: list[dict] = [] + token_cost_tracker = TokenCostTracker() + from toc_page_policy import TocPagePolicy + + toc_policy = TocPagePolicy.from_anatomy(anatomy) + + # ── Build coarse scopes ── + coarse_scopes = build_debug_coarse_scopes( + skeletons=skeletons, + filename=filename, + page_count=page_count, + anatomy=anatomy, + ) + + if not coarse_scopes: + root_skel = SectionSkeleton( + section_path=f"{filename}/Root", + level=1, + start_page=1, + end_page=page_count, + title="Root", + parent_path=filename, + evidence={"source": "fallback_root", "confidence": 0.0}, + ) + coarse_scopes = [ + { + "scope_id": scope_id_for_pages(1, page_count), + "skeletons": [root_skel], + "start_page": 1, + "end_page": page_count, + "strategy": "fallback_root", + "processing_pages": toc_policy.filter_processing_pages( + list(range(1, page_count + 1)) + ), + "excluded_toc_pages": sorted(toc_policy.pure_toc_pages), + } + ] + logger.info(" no skeleton hierarchy → fallback Root scope p1-{}", page_count) + + # ── Scope selection ── + if args.fat_only: + selected_scopes = [ + max(coarse_scopes, key=lambda s: int(s["end_page"]) - int(s["start_page"])) + ] + logger.info( + "🎯 --fat-only: 1/{} scopes selected {} p{}-{}", + len(coarse_scopes), + selected_scopes[0]["scope_id"], + selected_scopes[0]["start_page"], + selected_scopes[0]["end_page"], + ) + elif args.page_range: + parts = args.page_range.split("-") + pr_start = int(parts[0]) + pr_end = int(parts[1]) if len(parts) > 1 else pr_start + requested_pages = list(range(pr_start, pr_end + 1)) + pr_skeletons = [ + s for s in skeletons + if s.start_page <= pr_end and s.end_page >= pr_start + ] + selected_scopes = [ + { + "scope_id": scope_id_for_pages(pr_start, pr_end), + "skeletons": pr_skeletons, + "start_page": pr_start, + "end_page": pr_end, + "strategy": "manual_page_range", + "processing_pages": toc_policy.filter_processing_pages( + requested_pages + ), + "excluded_toc_pages": sorted( + set(requested_pages) & toc_policy.pure_toc_pages + ), + } + ] + logger.info(f" --page-range: p{pr_start}-{pr_end} ({len(pr_skeletons)} skeletons)") + else: + selected_scopes = coarse_scopes + logger.info( + " default: all {} scopes selected", len(selected_scopes), + ) + + record_stage( + trace_stages, + "C4.coarse_scopes", + variables={ + "total_coarse_scopes": len(coarse_scopes), + "selected_scopes": len(selected_scopes), + "mode": ( + "fat_only" if args.fat_only + else "page_range" if args.page_range + else "all_scopes" + ), + "scopes": [ + { + "scope_id": s["scope_id"], + "start_page": s["start_page"], + "end_page": s["end_page"], + "strategy": s.get("strategy", ""), + "skeleton_count": len(s["skeletons"]), + "processing_pages": list(s.get("processing_pages") or []), + "excluded_toc_pages": list(s.get("excluded_toc_pages") or []), + } + for s in selected_scopes + ], + }, + ) + + # ── Create per-scope directories ── + scopes_dir = out_dir / "scopes" + scopes_dir.mkdir(parents=True, exist_ok=True) + for s in selected_scopes: + scope_dir = scopes_dir / s["scope_id"] + scope_dir.mkdir(parents=True, exist_ok=True) + write_debug_json( + scope_dir / "skeletons.json", + { + **_serialize_scope_skeletons( + scope_id=str(s["scope_id"]), + start_page=int(s["start_page"]), + end_page=int(s["end_page"]), + strategy=str(s.get("strategy") or ""), + skeletons=s["skeletons"], + ), + "processing_pages": list(s.get("processing_pages") or []), + "excluded_toc_pages": list(s.get("excluded_toc_pages") or []), + }, + ) + # Placeholders for later stages (explicit empty slots for viewing). + write_debug_json(scope_dir / "page_tags.json", []) + write_debug_json(scope_dir / "assets.json", []) + + scope_rows = [ + { + "scope_id": str(scope["scope_id"]), + "start_page": int(scope["start_page"]), + "end_page": int(scope["end_page"]), + "strategy": str(scope.get("strategy") or ""), + "skeleton_count": len(scope["skeletons"]), + "processing_pages": list(scope.get("processing_pages") or []), + "excluded_toc_pages": list(scope.get("excluded_toc_pages") or []), + "artifact_path": str( + scopes_dir / str(scope["scope_id"]) / "skeletons.json" + ), + } + for scope in selected_scopes + ] + update_pipeline_state( + state_path, + stage=3, + payload={ + "selection_mode": ( + "fat_only" + if args.fat_only + else "page_range" + if args.page_range + else "all_scopes" + ), + "total_scope_count": len(coarse_scopes), + "selected_scope_count": len(selected_scopes), + "scopes": scope_rows, + }, + ) + (out_dir / "coarse_scopes.json").unlink(missing_ok=True) + + elapsed = time.time() - t_start + logger.info(f"✅ Stage 3 done in {elapsed:.1f}s") + logger.info(f" {len(selected_scopes)} scope dirs created → {scopes_dir}/") + + return stop_with_trace( + out_dir=out_dir, + stages=trace_stages, + stop_at="scope", + page_count=page_count, + pipeline_stage=3, + elapsed_s=elapsed, + token_cost_tracker=token_cost_tracker, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage4_fine_hierarchy.py b/apps/worker/scripts/page_memory/debug_pm_stage4_fine_hierarchy.py new file mode 100644 index 000000000..497a657d6 --- /dev/null +++ b/apps/worker/scripts/page_memory/debug_pm_stage4_fine_hierarchy.py @@ -0,0 +1,458 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Stage 4: Document-level page tagging + per-scope fine hierarchy. + +Renders and tags each processing page once (global concurrency), then fans +tag subsets into scopes for fine hierarchy refinement. + +Requires Stage 3 output: scopes//skeletons.json +Uses Stage 2 skeletons in pipeline_state for ``next_title_by_path``. + +Usage: + cd apps/worker + uv run python scripts/page_memory/debug_pm_stage4_fine_hierarchy.py \\ + --file /path/to/doc.pdf --scope-id p14-23 --out-suffix boundary_clip + uv run python scripts/page_memory/debug_pm_stage4_fine_hierarchy.py --file ... --all-scopes +""" + +import sys +from pathlib import Path as _Path + +sys.path.insert(0, str(_Path(__file__).resolve().parent)) + +import os +import time +from pathlib import Path +from typing import Any, cast + +from loguru import logger + +from _debug_pm_shared import ( + ScopeResult, + TokenCostTracker, + TraceStageAdapter, + add_scope_selection_args, + base_argparser, + list_scope_dirs, + load_anatomy_cache, + resolve_anatomy_cache_path, + load_pipeline_skeletons, + load_scope_skeletons_artifact, + pipeline_state_path, + record_stage, + require_file, + resolve_debug_scope_ids, + resolve_paths, + sort_skeletons, + stop_with_trace, + update_pipeline_state, + write_scope_artifacts, + write_top_level_artifacts, + page_scope_info, + _derive_hierarchy_page_scope, + _scope_manifest, + _serialize_skeletons, +) + + +def _resolve_scope_processing_pages( + *, + scope_meta: dict[str, Any], + skeletons: list[Any], + page_count: int, + toc_policy: Any, +) -> tuple[list[int], list[int]]: + processing_pages = [ + int(page) for page in (scope_meta.get("processing_pages") or []) + ] + if not processing_pages: + processing_pages = toc_policy.filter_processing_pages( + _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ) + ) + excluded_toc_pages = [ + int(page) for page in (scope_meta.get("excluded_toc_pages") or []) + ] + if not excluded_toc_pages: + excluded_toc_pages = sorted( + set(range( + int(scope_meta.get("start_page") or 1), + int(scope_meta.get("end_page") or page_count) + 1, + )) + & toc_policy.pure_toc_pages + ) + return processing_pages, excluded_toc_pages + + +def _run_fine_hierarchy_for_scope( + *, + scope_id: str, + scope_dir: Path, + out_dir: Path, + page_count: int, + rendered_by_page: dict[int, Any], + tags_by_page: dict[int, Any], + next_title_by_path: dict[str, str | None], + toc_policy: Any, + page_memory_config: Any, + token_cost_tracker: TokenCostTracker | None = None, +) -> ScopeResult: + """Consume shared page tags and refine one coarse scope.""" + from app.services.page_memory.fine_hierarchy import ( + compute_fat_leaf_pages, + refine_fat_leaf_skeletons, + ) + from app.services.page_memory.memory_service import _resolve_hierarchy_model + + scope_stages: list[dict[str, Any]] = [] + if token_cost_tracker is not None: + token_cost_tracker.register_child_thread() + + skel_path = scope_dir / "skeletons.json" + require_file(skel_path, hint=f"Run Stage 3 first to create {skel_path}") + scope_meta, active_skeletons = load_scope_skeletons_artifact(skel_path) + strategy = str(scope_meta.get("strategy") or "coarse_scope") + processing_pages, excluded_toc_pages = _resolve_scope_processing_pages( + scope_meta=scope_meta, + skeletons=active_skeletons, + page_count=page_count, + toc_policy=toc_policy, + ) + + scope_manifest = _scope_manifest( + scope_id=scope_id, + skeletons=active_skeletons, + page_count=page_count, + strategy=strategy, + processing_pages=processing_pages, + excluded_toc_pages=excluded_toc_pages, + ) + + logger.info( + "🔬 [scope {}] {} skeletons p{}-{} processing={}", + scope_id, + len(active_skeletons), + scope_meta.get("start_page", "?"), + scope_meta.get("end_page", "?"), + processing_pages, + ) + + if not processing_pages: + logger.info(" [scope {}] no processing pages after TOC exclusion", scope_id) + return ScopeResult( + scope_id=scope_id, + skeletons=active_skeletons, + tags=[], + assets_by_page={}, + rendered=[], + final_pages=[], + scope_manifest=scope_manifest, + trace_stages=scope_stages, + ) + + rendered = [ + rendered_by_page[page] + for page in processing_pages + if page in rendered_by_page + ] + tags = [ + tags_by_page[page] + for page in processing_pages + if page in tags_by_page + ] + + fine_min = page_memory_config.fine_min_pages + fat_leaf_pages = compute_fat_leaf_pages( + active_skeletons, + min_pages=fine_min, + exclude_pages=toc_policy.pure_toc_pages, + ) + if fat_leaf_pages: + active_skeletons = refine_fat_leaf_skeletons( + coarse_skeletons=active_skeletons, + tag_results=tags, + fat_leaf_pages=fat_leaf_pages, + next_title_by_path=next_title_by_path, + model_name=_resolve_hierarchy_model(page_memory_config), + max_tokens=page_memory_config.hierarchy_max_tokens, + max_depth=page_memory_config.max_heading_depth, + trace_recorder=TraceStageAdapter(scope_stages), + ) + logger.info( + " [scope {}] C4b: {} sections after fine hierarchy", + scope_id, len(active_skeletons), + ) + else: + logger.info( + " [scope {}] no fat-leaf pages (min={}); skip fine hierarchy", + scope_id, fine_min, + ) + + scope_manifest = _scope_manifest( + scope_id=scope_id, + skeletons=active_skeletons, + page_count=page_count, + strategy=f"{strategy}:refined", + processing_pages=processing_pages, + excluded_toc_pages=excluded_toc_pages, + ) + record_stage( + scope_stages, "C4b.fine_hierarchy", + page_info={"fat_leaf": page_scope_info(sorted(fat_leaf_pages))}, + variables={ + "scope_id": scope_id, + "scope": scope_manifest, + "sections": _serialize_skeletons(active_skeletons), + }, + ) + if token_cost_tracker is not None: + token_cost_tracker.snapshot_stage(f"C4b.fine_hierarchy:{scope_id}") + + write_scope_artifacts( + out_dir=out_dir, + scope_id=scope_id, + scope_manifest=scope_manifest, + hierarchy=active_skeletons, + tags=tags, + ) + + return ScopeResult( + scope_id=scope_id, + skeletons=active_skeletons, + tags=tags, + assets_by_page={}, + rendered=rendered, + final_pages=processing_pages, + scope_manifest=scope_manifest, + trace_stages=scope_stages, + ) + + +def main() -> int: + parser = base_argparser("Stage 4: Combined page tagging + fine hierarchy") + add_scope_selection_args(parser) + parser.add_argument( + "--max-workers", type=int, default=5, + help="Concurrent workers for scope fine hierarchy (default=5)", + ) + args = parser.parse_args() + + from app.services.document_agent.pdf_text import read_page_texts + from app.services.page_memory.fine_hierarchy import build_next_title_by_path + from app.services.page_memory.memory_service import _render_and_tag_document_pages + from toc_page_policy import TocPagePolicy + from shared.models.schemas.page_memory_config import PageMemoryConfig + + pdf_path, filename, out_dir = resolve_paths(args) + doc_agent_dir = out_dir / "_doc_agent" + anatomy_cache = resolve_anatomy_cache_path(out_dir) + state_path = pipeline_state_path(out_dir) + legacy_locate_cache = doc_agent_dir / "locate_cache.json" + scopes_dir = out_dir / "scopes" + + require_file( + anatomy_cache, + hint="Run Stage 1 first: uv run python scripts/page_memory/debug_pm_stage1_hierarchy.py --file ...", + ) + if not state_path.exists() and not legacy_locate_cache.exists(): + require_file( + state_path, + hint="Run Stage 2 first: uv run python scripts/page_memory/debug_pm_stage2_calibration.py --file ...", + ) + + anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) + page_count = anatomy.page_count + page_features = anatomy.page_features if anatomy else [] + page_labels = anatomy.page_labels if anatomy else [] + toc_policy = TocPagePolicy.from_anatomy(anatomy) + page_memory_config = PageMemoryConfig.default() + + all_skeletons = load_pipeline_skeletons( + state_path, + legacy_locate_cache=legacy_locate_cache, + ) + next_title_by_path = build_next_title_by_path(all_skeletons) + logger.info( + " next_title_by_path: {} paths ({} with tail anchor)", + len(next_title_by_path), + sum(1 for title in next_title_by_path.values() if title), + ) + + scope_ids = resolve_debug_scope_ids( + scopes_dir=scopes_dir, + scope_id=args.scope_id, + page_range=args.page_range, + fat_only=args.fat_only, + all_scopes=args.all_scopes, + list_scopes=args.list_scopes, + require_file="skeletons.json", + ) + partial_run = len(scope_ids) < len(list_scope_dirs(scopes_dir)) + + logger.info("█" * 70) + logger.info(f" STAGE 4: FINE HIERARCHY — {filename}") + logger.info(f" OUTPUT: {out_dir}") + logger.info(f" SCOPES ({len(scope_ids)}): {scope_ids}") + if partial_run: + logger.info(" MODE: partial — will not overwrite top-level hierarchy.json") + logger.info("█" * 70) + + t_start = time.time() + trace_stages: list[dict] = [] + token_cost_tracker = TokenCostTracker() + + selected_processing_pages: set[int] = set() + scope_payloads: list[tuple[str, Path, list[int]]] = [] + for sid in scope_ids: + scope_dir = scopes_dir / sid + scope_meta, skeletons = load_scope_skeletons_artifact(scope_dir / "skeletons.json") + processing_pages, _excluded = _resolve_scope_processing_pages( + scope_meta=scope_meta, + skeletons=skeletons, + page_count=page_count, + toc_policy=toc_policy, + ) + selected_processing_pages.update(processing_pages) + scope_payloads.append((sid, scope_dir, processing_pages)) + + processing_pages = sorted(selected_processing_pages) + page_texts = read_page_texts( + pdf_path, + processing_pages or list(range(1, page_count + 1)), + timeout=300, + ) + logger.info( + " document page stage: {} processing pages (of {}), tag_concurrency={}", + len(processing_pages), + page_count, + page_memory_config.tag_concurrency, + ) + + vlm_model = getattr(args, "vlm_model", None) or os.environ.get("IMAGE_MODEL") + rendered_by_page, tags_by_page = _render_and_tag_document_pages( + pdf_path=pdf_path, + output_dir=str(out_dir), + page_count=page_count, + processing_pages=processing_pages, + page_texts=page_texts, + page_features=page_features, + page_labels=page_labels, + vlm_model=vlm_model, + toc_policy=toc_policy, + page_memory_config=page_memory_config, + trace_recorder=TraceStageAdapter(trace_stages), + ) + token_cost_tracker.snapshot_stage("C3.page_tagger") + logger.info( + " tagged {} unique pages; refining {} scopes", + len(tags_by_page), + len(scope_ids), + ) + + def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: + return _run_fine_hierarchy_for_scope( + scope_id=scope_id, + scope_dir=scope_dir, + out_dir=out_dir, + page_count=page_count, + rendered_by_page=rendered_by_page, + tags_by_page=tags_by_page, + next_title_by_path=next_title_by_path, + toc_policy=toc_policy, + page_memory_config=page_memory_config, + token_cost_tracker=token_cost_tracker, + ) + + if args.max_workers > 1 and len(scope_ids) > 1: + import gevent + from gevent.pool import Pool as GeventPool + + logger.info( + " scope fine-hierarchy concurrency: {} workers × {} scopes", + args.max_workers, len(scope_ids), + ) + gpool = GeventPool(size=min(args.max_workers, len(scope_ids))) + greenlets = [ + gpool.spawn( + _run_selected_scope, + sid, + scope_dir, + ) + for sid, scope_dir, _pages in scope_payloads + ] + gevent.joinall(greenlets, raise_error=True) + scope_results = [cast(ScopeResult, g.value) for g in greenlets] + else: + logger.info(" serial fine hierarchy: {} scope(s)", len(scope_ids)) + scope_results = [ + _run_selected_scope(sid, scope_dir) + for sid, scope_dir, _pages in scope_payloads + ] + + for sr in scope_results: + trace_stages.extend(sr.trace_stages) + + merged_skeletons = sort_skeletons( + [skel for sr in scope_results for skel in sr.skeletons] + ) + merged_tags = [tags_by_page[page] for page in sorted(tags_by_page)] + if not partial_run: + write_top_level_artifacts( + out_dir=out_dir, + hierarchy=merged_skeletons, + tags=merged_tags, + ) + else: + logger.info( + " skipped top-level hierarchy.json merge (partial {}/{} scopes)", + len(scope_ids), + len(list_scope_dirs(scopes_dir)), + ) + + elapsed = time.time() - t_start + logger.info(f"✅ Stage 4 done in {elapsed:.1f}s") + logger.info( + f" {len(scope_results)} scopes processed, " + f"{len(merged_skeletons)} skeletons this run, " + f"{len(merged_tags)} unique page tags" + ) + for sid in scope_ids: + logger.info(f" → {scopes_dir / sid / 'fine_hierarchy.json'}") + + update_pipeline_state( + state_path, + stage=4, + payload={ + "partial_run": partial_run, + "processed_scope_ids": scope_ids, + "processed_scope_count": len(scope_results), + "skeleton_count": len(merged_skeletons), + "tagged_page_count": len(merged_tags), + "scope_artifacts": [ + str(scopes_dir / scope_id / "fine_hierarchy.json") + for scope_id in scope_ids + ], + }, + ) + + return stop_with_trace( + out_dir=out_dir, + stages=trace_stages, + stop_at="fine_hierarchy", + page_count=page_count, + pipeline_stage=4, + elapsed_s=elapsed, + scope_id=scope_ids[0] if len(scope_ids) == 1 else None, + token_cost_tracker=token_cost_tracker, + extra_summary={ + "scope_count": len(scope_results), + "skeleton_count": len(merged_skeletons), + "tagged_page_count": len(merged_tags), + }, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage5_assets.py b/apps/worker/scripts/page_memory/debug_pm_stage5_assets.py new file mode 100644 index 000000000..7241fda04 --- /dev/null +++ b/apps/worker/scripts/page_memory/debug_pm_stage5_assets.py @@ -0,0 +1,314 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Stage 5: Document-level page asset extraction (C5) — NO page tagging. + +Unions processing pages from selected scopes, renders/extracts each unique page +once, writes top-level ``assets.json``, and projects references into scope dirs. +Does NOT run page tagging (C3) — Stage 4 already produced shared document-level tags. + +Requires Stage 4 output: scopes//fine_hierarchy.json + +Usage: + cd apps/worker + uv run python scripts/page_memory/debug_pm_stage5_assets.py --file /path/to/doc.pdf + uv run python scripts/page_memory/debug_pm_stage5_assets.py --scope-id p1-100 + uv run python scripts/page_memory/debug_pm_stage5_assets.py --all-scopes +""" + +import sys +from pathlib import Path as _Path + +sys.path.insert(0, str(_Path(__file__).resolve().parent)) + +import os +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from loguru import logger + +from _debug_pm_shared import ( + TokenCostTracker, + add_scope_selection_args, + base_argparser, + load_anatomy_cache, + resolve_anatomy_cache_path, + load_hierarchy_artifact, + pipeline_state_path, + record_stage, + require_file, + resolve_debug_scope_ids, + resolve_paths, + stop_with_trace, + update_pipeline_state, + write_scope_artifacts, + page_scope_info, + _scope_manifest, + _derive_hierarchy_page_scope, + _serialize_assets, + write_debug_json, +) + + +@dataclass(frozen=True) +class _ScopeAssetContext: + scope_id: str + pages: list[int] + skeletons: list[Any] + scope_manifest: dict[str, Any] + + +def _load_scope_asset_context( + *, + scope_id: str, + scope_dir: Path, + page_count: int, + toc_policy: Any, +) -> _ScopeAssetContext | None: + fine_hierarchy_path = scope_dir / "fine_hierarchy.json" + require_file( + fine_hierarchy_path, + hint=f"Run Stage 4 first to produce {fine_hierarchy_path}", + ) + prior_scope, active_skeletons = load_hierarchy_artifact(fine_hierarchy_path) + if not active_skeletons: + logger.warning( + " [scope {}] no skeletons in fine_hierarchy.json — skipping", + scope_id, + ) + return None + + recorded_pages = prior_scope.get("processing_pages") + final_pages = ( + [int(page) for page in recorded_pages] + if isinstance(recorded_pages, list) + else toc_policy.filter_processing_pages( + _derive_hierarchy_page_scope( + skeletons=active_skeletons, + page_count=page_count, + ) + ) + ) + recorded_excluded = prior_scope.get("excluded_toc_pages") + excluded_toc_pages = ( + [int(page) for page in recorded_excluded] + if isinstance(recorded_excluded, list) + else sorted(toc_policy.pure_toc_pages) + ) + scope_manifest = _scope_manifest( + scope_id=scope_id, + skeletons=active_skeletons, + page_count=page_count, + strategy="fine:assets", + processing_pages=final_pages, + excluded_toc_pages=excluded_toc_pages, + ) + return _ScopeAssetContext( + scope_id=scope_id, + pages=final_pages, + skeletons=active_skeletons, + scope_manifest=scope_manifest, + ) + + +def main() -> int: + parser = base_argparser("Stage 5: Document-level asset extraction (C5)") + add_scope_selection_args(parser) + parser.add_argument( + "--max-workers", + type=int, + default=5, + help="Kept for CLI compatibility; Stage 5 extracts once at document level", + ) + args = parser.parse_args() + + from app.services.document_agent.pdf_text import read_page_texts + from app.services.page_memory.memory_service import ( + _project_assets_for_pages, + _resolve_asset_max_pages, + _select_rendered_pages_with_assets, + ) + from app.services.page_memory.page_assets import ( + extract_page_assets_from_renders, + get_asset_confidence_threshold, + page_asset_summary_enabled, + ) + from app.services.page_memory.page_renderer import render_document_pages + from toc_page_policy import TocPagePolicy + from shared.models.schemas.page_memory_config import PageMemoryConfig + + pdf_path, filename, out_dir = resolve_paths(args) + anatomy_cache = resolve_anatomy_cache_path(out_dir) + state_path = pipeline_state_path(out_dir) + scopes_dir = out_dir / "scopes" + + require_file( + anatomy_cache, + hint="Run Stage 1 first: uv run python scripts/page_memory/debug_pm_stage1_hierarchy.py --file ...", + ) + + anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) + page_count = anatomy.page_count + page_features = anatomy.page_features if anatomy else [] + toc_policy = TocPagePolicy.from_anatomy(anatomy) + + scope_ids = resolve_debug_scope_ids( + scopes_dir=scopes_dir, + scope_id=args.scope_id, + page_range=args.page_range, + fat_only=args.fat_only, + all_scopes=args.all_scopes, + list_scopes=args.list_scopes, + require_file="fine_hierarchy.json", + nonempty_json=True, + ) + logger.info("█" * 70) + logger.info(f" STAGE 5: DOCUMENT ASSET EXTRACTION — {filename}") + logger.info(f" OUTPUT: {out_dir}") + logger.info(f" SCOPES: {scope_ids}") + logger.info("█" * 70) + + t_start = time.time() + trace_stages: list[dict] = [] + token_cost_tracker = TokenCostTracker() + + pages = list(range(1, page_count + 1)) + page_texts = read_page_texts(pdf_path, pages, timeout=300) + logger.info(f" read {len(page_texts)} pages") + + scope_contexts: list[_ScopeAssetContext] = [] + for scope_id in scope_ids: + context = _load_scope_asset_context( + scope_id=scope_id, + scope_dir=scopes_dir / scope_id, + page_count=page_count, + toc_policy=toc_policy, + ) + if context is not None: + scope_contexts.append(context) + + union_pages = sorted( + {page for context in scope_contexts for page in context.pages} + ) + logger.info( + "🔬 document C5: {} unique pages across {} scopes", + len(union_pages), + len(scope_contexts), + ) + + rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=str(out_dir), + pages=union_pages, + page_features=page_features, + page_texts=page_texts, + ) + record_stage( + trace_stages, + "C1.render_pages", + page_info=page_scope_info([item.page_index for item in rendered]), + variables={"rendered_count": len(rendered)}, + ) + + asset_model = ( + getattr(args, "vlm_model", None) + or os.environ.get("PAGE_MEMORY_ASSET_MODEL") + or os.environ.get("IMAGE_MODEL") + ) + pm_config = PageMemoryConfig.default() + asset_max_pages = _resolve_asset_max_pages(page_count, pm_config) + asset_rendered = _select_rendered_pages_with_assets(rendered, page_features) + summary_enabled = page_asset_summary_enabled() + logger.info( + " C5: {}/{} rendered pages have coarse has_asset " + "(asset_max_pages={} summary_enabled={} model={})", + len(asset_rendered), + len(rendered), + asset_max_pages, + summary_enabled, + asset_model, + ) + + assets_by_page = extract_page_assets_from_renders( + pdf_path=pdf_path, + rendered_pages=asset_rendered, + output_dir=str(out_dir), + model_name=asset_model, + budget=None, + max_pages=asset_max_pages, + confidence_threshold=get_asset_confidence_threshold(), + summary_enabled=summary_enabled, + summary_concurrency=pm_config.asset_summary_concurrency, + table_engine=pm_config.table_engine, + table_merge_enabled=pm_config.table_merge_enabled, + ) + asset_count = sum(len(items) for items in assets_by_page.values()) + logger.info( + " C5: {} assets on {} pages (single document extraction)", + asset_count, + len(assets_by_page), + ) + record_stage( + trace_stages, + "C5.page_assets", + page_info=page_scope_info(sorted(assets_by_page)), + variables={ + "asset_count": asset_count, + "assets_by_page": { + page: [asset.asset_id for asset in assets] + for page, assets in assets_by_page.items() + }, + }, + ) + token_cost_tracker.snapshot_stage("C5.page_assets") + + write_debug_json(out_dir / "assets.json", _serialize_assets(assets_by_page)) + for context in scope_contexts: + projected = _project_assets_for_pages(assets_by_page, set(context.pages)) + write_scope_artifacts( + out_dir=out_dir, + scope_id=context.scope_id, + scope_manifest=context.scope_manifest, + hierarchy=context.skeletons, + tags=None, + assets_by_page=projected, + ) + + elapsed = time.time() - t_start + logger.info(f"✅ Stage 5 done in {elapsed:.1f}s") + logger.info( + f" {len(scope_contexts)} scopes, {asset_count} assets, " + f"{len(union_pages)} unique pages" + ) + update_pipeline_state( + state_path, + stage=5, + payload={ + "processed_scope_ids": [context.scope_id for context in scope_contexts], + "processed_scope_count": len(scope_contexts), + "asset_count": asset_count, + "asset_pages": sorted(assets_by_page), + "unique_pages": union_pages, + "document_assets": str(out_dir / "assets.json"), + "scope_artifacts": [ + str(scopes_dir / context.scope_id / "assets.json") + for context in scope_contexts + ], + }, + ) + + return stop_with_trace( + out_dir=out_dir, + stages=trace_stages, + stop_at="assets", + page_count=page_count, + pipeline_stage=5, + elapsed_s=elapsed, + token_cost_tracker=token_cost_tracker, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage6_tagging_finalize.py b/apps/worker/scripts/page_memory/debug_pm_stage6_tagging_finalize.py new file mode 100644 index 000000000..cec8ccb00 --- /dev/null +++ b/apps/worker/scripts/page_memory/debug_pm_stage6_tagging_finalize.py @@ -0,0 +1,697 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Stage 6: Canonical chunk assembly (C7) + finalize (C9). + +Loads the combined page tags produced by Stage 4, assembles canonical chunks, +and optionally produces chunks.json / doc_nav.json / manifest.json. + +Requires Stage 4 output: scopes//fine_hierarchy.json +Prefer Stage 5 document assets: assets.json +Legacy fallback: scopes//assets.json (deduped by asset_id) + +Usage: + cd apps/worker + uv run python scripts/page_memory/debug_pm_stage6_tagging_finalize.py --file /path/to/doc.pdf + uv run python scripts/page_memory/debug_pm_stage6_tagging_finalize.py --all-scopes --finalize + uv run python scripts/page_memory/debug_pm_stage6_tagging_finalize.py --scope-id p1-100 --finalize --run-db +""" + +import sys +from pathlib import Path as _Path + +sys.path.insert(0, str(_Path(__file__).resolve().parent)) + +import json +import os +import time +from collections import Counter +from pathlib import Path +from typing import Any, cast + +from loguru import logger + +from _debug_pm_shared import ( + ScopeResult, + TokenCostTracker, + add_scope_selection_args, + aggregate_stage_costs, + base_argparser, + build_production_job_metadata_from_stage_costs, + load_anatomy_cache, + resolve_anatomy_cache_path, + load_assets_artifact, + load_page_tags_artifact, + load_hierarchy_artifact, + load_stage_costs, + pipeline_state_path, + record_stage, + require_file, + resolve_debug_scope_ids, + resolve_paths, + sort_skeletons, + walk, + write_scope_artifacts, + write_top_level_artifacts, + stop_with_trace, + update_pipeline_state, + page_scope_info, + _scope_manifest, + _derive_hierarchy_page_scope, + _build_hierarchy_from_skeletons, +) + + +def _run_tagging_for_scope( + *, + scope_id: str, + scope_dir: Path, + pdf_path: str, + filename: str, + out_dir: Path, + page_count: int, + page_texts: dict[int, str], + page_features: list[Any], + toc_policy: Any, + args: Any, + token_cost_tracker: TokenCostTracker | None = None, +) -> ScopeResult: + """Load Stage-4 combined tags and rehydrate renders for final assembly.""" + from app.services.page_memory.page_renderer import render_document_pages + + scope_stages: list[dict[str, Any]] = [] + if token_cost_tracker is not None: + token_cost_tracker.register_child_thread() + + fine_hierarchy_path = scope_dir / "fine_hierarchy.json" + require_file(fine_hierarchy_path, hint=f"Run Stage 4 to produce {fine_hierarchy_path}") + prior_scope, active_skeletons = load_hierarchy_artifact(fine_hierarchy_path) + if not active_skeletons: + logger.warning(" [scope {}] no skeletons — skipping", scope_id) + return ScopeResult( + scope_id=scope_id, skeletons=[], tags=[], assets_by_page={}, + rendered=[], final_pages=[], scope_manifest={}, trace_stages=scope_stages, + ) + + tags_path = scope_dir / "page_tags.json" + require_file(tags_path, hint=f"Run Stage 4 to produce {tags_path}") + tags = load_page_tags_artifact(tags_path) + + # Load existing assets if available + assets_path = scope_dir / "assets.json" + assets_by_page: dict[int, list[Any]] = {} + if assets_path.exists(): + try: + assets_by_page = load_assets_artifact(assets_path) + except Exception as exc: + logger.warning( + " failed to load {}; continuing without assets: {}", + assets_path, + exc, + ) + + # Reuse Stage-4's exact scope contract. Fall back only for old artifacts. + recorded_pages = prior_scope.get("processing_pages") + final_pages = ( + [int(page) for page in recorded_pages] + if isinstance(recorded_pages, list) + else toc_policy.filter_processing_pages( + _derive_hierarchy_page_scope( + skeletons=active_skeletons, + page_count=page_count, + ) + ) + ) + recorded_excluded = prior_scope.get("excluded_toc_pages") + excluded_toc_pages = ( + [int(page) for page in recorded_excluded] + if isinstance(recorded_excluded, list) + else sorted(toc_policy.pure_toc_pages) + ) + scope_manifest = _scope_manifest( + scope_id=scope_id, + skeletons=active_skeletons, + page_count=page_count, + strategy="fine:finalize", + processing_pages=final_pages, + excluded_toc_pages=excluded_toc_pages, + ) + logger.info( + "🔬 [scope {}] loaded {} combined tags for {} processing pages", + scope_id, + len(tags), + len(final_pages), + ) + + # A separate debug process rehydrates the deterministic production renders. + rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=str(out_dir), + pages=final_pages, + page_features=page_features, + page_texts=page_texts, + ) + record_stage( + scope_stages, "C1.render_pages_rehydrated", + page_info=page_scope_info([r.page_index for r in rendered]), + variables={ + "scope_id": scope_id, + "rendered_count": len(rendered), + "tag_count": len(tags), + }, + ) + + # Preserve Stage-4 tags while attaching Stage-5 assets. + write_scope_artifacts( + out_dir=out_dir, + scope_id=scope_id, + scope_manifest=scope_manifest, + hierarchy=active_skeletons, + tags=tags, + assets_by_page=assets_by_page if assets_by_page else None, + ) + + return ScopeResult( + scope_id=scope_id, + skeletons=active_skeletons, + tags=tags, + assets_by_page=assets_by_page, + rendered=rendered, + final_pages=final_pages, + scope_manifest=scope_manifest, + trace_stages=scope_stages, + ) + + +def _build_report( + *, + filename: str, + anatomy, + toc_nodes: list, + skeletons: list, + tags: list, + chunks: list, + rendered: list, + elapsed: float, + token_cost_stages: list[dict[str, Any]] | None = None, +) -> str: + lines: list[str] = [] + ap = lines.append + + ap(f"# Page-Memory E2E Report — {filename}\n") + ap(f"- elapsed: **{elapsed:.1f}s**") + ap(f"- page_count: **{anatomy.page_count}**") + ap(f"- toc_pages: `{anatomy.toc_result.toc_pages}`") + + toc_node_rows = walk(toc_nodes) + toc_leaf_count = sum(1 for _, n in toc_node_rows if not n.children) + ap(f"- TitleNode: total **{len(toc_node_rows)}**, leaves **{toc_leaf_count}**\n") + + ap("## 1. Skeleton\n") + located = [s for s in skeletons if s.evidence.get("source") not in (None, "unlocated", "fallback_root")] + source_dist = Counter(s.evidence.get("source", "?") for s in skeletons) + ap(f"- leaf skeletons: **{len(skeletons)}**") + ap(f"- located: **{len(located)}** ({len(located)*100//max(len(skeletons),1)}%)") + ap(f"- source dist: `{dict(source_dist)}`\n") + + ap("## 2. Page Tags\n") + strategy_dist = Counter(t.strategy_used for t in tags) + ap(f"- total tagged: **{len(tags)}**") + ap(f"- strategy dist: `{dict(strategy_dist)}`\n") + + ap("## 3. Node Assembly\n") + type_dist = Counter(str(chunk.get("type", "?")) for chunk in chunks) + ap(f"- total chunks: **{len(chunks)}**") + ap(f"- type dist: `{dict(type_dist)}`\n") + + if token_cost_stages: + ap("\n## Token Cost\n") + ap("| stage | prompt | completion | calls | cost(USD) |") + ap("|-------|-------:|----------:|------:|----------:|") + t_pt = t_ct = t_calls = 0 + t_cost = 0.0 + for s in token_cost_stages: + pt = s.get("prompt_tokens", 0) + ct = s.get("completion_tokens", 0) + calls = s.get("calls", 0) + cost = float((s.get("cost") or {}).get("total_cost", 0)) + t_pt += pt + t_ct += ct + t_calls += calls + t_cost += cost + if calls: + ap(f"| {s['stage']} | {pt:,} | {ct:,} | {calls} | ${cost:.6f} |") + ap(f"| **total** | **{t_pt:,}** | **{t_ct:,}** | **{t_calls}** | **${t_cost:.6f}** |\n") + + return "\n".join(lines) + + +def main() -> int: + parser = base_argparser("Stage 6: Node assembly + finalize") + add_scope_selection_args(parser) + parser.add_argument( + "--max-workers", type=int, default=5, + help="Concurrent workers (default=5)", + ) + parser.add_argument( + "--finalize", action="store_true", + help="Run C9: chunks.json / doc_nav.json / manifest.json", + ) + parser.add_argument( + "--run-db", action="store_true", + help="Publish to local DB (implies --finalize)", + ) + parser.add_argument( + "--publish-job-id", default=None, + help="Explicit job_id for --run-db", + ) + parser.add_argument( + "--skip-assets", action="store_true", + help="With --run-db, skip asset upload", + ) + args = parser.parse_args() + if args.run_db and not args.finalize: + args.finalize = True + + from app.services.document_agent.pdf_text import read_page_texts + from app.services.document_agent.structure.hierarchy_locator import extract_toc_nodes + from app.services.page_memory.memory_service import ( + _append_toc_nav_skeletons, + _merge_static_toc_tags, + ) + from app.services.page_memory.node_assembler import ( + build_node_chunks, + build_toc_node_chunks, + merge_chunks_by_first_page, + ) + from app.services.page_memory.skeleton_extractor import collapse_single_child_chains + from shared.models.schemas.page_memory_config import PageMemoryConfig + + pdf_path, filename, out_dir = resolve_paths(args) + anatomy_cache = resolve_anatomy_cache_path(out_dir) + state_path = pipeline_state_path(out_dir) + scopes_dir = out_dir / "scopes" + + require_file( + anatomy_cache, + hint="Run Stage 1 first.", + ) + + anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) + page_count = anatomy.page_count + page_features = anatomy.page_features if anatomy else [] + from toc_page_policy import TocPagePolicy + + toc_policy = TocPagePolicy.from_anatomy(anatomy) + + scope_ids = resolve_debug_scope_ids( + scopes_dir=scopes_dir, + scope_id=args.scope_id, + page_range=args.page_range, + fat_only=args.fat_only, + all_scopes=args.all_scopes, + list_scopes=args.list_scopes, + require_file="fine_hierarchy.json", + nonempty_json=True, + ) + logger.info("█" * 70) + logger.info(f" STAGE 6: ASSEMBLY + FINALIZE — {filename}") + logger.info(f" OUTPUT: {out_dir}") + logger.info(f" SCOPES: {scope_ids}") + logger.info("█" * 70) + + t_start = time.time() + trace_stages: list[dict] = [] + token_cost_tracker = TokenCostTracker() + + # Read page texts + pages = list(range(1, page_count + 1)) + page_texts = read_page_texts(pdf_path, pages, timeout=300) + logger.info(f" read {len(page_texts)} pages") + + # ── Load combined tags and renders per scope ── + def _load_selected_scope(scope_id: str) -> ScopeResult: + return _run_tagging_for_scope( + scope_id=scope_id, + scope_dir=scopes_dir / scope_id, + pdf_path=pdf_path, + filename=filename, + out_dir=out_dir, + page_count=page_count, + page_texts=page_texts, + page_features=page_features, + toc_policy=toc_policy, + args=args, + token_cost_tracker=token_cost_tracker, + ) + + if args.max_workers > 1 and len(scope_ids) > 1: + import gevent + from gevent.pool import Pool as GeventPool + + logger.info( + " concurrent mode: {} workers × {} scopes", + args.max_workers, len(scope_ids), + ) + gpool = GeventPool(size=min(args.max_workers, len(scope_ids))) + greenlets = [ + gpool.spawn( + _load_selected_scope, + sid, + ) + for sid in scope_ids + ] + gevent.joinall(greenlets, raise_error=True) + scope_results = [cast(ScopeResult, g.value) for g in greenlets] + else: + logger.info(" serial mode: {} scope(s)", len(scope_ids)) + scope_results = [ + _load_selected_scope(sid) + for sid in scope_ids + ] + + # Merge trace + for sr in scope_results: + trace_stages.extend(sr.trace_stages) + + # ── Global merge ── + all_skeletons = sort_skeletons( + [skel for sr in scope_results for skel in sr.skeletons] + ) + tag_by_page: dict[int, Any] = {} + for sr in scope_results: + for t in sr.tags: + tag_by_page[t.page_index] = t + all_tags = sorted(tag_by_page.values(), key=lambda t: t.page_index) + + from app.services.page_memory.memory_service import _merge_assets_by_page + + document_assets_path = out_dir / "assets.json" + if document_assets_path.exists(): + try: + all_assets = load_assets_artifact(document_assets_path) + logger.info( + " loaded document assets.json: {} assets on {} pages", + sum(len(items) for items in all_assets.values()), + len(all_assets), + ) + except Exception as exc: + logger.warning( + " failed to load document assets.json ({}); falling back to scope assets", + exc, + ) + all_assets = _merge_assets_by_page( + sr.assets_by_page for sr in scope_results + ) + else: + all_assets = _merge_assets_by_page(sr.assets_by_page for sr in scope_results) + if all_assets: + logger.info( + " legacy fallback: merged {} scope asset groups → {} pages", + len(scope_results), + len(all_assets), + ) + + rendered_by_page: dict[int, Any] = {} + for sr in scope_results: + for rp in sr.rendered: + rendered_by_page[rp.page_index] = rp + all_rendered = sorted(rendered_by_page.values(), key=lambda rp: rp.page_index) + + active_pages = sorted({p for sr in scope_results for p in sr.final_pages}) + + # ── C4c: global collapse ── + pre_collapse = len(all_skeletons) + all_skeletons = collapse_single_child_chains(all_skeletons) + absorbed = pre_collapse - len(all_skeletons) + if absorbed: + logger.info( + "🪢 C4c collapse: {} → {} ({} absorbed)", + pre_collapse, len(all_skeletons), absorbed, + ) + record_stage( + trace_stages, + "C4c.collapse_single_child_chains", + variables={ + "pre_count": pre_collapse, + "post_count": len(all_skeletons), + "absorbed": absorbed, + }, + ) + + active_skeletons = all_skeletons + tags = _merge_static_toc_tags(all_tags, toc_policy) + nav_skeletons = _append_toc_nav_skeletons( + body_skeletons=active_skeletons, + anatomy=anatomy, + filename=filename, + ) + + # Write top-level artifacts + write_top_level_artifacts( + out_dir=out_dir, + hierarchy=nav_skeletons, + tags=tags, + assets_by_page=all_assets if all_assets else None, + ) + + # ── C7: Node assembly ── + logger.info("=" * 70) + logger.info("🧱 C7: assemble canonical chunks") + logger.info("=" * 70) + + tag_map = {t.page_index: t for t in tags} + render_map = {r.page_index: r for r in all_rendered} + raw_text_by_page: dict[int, str] = {} + image_path_by_page: dict[int, str] = {} + for page in active_pages: + rend = render_map.get(page) + raw_text_by_page[page] = ( + rend.raw_text if rend else page_texts.get(page, "") + ) or "" + if rend and rend.image_path and os.path.exists(rend.image_path): + image_path_by_page[page] = rend.image_path + + page_memory_config = PageMemoryConfig.default() + body_chunks = build_node_chunks( + skeletons=active_skeletons, + raw_text_by_page=raw_text_by_page, + image_path_by_page=image_path_by_page, + tag_by_page=tag_map, + filename=filename, + vlm_model=args.vlm_model or os.environ.get("IMAGE_MODEL"), + page_assets_by_page=all_assets if all_assets else None, + node_assembly_concurrency=page_memory_config.node_assembly_concurrency, + body_start_by_page=toc_policy.body_start_by_page(), + ) + toc_chunks = build_toc_node_chunks(anatomy=anatomy, filename=filename) + canonical_chunks = merge_chunks_by_first_page(toc_chunks, body_chunks) + from shared.services.chunks.canonical_chunk_builder import chunks_as_json + + chunks = cast( + list[dict[str, Any]], + chunks_as_json(canonical_chunks), + ) + logger.info(f" C7: {len(chunks)} canonical chunks") + record_stage( + trace_stages, + "C7.node_assembly", + page_info=page_scope_info(active_pages), + variables={"chunk_count": len(chunks)}, + ) + token_cost_tracker.snapshot_stage("C7.node_assembly") + + # ── C9: finalize ── + doc_nav: dict[str, Any] = {} + hierarchy_dict: dict[str, Any] = {} + if args.finalize: + logger.info("=" * 70) + logger.info("📦 C9: finalize (chunks → doc_nav → manifest)") + logger.info("=" * 70) + + from shared.services.storage.zip_doc_navigation import ZipDocNavigationBuilder + + t_fin = time.time() + type_dist = Counter(c.get("type", "?") for c in chunks) + logger.info(f" chunks: {len(chunks)} ({dict(type_dist)})") + + chunks_path = out_dir / "chunks.json" + chunks_path.write_text( + json.dumps({"chunks": chunks}, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + doc_nav = ZipDocNavigationBuilder().build_doc_nav(chunks, filename) + doc_nav_path = out_dir / "doc_nav.json" + doc_nav_path.write_text( + json.dumps(doc_nav, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + hierarchy_dict = ( + _build_hierarchy_from_skeletons(active_skeletons) + if active_skeletons + else {} + ) + try: + from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries + enrich_doc_nav_summaries(str(out_dir.parent), source_file=filename, use_llm=False) + except Exception as exc: + logger.warning(f" enrich failed (non-fatal): {exc}") + + logger.info(f" finalize done in {time.time() - t_fin:.1f}s") + record_stage( + trace_stages, + "C9.finalize", + page_info=page_scope_info(active_pages), + variables={ + "chunk_count": len(chunks), + "doc_nav_sections": len(doc_nav.get("sections", [])), + }, + ) + token_cost_tracker.snapshot_stage("C9.finalize") + + if args.run_db: + from scripts._debug_publish import publish_debug_result_dir + + publish_result = publish_debug_result_dir( + result_dir=out_dir, + source_file_name=filename, + chunks=chunks, + job_id=args.publish_job_id, + parse_track="page_memory", + upload_assets=not args.skip_assets, + ) + record_stage( + trace_stages, + "C10.debug_publish", + variables={"publish_result": publish_result.to_dict()}, + ) + + # ── Final trace + cross-stage cost rollup ── + toc_nodes = ( + extract_toc_nodes(anatomy.toc_hierarchies) if anatomy.toc_hierarchies else [] + ) + elapsed = time.time() - t_start + stop_with_trace( + out_dir=out_dir, + stages=trace_stages, + stop_at="finalize" if args.finalize else "assembly", + page_count=page_count, + pipeline_stage=6, + elapsed_s=elapsed, + token_cost_tracker=token_cost_tracker, + final_status="success", + extra_summary={ + "scope_ids": [sr.scope_id for sr in scope_results], + "chunk_count": len(chunks), + }, + ) + aggregated = aggregate_stage_costs(load_stage_costs(out_dir)) + report = _build_report( + filename=filename, + anatomy=anatomy, + toc_nodes=toc_nodes, + skeletons=active_skeletons, + tags=tags, + chunks=chunks, + rendered=all_rendered, + elapsed=float(aggregated.get("elapsed_s") or elapsed), + token_cost_stages=list( + ((aggregated.get("token_cost") or {}).get("by_stage") or []) + ), + ) + report_path = out_dir / "debug" / "report.md" + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(report, encoding="utf-8") + + if args.finalize: + from shared.services.storage.zip_manifest_schema import ZipManifestBuilder + + manifest = ZipManifestBuilder().generate_manifest( + job_id=filename, + data_id=None, + source_file_name=filename, + statistics=doc_nav.get("stats", {}), + job_metadata=build_production_job_metadata_from_stage_costs( + page_count=page_count, + ledger=load_stage_costs(out_dir), + ), + hierarchy=hierarchy_dict, + ) + (out_dir / "manifest.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + # Extra production ZIP: same members as the user-facing page_memory package. + # Leaves the full debug workspace untouched. + from shared.services.storage.zip_result_service import ZipResultService + + production_zip_name = "production_result.zip" + zip_path, checksum, statistics, zip_size = ZipResultService().generate_zip_package( + job_id=filename, + chunks=chunks, + add_dir=str(out_dir), + source_file_name=filename, + data_id=None, + job_metadata=build_production_job_metadata_from_stage_costs( + page_count=page_count, + ledger=load_stage_costs(out_dir), + ), + temp_dir=str(out_dir), + ) + desired_zip = out_dir / production_zip_name + generated_zip = Path(zip_path) + if generated_zip.resolve() != desired_zip.resolve(): + if desired_zip.exists(): + desired_zip.unlink() + generated_zip.replace(desired_zip) + zip_path = str(desired_zip) + record_stage( + trace_stages, + "C9.production_zip", + variables={ + "zip_path": zip_path, + "zip_size": zip_size, + "checksum": checksum, + "statistics": statistics, + }, + ) + logger.info(f" production ZIP → {zip_path} ({zip_size} bytes)") + + update_pipeline_state( + state_path, + stage=6, + payload={ + "processed_scope_ids": [sr.scope_id for sr in scope_results], + "finalized": bool(args.finalize), + "chunk_count": len(chunks) if chunks is not None else None, + "published": bool(args.run_db), + "elapsed_s": elapsed, + }, + ) + + logger.info("") + logger.info("═" * 70) + logger.info(f" ✅ DONE in {elapsed:.1f}s → {out_dir}") + logger.info( + " total pipeline: {:.1f}s / ${:.6f}", + float(aggregated.get("elapsed_s") or elapsed), + float( + ((aggregated.get("token_cost") or {}).get("total") or {}).get("total_cost") + or 0 + ), + ) + logger.info(f" report → {report_path}") + if args.finalize: + logger.info(f" chunks → {out_dir / 'chunks.json'}") + logger.info(f" doc_nav → {out_dir / 'doc_nav.json'}") + logger.info(f" prod ZIP → {out_dir / 'production_result.zip'}") + logger.info("═" * 70) + print(report) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/toc_page_policy.py b/apps/worker/scripts/page_memory/toc_page_policy.py new file mode 100644 index 000000000..4a15f5f50 --- /dev/null +++ b/apps/worker/scripts/page_memory/toc_page_policy.py @@ -0,0 +1,33 @@ +"""Debug-only TOC page policy for staged page_memory scripts. + +Production page_memory no longer ships this helper; debug stages still need a +single place to read ``toc_result.toc_pages`` and filter processing ranges. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True) +class TocPagePolicy: + pure_toc_pages: frozenset[int] = field(default_factory=frozenset) + + @classmethod + def from_anatomy(cls, anatomy: Any | None) -> TocPagePolicy: + toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None + pages = getattr(toc_result, "toc_pages", None) or [] + pure: set[int] = set() + for page in pages: + try: + pure.add(int(page)) + except (TypeError, ValueError): + continue + return cls(pure_toc_pages=frozenset(pure)) + + def filter_processing_pages(self, pages: list[int]) -> list[int]: + return [page for page in pages if page not in self.pure_toc_pages] + + def body_start_by_page(self) -> dict[int, float]: + return {} diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 66267f357..398da4e31 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -121,6 +121,105 @@ def test_parse_task_should_process_uploaded_file_through_real_contract_boundarie assert contract.find_task_workspaces(tmp_path, job["job_id"]) == [] +def test_parse_task_should_publish_each_non_null_source_path_once( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + contract = WorkerParseContract.create() + contract.use_workspace_root(monkeypatch, tmp_path) + contract.use_billing(monkeypatch, is_enabled=False) + + source_file_name = "contract-duplicate-path.xlsx" + job = contract.create_file_job( + source_file_name=source_file_name, + job_id_prefix="job_duplicate_path", + ) + contract.upload_source_file( + local_file_path=_SAMPLE_XLSX_PATH, + s3_key=job["s3_key"], + ) + + def fake_execute_document_parse( + *, + job_id: str, + job_context: object, + prepared_source: object, + output_dir: str, + ) -> ParseOutput: + del job_id, job_context, prepared_source + parsed_df = pd.DataFrame( + [ + { + "know_id": "first-shared-path", + "type": "text", + "content": "first shared-path chunk", + "path": f"{source_file_name}/Root/Shared", + }, + { + "know_id": "duplicate-shared-path", + "type": "text", + "content": "duplicate shared-path chunk", + "path": f"{source_file_name}/Root/Shared", + }, + { + "know_id": "nul-duplicate-shared-path", + "type": "text", + "content": "NUL duplicate shared-path chunk", + "path": f"{source_file_name}/Root\x00/Shared", + }, + { + "know_id": "first-null-path", + "type": "text", + "content": "first null-path chunk", + "path": "", + }, + { + "know_id": "second-null-path", + "type": "text", + "content": "second null-path chunk", + "path": "", + }, + ] + ) + return ParseOutput(output_dir=output_dir, parsed_df=parsed_df) + + monkeypatch.setattr( + "app.services.document_ingestion.processing_run.execute_document_parse", + fake_execute_document_parse, + ) + + celery_result = contract.enqueue_parse_task( + job_id=job["job_id"], + user_id=job["user_id"], + ) + + assert celery_result.successful() + assert celery_result.result["status"] == "success" + + observed = contract.observe_successful_job(job["job_id"]) + job_chunks = observed["job_chunks"] + document_chunks = observed["document_chunks"] + + assert [row["chunk_id"] for row in job_chunks] == [ + "first-shared-path", + "duplicate-shared-path", + "nul-duplicate-shared-path", + "first-null-path", + "second-null-path", + ] + assert [row["chunk_id"] for row in document_chunks] == [ + "first-shared-path", + "first-null-path", + "second-null-path", + ] + assert [row["source_chunk_path"] for row in document_chunks] == [ + f"{source_file_name}/Root/Shared", + None, + None, + ] + + def test_parse_task_result_zip_includes_page_citation_assets( worker_contract_environment: None, monkeypatch: pytest.MonkeyPatch, @@ -223,6 +322,87 @@ def fake_execute_document_parse( ] +def test_parse_task_sanitizes_nul_characters_at_database_boundary( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + contract = WorkerParseContract.create() + contract.use_workspace_root(monkeypatch, tmp_path) + contract.use_billing(monkeypatch, is_enabled=False) + + source_file = tmp_path / "nul-source.txt" + source_file.write_text("source", encoding="utf-8") + job = contract.create_file_job( + source_file_name=source_file.name, + job_id_prefix="job_parse_nul", + ) + contract.upload_source_file( + local_file_path=source_file, + s3_key=job["s3_key"], + ) + + def fake_execute_document_parse( + *, + job_id: str, + job_context: object, + prepared_source: object, + output_dir: str, + ) -> ParseOutput: + del job_id, job_context, prepared_source + parsed_df = pd.DataFrame( + [ + { + "know_id": "chunk-with-nul", + "type": "text", + "content": "Text\x00content", + "path": f"{source_file.name}/Root\x00/Introduction", + "length": 12, + "keywords": "", + "summary": "Summary\x00value", + "tokens": "", + "connectto": "", + "page_nums": "1", + "extra_metadata": json.dumps({"note": "Nested\x00value"}), + } + ] + ) + return ParseOutput(output_dir=output_dir, parsed_df=parsed_df) + + monkeypatch.setattr( + "app.services.document_ingestion.processing_run.execute_document_parse", + fake_execute_document_parse, + ) + + celery_result = contract.enqueue_parse_task( + job_id=job["job_id"], + user_id=job["user_id"], + ) + + assert celery_result.successful() + assert celery_result.result["status"] == "success" + + observed = contract.observe_successful_job(job["job_id"]) + job_chunk = observed["job_chunks"][0] + document_chunk = observed["document_chunks"][0] + + assert job_chunk["text"] == "Textcontent" + assert job_chunk["path"] == f"{source_file.name}/Root/Introduction" + assert job_chunk["chunk_metadata"]["summary"] == "Summaryvalue" + assert job_chunk["chunk_metadata"]["note"] == "Nestedvalue" + assert document_chunk["content"] == "Textcontent" + assert document_chunk["source_chunk_path"] == ( + f"{source_file.name}/Root/Introduction" + ) + + result_zip = contract.read_result_zip( + result_s3_key=observed["result"]["result_s3_key"], + tmp_path=tmp_path, + ) + archived_chunk = result_zip["chunks"]["chunks"][0] + assert archived_chunk["content"] == "Text\x00content" + + def test_parse_task_should_charge_user_when_billing_is_enabled( worker_contract_environment: None, monkeypatch: pytest.MonkeyPatch, diff --git a/apps/worker/tests/contract/test_processing_run_contract.py b/apps/worker/tests/contract/test_processing_run_contract.py new file mode 100644 index 000000000..d219627fb --- /dev/null +++ b/apps/worker/tests/contract/test_processing_run_contract.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from shared.core.exceptions.domain_exceptions import UnavailableException + + +def test_should_skip_duplicate_delivery_when_processing_lock_is_held( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import app.services.document_ingestion.processing_run as processing_run + + class FakeLock: + def __init__(self, _redis_service: object, _job_id: str) -> None: + pass + + def __enter__(self) -> "FakeLock": + raise UnavailableException( + internal_message="Could not acquire processing lock for job job-lock", + retry_after=120, + ) + + def __exit__(self, *_args: object) -> bool: + return False + + monkeypatch.setattr( + processing_run, + "load_parse_job_context", + lambda *args, **kwargs: SimpleNamespace(redis_service=object()), + ) + monkeypatch.setattr(processing_run, "mark_job_running", lambda *args: True) + monkeypatch.setattr( + processing_run, + "get_sync_job_lifecycle_service", + object, + ) + monkeypatch.setattr(processing_run, "RedisJobLock", FakeLock) + + result = processing_run.DocumentProcessingRun().execute( + job_id="job-lock", + user_id="contract-user", + ) + + assert result == { + "status": "skipped", + "job_id": "job-lock", + "reason": "job_already_processing", + } + + +def test_should_not_treat_other_unavailable_errors_as_lock_contention( + worker_contract_environment: None, +) -> None: + import app.services.document_ingestion.processing_run as processing_run + + error = UnavailableException( + internal_message="Job state is still settling", + retry_after=120, + ) + + assert processing_run._is_processing_lock_contention(error) is False diff --git a/apps/worker/tests/contract/test_profile_agent_protocol_contract.py b/apps/worker/tests/contract/test_profile_agent_protocol_contract.py index 60a4c07aa..63452c1ef 100644 --- a/apps/worker/tests/contract/test_profile_agent_protocol_contract.py +++ b/apps/worker/tests/contract/test_profile_agent_protocol_contract.py @@ -33,7 +33,6 @@ def test_planner_verdict_now_falls_through_to_ready_to_shard() -> None: "is_scanned": True, "category": "Feasibility Study Report", "routing_category": "generic", - "category_rationale": "scanned prose", "language": "zh", "rationale": "scanned PDF not atlas", "header_y": None, diff --git a/apps/worker/tests/contract/test_structure_anchoring_contract.py b/apps/worker/tests/contract/test_structure_anchoring_contract.py index d81465f65..11008b5cd 100644 --- a/apps/worker/tests/contract/test_structure_anchoring_contract.py +++ b/apps/worker/tests/contract/test_structure_anchoring_contract.py @@ -190,9 +190,7 @@ def test_anchor_hierarchy_uses_calibration_phase1() -> None: offset_status="ok", entry_indices=[0], samples=[ - CalibrationSample( - title="Only", printed_label=2, physical=4 - ) + CalibrationSample(title="Only", physical=4) ], ) ], @@ -276,7 +274,7 @@ def test_phase2_recalibrate_miss_drops_suffix_from_tree() -> None: offset_status="ok", entry_indices=[0, 1, 2, 3], samples=[ - CalibrationSample(title="Ch1", printed_label=1, physical=11) + CalibrationSample(title="Ch1", physical=11) ], ) ], @@ -358,7 +356,7 @@ def test_multi_regime_phase2_merges_physical_overrides() -> None: offset_status="ok", entry_indices=[0], samples=[ - CalibrationSample(title="Glossary", printed_label="iv", physical=20) + CalibrationSample(title="Glossary", physical=20) ], ), CalibrationRegime( @@ -367,7 +365,7 @@ def test_multi_regime_phase2_merges_physical_overrides() -> None: offset_status="ok", entry_indices=[1, 2], samples=[ - CalibrationSample(title="Summary", printed_label=1, physical=21) + CalibrationSample(title="Summary", physical=21) ], ), CalibrationRegime( @@ -376,9 +374,7 @@ def test_multi_regime_phase2_merges_physical_overrides() -> None: offset_status="ok", entry_indices=[3], samples=[ - CalibrationSample( - title="Financials", printed_label="F-1", physical=301 - ) + CalibrationSample(title="Financials", physical=301) ], ), ], diff --git a/deploy/ecs/README.md b/deploy/ecs/README.md index 564611fb3..e19200afe 100644 --- a/deploy/ecs/README.md +++ b/deploy/ecs/README.md @@ -2,9 +2,9 @@ These files are deployment templates for the shared `knowhere-fargate` cluster. They do not contain secret values and are not registered automatically. -The staging templates intentionally omit `S3_ACCESS_KEY_ID` and `S3_SECRET_ACCESS_KEY`. With `S3_TYPE=s3`, boto3 obtains temporary authenticated credentials from the ECS task role. +The templates intentionally omit `S3_ACCESS_KEY_ID` and `S3_SECRET_ACCESS_KEY`. With `S3_TYPE=s3`, boto3 obtains temporary authenticated credentials from the ECS task role. The same templates render staging or production values; the release workflow selects production only for a published release whose tag points to `main`. -The single staging Secrets Manager secret supplied to the renderer must be a JSON secret with these keys: +The environment-specific Secrets Manager secret supplied to the renderer must be a JSON secret with these keys: - API: `DATABASE_URL`, `REDIS_HOST`, `REDIS_PORT`, `REDIS_PASSWORD`, `CELERY_REDIS_URL`, `SECRET_KEY`, `DS_KEY`, `ALI_API_KEYS`, `ARK_API_KEY`, `GPT_API_KEY`, `MINERU_API_KEYS`, `STRIPE_SECRET_KEY`, `STRIPE_WEBHOOK_SECRET`, `WEBHOOK_MASTER_KEY`, `LOGFIRE_TOKEN`, `QSTASH_TOKEN`, `QSTASH_CURRENT_SIGNING_KEY`, `QSTASH_NEXT_SIGNING_KEY` - Worker: the API keys above plus `CELERY_REDIS_PASSWORD` and `ILOVEAPI_KEYS` @@ -17,8 +17,24 @@ WORKER_IMAGE=107424103509.dkr.ecr.us-east-1.amazonaws.com/knowhere/knowhere-work EXECUTION_ROLE_ARN=arn:aws:iam::107424103509:role/knowhere-fargate-staging-execution-role \ API_TASK_ROLE_ARN=arn:aws:iam::107424103509:role/knowhere-api-staging-task-role \ WORKER_TASK_ROLE_ARN=arn:aws:iam::107424103509:role/knowhere-worker-staging-task-role \ -STAGING_SECRETS_ARN=arn:aws:secretsmanager:us-east-1:107424103509:secret:knowhere/staging/runtime-... \ -python deploy/ecs/render_task_definitions.py --output-dir /tmp/knowhere-ecs-rendered +SECRETS_ARN=arn:aws:secretsmanager:us-east-1:107424103509:secret:knowhere/staging/runtime-... \ +DEPLOYMENT_ENVIRONMENT=staging \ +RUNTIME_ENVIRONMENT=staging \ +APP_ENV=staging \ +DB_SSL_MODE=require \ +API_DB_POOL_SIZE=5 \ +API_DB_MAX_OVERFLOW=5 \ +WORKER_DB_SYNC_POOL_SIZE=2 \ +WORKER_DB_SYNC_MAX_OVERFLOW=2 \ +S3_BUCKET_NAME=knowhere-storage-staging \ +INTERNAL_DASHBOARD_ENDPOINT=https://staging.knowhereto.ai \ +FRONTEND_URL=https://staging.knowhereto.ai \ +API_WEBHOOK_ENDPOINT=https://api-staging.knowhereto.ai/v1/internal/s3-events \ +SNS_TOPIC_ARN=arn:aws:sns:us-east-1:107424103509:knowhere-staging-s3-events \ +QSTASH_CALLBACK_BASE_URL=https://api-staging.knowhereto.ai/api/v1 \ +WORKER_CPU=2048 \ +WORKER_MEMORY=4096 \ +python deploy/ecs/render_task_definitions.py --environment staging --output-dir /tmp/knowhere-ecs-rendered ``` The output directory is deployment-only and must not be committed. The renderer fails on missing inputs, unresolved placeholders, or either long-lived S3 credential variable. @@ -33,6 +49,26 @@ The staging workflow in `.github/workflows/build-images.yml` expects these GitHu Before an ECS staging deployment, an operator must create and verify the ECS services, network configuration, API load-balancer target, CloudWatch log groups, and runtime secret. The workflow validates those resources and fails without registering or updating a service when any prerequisite is missing. It does not create or delete AWS resources. +## Production workflow prerequisites + +The production release path uses the same ECS cluster and renders the approved +2 vCPU / 4 GiB worker task size with `WORKER_CONCURRENCY=10`. It requires these +additional GitHub Actions secrets: + +- `AWS_ECS_PROD_EXECUTION_ROLE_ARN` +- `AWS_ECS_PROD_API_TASK_ROLE_ARN` +- `AWS_ECS_PROD_WORKER_TASK_ROLE_ARN` +- `AWS_ECS_PROD_SECRETS_ARN` +- `PRODUCTION_MIGRATION_DATABASE_URL` (direct migration-only database URL) + +Before publishing the production release, an operator must create and verify +the production runtime secret, IAM roles, API and worker ECS services, API +target group/listener rule, production ACM certificate attachment, and +`/ecs/knowhere-api-prod` and `/ecs/knowhere-worker-prod` log groups. The release +workflow validates these resources, registers immutable image-digest task +definitions, runs the production migration first, and then updates the ECS +services. It does not create or delete AWS resources. + ## Manual staging availability `.github/workflows/manage-staging.yml` exposes three manually dispatched diff --git a/deploy/ecs/render_task_definitions.py b/deploy/ecs/render_task_definitions.py index 4312533ce..458cb8a76 100644 --- a/deploy/ecs/render_task_definitions.py +++ b/deploy/ecs/render_task_definitions.py @@ -1,4 +1,4 @@ -"""Render staging ECS task-definition templates without storing secrets.""" +"""Render ECS task-definition templates without storing secrets.""" from __future__ import annotations @@ -17,7 +17,23 @@ "EXECUTION_ROLE_ARN", "API_TASK_ROLE_ARN", "WORKER_TASK_ROLE_ARN", - "STAGING_SECRETS_ARN", + "SECRETS_ARN", + "DEPLOYMENT_ENVIRONMENT", + "RUNTIME_ENVIRONMENT", + "APP_ENV", + "DB_SSL_MODE", + "API_DB_POOL_SIZE", + "API_DB_MAX_OVERFLOW", + "WORKER_DB_SYNC_POOL_SIZE", + "WORKER_DB_SYNC_MAX_OVERFLOW", + "S3_BUCKET_NAME", + "INTERNAL_DASHBOARD_ENDPOINT", + "FRONTEND_URL", + "API_WEBHOOK_ENDPOINT", + "SNS_TOPIC_ARN", + "QSTASH_CALLBACK_BASE_URL", + "WORKER_CPU", + "WORKER_MEMORY", ) FORBIDDEN_ENVIRONMENT_NAMES: Final[frozenset[str]] = frozenset( {"S3_ACCESS_KEY_ID", "S3_SECRET_ACCESS_KEY"} @@ -86,6 +102,12 @@ def parse_arguments() -> argparse.Namespace: """Parse renderer CLI arguments.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument( + "--environment", + choices=("staging", "production"), + required=True, + help="Deployment environment represented by the rendered definitions.", + ) return parser.parse_args() @@ -99,18 +121,25 @@ def load_variables() -> dict[str, str]: def main() -> None: - """Render both staging task definitions.""" + """Render both API and worker task definitions.""" arguments = parse_arguments() variables = load_variables() + expected_environment = "staging" if arguments.environment == "staging" else "prod" + if variables["DEPLOYMENT_ENVIRONMENT"] != expected_environment: + raise ValueError( + "DEPLOYMENT_ENVIRONMENT must match the selected environment: " + f"expected {expected_environment}, got {variables['DEPLOYMENT_ENVIRONMENT']}" + ) template_directory = Path(__file__).parent + output_suffix = variables["DEPLOYMENT_ENVIRONMENT"] render_template( template_directory / "task-definition-api.staging.json", - arguments.output_dir / "knowhere-api-staging.json", + arguments.output_dir / f"knowhere-api-{output_suffix}.json", variables, ) render_template( template_directory / "task-definition-worker.staging.json", - arguments.output_dir / "knowhere-worker-staging.json", + arguments.output_dir / f"knowhere-worker-{output_suffix}.json", variables, ) diff --git a/deploy/ecs/task-definition-api.staging.json b/deploy/ecs/task-definition-api.staging.json index 71025dd3d..0e1013c99 100644 --- a/deploy/ecs/task-definition-api.staging.json +++ b/deploy/ecs/task-definition-api.staging.json @@ -1,5 +1,5 @@ { - "family": "knowhere-api-staging", + "family": "knowhere-api-${DEPLOYMENT_ENVIRONMENT}", "taskRoleArn": "${API_TASK_ROLE_ARN}", "executionRoleArn": "${EXECUTION_ROLE_ARN}", "networkMode": "awsvpc", @@ -24,15 +24,15 @@ } ], "environment": [ - {"name": "ENVIRONMENT", "value": "staging"}, - {"name": "APP_ENV", "value": "staging"}, - {"name": "DB_SSL_MODE", "value": "require"}, - {"name": "DB_POOL_SIZE", "value": "5"}, - {"name": "DB_MAX_OVERFLOW", "value": "5"}, + {"name": "ENVIRONMENT", "value": "${RUNTIME_ENVIRONMENT}"}, + {"name": "APP_ENV", "value": "${APP_ENV}"}, + {"name": "DB_SSL_MODE", "value": "${DB_SSL_MODE}"}, + {"name": "DB_POOL_SIZE", "value": "${API_DB_POOL_SIZE}"}, + {"name": "DB_MAX_OVERFLOW", "value": "${API_DB_MAX_OVERFLOW}"}, {"name": "TMP_PATH", "value": "/tmp/aismart_bid"}, {"name": "S3_TYPE", "value": "s3"}, - {"name": "S3_BUCKET_NAME", "value": "knowhere-storage-staging"}, - {"name": "S3_RESULTS_BUCKET", "value": "knowhere-storage-staging"}, + {"name": "S3_BUCKET_NAME", "value": "${S3_BUCKET_NAME}"}, + {"name": "S3_RESULTS_BUCKET", "value": "${S3_BUCKET_NAME}"}, {"name": "S3_TEMP_PATH", "value": "/tmp"}, {"name": "S3_REGION", "value": "us-east-1"}, {"name": "S3_USE_SSL", "value": "true"}, @@ -48,32 +48,34 @@ {"name": "RATE_LIMIT_ENABLED", "value": "true"}, {"name": "TELEMETRY_ENABLED", "value": "false"}, {"name": "API_STANDALONE_MODE_ENABLED", "value": "true"}, - {"name": "INTERNAL_DASHBOARD_ENDPOINT", "value": "https://staging.knowhereto.ai"}, - {"name": "FRONTEND_URL", "value": "https://staging.knowhereto.ai"}, + {"name": "INTERNAL_DASHBOARD_ENDPOINT", "value": "${INTERNAL_DASHBOARD_ENDPOINT}"}, + {"name": "FRONTEND_URL", "value": "${FRONTEND_URL}"}, {"name": "BILLING_ENABLED", "value": "true"}, - {"name": "QSTASH_CALLBACK_BASE_URL", "value": "https://api-staging.knowhereto.ai/api/v1"}, + {"name": "QSTASH_CALLBACK_BASE_URL", "value": "${QSTASH_CALLBACK_BASE_URL}"}, {"name": "AWS_REGION", "value": "us-east-1"}, - {"name": "AWS_ACCOUNT_ID", "value": "107424103509"} + {"name": "AWS_ACCOUNT_ID", "value": "107424103509"}, + {"name": "API_WEBHOOK_ENDPOINT", "value": "${API_WEBHOOK_ENDPOINT}"}, + {"name": "SNS_TOPIC_ARN", "value": "${SNS_TOPIC_ARN}"} ], "secrets": [ - {"name": "DATABASE_URL", "valueFrom": "${STAGING_SECRETS_ARN}:DATABASE_URL::"}, - {"name": "REDIS_HOST", "valueFrom": "${STAGING_SECRETS_ARN}:REDIS_HOST::"}, - {"name": "REDIS_PORT", "valueFrom": "${STAGING_SECRETS_ARN}:REDIS_PORT::"}, - {"name": "REDIS_PASSWORD", "valueFrom": "${STAGING_SECRETS_ARN}:REDIS_PASSWORD::"}, - {"name": "CELERY_REDIS_URL", "valueFrom": "${STAGING_SECRETS_ARN}:CELERY_REDIS_URL::"}, - {"name": "SECRET_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:SECRET_KEY::"}, - {"name": "DS_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:DS_KEY::"}, - {"name": "ALI_API_KEYS", "valueFrom": "${STAGING_SECRETS_ARN}:ALI_API_KEYS::"}, - {"name": "ARK_API_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:ARK_API_KEY::"}, - {"name": "GPT_API_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:GPT_API_KEY::"}, - {"name": "MINERU_API_KEYS", "valueFrom": "${STAGING_SECRETS_ARN}:MINERU_API_KEYS::"}, - {"name": "STRIPE_SECRET_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:STRIPE_SECRET_KEY::"}, - {"name": "STRIPE_WEBHOOK_SECRET", "valueFrom": "${STAGING_SECRETS_ARN}:STRIPE_WEBHOOK_SECRET::"}, - {"name": "WEBHOOK_MASTER_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:WEBHOOK_MASTER_KEY::"}, - {"name": "LOGFIRE_TOKEN", "valueFrom": "${STAGING_SECRETS_ARN}:LOGFIRE_TOKEN::"}, - {"name": "QSTASH_TOKEN", "valueFrom": "${STAGING_SECRETS_ARN}:QSTASH_TOKEN::"}, - {"name": "QSTASH_CURRENT_SIGNING_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:QSTASH_CURRENT_SIGNING_KEY::"}, - {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"} + {"name": "DATABASE_URL", "valueFrom": "${SECRETS_ARN}:DATABASE_URL::"}, + {"name": "REDIS_HOST", "valueFrom": "${SECRETS_ARN}:REDIS_HOST::"}, + {"name": "REDIS_PORT", "valueFrom": "${SECRETS_ARN}:REDIS_PORT::"}, + {"name": "REDIS_PASSWORD", "valueFrom": "${SECRETS_ARN}:REDIS_PASSWORD::"}, + {"name": "CELERY_REDIS_URL", "valueFrom": "${SECRETS_ARN}:CELERY_REDIS_URL::"}, + {"name": "SECRET_KEY", "valueFrom": "${SECRETS_ARN}:SECRET_KEY::"}, + {"name": "DS_KEY", "valueFrom": "${SECRETS_ARN}:DS_KEY::"}, + {"name": "ALI_API_KEYS", "valueFrom": "${SECRETS_ARN}:ALI_API_KEYS::"}, + {"name": "ARK_API_KEY", "valueFrom": "${SECRETS_ARN}:ARK_API_KEY::"}, + {"name": "GPT_API_KEY", "valueFrom": "${SECRETS_ARN}:GPT_API_KEY::"}, + {"name": "MINERU_API_KEYS", "valueFrom": "${SECRETS_ARN}:MINERU_API_KEYS::"}, + {"name": "STRIPE_SECRET_KEY", "valueFrom": "${SECRETS_ARN}:STRIPE_SECRET_KEY::"}, + {"name": "STRIPE_WEBHOOK_SECRET", "valueFrom": "${SECRETS_ARN}:STRIPE_WEBHOOK_SECRET::"}, + {"name": "WEBHOOK_MASTER_KEY", "valueFrom": "${SECRETS_ARN}:WEBHOOK_MASTER_KEY::"}, + {"name": "LOGFIRE_TOKEN", "valueFrom": "${SECRETS_ARN}:LOGFIRE_TOKEN::"}, + {"name": "QSTASH_TOKEN", "valueFrom": "${SECRETS_ARN}:QSTASH_TOKEN::"}, + {"name": "QSTASH_CURRENT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_CURRENT_SIGNING_KEY::"}, + {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"} ], "healthCheck": { "command": ["CMD-SHELL", "curl -f http://localhost:5005/health || exit 1"], @@ -85,7 +87,7 @@ "logConfiguration": { "logDriver": "awslogs", "options": { - "awslogs-group": "/ecs/knowhere-api-staging", + "awslogs-group": "/ecs/knowhere-api-${DEPLOYMENT_ENVIRONMENT}", "awslogs-region": "us-east-1", "awslogs-stream-prefix": "api" } @@ -94,7 +96,7 @@ ], "tags": [ {"key": "Project", "value": "knowhere"}, - {"key": "Environment", "value": "staging"}, + {"key": "Environment", "value": "${DEPLOYMENT_ENVIRONMENT}"}, {"key": "Service", "value": "api"}, {"key": "ManagedBy", "value": "knowhere-api-infra"} ] diff --git a/deploy/ecs/task-definition-worker.staging.json b/deploy/ecs/task-definition-worker.staging.json index f0c2a709d..202339d1a 100644 --- a/deploy/ecs/task-definition-worker.staging.json +++ b/deploy/ecs/task-definition-worker.staging.json @@ -1,11 +1,11 @@ { - "family": "knowhere-worker-staging", + "family": "knowhere-worker-${DEPLOYMENT_ENVIRONMENT}", "taskRoleArn": "${WORKER_TASK_ROLE_ARN}", "executionRoleArn": "${EXECUTION_ROLE_ARN}", "networkMode": "awsvpc", "requiresCompatibilities": ["FARGATE"], - "cpu": "2048", - "memory": "4096", + "cpu": "${WORKER_CPU}", + "memory": "${WORKER_MEMORY}", "ephemeralStorage": {"sizeInGiB": 21}, "runtimePlatform": { "cpuArchitecture": "X86_64", @@ -18,16 +18,16 @@ "essential": true, "stopTimeout": 120, "environment": [ - {"name": "ENVIRONMENT", "value": "staging"}, - {"name": "APP_ENV", "value": "staging"}, - {"name": "DB_SSL_MODE", "value": "require"}, - {"name": "DB_SYNC_POOL_SIZE", "value": "2"}, - {"name": "DB_SYNC_MAX_OVERFLOW", "value": "2"}, + {"name": "ENVIRONMENT", "value": "${RUNTIME_ENVIRONMENT}"}, + {"name": "APP_ENV", "value": "${APP_ENV}"}, + {"name": "DB_SSL_MODE", "value": "${DB_SSL_MODE}"}, + {"name": "DB_SYNC_POOL_SIZE", "value": "${WORKER_DB_SYNC_POOL_SIZE}"}, + {"name": "DB_SYNC_MAX_OVERFLOW", "value": "${WORKER_DB_SYNC_MAX_OVERFLOW}"}, {"name": "WORKER_CONCURRENCY", "value": "10"}, {"name": "TMP_PATH", "value": "/tmp/aismart_bid"}, {"name": "S3_TYPE", "value": "s3"}, - {"name": "S3_BUCKET_NAME", "value": "knowhere-storage-staging"}, - {"name": "S3_RESULTS_BUCKET", "value": "knowhere-storage-staging"}, + {"name": "S3_BUCKET_NAME", "value": "${S3_BUCKET_NAME}"}, + {"name": "S3_RESULTS_BUCKET", "value": "${S3_BUCKET_NAME}"}, {"name": "S3_TEMP_PATH", "value": "/tmp"}, {"name": "S3_REGION", "value": "us-east-1"}, {"name": "S3_USE_SSL", "value": "true"}, @@ -45,10 +45,10 @@ {"name": "AWS_ACCOUNT_ID", "value": "107424103509"}, {"name": "RATE_LIMIT_ENABLED", "value": "false"}, {"name": "TELEMETRY_ENABLED", "value": "false"}, - {"name": "INTERNAL_DASHBOARD_ENDPOINT", "value": "https://staging.knowhereto.ai"}, - {"name": "API_WEBHOOK_ENDPOINT", "value": "https://api-staging.knowhereto.ai/v1/internal/s3-events"}, - {"name": "SNS_TOPIC_ARN", "value": "arn:aws:sns:us-east-1:107424103509:knowhere-staging-s3-events"}, - {"name": "QSTASH_CALLBACK_BASE_URL", "value": "https://api-staging.knowhereto.ai/api/v1"}, + {"name": "INTERNAL_DASHBOARD_ENDPOINT", "value": "${INTERNAL_DASHBOARD_ENDPOINT}"}, + {"name": "API_WEBHOOK_ENDPOINT", "value": "${API_WEBHOOK_ENDPOINT}"}, + {"name": "SNS_TOPIC_ARN", "value": "${SNS_TOPIC_ARN}"}, + {"name": "QSTASH_CALLBACK_BASE_URL", "value": "${QSTASH_CALLBACK_BASE_URL}"}, {"name": "HF_HOME", "value": "/mnt/models/huggingface"}, {"name": "TRANSFORMERS_CACHE", "value": "/mnt/models/huggingface"}, {"name": "BILLING_ENABLED", "value": "true"}, @@ -64,24 +64,24 @@ {"name": "ALL_DF_COLS", "value": "content,path,type,length,keywords,summary,know_id,tokens,connectto,addtime,page_nums"} ], "secrets": [ - {"name": "DATABASE_URL", "valueFrom": "${STAGING_SECRETS_ARN}:DATABASE_URL::"}, - {"name": "REDIS_HOST", "valueFrom": "${STAGING_SECRETS_ARN}:REDIS_HOST::"}, - {"name": "REDIS_PORT", "valueFrom": "${STAGING_SECRETS_ARN}:REDIS_PORT::"}, - {"name": "REDIS_PASSWORD", "valueFrom": "${STAGING_SECRETS_ARN}:REDIS_PASSWORD::"}, - {"name": "CELERY_REDIS_URL", "valueFrom": "${STAGING_SECRETS_ARN}:CELERY_REDIS_URL::"}, - {"name": "CELERY_REDIS_PASSWORD", "valueFrom": "${STAGING_SECRETS_ARN}:CELERY_REDIS_PASSWORD::"}, - {"name": "SECRET_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:SECRET_KEY::"}, - {"name": "WEBHOOK_MASTER_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:WEBHOOK_MASTER_KEY::"}, - {"name": "DS_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:DS_KEY::"}, - {"name": "ALI_API_KEYS", "valueFrom": "${STAGING_SECRETS_ARN}:ALI_API_KEYS::"}, - {"name": "ARK_API_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:ARK_API_KEY::"}, - {"name": "GPT_API_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:GPT_API_KEY::"}, - {"name": "MINERU_API_KEYS", "valueFrom": "${STAGING_SECRETS_ARN}:MINERU_API_KEYS::"}, - {"name": "ILOVEAPI_KEYS", "valueFrom": "${STAGING_SECRETS_ARN}:ILOVEAPI_KEYS::"}, - {"name": "LOGFIRE_TOKEN", "valueFrom": "${STAGING_SECRETS_ARN}:LOGFIRE_TOKEN::"}, - {"name": "QSTASH_TOKEN", "valueFrom": "${STAGING_SECRETS_ARN}:QSTASH_TOKEN::"}, - {"name": "QSTASH_CURRENT_SIGNING_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:QSTASH_CURRENT_SIGNING_KEY::"}, - {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${STAGING_SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"} + {"name": "DATABASE_URL", "valueFrom": "${SECRETS_ARN}:DATABASE_URL::"}, + {"name": "REDIS_HOST", "valueFrom": "${SECRETS_ARN}:REDIS_HOST::"}, + {"name": "REDIS_PORT", "valueFrom": "${SECRETS_ARN}:REDIS_PORT::"}, + {"name": "REDIS_PASSWORD", "valueFrom": "${SECRETS_ARN}:REDIS_PASSWORD::"}, + {"name": "CELERY_REDIS_URL", "valueFrom": "${SECRETS_ARN}:CELERY_REDIS_URL::"}, + {"name": "CELERY_REDIS_PASSWORD", "valueFrom": "${SECRETS_ARN}:CELERY_REDIS_PASSWORD::"}, + {"name": "SECRET_KEY", "valueFrom": "${SECRETS_ARN}:SECRET_KEY::"}, + {"name": "WEBHOOK_MASTER_KEY", "valueFrom": "${SECRETS_ARN}:WEBHOOK_MASTER_KEY::"}, + {"name": "DS_KEY", "valueFrom": "${SECRETS_ARN}:DS_KEY::"}, + {"name": "ALI_API_KEYS", "valueFrom": "${SECRETS_ARN}:ALI_API_KEYS::"}, + {"name": "ARK_API_KEY", "valueFrom": "${SECRETS_ARN}:ARK_API_KEY::"}, + {"name": "GPT_API_KEY", "valueFrom": "${SECRETS_ARN}:GPT_API_KEY::"}, + {"name": "MINERU_API_KEYS", "valueFrom": "${SECRETS_ARN}:MINERU_API_KEYS::"}, + {"name": "ILOVEAPI_KEYS", "valueFrom": "${SECRETS_ARN}:ILOVEAPI_KEYS::"}, + {"name": "LOGFIRE_TOKEN", "valueFrom": "${SECRETS_ARN}:LOGFIRE_TOKEN::"}, + {"name": "QSTASH_TOKEN", "valueFrom": "${SECRETS_ARN}:QSTASH_TOKEN::"}, + {"name": "QSTASH_CURRENT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_CURRENT_SIGNING_KEY::"}, + {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"} ], "healthCheck": { "command": ["CMD-SHELL", "python -c \"from shared.services.worker_health import assert_worker_healthy; assert_worker_healthy()\""], @@ -93,7 +93,7 @@ "logConfiguration": { "logDriver": "awslogs", "options": { - "awslogs-group": "/ecs/knowhere-worker-staging", + "awslogs-group": "/ecs/knowhere-worker-${DEPLOYMENT_ENVIRONMENT}", "awslogs-region": "us-east-1", "awslogs-stream-prefix": "worker" } @@ -102,7 +102,7 @@ ], "tags": [ {"key": "Project", "value": "knowhere"}, - {"key": "Environment", "value": "staging"}, + {"key": "Environment", "value": "${DEPLOYMENT_ENVIRONMENT}"}, {"key": "Service", "value": "worker"}, {"key": "ManagedBy", "value": "knowhere-api-infra"} ] diff --git a/deploy/ecs/test_manage_staging_workflow.py b/deploy/ecs/test_manage_staging_workflow.py index f51c2db98..c50f0a7a9 100644 --- a/deploy/ecs/test_manage_staging_workflow.py +++ b/deploy/ecs/test_manage_staging_workflow.py @@ -54,13 +54,48 @@ def test_start_restores_healthy_workers_before_api() -> None: start_block.index("wait_for_service knowhere-worker-staging") ) assert start_block.index("wait_for_service knowhere-worker-staging") < ( - start_block.index('if [ "$healthy_workers" -ne 2 ]') + start_block.index("wait_for_healthy_workers") ) - assert start_block.index('if [ "$healthy_workers" -ne 2 ]') < ( + assert start_block.index("wait_for_healthy_workers") < ( start_block.index("update_service knowhere-api-staging 1") ) assert start_block.index("update_service knowhere-api-staging 1") < ( - start_block.index("https://api-staging.knowhereto.ai/health") + start_block.index("wait_for_service knowhere-api-staging") + ) + assert start_block.index("wait_for_service knowhere-api-staging") < ( + start_block.index("wait_for_public_api_health") + ) + + +def test_worker_health_gate_polls_through_container_start_period() -> None: + """A newly running worker may remain health-unknown during startPeriod.""" + workflow: str = _read_workflow() + health_gate: str = workflow.split( + " wait_for_healthy_workers() {", maxsplit=1 + )[1].split("\n }", maxsplit=1)[0] + + assert 'worker_health_deadline="$((SECONDS + 300))"' in health_gate + assert "while (( SECONDS < worker_health_deadline )); do" in health_gate + assert "healthStatus==`HEALTHY`" in health_gate + assert 'if [ "$healthy_worker_count" -eq 2 ]; then' in health_gate + assert "sleep 15" in health_gate + assert 'echo "Timed out waiting for two healthy worker tasks"' in health_gate + + +def test_public_api_health_gate_retries_transient_alb_errors() -> None: + """A newly running API task may not yet have a healthy public ALB route.""" + workflow: str = _read_workflow() + health_gate: str = workflow.split( + " wait_for_public_api_health() {", maxsplit=1 + )[1].split("\n }", maxsplit=1)[0] + + assert 'api_health_deadline="$((SECONDS + 300))"' in health_gate + assert "while (( SECONDS < api_health_deadline )); do" in health_gate + assert "https://api-staging.knowhereto.ai/health" in health_gate + assert "return 0" in health_gate + assert "sleep 10" in health_gate + assert 'echo "Timed out waiting for the public staging API health endpoint"' in ( + health_gate ) @@ -95,7 +130,13 @@ def test_status_uses_the_jobs_ledger_for_backlog() -> None: def test_start_reports_readiness_time() -> None: """Operators can compare cold-start readiness with the schedule lead time.""" workflow: str = _read_workflow() + start_block: str = workflow.split(" start)", maxsplit=1)[1].split( + " ;;", maxsplit=1 + )[0] assert 'start_started_epoch="$(date +%s)"' in workflow + assert start_block.index("wait_for_public_api_health") < start_block.index( + 'startup_seconds="$(($(date +%s) - start_started_epoch))"' + ) assert 'startup_seconds="$(($(date +%s) - start_started_epoch))"' in workflow assert "startupSeconds" in workflow diff --git a/deploy/ecs/test_render_task_definitions.py b/deploy/ecs/test_render_task_definitions.py index 72b13457a..b9c1f5698 100644 --- a/deploy/ecs/test_render_task_definitions.py +++ b/deploy/ecs/test_render_task_definitions.py @@ -18,7 +18,23 @@ "EXECUTION_ROLE_ARN": "execution-role", "API_TASK_ROLE_ARN": "api-role", "WORKER_TASK_ROLE_ARN": "worker-role", - "STAGING_SECRETS_ARN": "secrets-arn", + "SECRETS_ARN": "secrets-arn", + "DEPLOYMENT_ENVIRONMENT": "staging", + "RUNTIME_ENVIRONMENT": "staging", + "APP_ENV": "staging", + "DB_SSL_MODE": "require", + "API_DB_POOL_SIZE": "5", + "API_DB_MAX_OVERFLOW": "5", + "WORKER_DB_SYNC_POOL_SIZE": "2", + "WORKER_DB_SYNC_MAX_OVERFLOW": "2", + "S3_BUCKET_NAME": "knowhere-storage-staging", + "INTERNAL_DASHBOARD_ENDPOINT": "https://staging.knowhereto.ai", + "FRONTEND_URL": "https://staging.knowhereto.ai", + "API_WEBHOOK_ENDPOINT": "https://api-staging.knowhereto.ai/v1/internal/s3-events", + "SNS_TOPIC_ARN": "arn:aws:sns:us-east-1:107424103509:knowhere-staging-s3-events", + "QSTASH_CALLBACK_BASE_URL": "https://api-staging.knowhereto.ai/api/v1", + "WORKER_CPU": "2048", + "WORKER_MEMORY": "4096", } SHARED_STAGING_ENVIRONMENT: dict[str, str] = { @@ -95,15 +111,19 @@ def test_staging_templates_render_without_long_lived_s3_keys( ], ) def test_staging_templates_preserve_expected_staging_configuration( + tmp_path: Path, template_name: str, container_name: str, expected_environment: dict[str, str], ) -> None: """Fargate preserves the verified staging settings captured on 2026-08-13.""" - definition_path: Path = TEMPLATE_DIRECTORY / template_name - definition: dict[str, object] = json.loads( - definition_path.read_text(encoding="utf-8") + output_path: Path = tmp_path / template_name + render_template( + TEMPLATE_DIRECTORY / template_name, + output_path, + RENDER_VARIABLES, ) + definition: dict[str, object] = json.loads(output_path.read_text(encoding="utf-8")) container_definitions: list[dict[str, object]] = definition[ "containerDefinitions" ] @@ -119,12 +139,15 @@ def test_staging_templates_preserve_expected_staging_configuration( assert environment_values[name] == value -def test_staging_worker_preserves_evidence_selected_capacity() -> None: +def test_staging_worker_preserves_evidence_selected_capacity(tmp_path: Path) -> None: """Worker capacity matches the staging load evidence recorded in issue 22.""" - definition_path: Path = TEMPLATE_DIRECTORY / "task-definition-worker.staging.json" - definition: dict[str, object] = json.loads( - definition_path.read_text(encoding="utf-8") + output_path: Path = tmp_path / "task-definition-worker.staging.json" + render_template( + TEMPLATE_DIRECTORY / "task-definition-worker.staging.json", + output_path, + RENDER_VARIABLES, ) + definition: dict[str, object] = json.loads(output_path.read_text(encoding="utf-8")) # The measured production-envelope replay rejected 1-vCPU tasks and passed # on two fixed 2-vCPU tasks, so CD must not restore the rejected task size. diff --git a/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py index 22983a1ec..36a60abf5 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py @@ -1,12 +1,13 @@ from __future__ import annotations -from typing import Any +from typing import Any, cast from uuid import uuid4 from sqlalchemy import delete, select from sqlalchemy.orm import Session from shared.models.database.job_result import JobChunk, JobResult +from shared.utils.json_utils import remove_nul_characters class SyncJobResultWriter: @@ -59,23 +60,24 @@ def replace_chunks( chunk_models = [] for index, chunk in enumerate(chunks): - chunk_identifier = chunk.get("chunk_id") or str(uuid4()) - metadata = chunk.get("metadata") - chunk_text = chunk.get("text") or chunk.get("content") + safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) + chunk_identifier = safe_chunk.get("chunk_id") or str(uuid4()) + metadata = safe_chunk.get("metadata") + chunk_text = safe_chunk.get("text") or safe_chunk.get("content") chunk_path = ( metadata.get("path") if isinstance(metadata, dict) and metadata.get("path") - else chunk.get("path") + else safe_chunk.get("path") ) chunk_models.append( JobChunk( job_result_id=job_result_id, - chunk_id=chunk_identifier, - chunk_type=chunk.get("type", "paragraph"), + chunk_id=str(chunk_identifier), + chunk_type=str(safe_chunk.get("type", "paragraph")), text=str(chunk_text) if chunk_text is not None else None, path=str(chunk_path) if chunk_path is not None else None, chunk_metadata=metadata, - sort_order=chunk.get("order", index), + sort_order=safe_chunk.get("order", index), ) ) db.add_all(chunk_models) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py index 2edfc2b66..3706ed998 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Literal +from typing import Any, Literal, cast from loguru import logger from sqlalchemy.orm import Session @@ -11,6 +11,7 @@ from shared.services.jobs.lifecycle.publication import SyncJobPublicationFinalizer from shared.services.jobs.lifecycle.result_writer import SyncJobResultWriter from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox +from shared.utils.json_utils import remove_nul_characters @dataclass(frozen=True) @@ -81,6 +82,17 @@ def finalize( section_summaries: dict[str, str] | None, document_top_summary: str | None = None, ) -> JobSuccessFinalization: + safe_chunks = cast( + list[dict[str, Any]], remove_nul_characters(chunks) + ) + safe_section_summaries = cast( + dict[str, str] | None, + remove_nul_characters(section_summaries), + ) + safe_document_top_summary = cast( + str | None, + remove_nul_characters(document_top_summary), + ) job_result = self._result_writer.upsert_job_result( db, job_id, @@ -89,14 +101,14 @@ def finalize( result_s3_key=result_s3_key, result_size=zip_size, ) - self._result_writer.replace_chunks(db, job_result.id, chunks) + self._result_writer.replace_chunks(db, job_result.id, safe_chunks) publication_outcome = self._publication_finalizer.publish_result( db, job_id=job_id, job_result_id=job_result.id, - chunks=chunks, - section_summaries=section_summaries, - document_top_summary=document_top_summary, + chunks=safe_chunks, + section_summaries=safe_section_summaries, + document_top_summary=safe_document_top_summary, ) transition_outcome = self._state_machine.mark_completed_outcome( diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index deed9135e..9ebbdd74a 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any +from typing import Any, cast from uuid import uuid4 from sqlalchemy import delete @@ -16,6 +16,35 @@ build_term_search_text, section_path_from_chunk_path, ) +from shared.utils.json_utils import remove_nul_characters + + +def deduplicate_chunks_by_source_path( + chunks: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Keep the first chunk for each non-null source path. + + The revision-path uniqueness constraint applies to non-null + ``source_chunk_path`` values. Text chunks without a source path are + intentionally retained because PostgreSQL permits multiple NULL values + for that constraint and those chunks can carry distinct content. + """ + seen_source_paths: set[str] = set() + deduplicated_chunks: list[dict[str, Any]] = [] + + for chunk in chunks: + safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) + source_path = _get_source_path( + chunk=safe_chunk, + chunk_metadata=_get_chunk_metadata(safe_chunk), + ) + if source_path is not None: + if source_path in seen_source_paths: + continue + seen_source_paths.add(source_path) + deduplicated_chunks.append(chunk) + + return deduplicated_chunks def replace_document_revision_content( @@ -31,8 +60,12 @@ def replace_document_revision_content( db=db, scope=scope, section_summaries=section_summaries, ) for index, chunk in enumerate(chunks): - chunk_metadata = _get_chunk_metadata(chunk) - source_path = _get_source_path(chunk=chunk, chunk_metadata=chunk_metadata) + safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) + chunk_metadata = _get_chunk_metadata(safe_chunk) + source_path = _get_source_path( + chunk=safe_chunk, + chunk_metadata=chunk_metadata, + ) section_path = section_path_from_chunk_path( source_path, source_file_name=scope.source_file_name, @@ -40,7 +73,7 @@ def replace_document_revision_content( section = section_publisher.ensure_section(section_path) db.add( _build_document_chunk( - chunk=chunk, + chunk=safe_chunk, chunk_metadata=chunk_metadata, source_path=source_path, section=section, @@ -60,7 +93,10 @@ def __init__( ) -> None: self._db = db self._scope = scope - self._section_summaries = section_summaries or {} + self._section_summaries = cast( + dict[str, str], + remove_nul_characters(section_summaries or {}), + ) self._sections_by_path: dict[str, DocumentSection] = {} def ensure_section(self, section_path: str) -> DocumentSection: @@ -179,7 +215,11 @@ def _get_source_path( chunk_metadata: dict[str, Any], ) -> str | None: source_path = chunk_metadata.get("path") or chunk.get("path") - return str(source_path) if source_path is not None else None + if source_path is None: + return None + + normalized_source_path = str(source_path) + return normalized_source_path or None def _get_sort_order(chunk: dict[str, Any], fallback_sort_order: int) -> int: diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 8911c3b7c..e201916e5 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -23,6 +23,7 @@ from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope from shared.services.retrieval.publication_content import ( + deduplicate_chunks_by_source_path, replace_document_revision_content, ) from shared.services.retrieval.publication_models import ( @@ -106,7 +107,7 @@ def _publish_document_state_for_job( ) document_metadata = JobMetadataHelper.get_document_metadata(job_metadata) - deduped_chunks = chunks + deduped_chunks = deduplicate_chunks_by_source_path(chunks) # If ALL chunks are duplicates → skip document creation entirely if not deduped_chunks: diff --git a/packages/shared-python/shared/utils/json_utils.py b/packages/shared-python/shared/utils/json_utils.py index fa2de8c73..275c6be32 100644 --- a/packages/shared-python/shared/utils/json_utils.py +++ b/packages/shared-python/shared/utils/json_utils.py @@ -7,6 +7,29 @@ from typing import Any, Mapping, MutableSet +def remove_nul_characters(value: object) -> object: + """Remove PostgreSQL-incompatible NUL characters from JSON-like values. + + Parser output is persisted both as JSON metadata and as text columns. The + PostgreSQL text encoders reject U+0000, so clean it at the persistence + boundary while preserving the shape and non-string scalar values of the + payload. + """ + if isinstance(value, str): + return value.replace("\x00", "") + if isinstance(value, Mapping): + return { + remove_nul_characters(key) if isinstance(key, str) else key: + remove_nul_characters(item) + for key, item in value.items() + } + if isinstance(value, list): + return [remove_nul_characters(item) for item in value] + if isinstance(value, tuple): + return tuple(remove_nul_characters(item) for item in value) + return value + + def make_json_safe( value: Any, *, max_preview_rows: int = 5, _visited: MutableSet[int] | None = None ) -> Any: