From 5a7763770d8fe186e4bd39696b1149ac1515c943 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 01:46:17 +0800 Subject: [PATCH 01/40] refactor: deepen service and retrieval architecture --- apps/api/app/api/v1/routes/billing.py | 479 +------ apps/api/app/api/v1/routes/jobs.py | 1026 +-------------- .../api/app/api/v1/routes/qstash_callbacks.py | 276 +--- apps/api/app/api/v1/routes/s3_events.py | 536 +------- .../services/billing/billing_app_service.py | 364 ++++++ apps/api/app/services/job_creation_service.py | 463 +++++++ apps/api/app/services/job_read_service.py | 173 +++ .../app/services/job_response_projection.py | 172 +++ .../job_upload_confirmation_service.py | 124 ++ .../app/services/qstash_callback_service.py | 234 ++++ apps/api/app/services/s3_events/__init__.py | 1 + .../app/services/s3_events/event_handlers.py | 192 +++ apps/api/app/services/s3_events/service.py | 52 + .../s3_events/signature_verification.py | 39 + .../s3_events/subscription_service.py | 47 + .../s3_events/upload_event_service.py | 96 ++ apps/api/main.py | 2 +- .../tests/contract/test_billing_contract.py | 20 +- .../contract/test_job_creation_contract.py | 13 +- .../contract/test_qstash_callback_contract.py | 24 +- .../tests/contract/test_retrieval_contract.py | 671 +++++++++- .../tests/contract/test_s3_event_contract.py | 20 +- apps/worker/app/core/tasks/kb_tasks.py | 618 +-------- .../services/workload/parse_job_service.py | 467 +++++++ .../services/workload/url_upload_service.py | 154 +++ .../contract/test_parse_task_contract.py | 95 +- .../contract/test_url_upload_contract.py | 30 +- .../services/retrieval/agent_navigate.py | 1137 ----------------- .../services/retrieval/agentic/asset_tools.py | 266 ++++ .../retrieval/agentic/discovery_tools.py | 294 +++++ .../services/retrieval/agentic/evidence.py | 345 +++++ .../retrieval/agentic/evidence_renderer.py | 182 +++ .../retrieval/agentic/knowledge_map.py | 93 ++ .../retrieval/agentic/navigation_tools.py | 564 ++++++++ .../retrieval/agentic/orchestrator.py | 374 +----- .../services/retrieval/agentic/prompts.py | 223 ++++ .../retrieval/agentic/section_tree.py | 454 +++++++ .../services/retrieval/agentic/tools.py | 966 +------------- .../shared/services/retrieval/app_service.py | 1046 ++------------- .../shared/services/retrieval/assets.py | 14 + .../services/retrieval/hit_stats_recorder.py | 68 + .../shared/services/retrieval/hydration.py | 549 ++++++++ .../shared/services/retrieval/ranking.py | 202 +++ .../services/retrieval/response_projection.py | 92 ++ .../services/retrieval/scoped_corpus.py | 102 ++ .../shared/services/retrieval/scoring.py | 94 ++ .../shared/services/retrieval/settings.py | 20 + .../retrieval/workflow/orchestrator.py | 34 +- 48 files changed, 7238 insertions(+), 6269 deletions(-) create mode 100644 apps/api/app/services/billing/billing_app_service.py create mode 100644 apps/api/app/services/job_creation_service.py create mode 100644 apps/api/app/services/job_read_service.py create mode 100644 apps/api/app/services/job_response_projection.py create mode 100644 apps/api/app/services/job_upload_confirmation_service.py create mode 100644 apps/api/app/services/qstash_callback_service.py create mode 100644 apps/api/app/services/s3_events/__init__.py create mode 100644 apps/api/app/services/s3_events/event_handlers.py create mode 100644 apps/api/app/services/s3_events/service.py create mode 100644 apps/api/app/services/s3_events/signature_verification.py create mode 100644 apps/api/app/services/s3_events/subscription_service.py create mode 100644 apps/api/app/services/s3_events/upload_event_service.py create mode 100644 apps/worker/app/services/workload/parse_job_service.py create mode 100644 apps/worker/app/services/workload/url_upload_service.py delete mode 100644 packages/shared-python/shared/services/retrieval/agent_navigate.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/asset_tools.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/evidence.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/prompts.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/section_tree.py create mode 100644 packages/shared-python/shared/services/retrieval/assets.py create mode 100644 packages/shared-python/shared/services/retrieval/hit_stats_recorder.py create mode 100644 packages/shared-python/shared/services/retrieval/hydration.py create mode 100644 packages/shared-python/shared/services/retrieval/ranking.py create mode 100644 packages/shared-python/shared/services/retrieval/response_projection.py create mode 100644 packages/shared-python/shared/services/retrieval/scoped_corpus.py create mode 100644 packages/shared-python/shared/services/retrieval/scoring.py create mode 100644 packages/shared-python/shared/services/retrieval/settings.py diff --git a/apps/api/app/api/v1/routes/billing.py b/apps/api/app/api/v1/routes/billing.py index 2ae4721bd..f9c2a2dd7 100644 --- a/apps/api/app/api/v1/routes/billing.py +++ b/apps/api/app/api/v1/routes/billing.py @@ -1,134 +1,74 @@ -""" -Billing API Routes -""" +"""Billing API routes.""" from typing import Optional -from app.services.billing.stripe_service import StripeService +from app.services.billing.billing_app_service import ( + ParseUsageResponse, + buy_credits_for_user, + buy_credits_package_for_user, + get_credits_balance_for_user, + get_parse_usage_overview_for_user, + get_price_configs_payload, + get_transaction_history_for_user, + get_usage_stats_for_user, + handle_stripe_webhook, +) from app.services.rate_limit.dependencies import CurrentUser, with_current_user from fastapi import APIRouter, Depends, Query, Request -from pydantic import BaseModel -from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.billing import MicroDollar -from shared.core.config import settings from shared.core.database import get_db -from shared.core.exceptions.domain_exceptions import StripeServiceException -from shared.models.database.credits_transaction import CreditsTransaction -from shared.models.database.job import Job -from shared.models.database.stripe_price_config import StripePriceConfig from shared.models.schemas.billing import ( BuyCreditsPackageRequest, BuyCreditsRequest, CheckoutSessionResponse, CreditsBalanceResponse, PaymentIntentResponse, - TransactionHistoryResponse, UsageStatsResponse, ) -from shared.services.billing import CreditsService router = APIRouter(tags=["Billing"]) -class ParseUsageResponse(BaseModel): - """Parse usage overview response""" - - request_total: int - mom_growth: float - credits_used: float - estimated_amount: Optional[float] - success_rate: float - avg_processing_time: float - - -@router.post("/buy-credits", summary="Buy Credits") +@router.post("/buy-credits", summary="Buy Credits", response_model=PaymentIntentResponse) async def buy_credits( request: BuyCreditsRequest, current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """Buy credits via Stripe payment intent""" - stripe_service = StripeService() - - try: - # Calculate amount (100 Credits = ¥2, i.e. 1 Credit = ¥0.02) - amount_cny = request.credits_amount * 0.02 # CNY amount - amount_cents = int(amount_cny * 100) # Convert to cents - - payment_intent = await stripe_service.create_payment_intent( - user_id=current_user.user_id, - amount=amount_cents, - credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, - currency="cny", - ) - - return PaymentIntentResponse( - client_secret=payment_intent["client_secret"], - payment_intent_id=payment_intent["payment_intent_id"], - ) - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to buy credits: {str(e)}" - ) +) -> PaymentIntentResponse: + return await buy_credits_for_user( + request=request, + user_id=current_user.user_id, + ) @router.get( - "/credits", summary="Get Credits Balance", response_model=CreditsBalanceResponse + "/credits", + summary="Get Credits Balance", + response_model=CreditsBalanceResponse, ) async def get_credits_balance( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """Get the current credits balance for the authenticated user""" - credits_service = CreditsService() - - try: - # Ensure user is initialized - await credits_service.ensure_user_initialized(db, current_user.user_id) - await db.commit() - - balance_micro_dollar = await credits_service.get_balance( - db, current_user.user_id - ) +) -> CreditsBalanceResponse: + return await get_credits_balance_for_user(db, user_id=current_user.user_id) - return CreditsBalanceResponse( - credits_balance=MicroDollar(balance_micro_dollar).to_credit() - ) - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get credits balance: {str(e)}" - ) - - -@router.get("/usage", summary="Get Usage Statistics") +@router.get( + "/usage", + summary="Get Usage Statistics", + response_model=UsageStatsResponse, +) async def get_usage_stats( period: str = "month", current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """Get usage statistics for the authenticated user""" - credits_service = CreditsService() - - try: - stats = await credits_service.get_usage_stats(db, current_user.user_id, period) - - return UsageStatsResponse( - period=stats["period"], - total_credits_used=MicroDollar(stats["total_used"]).to_credit(), - api_calls_count=stats["transaction_count"], - success_rate=95.0, # TODO: Calculate actual success rate from usage logs - average_response_time=stats.get("avg_response_time", 0), - top_endpoints=[], # TODO: Get top endpoints from usage logs - ) - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get usage statistics: {str(e)}" - ) +) -> UsageStatsResponse: + return await get_usage_stats_for_user( + db, + user_id=current_user.user_id, + period=period, + ) @router.get( @@ -139,85 +79,11 @@ async def get_usage_stats( async def parse_usage_overview( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """ - Returns a usage overview: - - Total request count (deprecated, always 0) - - Month-over-month growth (deprecated, always 0) - - Credits used (from credits_transactions table, including usage and refund types) - - Estimated amount (using the first credits_package unit price: amount_cents / (100 * credits_amount)) - - Success rate (jobs: done out of terminal-state jobs) - - Average processing time (jobs: updated_at - created_at, in seconds) - """ - try: - # request_total and mom_growth: UsageLog is deprecated, hardcoded to 0 - total_requests = 0 - mom_growth = 0.0 - - # Credits used: sum usage and refund types from credits_transactions - # Usage type is negative (deduction), refund type is positive (return) - # Net consumption = abs(sum(usage + refund)), then convert to display credits - credits_row = await db.execute( - select(func.coalesce(func.sum(CreditsTransaction.credits_amount), 0)) - .where(CreditsTransaction.user_id == current_user.user_id) - .where(CreditsTransaction.transaction_type.in_(["usage", "refund"])) - ) - # Cast Decimal to int is safe here because: - # 1. Source column is BigInteger (whole numbers only) - # 2. Postgres returns Decimal to avoid overflow - # 3. Sum of integers has no fractional part, so int() is lossless - total_micro_credits_used = int(abs(credits_row.scalar_one() or 0)) - - # Success rate & average processing time (terminal-state jobs only: done / failed) - job_row = await db.execute( - select( - func.count().filter(Job.status == "done").label("done_cnt"), - func.count() - .filter(Job.status.in_(["done", "failed"])) - .label("terminal_cnt"), - func.avg(func.extract("epoch", Job.updated_at - Job.created_at)) - .filter(Job.status.in_(["done", "failed"])) - .label("avg_secs"), - ).where(Job.user_id == current_user.user_id) - ) - job_stats = job_row.first() or (0, 0, 0.0) - done_cnt = getattr(job_stats, "done_cnt", 0) or 0 - terminal_cnt = getattr(job_stats, "terminal_cnt", 0) or 0 - success_rate = (done_cnt / terminal_cnt * 100) if terminal_cnt > 0 else 0.0 - avg_processing_time = round( - float(getattr(job_stats, "avg_secs", 0.0) or 0.0), 2 - ) - - # Estimated amount: use the first credits_package price config - price_row = await db.execute( - select(StripePriceConfig) - .where(StripePriceConfig.product_type == "credits_package") - .where(StripePriceConfig.is_active.is_(True)) - .order_by(StripePriceConfig.created_at) - .limit(1) - ) - price_cfg = price_row.scalar_one_or_none() - estimated_amount = None - if price_cfg and price_cfg.credits_amount and price_cfg.credits_amount > 0: - estimated_amount = round( - price_cfg.amount_cents - * total_micro_credits_used - / (100 * price_cfg.credits_amount), - 4, - ) - - return ParseUsageResponse( - request_total=total_requests or 0, - mom_growth=round(mom_growth, 2), - credits_used=MicroDollar(total_micro_credits_used).to_credit() or 0, - estimated_amount=estimated_amount, # in dollar - success_rate=round(success_rate, 2), - avg_processing_time=avg_processing_time, - ) - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get parse usage overview: {str(e)}" - ) +) -> ParseUsageResponse: + return await get_parse_usage_overview_for_user( + db, + user_id=current_user.user_id, + ) @router.get("/history", summary="Get Transaction History") @@ -226,262 +92,45 @@ async def get_transaction_history( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ): - """Get credits transaction history for the authenticated user""" - credits_service = CreditsService() - - try: - transactions = await credits_service.get_transaction_history( - db, current_user.user_id, limit - ) - - transaction_list = [ - TransactionHistoryResponse( - id=tx.id, - credits_amount=MicroDollar(tx.credits_amount).to_credit(), - transaction_type=tx.transaction_type, - description=tx.description, - created_at=tx.created_at, - ) - for tx in transactions - ] - - return transaction_list - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get transaction history: {str(e)}" - ) + return await get_transaction_history_for_user( + db, + user_id=current_user.user_id, + limit=limit, + ) @router.get("/price-configs", summary="Get Price Configurations") async def get_price_configs( product_type: Optional[str] = Query( - None, description="Product type: subscription or credits_package" + None, + description="Product type: subscription or credits_package", ), db: AsyncSession = Depends(get_db), -): - """Get price configuration list (subscriptions or credits packages)""" - try: - from app.services.billing.price_config_service import PriceConfigService +) -> dict[str, list[dict]]: + return await get_price_configs_payload(db, product_type=product_type) - price_config_service = PriceConfigService() - if product_type == "subscription": - # Get all subscription type configs - configs = await price_config_service.repository.get_all_active(db) - subscription_configs = [ - c for c in configs if c.product_type == "subscription" - ] - return { - "subscriptions": [ - { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": ( - config.extra_metadata.get( - "display_name", config.plan_id.upper() - ) - if config.extra_metadata - else config.plan_id.upper() - ), - "description": ( - config.extra_metadata.get("description", "") - if config.extra_metadata - else "" - ), - "features": ( - config.extra_metadata.get("features", []) - if config.extra_metadata - else [] - ), - "popular": ( - config.extra_metadata.get("frontend_config", {}).get( - "popular", False - ) - if config.extra_metadata - else False - ), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": config.extra_metadata or {}, - } - for config in subscription_configs - ], - "credits_packages": [], - } - elif product_type == "credits_package": - # Get all credits package configs - credits_configs = await price_config_service.get_all_credits_packages(db) - return { - "subscriptions": [], - "credits_packages": [ - { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": ( - config.extra_metadata.get( - "display_name", - f"{MicroDollar(config.credits_amount).to_credit()} Credits", - ) - if config.extra_metadata - else f"{MicroDollar(config.credits_amount).to_credit()} Credits" - ), - "description": ( - config.extra_metadata.get("description", "") - if config.extra_metadata - else "" - ), - "credits_amount": MicroDollar( - config.credits_amount - ).to_credit(), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": config.extra_metadata or {}, - } - for config in credits_configs - ], - } - else: - # Get all configs - configs = await price_config_service.repository.get_all_active(db) - subscriptions = [c for c in configs if c.product_type == "subscription"] - credits_packages = [ - c for c in configs if c.product_type == "credits_package" - ] - - return { - "subscriptions": [ - { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": ( - config.extra_metadata.get( - "display_name", config.plan_id.upper() - ) - if config.extra_metadata - else config.plan_id.upper() - ), - "description": ( - config.extra_metadata.get("description", "") - if config.extra_metadata - else "" - ), - "features": ( - config.extra_metadata.get("features", []) - if config.extra_metadata - else [] - ), - "popular": ( - config.extra_metadata.get("frontend_config", {}).get( - "popular", False - ) - if config.extra_metadata - else False - ), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": config.extra_metadata or {}, - } - for config in subscriptions - ], - "credits_packages": [ - { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": ( - config.extra_metadata.get( - "display_name", - f"{MicroDollar(config.credits_amount).to_credit()} Credits", - ) - if config.extra_metadata - else f"{MicroDollar(config.credits_amount).to_credit()} Credits" - ), - "description": ( - config.extra_metadata.get("description", "") - if config.extra_metadata - else "" - ), - "credits_amount": MicroDollar( - config.credits_amount - ).to_credit(), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": config.extra_metadata or {}, - } - for config in credits_packages - ], - } - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to get price configurations: {str(e)}" - ) - - -@router.post("/buy-credits-package", summary="Buy Credits Package by Price ID") +@router.post( + "/buy-credits-package", + summary="Buy Credits Package by Price ID", + response_model=CheckoutSessionResponse, +) async def buy_credits_package( request: BuyCreditsPackageRequest, current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), -): - """Buy a credits package by its Stripe price ID""" - from sqlalchemy import select - - from shared.models.database.user import User - - stripe_service = StripeService() - - try: - # Query user email from database - result = await db.execute( - select(User.email).where(User.id == current_user.user_id) - ) - user_email = result.scalar_one_or_none() - - frontend_url = settings.FRONTEND_URL - success_url = f"{frontend_url}/billing?success=true&type=credits_package" - cancel_url = f"{frontend_url}/billing?canceled=true" - - checkout_url = await stripe_service.create_checkout_session_for_credits_package( - db=db, - user_id=current_user.user_id, - price_id=request.price_id, - success_url=success_url, - cancel_url=cancel_url, - quantity=request.quantity, - email=user_email, - ) - - return CheckoutSessionResponse(checkout_url=checkout_url, session_id="") - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to create credits package purchase: {str(e)}" - ) +) -> CheckoutSessionResponse: + return await buy_credits_package_for_user( + db, + request=request, + user_id=current_user.user_id, + ) @router.post("/webhook", summary="Stripe Webhook") async def stripe_webhook(request: Request, db: AsyncSession = Depends(get_db)): - """Handle Stripe webhook events""" - stripe_service = StripeService() - - try: - payload = await request.body() - sig_header = request.headers.get("stripe-signature") - if not sig_header: - raise StripeServiceException( - internal_message="Missing stripe-signature header" - ) - - result = await stripe_service.handle_webhook(db, payload, sig_header) - - return result - - except Exception as e: - raise StripeServiceException( - internal_message=f"Failed to handle webhook: {str(e)}" - ) + return await handle_stripe_webhook( + db, + payload=await request.body(), + stripe_signature=request.headers.get("stripe-signature"), + ) diff --git a/apps/api/app/api/v1/routes/jobs.py b/apps/api/app/api/v1/routes/jobs.py index d7c6972a2..b6afb84f6 100644 --- a/apps/api/app/api/v1/routes/jobs.py +++ b/apps/api/app/api/v1/routes/jobs.py @@ -4,295 +4,39 @@ from __future__ import annotations -import os -import uuid -from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Literal, Optional, cast -from urllib.parse import urlparse +from datetime import datetime +from typing import Optional -from app.repositories.job_repository import JobRepository -from app.services.job_document_scope_service import ( - find_active_job_for_document, - is_active_document_job_unique_violation, - raise_document_ingestion_conflict, - resolve_effective_document_scope, +from app.services.job_creation_service import create_job_from_request +from app.services.job_read_service import ( + get_job_result_for_user, + list_jobs_for_user, ) -from app.services.knowledge.kb_orchestrator import KBOrchestrator from app.services.rate_limit.dependencies import ( CurrentUser, enforce_job_creation_capacity, require_billing_limits, with_current_user, ) -from app.services.state_machine import JobStateMachine +from app.services.job_upload_confirmation_service import confirm_job_upload from fastapi import APIRouter, Depends, Query, Request -from loguru import logger -from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.billing import MicroDollar -from shared.core.config import settings from shared.core.database import get_db -from shared.core.exceptions.domain_exceptions import ( - ConflictException, - JobOperationException, - NotFoundException, - PermissionDeniedException, - RateLimitException, - UnavailableException, - ValidationException, -) -from shared.core.exceptions.webhook_exceptions import WebhookConfigException -from shared.core.state_machine.states import JobStatus from shared.models.schemas.job import ( ConfirmUploadRequest, JobCreate, JobList, JobResponse, JobResultResponse, - StandardErrorObject, -) -from shared.services.storage.file_upload_service import FileUploadService -from shared.utils.utc_now import utc_now_naive -from shared.utils.url_security import ( - validate_http_url_and_resolve_ip_async, ) -from shared.utils.error_details import normalize_error_details -from shared.utils.url_file_type import resolve_file_extension_async router = APIRouter(tags=["Jobs"]) -JobStatusValue = Literal[ - "pending", "waiting-file", "running", "converting", "done", "failed" -] # ==================== Shared Helpers ==================== -def get_supported_formats() -> str: - """Return the supported file extensions as a comma-separated string.""" - return ", ".join(sorted(settings.get_supported_extensions())) - - -async def transition_to_uploaded( - db: AsyncSession, - job_id: str, - job_type: str, - trigger: str = "manual_upload_completed", -): - """ - Move the job into the uploaded flow. - - Args: - db: Database session. - job_id: Job identifier. - job_type: Job type. - trigger: Transition trigger. - """ - state_machine = JobStateMachine() - - # Once the upload is confirmed, transition the job to pending. - await state_machine.transition( - db, job_id, JobStatus.PENDING.value, trigger, None, "system" - ) - - -async def start_workflow_for_job( - db: AsyncSession, - job_id: str, - job_type: str, - source_type: str, - user_id: str, - file_path: Optional[str] = None, - file_url: Optional[str] = None, -): - """ - Start the workflow for a job. - - Args: - db: Database session. - job_id: Job identifier. - job_type: Job type. - source_type: Source type. - user_id: User identifier. - file_path: File path. - file_url: File URL. - """ - if job_type == "kb_management": - orchestrator = KBOrchestrator() - await orchestrator.start_workflow( - db=db, - job_id=job_id, - source_type=source_type, - file_path=file_path, - file_url=file_url, - user_id=user_id, - ) - else: - raise ValidationException( - user_message="Unsupported job type", - violations=[ - { - "field": "job_type", - "description": f"Job type '{job_type}' is not supported", - } - ], - ) - - -def check_job_permission(job, user_id: str, job_id: str) -> None: - """ - Verify that the job belongs to the current user. - - Args: - job: Job object. - user_id: Current user ID. - job_id: Requested job ID. - - Raises: - HTTPException: Raised when the user does not own the job. - """ - if not job: - raise NotFoundException( - resource="Job", resource_id=job_id, internal_message="Job not found" - ) - - if str(job.user_id) != user_id: - raise PermissionDeniedException( - user_message="You don't have permission to access this job", - ) - - -def _build_error_response( - job: Any, job_metadata: Optional[dict] = None -) -> Optional[StandardErrorObject]: - """ - Build StandardErrorObject for embedded error pattern. - - Args: - job: Job object with job_id, error_code, and error_message - job_metadata: Job metadata dict that may contain error_details - - Returns: - StandardErrorObject or None - """ - if not job.error_message: - return None - - # Extract error_details from job_metadata if present - error_details = None - if job_metadata and isinstance(job_metadata, dict): - error_details = normalize_error_details(job_metadata.get("error_details")) - - return StandardErrorObject( - code=job.error_code or "UNKNOWN", - message=job.error_message, - request_id=job.job_id, - details=error_details, - ) - - -def create_job_response( - job_id: str, - job, - source_type: str, - data_id: Optional[str], - namespace: Optional[str] = None, - document_id: Optional[str] = None, - upload_url: Optional[str] = None, - upload_headers: Optional[dict] = None, - expires_in: Optional[int] = None, -) -> JobResponse: - """ - Build a JobResponse object. - - Args: - job_id: Job identifier. - job: Job object. - source_type: Source type. - data_id: Data identifier. - upload_url: Upload URL in file mode. - upload_headers: Upload headers in file mode. - expires_in: Upload expiry in file mode. - - Returns: - JobResponse: Serialized job response payload. - """ - return JobResponse( - job_id=job_id, - status=job.status, - source_type=source_type, - data_id=data_id, - namespace=namespace, - document_id=document_id, - created_at=job.created_at, - upload_url=upload_url, - upload_headers=upload_headers, - expires_in=expires_in, - ) - - -def resolve_public_document_id(job) -> Optional[str]: - """Expose document_id only after it is published in the persisted job result.""" - job_result = getattr(job, "job_result", None) - published_document_id = getattr(job_result, "document_id", None) - if isinstance(published_document_id, str) and published_document_id: - return published_document_id - - return None - - -def validate_file_type(file_name: str) -> bool: - """ - Return whether the file extension is supported. - - Args: - file_name: File name. - - Returns: - bool: Whether the file type is supported. - """ - if not file_name: - return False - - file_extension = os.path.splitext(file_name)[1].lower() - - return file_extension in settings.get_supported_extensions() - - -def ensure_utc(dt: Optional[datetime]) -> Optional[datetime]: - """Normalize a datetime to UTC.""" - if not dt: - return None - if dt.tzinfo: - return dt.astimezone(timezone.utc) - return dt.replace(tzinfo=timezone.utc) - - -def normalize_naive_utc_filter_datetime(dt: Optional[datetime]) -> Optional[datetime]: - """Convert a query datetime into the naive UTC form used by database columns.""" - if dt is None: - return None - if dt.tzinfo is None or dt.utcoffset() is None: - return dt - return dt.astimezone(timezone.utc).replace(tzinfo=None) - - -def require_utc(dt: Optional[datetime], *, field_name: str) -> datetime: - """Normalize a required datetime to UTC.""" - normalized_dt = ensure_utc(dt) - if normalized_dt is None: - raise JobOperationException( - internal_message=f"Job is missing required datetime field: {field_name}" - ) - return normalized_dt - - -def to_job_status_value(status: str) -> JobStatusValue: - """Cast a persisted job status into the response literal type.""" - return cast(JobStatusValue, status) - - @router.post("", response_model=JobResponse, summary="Create a parsing job") @router.post("/", include_in_schema=False) async def create_job( # pyright: ignore[reportGeneralTypeIssues] @@ -304,358 +48,13 @@ async def create_job( # pyright: ignore[reportGeneralTypeIssues] """ Create a parsing job. """ - try: - job_id = f"job_{uuid.uuid4().hex[:12]}" - # Validate input parameters. - if payload.source_type == "file" and not payload.file_name: - raise ValidationException( - user_message="file_name is required when source_type is 'file'", - violations=[ - { - "field": "file_name", - "description": "Required for file source type", - } - ], - ) - if payload.source_type == "url" and not payload.source_url: - raise ValidationException( - user_message="source_url is required when source_type is 'url'", - violations=[ - { - "field": "source_url", - "description": "Required for url source type", - } - ], - ) - - # Validate webhook config if present - if payload.webhook: - # Check for URL validity - if payload.webhook.url: - validation_result = await validate_http_url_and_resolve_ip_async( - payload.webhook.url, - ) - if not validation_result.is_valid: - raise WebhookConfigException( - user_message="Invalid webhook URL", - internal_message=f"Webhook validation failed: {validation_result.error_message}", - ) - - # Validate the source file type. - if ( - payload.source_type == "file" - and payload.file_name - and not validate_file_type(payload.file_name) - ): - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported file type. Supported formats: {supported_formats}", - violations=[ - {"field": "file_name", "description": "File type not supported"} - ], - ) - elif payload.source_type == "url": - assert payload.source_url is not None - # Resolve file type from URL path or Content-Type header - file_ext = await resolve_file_extension_async(payload.source_url) - if not file_ext: - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported URL file type. Supported formats: {supported_formats}", - violations=[ - { - "field": "source_url", - "description": "URL file type not supported", - } - ], - ) - - job_type = "kb_management" - - # Keep job creation lightweight. - from shared.services.redis import RedisServiceFactory - - redis_service = RedisServiceFactory.get_service() - - # Build job_metadata without embedding user_config. - from shared.models.schemas.job_metadata import JobMetadataHelper - - job_metadata = JobMetadataHelper.create_from_request(payload) - requested_document_id = cast(Optional[str], job_metadata.get("document_id")) - if requested_document_id: - active_job = await find_active_job_for_document( - db, - user_id=current_user.user_id, - document_id=requested_document_id, - ) - if active_job is not None: - raise_document_ingestion_conflict( - document_id=requested_document_id, - active_job_id=active_job.job_id, - ) - ( - effective_document_id, - effective_namespace, - ) = await resolve_effective_document_scope( - db, - user_id=current_user.user_id, - document_id=requested_document_id, - requested_namespace=cast(Optional[str], payload.namespace), - ) - if not requested_document_id: - active_job = await find_active_job_for_document( - db, - user_id=current_user.user_id, - document_id=effective_document_id, - ) - if active_job is not None: - raise_document_ingestion_conflict( - document_id=effective_document_id, - active_job_id=active_job.job_id, - ) - job_metadata["document_id"] = effective_document_id - job_metadata["namespace"] = effective_namespace - - # Enforce Layers 2-3 immediately before DB insert so the row lock - # lifetime is limited to capacity check + create_job commit. - await enforce_job_creation_capacity( - request=http_request, - db=db, - current_user=current_user, - ) - - if payload.source_type == "file": - # File-upload mode: reserve the job row first. - assert payload.file_name is not None - file_extension = os.path.splitext(payload.file_name)[1] - s3_key = f"uploads/{job_id}{file_extension}" - job_metadata["source_file_name"] = payload.file_name - job_metadata["source_type"] = "file" - - # Create the waiting-file job row with the final S3 key in one insert. - job_repo = JobRepository() - try: - job = await job_repo.create_job( - db=db, - job_id=job_id, - user_id=current_user.user_id, - job_type=job_type, - source_type="file", - file_path=None, # The file has not been uploaded yet. - webhook_url=payload.webhook.url if payload.webhook else None, - metadata=job_metadata, - initial_state="waiting-file", - s3_key=s3_key, - ) - except IntegrityError as exc: - if is_active_document_job_unique_violation(exc): - raise_document_ingestion_conflict(document_id=effective_document_id) - raise - - if not job: - raise JobOperationException( - internal_message="Failed to create job in database" - ) - - # Generate the presigned upload URL. - upload_service = FileUploadService() - upload_info = await upload_service.generate_upload_url( - job_id, file_extension - ) - - # 3. Cache job_metadata in Redis for two hours. - from shared.services.redis.job_metadata_service import JobMetadataService - - metadata_service = JobMetadataService(redis_service) - await metadata_service.save_metadata(job_id, job_metadata) - - # 4. Cache the basic job info in Redis for two hours. - from datetime import datetime - - from shared.services.redis import JobInfoRedisService - - job_info_service = JobInfoRedisService(redis_service) - job_info = { - "job_id": job_id, - "s3_key": s3_key, - "user_id": current_user.user_id, - "webhook_enabled": bool(payload.webhook and payload.webhook.url), - "job_type": job_type, - "source_type": "file", - "created_at": datetime.now(timezone.utc).isoformat(), - } - await job_info_service.save_job_info(job_id, job_info) - - logger.info( - f"Job {job_id} upload_url returned to client: {upload_info['upload_url']}" - ) - - # Build the response payload. - response = create_job_response( - job_id=job_id, - job=job, - source_type="file", - data_id=payload.data_id, - namespace=effective_namespace, - upload_url=upload_info["upload_url"], - upload_headers=upload_info["upload_headers"], - expires_in=upload_info["expires_in"], - ) - - return response - - else: - # URL mode: create the job first, then download and upload asynchronously. - try: - assert payload.source_url is not None - # Resolve file extension (URL path first, then Content-Type header) - file_extension = await resolve_file_extension_async(payload.source_url) - if not file_extension: - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported URL file type. Supported formats: {supported_formats}", - violations=[ - { - "field": "source_url", - "description": "URL file type not supported", - } - ], - ) - - parsed_url = urlparse(payload.source_url) - url_basename = str(os.path.basename(parsed_url.path)) - # Ensure source_file_name carries the correct extension. - # URLs like arxiv.org/pdf/1706.03762 have no real extension in the path. - if ( - url_basename - and os.path.splitext(url_basename)[1].lower() == file_extension - ): - source_file_name = url_basename - elif url_basename: - source_file_name = f"{url_basename}{file_extension}" - else: - source_file_name = ( - f"url_file_{uuid.uuid4().hex[:8]}{file_extension}" - ) - - s3_key = f"uploads/{job_id}{file_extension}" - - job_metadata.update( - { - "source_file_name": source_file_name, - "source_url": payload.source_url, - "source_type": "url", - } - ) - - # Create the waiting-file job row; the file will arrive asynchronously. - job_repo = JobRepository() - try: - job = await job_repo.create_job( - db=db, - job_id=job_id, - user_id=current_user.user_id, - job_type=job_type, - source_type="url", - file_path=None, - webhook_url=payload.webhook.url if payload.webhook else None, - metadata=job_metadata, - initial_state=JobStatus.WAITING_FILE.value, # Reuse waiting-file for URL uploads. - s3_key=s3_key, # Precomputed target S3 key. - ) - except IntegrityError as exc: - if is_active_document_job_unique_violation(exc): - raise_document_ingestion_conflict( - document_id=effective_document_id - ) - raise - - if not job: - raise JobOperationException( - internal_message="Failed to create URL job in database" - ) - - # Cache job_metadata in Redis for two hours. - from shared.services.redis.job_metadata_service import ( - JobMetadataService, - ) - - metadata_service = JobMetadataService(redis_service) - await metadata_service.save_metadata(job_id, job_metadata) - - # Cache the basic job info in Redis for two hours. - from datetime import datetime - - from shared.services.redis import JobInfoRedisService - - job_info_service = JobInfoRedisService(redis_service) - job_info = { - "job_id": job_id, - "s3_key": s3_key, - "user_id": current_user.user_id, - "webhook_enabled": bool(payload.webhook and payload.webhook.url), - "job_type": job_type, - "source_type": "url", - "created_at": datetime.now(timezone.utc).isoformat(), - } - await job_info_service.save_job_info(job_id, job_info) - - # Start the URL download/upload task asynchronously in the worker. - from shared.core.celery_app import get_celery_app - - celery_app = get_celery_app() - upload_url_file_task = celery_app.signature( - "app.core.tasks.kb_tasks.upload_url_file_task" - ) - upload_url_file_task.apply_async( - args=[job_id, payload.source_url, current_user.user_id], - kwargs={ - "job_type": job_type, - }, - ) - - # Build the response payload. - response = create_job_response( - job_id=job_id, - job=job, - source_type="url", - data_id=payload.data_id, - namespace=effective_namespace, - ) - - return response - - except ValidationException: - raise - except WebhookConfigException: - raise - except ConflictException: - raise - except (RateLimitException, UnavailableException): - raise - except JobOperationException: - raise - except Exception as e: - logger.error(f"Failed to create URL job: {e}") - raise JobOperationException( - internal_message=f"URL job creation failed: {str(e)}" - ) - - except NotFoundException: - raise - except ValidationException: - raise - except ConflictException: - raise - except WebhookConfigException: - raise - except (RateLimitException, UnavailableException): - raise - except JobOperationException: - raise - except Exception as e: - logger.error(f"Failed to create job: {e}") - raise JobOperationException(internal_message=f"Job creation failed: {str(e)}") + return await create_job_from_request( + db, + payload=payload, + current_user=current_user, + enforce_capacity=enforce_job_creation_capacity, + request=http_request, + ) @router.get("", response_model=JobList, summary="List jobs") @@ -680,189 +79,17 @@ async def list_jobs( """ List jobs for the current user. """ - try: - job_repo = JobRepository() - - if recent_days not in (None, 1, 7, 30): - raise ValidationException( - user_message="recent_days only supports 1, 7, or 30", - violations=[{"field": "recent_days", "description": "Invalid value"}], - ) - created_after: Optional[datetime] = None - if recent_days: - created_after = utc_now_naive() - timedelta(days=recent_days) - - normalized_start_time = normalize_naive_utc_filter_datetime(start_time) - normalized_end_time = normalize_naive_utc_filter_datetime(end_time) - - if ( - normalized_start_time - and normalized_end_time - and normalized_start_time > normalized_end_time - ): - raise ValidationException( - user_message="start_time cannot be later than end_time", - violations=[ - {"field": "start_time", "description": "Must be before end_time"} - ], - ) - # start_time / end_time take priority over recent_days. - if normalized_start_time: - created_after = normalized_start_time - created_before = normalized_end_time - - # Count matching rows. - total_count = await job_repo.count_jobs_by_user( - db=db, - user_id=current_user.user_id, - created_after=created_after, - created_before=created_before, - job_type=job_type, - job_status=job_status, - ) - - # Fetch the matching jobs. - jobs = await job_repo.get_jobs_by_user( - db=db, - user_id=current_user.user_id, - limit=page_size, - offset=(page - 1) * page_size, - created_after=created_after, - created_before=created_before, - job_type=job_type, - job_status=job_status, - ) - - # Build the response payload. - job_responses = [] - upload_service = FileUploadService() - from shared.models.schemas.job_metadata import JobMetadataHelper - from shared.services.redis import RedisServiceFactory - - redis_service = RedisServiceFactory.get_service() - for job in jobs: - # Load job_metadata through the shared access path. - job_metadata = await job_repo.get_job_metadata( - db, job.job_id, redis_service - ) - job_result = job.job_result - status_for_api = to_job_status_value(job.status) - - result_url = None - result = None - result_url_expires_at = job.created_at # Default to created_at. - - if job_result and job_result.result_s3_key: - result_url_info = cast( - Dict[str, Any], - await upload_service.generate_download_url( - job_result.result_s3_key - ), - ) - result_url = result_url_info["download_url"] - - # Read checksum-only data from inline_payload when present. - if job_result.inline_payload: - result = job_result.inline_payload - - # Compute result_url_expires_at when a download URL was issued. - if result_url: - expires_in = int(result_url_info.get("expires_in", 3600)) - result_url_expires_at = utc_now_naive() + timedelta( - seconds=expires_in - ) - - original_request = ( - job_metadata.get("original_request") - if isinstance(job_metadata, dict) - else {} - ) - source_url = ( - original_request.get("source_url") - if isinstance(original_request, dict) - else None - ) - file_name = None - if source_url: - parsed_source = urlparse(source_url) - file_name = os.path.basename(parsed_source.path) or None - if not file_name and isinstance(original_request, dict): - file_name = original_request.get("file_name") - file_extension = None - if file_name: - ext = os.path.splitext(file_name)[1] - file_extension = ext[1:].upper() if ext else None - - parsing_params = {} - if isinstance(original_request, dict): - parsing_params = original_request.get("parsing_params") or {} - if not parsing_params and isinstance(job_metadata, dict): - parsing_params = job_metadata.get("parsing_params") or {} - model = ( - parsing_params.get("model") - if isinstance(parsing_params, dict) - else None - ) - ocr_enabled = ( - parsing_params.get("ocr_enabled") - if isinstance(parsing_params, dict) - else None - ) - - duration_seconds = None - if job.updated_at and job.created_at: - duration_seconds = (job.updated_at - job.created_at).total_seconds() - - job_responses.append( - JobResultResponse( - job_id=job.job_id, - namespace=JobMetadataHelper.get_field(job_metadata, "namespace"), - document_id=resolve_public_document_id(job), - status=status_for_api, - source_type=job.source_type, - data_id=JobMetadataHelper.get_field(job_metadata, "data_id"), - created_at=require_utc(job.created_at, field_name="created_at"), - progress=None, # The list view does not expose detailed progress. - error=_build_error_response(job, job_metadata), - result=result, - result_url=result_url, - result_url_expires_at=require_utc( - result_url_expires_at, - field_name="result_url_expires_at", - ), - file_name=file_name, - file_extension=file_extension, - model=model, - ocr_enabled=ocr_enabled, - duration_seconds=duration_seconds, - credits_spent=( - MicroDollar(job.credits_charged).to_credit() - if hasattr(job, "credits_charged") - else 0 - ), - ) - ) - - # Compute the total page count. - import math - - total_pages = math.ceil(total_count / page_size) if total_count > 0 else 0 - - response = JobList( - jobs=job_responses, - total=total_count, - page=page, - page_size=page_size, - total_pages=total_pages, - ) - - return response - - except Exception as e: - logger.error(f"Failed to list jobs: {e}") - raise JobOperationException( - internal_message=f"Failed to get job list: {str(e)}" - ) + return await list_jobs_for_user( + db, + user_id=current_user.user_id, + page=page, + page_size=page_size, + job_status=job_status, + job_type=job_type, + recent_days=recent_days, + start_time=start_time, + end_time=end_time, + ) @router.get("/{job_id}", response_model=JobResultResponse, summary="Get a job result") @@ -874,138 +101,11 @@ async def get_job_result( """ Return the result payload for one job. """ - try: - job_repo = JobRepository() - - # Load the job and verify access. - job = await job_repo.get_job_by_id(db, job_id) - check_job_permission(job, current_user.user_id, job_id) - assert job is not None - - status_for_api = to_job_status_value(job.status) - - # Load detailed progress from Redis while the job is running. - progress = None - if status_for_api == "running": - # TODO: Load detailed progress from Redis and convert it to the progress schema. - # from shared.services.redis import RedisServiceFactory - # redis_service = RedisServiceFactory.get_service() - # from shared.utils.redis_key_builder import redis_key_builder - - # progress_key = redis_key_builder.task_progress(job_id) - # progress = await redis_service.hgetall(progress_key) - progress = {"total_pages": 10, "processed_pages": 5} - - # Load job_metadata through the shared access path. - from shared.models.schemas.job_metadata import JobMetadataHelper - from shared.services.redis import RedisServiceFactory - - redis_service = RedisServiceFactory.get_service() - job_metadata = await job_repo.get_job_metadata(db, job_id, redis_service) - - # Result delivery fields. - job_result = job.job_result - result_url = None - result = None - result_url_expires_at = job.created_at # Default to created_at. - - if job_result and job_result.result_s3_key: - upload_service = FileUploadService() - result_url_info = cast( - Dict[str, Any], - await upload_service.generate_download_url(job_result.result_s3_key), - ) - result_url = result_url_info["download_url"] - expires_in = int(result_url_info["expires_in"]) - - # Read checksum/statistics data from inline_payload when present. - if job_result.inline_payload: - result = job_result.inline_payload - - # Compute result_url_expires_at when a download URL was issued. - if result_url: - from datetime import datetime, timedelta - - result_url_expires_at = datetime.now() + timedelta(seconds=expires_in) - - original_request = ( - job_metadata.get("original_request") - if isinstance(job_metadata, dict) - else {} - ) - source_url = ( - original_request.get("source_url") - if isinstance(original_request, dict) - else None - ) - file_name = None - if source_url: - parsed_source = urlparse(source_url) - file_name = os.path.basename(parsed_source.path) or None - if not file_name and isinstance(original_request, dict): - file_name = original_request.get("file_name") - file_extension = None - if file_name: - ext = os.path.splitext(file_name)[1] - file_extension = ext[1:].upper() if ext else None - - parsing_params = {} - if isinstance(original_request, dict): - parsing_params = original_request.get("parsing_params") or {} - if not parsing_params and isinstance(job_metadata, dict): - parsing_params = job_metadata.get("parsing_params") or {} - model = ( - parsing_params.get("model") if isinstance(parsing_params, dict) else None - ) - ocr_enabled = ( - parsing_params.get("ocr_enabled") - if isinstance(parsing_params, dict) - else None - ) - - response_data = JobResultResponse( - job_id=job.job_id, - namespace=JobMetadataHelper.get_field(job_metadata, "namespace"), - document_id=resolve_public_document_id(job), - status=status_for_api, - source_type=job.source_type, - data_id=JobMetadataHelper.get_field(job_metadata, "data_id"), - created_at=require_utc(job.created_at, field_name="created_at"), - progress=progress, - error=_build_error_response(job, job_metadata), - result=result, - result_url=result_url, - result_url_expires_at=require_utc( - result_url_expires_at, - field_name="result_url_expires_at", - ), - file_name=file_name, - file_extension=file_extension, - model=model, - ocr_enabled=ocr_enabled, - duration_seconds=( - (job.updated_at - job.created_at).total_seconds() - if job.updated_at and job.created_at - else None - ), - credits_spent=( - MicroDollar(job.credits_charged).to_credit() - if hasattr(job, "credits_charged") - else 0 - ), - ) - - return response_data - - except NotFoundException: - raise - except PermissionDeniedException: - raise - except Exception as e: - logger.error(f"Failed to get job result: {e}") - raise JobOperationException( - internal_message=f"Failed to get job result: {str(e)}" - ) + return await get_job_result_for_user( + db, + job_id=job_id, + user_id=current_user.user_id, + ) @router.post( @@ -1022,63 +122,9 @@ async def confirm_upload( """ Confirm a completed file upload as a fallback path. """ - try: - job_repo = JobRepository() - - # Load the job and verify access. - job = await job_repo.get_job_by_id(db, job_id) - check_job_permission(job, current_user.user_id, job_id) - assert job is not None - - # Check the current job state. - logger.info(f"Confirm upload - Job {job_id} current status: {job.status}") - if job.status not in [JobStatus.PENDING.value, JobStatus.WAITING_FILE.value]: - # If the webhook already advanced the job, return success idempotently. - logger.info(f"Job {job_id} already processed, status: {job.status}") - return {"message": "Job status already updated"} - - # Verify that the S3 object exists. - if not job.s3_key: - raise ValidationException( - user_message="Job is missing S3 key information", - violations=[ - {"field": "s3_key", "description": "S3 key not set for this job"} - ], - ) - - upload_service = FileUploadService() - file_info = await upload_service.verify_s3_file_exists(job.s3_key) - - if not file_info.get("exists"): - raise ValidationException( - user_message="S3 file does not exist, please upload the file first", - violations=[{"field": "file", "description": "File not found in S3"}], - ) - - # Advance the job state. - await transition_to_uploaded( - db, job_id, job.job_type, "manual_upload_completed" - ) - - # Start job processing. - await start_workflow_for_job( - db=db, - job_id=job_id, - job_type=job.job_type, - source_type="file", - user_id=current_user.user_id, - ) - - return {"message": "File upload confirmed; processing started"} - - except NotFoundException: - raise - except PermissionDeniedException: - raise - except ValidationException: - raise - except Exception as e: - logger.error(f"Failed to confirm upload: {e}") - raise JobOperationException( - internal_message=f"Failed to confirm upload: {str(e)}" - ) + return await confirm_job_upload( + db, + job_id=job_id, + request=request, + user_id=current_user.user_id, + ) diff --git a/apps/api/app/api/v1/routes/qstash_callbacks.py b/apps/api/app/api/v1/routes/qstash_callbacks.py index 84cae94f2..495cb22ea 100644 --- a/apps/api/app/api/v1/routes/qstash_callbacks.py +++ b/apps/api/app/api/v1/routes/qstash_callbacks.py @@ -1,263 +1,31 @@ -""" -QStash callback endpoints. - -These endpoints receive delivery status from Upstash QStash after it -delivers (or fails to deliver) a webhook to the customer's endpoint. - -Both endpoints verify the QStash JWT signature before processing. -""" +"""QStash callback endpoints.""" from __future__ import annotations -import json -from datetime import datetime, timezone -from typing import Any, Dict, Optional -from uuid import NAMESPACE_URL, uuid5 - +from app.services import qstash_callback_service from fastapi import APIRouter, Request, Response -from loguru import logger -from sqlalchemy import select - -from shared.core.config import app_config -from shared.core.database_sync import get_sync_db_context -from shared.models.database.webhook import WebhookEvent, WebhookEventStatus -from shared.models.database.webhook_log import WebhookLog router = APIRouter(tags=["QStash Callbacks"]) -def _get_qstash_verification_url(callback_path: str, request_url: str) -> str: - """Build the URL used for QStash signature verification. - - Prefer the configured public callback URL because ingress/TLS termination - can make ``request.url`` appear as an internal ``http://`` URL. - """ - callback_base_url = app_config.QSTASH_CALLBACK_BASE_URL - - if callback_base_url: - return f"{callback_base_url.rstrip('/')}{callback_path}" - - return request_url - - -def _verify_qstash_signature(raw_body: bytes, signature: str, url: str) -> bool: - """Verify the QStash JWT signature on an inbound callback.""" - current_key = app_config.QSTASH_CURRENT_SIGNING_KEY - next_key = app_config.QSTASH_NEXT_SIGNING_KEY - - if not current_key or not next_key: - logger.error("QStash signing keys not configured — rejecting callback") - return False - - try: - from qstash import Receiver - - receiver = Receiver( - current_signing_key=current_key, - next_signing_key=next_key, - ) - receiver.verify( - body=raw_body.decode("utf-8"), - signature=signature, - url=url, - ) - return True - except Exception as exc: - logger.warning( - "QStash signature verification failed: error_type={error_type}, url={url}", - error_type=type(exc).__name__, - url=url, - ) - return False - - -def _extract_callback_data(body: bytes) -> Dict[str, Any]: - """Parse the QStash callback body.""" - try: - return json.loads(body) - except (json.JSONDecodeError, ValueError): - return {"raw": body.decode("utf-8", errors="replace")} - - -def _normalize_header_value(value: Any) -> Optional[str]: - """Normalize a callback header value to a single string.""" - if isinstance(value, list): - if not value: - return None - first_value = value[0] - return first_value if isinstance(first_value, str) else str(first_value) - - if isinstance(value, str): - return value - - if value is None: - return None - - return str(value) - - -def _find_event_id(data: Dict[str, Any]) -> Optional[str]: - """Extract the Knowhere event ID from QStash sourceHeader.""" - source_header = data.get("sourceHeader", {}) or {} - event_id = _normalize_header_value( - source_header.get("X-Knowhere-Event-Id") - or source_header.get("x-knowhere-event-id") - ) - if not event_id: - for key, value in source_header.items(): - if key.lower() == "x-knowhere-event-id": - event_id = _normalize_header_value(value) - break - return event_id - - -def _build_callback_log_idempotency_key( - qstash_message_id: Optional[str], - event_id: str, -) -> str: - """Build a fixed-width idempotency key for webhook_logs. - - ``webhook_logs.idempotency_key`` is limited to 36 characters. QStash - ``sourceMessageId`` is longer, so store the raw value in - ``qstash_message_id`` and derive a stable UUID from it for the - idempotency key column. - """ - if qstash_message_id: - return str(uuid5(NAMESPACE_URL, qstash_message_id)) - - return event_id - - -def _get_response_status_code(value: Any) -> Optional[int]: - """Return the destination response status reported by QStash.""" - if value is None: - return None - - try: - return int(value) - except (TypeError, ValueError): - return None - - -def _is_success_response_status(status_code: Optional[int]) -> bool: - """Return whether a destination response status is successful.""" - return status_code is not None and 200 <= status_code < 300 - - -def _get_callback_event_status(data: Dict[str, Any]) -> str: - """Map a normal QStash callback to the current webhook event status.""" - response_status = _get_response_status_code(data.get("status")) - if _is_success_response_status(response_status): - return WebhookEventStatus.DELIVERED - - return WebhookEventStatus.DELIVERING - - -def _resolve_event_status(current_status: str, callback_status: str) -> str: - """Apply callback status without downgrading terminal delivery state.""" - if current_status in ( - WebhookEventStatus.DELIVERED, - WebhookEventStatus.FAILED, - WebhookEventStatus.CANCELED, - ): - return current_status - - return callback_status - - -def _process_qstash_callback( - data: Dict[str, Any], - event_id: str, - callback_status: str, - log_label: str, -) -> Response: - """Shared logic for both success and failure QStash callbacks. - - Fetches the WebhookEvent, updates its status, and writes a WebhookLog entry. - """ - response_status_code = _get_response_status_code(data.get("status")) - response_body = data.get("body", "") - qstash_message_id = data.get("sourceMessageId") - retried = data.get("retried", 0) - is_failed_delivery_attempt = ( - callback_status == WebhookEventStatus.FAILED - or ( - callback_status == WebhookEventStatus.DELIVERING - and not _is_success_response_status(response_status_code) - ) - ) - error_message = None - if is_failed_delivery_attempt: - error_message = data.get("error") or response_body - - with get_sync_db_context() as db: - event = db.execute( - select(WebhookEvent).where(WebhookEvent.id == event_id) - ).scalar_one_or_none() - - if not event: - logger.warning(f"QStash {log_label}: event {event_id} not found in DB") - return Response(status_code=200, content="OK (event not found)") - - now = datetime.now(timezone.utc).replace(tzinfo=None) - event_status = _resolve_event_status(event.status, callback_status) - attempt_number = retried + 1 - event.status = event_status - event.attempts = max(event.attempts, attempt_number) - event.updated_at = now - - log = WebhookLog( - job_id=event.job_id, - event_id=event.id, - webhook_url=event.target_url, - attempt_number=attempt_number, - request_payload=event.payload, - signature="", - idempotency_key=_build_callback_log_idempotency_key( - qstash_message_id, event.id - ), - response_status_code=response_status_code, - response_body=response_body[:4096] if response_body else None, - error_message=str(error_message)[:4096] if error_message else None, - duration_ms=0, - qstash_message_id=qstash_message_id, - ) - db.add(log) - db.commit() - - return Response(status_code=200, content="OK") - - @router.post("/qstash/callback") async def handle_qstash_callback(request: Request) -> Response: """Handle QStash success callback after webhook delivery.""" raw_body = await request.body() signature = request.headers.get("upstash-signature", "") - verification_url = _get_qstash_verification_url( + verification_url = qstash_callback_service.get_qstash_verification_url( "/webhooks/qstash/callback", str(request.url), ) - if not _verify_qstash_signature(raw_body, signature, verification_url): + if not qstash_callback_service.verify_qstash_signature( + raw_body, + signature, + verification_url, + ): return Response(status_code=401, content="Invalid signature") - data = _extract_callback_data(raw_body) - event_id = _find_event_id(data) - - if not event_id: - logger.warning("QStash callback: missing event_id, cannot correlate") - return Response(status_code=200, content="OK (no event_id)") - - retried = data.get("retried", 0) - logger.info( - f"QStash callback: event_id={event_id}, status={data.get('status')}, " - f"retried={retried}, qstash_message_id={data.get('sourceMessageId')}" - ) - - event_status = _get_callback_event_status(data) - return _process_qstash_callback( - data, event_id, event_status, "callback" - ) + return qstash_callback_service.handle_qstash_success_callback(raw_body) @router.post("/qstash/failure") @@ -265,28 +33,16 @@ async def handle_qstash_failure(request: Request) -> Response: """Handle QStash failure callback after all retries exhausted.""" raw_body = await request.body() signature = request.headers.get("upstash-signature", "") - verification_url = _get_qstash_verification_url( + verification_url = qstash_callback_service.get_qstash_verification_url( "/webhooks/qstash/failure", str(request.url), ) - if not _verify_qstash_signature(raw_body, signature, verification_url): + if not qstash_callback_service.verify_qstash_signature( + raw_body, + signature, + verification_url, + ): return Response(status_code=401, content="Invalid signature") - data = _extract_callback_data(raw_body) - event_id = _find_event_id(data) - - if not event_id: - logger.warning("QStash failure callback: missing event_id, cannot correlate") - return Response(status_code=200, content="OK (no event_id)") - - retried = data.get("retried", 0) - max_retries = data.get("maxRetries", 0) - logger.warning( - f"QStash failure: event_id={event_id}, status={data.get('status')}, " - f"retried={retried}/{max_retries}, qstash_message_id={data.get('sourceMessageId')}" - ) - - return _process_qstash_callback( - data, event_id, WebhookEventStatus.FAILED, "failure" - ) + return qstash_callback_service.handle_qstash_failure_callback(raw_body) diff --git a/apps/api/app/api/v1/routes/s3_events.py b/apps/api/app/api/v1/routes/s3_events.py index ee92dec6a..56d4bdf58 100644 --- a/apps/api/app/api/v1/routes/s3_events.py +++ b/apps/api/app/api/v1/routes/s3_events.py @@ -1,130 +1,13 @@ -""" -S3 event webhook routes. -""" +"""S3 event webhook routes.""" -import base64 -import json -import os -from typing import Any, Dict - -from app.repositories.job_repository import JobRepository -from app.services.knowledge.kb_orchestrator import KBOrchestrator -from app.services.state_machine import JobStateMachine +from app.services.s3_events.service import safely_handle_s3_event_post from fastapi import APIRouter, Header, Request from loguru import logger -from shared.core.database import get_db_context from shared.core.logging import LogEvent -from shared.core.state_machine.states import JobStatus -from shared.models.schemas.oss_event import OSSEvent -from shared.models.schemas.s3_event import S3Event -from shared.utils.pinned_outbound_http import ( - send_pinned_outbound_request, -) -from shared.utils.url_security import ( - validate_http_url_and_resolve_ip_async, -) router = APIRouter(tags=["Internal"]) -SNS_SUBSCRIPTION_TIMEOUT_SECONDS = 10 - - -def verify_sns_signature(request_body: bytes, signature: str, message: str) -> bool: - """ - Validate an SNS message signature. - - Args: - request_body: Request body. - signature: Signature header value. - message: Message payload. - - Returns: - bool: Whether validation succeeded. - """ - try: - # This is intentionally simplified. Production code should use the AWS SDK. - return True - except Exception as e: - logger.error(f"SNS signature verification failed: {e}") - return False - - -def verify_minio_signature(auth_token: str, expected_token: str) -> bool: - """ - Validate the MinIO webhook token. - - Args: - auth_token: Token supplied by the request. - expected_token: Token configured on the server. - - Returns: - bool: Whether validation succeeded. - """ - if not expected_token: - return True # Skip verification when no token is configured. - - return auth_token == expected_token - - -def verify_oss_signature(request_body: bytes, headers: Dict[str, str]) -> bool: - """ - Validate an OSS callback signature. - - Args: - request_body: Request body. - headers: Request headers. - - Returns: - bool: Whether validation succeeded. - """ - try: - from shared.core.config import settings - - # Allow an opt-out for local development and controlled environments. - if not getattr(settings, "OSS_EVENT_VERIFY_SIGNATURE", True): - return True - - # OSS callback verification is simplified here. Production code should - # follow the official OSS callback verification flow. - callback_key = getattr(settings, "OSS_EVENT_CALLBACK_KEY", "") - if not callback_key: - logger.warning( - "OSS_EVENT_CALLBACK_KEY is not configured; skipping signature verification" - ) - return True - - # TODO: Implement OSS callback signature verification. - # Expected steps: - # 1. Read the signature metadata from the headers. - # 2. Compute the signature with callback_key. - # 3. Compare the computed and provided signatures. - - return True - except Exception as e: - logger.error(f"OSS signature verification failed: {e}") - return False - - -def extract_job_id_from_s3_key(s3_key: str) -> str | None: - """ - Extract the job_id from an S3 object key. - - Args: - s3_key: S3 key in the format uploads/{job_id}.ext. - - Returns: - str: Job identifier. - """ - if not s3_key.startswith("uploads/"): - return None - - # Strip the uploads/ prefix and remove the file extension. - filename = s3_key[8:] # Remove the "uploads/" prefix. - job_id = os.path.splitext(filename)[0] - - return job_id - @router.get("/s3-events", response_model=dict, summary="Handle S3 webhook GET requests") async def handle_s3_events_get( @@ -132,16 +15,13 @@ async def handle_s3_events_get( x_amz_sns_message_type: str = Header(None, alias="x-amz-sns-message-type"), x_minio_auth_token: str = Header(None, alias="x-minio-auth-token"), authorization: str = Header(None), -): - """ - Handle S3-event GET requests, primarily for SNS subscription confirmation. - """ +) -> dict[str, str]: + """Handle S3-event GET requests, primarily for SNS subscription confirmation.""" logger.info("======== S3 event GET request ========") logger.info(f"Headers: {dict(request.headers)}") if request.client: logger.info(f"Client IP: {request.client.host}") - # Handle SNS subscription confirmation requests. if x_amz_sns_message_type == "SubscriptionConfirmation": logger.info("Received an SNS subscription confirmation request") return {"message": "SNS subscription confirmed"} @@ -157,411 +37,17 @@ async def handle_s3_events( x_amz_sns_message_type: str = Header(None, alias="x-amz-sns-message-type"), x_minio_auth_token: str = Header(None, alias="x-minio-auth-token"), authorization: str = Header(None), -): - """ - Handle S3-event POST requests from AWS SNS, MinIO, or OSS. - """ +) -> dict[str, str]: + """Handle S3-event POST requests from AWS SNS, MinIO, OSS, or tests.""" logger.bind(event=LogEvent.S3_WEBHOOK_EVENT).info( f"S3 event Headers: {dict(request.headers)}" ) if request.client: logger.info(f"Client IP: {request.client.host}") - try: - # Read the request body. - body = await request.body() - headers = dict(request.headers) - - # Determine the event source. - if x_amz_sns_message_type: - # AWS SNS event. - result = await handle_sns_event(body) - if result: - return result - elif _is_oss_event(headers): - # OSS event, including Aliyun MNS proxy notifications. - await handle_oss_event(body, headers) - elif x_minio_auth_token: - # MinIO event, identified by the dedicated x-minio-auth-token header. - await handle_minio_event(body, x_minio_auth_token) - else: - # Direct S3 event payload used in tests. - await handle_direct_s3_event(body) - - return {"message": "Event handled successfully"} - - except Exception as e: - logger.error(f"Failed to handle S3 event: {e}") - # Return 200 even on failure so the upstream storage service does not retry blindly. - return {"message": "Event handling completed"} - - -async def handle_sns_event(body: bytes): - """ - Handle an AWS SNS event payload. - """ - try: - # Parse the SNS message envelope. - sns_message = json.loads(body.decode("utf-8")) - - # Branch on the SNS message type. - message_type = sns_message.get("Type") - logger.info(f"SNS message type: {message_type}") - - if message_type == "SubscriptionConfirmation": - # Handle subscription confirmation. - logger.info("Received an SNS subscription confirmation request") - subscribe_url = sns_message.get("SubscribeURL") - if subscribe_url: - logger.info(f"SNS subscription confirmation URL: {subscribe_url}") - # Visit the URL to confirm the subscription. - return await confirm_sns_subscription(subscribe_url) - else: - logger.warning( - "SNS subscription confirmation did not include SubscribeURL" - ) - return {"message": "SNS subscription confirmation failed"} - - elif message_type == "Notification": - # Handle notification messages. - logger.info("Received an SNS notification") - logger.info(f"SNS message payload: {sns_message}") - - # Parse the embedded S3 event. - try: - s3_event_data = json.loads(sns_message["Message"]) - logger.info(f"S3 event payload: {s3_event_data}") - - # Skip S3 test events — AWS/LocalStack sends these when - # bucket notification configuration is first applied. - # They lack the standard Records[] structure. - if ( - isinstance(s3_event_data, dict) - and s3_event_data.get("Event") == "s3:TestEvent" - ): - logger.info("Skip S3 test event") - return {"message": "S3 test event confirmed and skipped"} - s3_event = S3Event(**s3_event_data) - - # Process the upload events. - await process_upload_events(s3_event) - except Exception as e: - logger.error(f"Failed to parse the S3 event payload: {e}") - logger.error(f"SNS payload: {sns_message}") - # Fall back to treating the SNS payload itself as an S3 event. - try: - s3_event = S3Event(**sns_message) - await process_upload_events(s3_event) - except Exception as e2: - logger.error( - f"Fallback parsing of the SNS payload as an S3 event also failed: {e2}" - ) - raise - else: - logger.warning(f"Unknown SNS message type: {message_type}") - return {"message": f"Unknown SNS message type: {message_type}"} - - except Exception as e: - logger.error(f"Failed to handle SNS event: {e}") - raise - - -async def confirm_sns_subscription(subscribe_url: str) -> dict[str, str]: - """Confirm an SNS subscription after SSRF validation and IP pinning.""" - validation = await validate_http_url_and_resolve_ip_async( - subscribe_url, + return await safely_handle_s3_event_post( + body=await request.body(), + headers=dict(request.headers), + sns_message_type=x_amz_sns_message_type, + minio_auth_token=x_minio_auth_token, ) - if not validation.is_valid: - logger.warning( - f"SNS subscription confirmation URL failed validation: {validation.error_message}" - ) - return {"message": "SNS subscription confirmation failed"} - - if not validation.validated_ip: - logger.warning("SNS subscription confirmation URL validation returned no IP") - return {"message": "SNS subscription confirmation failed"} - - try: - response = await send_pinned_outbound_request( - method="GET", - url=subscribe_url, - pinned_ip=validation.validated_ip, - timeout_seconds=SNS_SUBSCRIPTION_TIMEOUT_SECONDS, - ) - if response.status == 200: - logger.info("SNS subscription confirmed successfully") - return {"message": "SNS subscription confirmed"} - - if 300 <= response.status < 400: - logger.warning( - f"SNS subscription confirmation redirect blocked, status={response.status}" - ) - else: - logger.error( - f"SNS subscription confirmation failed, status={response.status}" - ) - return {"message": "SNS subscription confirmation failed"} - except Exception as e: - logger.error(f"Failed to reach the SNS confirmation URL: {e}") - return {"message": "SNS subscription confirmation failed"} - - -async def handle_minio_event(body: bytes, auth_token: str): - """ - Handle a MinIO webhook event. - """ - try: - # Validate the webhook token. - from shared.core.config import settings - - expected_token = getattr(settings, "S3_WEBHOOK_AUTH_TOKEN", "") - - if not verify_minio_signature(auth_token, expected_token): - logger.warning("MinIO webhook authentication failed") - return - - # Parse the S3 event payload. - s3_event_data = json.loads(body.decode("utf-8")) - s3_event = S3Event(**s3_event_data) - - # Process the upload events. - await process_upload_events(s3_event) - - except Exception as e: - logger.error(f"Failed to handle MinIO event: {e}") - - -async def handle_direct_s3_event(body: bytes): - """ - Handle a direct S3 event payload used in tests. - """ - try: - # Parse the S3 event payload. - s3_event_data = json.loads(body.decode("utf-8")) - s3_event = S3Event(**s3_event_data) - - # Process the upload events. - await process_upload_events(s3_event) - - except Exception as e: - logger.error(f"Failed to handle direct S3 event: {e}") - - -def _is_oss_event(headers: Dict[str, str]) -> bool: - """ - Return whether the incoming request looks like an OSS event. - - Args: - headers: Request headers. - - Returns: - bool: Whether the request matches OSS event heuristics. - """ - # Identify OSS events by storage type, known headers, or request shape. - - storage_type = os.getenv("S3_TYPE", "s3").lower() - if storage_type == "oss": - return True - - # Also look for OSS-specific headers such as x-oss-pub-key-url. - if "x-oss-pub-key-url" in headers: - return True - - # Recognize Aliyun MNS proxy headers and user agents. - if "x-mns-version" in headers or "x-mns-signing-cert-url" in headers: - return True - user_agent = headers.get("user-agent") or headers.get("User-Agent") - if user_agent and "Aliyun Notification Service Agent" in user_agent: - return True - - return False - - -async def handle_oss_event(body: bytes, headers: Dict[str, str]): - """ - Handle an OSS event payload. - """ - try: - # Verify the callback signature. - if not verify_oss_signature(body, headers): - logger.warning("OSS event signature verification failed") - return - - # Parse the OSS payload, including MNS wrapper envelopes. - event_data = json.loads(body.decode("utf-8")) - logger.info(f"OSS event payload: {event_data}") - # MNS may place the real event inside Message as base64 or raw JSON. - if isinstance(event_data, dict) and "Message" in event_data: - inner = event_data.get("Message") - if isinstance(inner, str): - decoded = None - # Prefer base64 decoding first. - try: - decoded_bytes = base64.b64decode(inner, validate=True) - decoded_str = decoded_bytes.decode("utf-8") - decoded = json.loads(decoded_str) - except Exception: - decoded = None - - if decoded is None: - # Fall back to parsing the raw JSON string directly. - try: - decoded = json.loads(inner) - except Exception: - decoded = None - - if decoded is not None: - event_data = decoded - logger.info(f"Decoded MNS Message payload: {event_data}") - elif isinstance(inner, dict): - event_data = inner - - # Detect the payload shape. - if "events" in event_data: - # Standard OSS event format. - oss_event = OSSEvent(**event_data) - elif "Records" in event_data: - # Compatibility path for S3-like payloads emitted by OSS. - oss_event = _convert_s3_format_to_oss(event_data) - else: - logger.error(f"Unknown OSS event format: {event_data}") - return - - # Convert to S3Event so the existing upload flow can be reused. - s3_event = oss_event.to_s3_event() - - # Process the upload events. - await process_upload_events(s3_event) - - except Exception as e: - logger.error(f"Failed to handle OSS event: {e}") - raise - - -def _convert_s3_format_to_oss(event_data: Dict[str, Any]) -> OSSEvent: - """ - Convert an S3-style event payload into an OSS event payload. - - Args: - event_data: S3-format event data. - - Returns: - OSSEvent: OSS event object. - """ - from shared.models.schemas.oss_event import OSSEventRecord - - # Convert each S3-style record into the OSS schema. - records = event_data.get("Records", []) - oss_records = [] - - for record in records: - oss_record = OSSEventRecord( - eventName=record.get("eventName", "").replace("s3:", ""), - eventSource="acs:oss", - eventTime=record.get("eventTime", ""), - region=record.get("awsRegion", ""), - oss={ - "bucket": record.get("s3", {}).get("bucket", {}), - "object": record.get("s3", {}).get("object", {}), - }, - ) - oss_records.append(oss_record) - - return OSSEvent(events=oss_records) - - -async def process_upload_events(s3_event: S3Event): - """ - Process upload events delivered by S3-compatible storage. - - Args: - s3_event: S3 event object. - """ - try: - # Gather only upload-related records. - upload_events = s3_event.get_upload_events() - - # Instantiate services once outside the loop - job_repo = JobRepository() - - for event in upload_events: - # Read the object key from the event record. - s3_key = event.object_key or event.s3.get("object", {}).get("key") - if not s3_key: - continue - - # Extract the job_id from the object key. - job_id = extract_job_id_from_s3_key(s3_key) - if not job_id: - logger.warning(f"Could not extract job_id from S3 key: {s3_key}") - continue - - logger.info(f"Processing S3 upload event: {s3_key} -> job_id={job_id}") - - # Load the matching job. - async with get_db_context() as db: - job = await job_repo.get_job_by_id(db, job_id) - - if not job: - logger.warning(f"No job found for upload event: {job_id}") - continue - - # Only react while the job is still waiting for file upload. - if job.status != "waiting-file": - logger.info( - f"Job {job_id} is not in waiting-file status: {job.status}" - ) - continue - - # Check if upload window has expired (race-condition safe via optimistic lock) - from shared.core.config import settings - from shared.core.state_machine.states import is_job_expired - - if is_job_expired(job.updated_at, settings.JOB_WAITING_EXPIRE_SECONDS): - logger.warning(f"Job {job_id} upload expired, marking failed") - state_machine = JobStateMachine() - await state_machine.mark_failed( - db, - job_id, - "Upload expired: file was not uploaded within the allowed time window", - error_code="UPLOAD_EXPIRED", - ) - continue - - # Skip S3 file verification — we are processing the upload - # notification itself, so the file is guaranteed to exist. - - # Advance the job state. - state_machine = JobStateMachine() - - # Once upload is complete, move the job to pending. - await state_machine.transition( - db, - job_id, - JobStatus.PENDING.value, - "s3_upload_completed", - None, - "system", - ) - - # Start job processing. - if job.job_type == "kb_management": - orchestrator = KBOrchestrator() - await orchestrator.start_workflow( - db=db, - job_id=job_id, - source_type="file", - file_path=None, - file_url=None, - user_id=str(job.user_id), - ) - else: - logger.warning( - f"Unsupported job type for upload event: {job.job_type}, job_id={job_id}" - ) - - logger.info(f"Triggered processing for job {job_id}") - - except Exception as e: - logger.error(f"Failed to process upload events: {e}") - raise diff --git a/apps/api/app/services/billing/billing_app_service.py b/apps/api/app/services/billing/billing_app_service.py new file mode 100644 index 000000000..c7c3e6b5b --- /dev/null +++ b/apps/api/app/services/billing/billing_app_service.py @@ -0,0 +1,364 @@ +from __future__ import annotations + +from typing import Optional + +from app.services.billing.price_config_service import PriceConfigService +from app.services.billing.stripe_service import StripeService +from pydantic import BaseModel +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.billing import MicroDollar +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import StripeServiceException +from shared.models.database.credits_transaction import CreditsTransaction +from shared.models.database.job import Job +from shared.models.database.stripe_price_config import StripePriceConfig +from shared.models.database.user import User +from shared.models.schemas.billing import ( + BuyCreditsPackageRequest, + BuyCreditsRequest, + CheckoutSessionResponse, + CreditsBalanceResponse, + PaymentIntentResponse, + TransactionHistoryResponse, + UsageStatsResponse, +) +from shared.services.billing import CreditsService + + +class ParseUsageResponse(BaseModel): + request_total: int + mom_growth: float + credits_used: float + estimated_amount: Optional[float] + success_rate: float + avg_processing_time: float + + +async def buy_credits_for_user( + *, + request: BuyCreditsRequest, + user_id: str, +) -> PaymentIntentResponse: + stripe_service = StripeService() + try: + amount_cny = request.credits_amount * 0.02 + amount_cents = int(amount_cny * 100) + payment_intent = await stripe_service.create_payment_intent( + user_id=user_id, + amount=amount_cents, + credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, + currency="cny", + ) + + return PaymentIntentResponse( + client_secret=payment_intent["client_secret"], + payment_intent_id=payment_intent["payment_intent_id"], + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to buy credits: {str(exc)}" + ) + + +async def get_credits_balance_for_user( + db: AsyncSession, + *, + user_id: str, +) -> CreditsBalanceResponse: + credits_service = CreditsService() + try: + await credits_service.ensure_user_initialized(db, user_id) + await db.commit() + + balance_micro_dollar = await credits_service.get_balance(db, user_id) + return CreditsBalanceResponse( + credits_balance=MicroDollar(balance_micro_dollar).to_credit() + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get credits balance: {str(exc)}" + ) + + +async def get_usage_stats_for_user( + db: AsyncSession, + *, + user_id: str, + period: str, +) -> UsageStatsResponse: + credits_service = CreditsService() + try: + stats = await credits_service.get_usage_stats(db, user_id, period) + return UsageStatsResponse( + period=stats["period"], + total_credits_used=MicroDollar(stats["total_used"]).to_credit(), + api_calls_count=stats["transaction_count"], + success_rate=95.0, + average_response_time=stats.get("avg_response_time", 0), + top_endpoints=[], + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get usage statistics: {str(exc)}" + ) + + +async def get_parse_usage_overview_for_user( + db: AsyncSession, + *, + user_id: str, +) -> ParseUsageResponse: + try: + total_micro_credits_used = await _load_total_parse_micro_credits_used( + db, + user_id=user_id, + ) + success_rate, avg_processing_time = await _load_parse_job_usage_stats( + db, + user_id=user_id, + ) + estimated_amount = await _estimate_parse_usage_amount( + db, + total_micro_credits_used=total_micro_credits_used, + ) + + return ParseUsageResponse( + request_total=0, + mom_growth=0.0, + credits_used=MicroDollar(total_micro_credits_used).to_credit() or 0, + estimated_amount=estimated_amount, + success_rate=round(success_rate, 2), + avg_processing_time=avg_processing_time, + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get parse usage overview: {str(exc)}" + ) + + +async def get_transaction_history_for_user( + db: AsyncSession, + *, + user_id: str, + limit: int, +) -> list[TransactionHistoryResponse]: + credits_service = CreditsService() + try: + transactions = await credits_service.get_transaction_history( + db, + user_id, + limit, + ) + return [ + TransactionHistoryResponse( + id=transaction.id, + credits_amount=MicroDollar(transaction.credits_amount).to_credit(), + transaction_type=transaction.transaction_type, + description=transaction.description, + created_at=transaction.created_at, + ) + for transaction in transactions + ] + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get transaction history: {str(exc)}" + ) + + +async def get_price_configs_payload( + db: AsyncSession, + *, + product_type: str | None, +) -> dict[str, list[dict]]: + try: + price_config_service = PriceConfigService() + + if product_type == "subscription": + configs = await price_config_service.repository.get_all_active(db) + return { + "subscriptions": [ + _subscription_config_payload(config) + for config in configs + if config.product_type == "subscription" + ], + "credits_packages": [], + } + + if product_type == "credits_package": + credits_configs = await price_config_service.get_all_credits_packages(db) + return { + "subscriptions": [], + "credits_packages": [ + _credits_package_config_payload(config) + for config in credits_configs + ], + } + + configs = await price_config_service.repository.get_all_active(db) + return { + "subscriptions": [ + _subscription_config_payload(config) + for config in configs + if config.product_type == "subscription" + ], + "credits_packages": [ + _credits_package_config_payload(config) + for config in configs + if config.product_type == "credits_package" + ], + } + + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get price configurations: {str(exc)}" + ) + + +async def buy_credits_package_for_user( + db: AsyncSession, + *, + request: BuyCreditsPackageRequest, + user_id: str, +) -> CheckoutSessionResponse: + stripe_service = StripeService() + try: + result = await db.execute(select(User.email).where(User.id == user_id)) + user_email = result.scalar_one_or_none() + + frontend_url = settings.FRONTEND_URL + success_url = f"{frontend_url}/billing?success=true&type=credits_package" + cancel_url = f"{frontend_url}/billing?canceled=true" + + checkout_url = await stripe_service.create_checkout_session_for_credits_package( + db=db, + user_id=user_id, + price_id=request.price_id, + success_url=success_url, + cancel_url=cancel_url, + quantity=request.quantity, + email=user_email, + ) + + return CheckoutSessionResponse(checkout_url=checkout_url, session_id="") + + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to create credits package purchase: {str(exc)}" + ) + + +async def handle_stripe_webhook( + db: AsyncSession, + *, + payload: bytes, + stripe_signature: str | None, +): + stripe_service = StripeService() + try: + if not stripe_signature: + raise StripeServiceException( + internal_message="Missing stripe-signature header" + ) + return await stripe_service.handle_webhook(db, payload, stripe_signature) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to handle webhook: {str(exc)}" + ) + + +async def _load_total_parse_micro_credits_used( + db: AsyncSession, + *, + user_id: str, +) -> int: + credits_row = await db.execute( + select(func.coalesce(func.sum(CreditsTransaction.credits_amount), 0)) + .where(CreditsTransaction.user_id == user_id) + .where(CreditsTransaction.transaction_type.in_(["usage", "refund"])) + ) + return int(abs(credits_row.scalar_one() or 0)) + + +async def _load_parse_job_usage_stats( + db: AsyncSession, + *, + user_id: str, +) -> tuple[float, float]: + job_row = await db.execute( + select( + func.count().filter(Job.status == "done").label("done_cnt"), + func.count() + .filter(Job.status.in_(["done", "failed"])) + .label("terminal_cnt"), + func.avg(func.extract("epoch", Job.updated_at - Job.created_at)) + .filter(Job.status.in_(["done", "failed"])) + .label("avg_secs"), + ).where(Job.user_id == user_id) + ) + job_stats = job_row.first() or (0, 0, 0.0) + done_count = getattr(job_stats, "done_cnt", 0) or 0 + terminal_count = getattr(job_stats, "terminal_cnt", 0) or 0 + success_rate = (done_count / terminal_count * 100) if terminal_count > 0 else 0.0 + avg_processing_time = round( + float(getattr(job_stats, "avg_secs", 0.0) or 0.0), + 2, + ) + return success_rate, avg_processing_time + + +async def _estimate_parse_usage_amount( + db: AsyncSession, + *, + total_micro_credits_used: int, +) -> float | None: + price_row = await db.execute( + select(StripePriceConfig) + .where(StripePriceConfig.product_type == "credits_package") + .where(StripePriceConfig.is_active.is_(True)) + .order_by(StripePriceConfig.created_at) + .limit(1) + ) + price_cfg = price_row.scalar_one_or_none() + if not price_cfg or not price_cfg.credits_amount or price_cfg.credits_amount <= 0: + return None + + return round( + price_cfg.amount_cents + * total_micro_credits_used + / (100 * price_cfg.credits_amount), + 4, + ) + + +def _subscription_config_payload(config) -> dict: + metadata = config.extra_metadata or {} + return { + "id": config.plan_id, + "plan_id": config.plan_id, + "price_id": config.price_id, + "name": metadata.get("display_name", config.plan_id.upper()), + "description": metadata.get("description", ""), + "features": metadata.get("features", []), + "popular": metadata.get("frontend_config", {}).get("popular", False), + "amount_cents": config.amount_cents, + "currency": config.currency, + "metadata": metadata, + } + + +def _credits_package_config_payload(config) -> dict: + metadata = config.extra_metadata or {} + credit_amount = MicroDollar(config.credits_amount).to_credit() + return { + "id": config.plan_id, + "plan_id": config.plan_id, + "price_id": config.price_id, + "name": metadata.get("display_name", f"{credit_amount} Credits"), + "description": metadata.get("description", ""), + "credits_amount": credit_amount, + "amount_cents": config.amount_cents, + "currency": config.currency, + "metadata": metadata, + } diff --git a/apps/api/app/services/job_creation_service.py b/apps/api/app/services/job_creation_service.py new file mode 100644 index 000000000..c5ed2b284 --- /dev/null +++ b/apps/api/app/services/job_creation_service.py @@ -0,0 +1,463 @@ +from __future__ import annotations + +import os +import uuid +from datetime import datetime, timezone +from typing import Optional, cast +from urllib.parse import urlparse + +from app.repositories.job_repository import JobRepository +from app.services.job_document_scope_service import ( + find_active_job_for_document, + is_active_document_job_unique_violation, + raise_document_ingestion_conflict, + resolve_effective_document_scope, +) +from app.services.rate_limit.data_structures import CurrentUser +from fastapi import Request +from loguru import logger +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + ConflictException, + JobOperationException, + NotFoundException, + RateLimitException, + UnavailableException, + ValidationException, +) +from shared.core.exceptions.webhook_exceptions import WebhookConfigException +from shared.core.state_machine.states import JobStatus +from shared.models.schemas.job import JobCreate, JobResponse +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.redis import JobInfoRedisService, RedisServiceFactory +from shared.services.redis.job_metadata_service import JobMetadataService +from shared.services.storage.file_upload_service import FileUploadService +from shared.utils.url_file_type import resolve_file_extension_async +from shared.utils.url_security import validate_http_url_and_resolve_ip_async + +JOB_TYPE_KB_MANAGEMENT = "kb_management" + + +def get_supported_formats() -> str: + return ", ".join(sorted(settings.get_supported_extensions())) + + +def validate_file_type(file_name: str) -> bool: + if not file_name: + return False + file_extension = os.path.splitext(file_name)[1].lower() + return file_extension in settings.get_supported_extensions() + + +def create_job_response( + job_id: str, + job, + source_type: str, + data_id: Optional[str], + namespace: Optional[str] = None, + document_id: Optional[str] = None, + upload_url: Optional[str] = None, + upload_headers: Optional[dict] = None, + expires_in: Optional[int] = None, +) -> JobResponse: + return JobResponse( + job_id=job_id, + status=job.status, + source_type=source_type, + data_id=data_id, + namespace=namespace, + document_id=document_id, + created_at=job.created_at, + upload_url=upload_url, + upload_headers=upload_headers, + expires_in=expires_in, + ) + + +async def _validate_create_job_payload(payload: JobCreate) -> None: + if payload.source_type == "file" and not payload.file_name: + raise ValidationException( + user_message="file_name is required when source_type is 'file'", + violations=[ + { + "field": "file_name", + "description": "Required for file source type", + } + ], + ) + if payload.source_type == "url" and not payload.source_url: + raise ValidationException( + user_message="source_url is required when source_type is 'url'", + violations=[ + { + "field": "source_url", + "description": "Required for url source type", + } + ], + ) + + if payload.webhook and payload.webhook.url: + validation_result = await validate_http_url_and_resolve_ip_async( + payload.webhook.url, + ) + if not validation_result.is_valid: + raise WebhookConfigException( + user_message="Invalid webhook URL", + internal_message=f"Webhook validation failed: {validation_result.error_message}", + ) + + if ( + payload.source_type == "file" + and payload.file_name + and not validate_file_type(payload.file_name) + ): + supported_formats = get_supported_formats() + raise ValidationException( + user_message=f"Unsupported file type. Supported formats: {supported_formats}", + violations=[ + {"field": "file_name", "description": "File type not supported"} + ], + ) + + if payload.source_type == "url": + assert payload.source_url is not None + file_extension = await resolve_file_extension_async(payload.source_url) + if not file_extension: + supported_formats = get_supported_formats() + raise ValidationException( + user_message=f"Unsupported URL file type. Supported formats: {supported_formats}", + violations=[ + { + "field": "source_url", + "description": "URL file type not supported", + } + ], + ) + + +async def _resolve_job_metadata( + db: AsyncSession, + *, + payload: JobCreate, + current_user: CurrentUser, +) -> tuple[dict, str, str]: + job_metadata = JobMetadataHelper.create_from_request(payload) + requested_document_id = cast(Optional[str], job_metadata.get("document_id")) + if requested_document_id: + active_job = await find_active_job_for_document( + db, + user_id=current_user.user_id, + document_id=requested_document_id, + ) + if active_job is not None: + raise_document_ingestion_conflict( + document_id=requested_document_id, + active_job_id=active_job.job_id, + ) + ( + effective_document_id, + effective_namespace, + ) = await resolve_effective_document_scope( + db, + user_id=current_user.user_id, + document_id=requested_document_id, + requested_namespace=cast(Optional[str], payload.namespace), + ) + if not requested_document_id: + active_job = await find_active_job_for_document( + db, + user_id=current_user.user_id, + document_id=effective_document_id, + ) + if active_job is not None: + raise_document_ingestion_conflict( + document_id=effective_document_id, + active_job_id=active_job.job_id, + ) + job_metadata["document_id"] = effective_document_id + job_metadata["namespace"] = effective_namespace + return job_metadata, effective_document_id, effective_namespace + + +async def _cache_job_creation_state( + *, + job_id: str, + s3_key: str, + user_id: str, + webhook_enabled: bool, + source_type: str, + job_metadata: dict, +) -> None: + redis_service = RedisServiceFactory.get_service() + metadata_service = JobMetadataService(redis_service) + await metadata_service.save_metadata(job_id, job_metadata) + + job_info_service = JobInfoRedisService(redis_service) + job_info = { + "job_id": job_id, + "s3_key": s3_key, + "user_id": user_id, + "webhook_enabled": webhook_enabled, + "job_type": JOB_TYPE_KB_MANAGEMENT, + "source_type": source_type, + "created_at": datetime.now(timezone.utc).isoformat(), + } + await job_info_service.save_job_info(job_id, job_info) + + +async def _create_waiting_job( + db: AsyncSession, + *, + job_id: str, + user_id: str, + source_type: str, + webhook_url: str | None, + job_metadata: dict, + s3_key: str, + effective_document_id: str, +): + job_repo = JobRepository() + try: + return await job_repo.create_job( + db=db, + job_id=job_id, + user_id=user_id, + job_type=JOB_TYPE_KB_MANAGEMENT, + source_type=source_type, + file_path=None, + webhook_url=webhook_url, + metadata=job_metadata, + initial_state=JobStatus.WAITING_FILE.value, + s3_key=s3_key, + ) + except IntegrityError as exc: + if is_active_document_job_unique_violation(exc): + raise_document_ingestion_conflict(document_id=effective_document_id) + raise + + +def _resolve_url_source_file_name(*, source_url: str, file_extension: str) -> str: + parsed_url = urlparse(source_url) + url_basename = str(os.path.basename(parsed_url.path)) + if url_basename and os.path.splitext(url_basename)[1].lower() == file_extension: + return url_basename + if url_basename: + return f"{url_basename}{file_extension}" + return f"url_file_{uuid.uuid4().hex[:8]}{file_extension}" + + +def _schedule_url_upload(*, job_id: str, source_url: str, user_id: str) -> None: + from shared.core.celery_app import get_celery_app + + celery_app = get_celery_app() + upload_url_file_task = celery_app.signature( + "app.core.tasks.kb_tasks.upload_url_file_task" + ) + upload_url_file_task.apply_async( + args=[job_id, source_url, user_id], + kwargs={ + "job_type": JOB_TYPE_KB_MANAGEMENT, + }, + ) + + +async def _create_file_job( + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + job_metadata: dict, + effective_document_id: str, + effective_namespace: str, +) -> JobResponse: + assert payload.file_name is not None + file_extension = os.path.splitext(payload.file_name)[1] + s3_key = f"uploads/{job_id}{file_extension}" + job_metadata["source_file_name"] = payload.file_name + job_metadata["source_type"] = "file" + + job = await _create_waiting_job( + db, + job_id=job_id, + user_id=current_user.user_id, + source_type="file", + webhook_url=payload.webhook.url if payload.webhook else None, + job_metadata=job_metadata, + s3_key=s3_key, + effective_document_id=effective_document_id, + ) + if not job: + raise JobOperationException( + internal_message="Failed to create job in database" + ) + + upload_service = FileUploadService() + upload_info = await upload_service.generate_upload_url(job_id, file_extension) + + await _cache_job_creation_state( + job_id=job_id, + s3_key=s3_key, + user_id=current_user.user_id, + webhook_enabled=bool(payload.webhook and payload.webhook.url), + source_type="file", + job_metadata=job_metadata, + ) + + logger.info( + f"Job {job_id} upload_url returned to client: {upload_info['upload_url']}" + ) + return create_job_response( + job_id=job_id, + job=job, + source_type="file", + data_id=payload.data_id, + namespace=effective_namespace, + upload_url=upload_info["upload_url"], + upload_headers=upload_info["upload_headers"], + expires_in=upload_info["expires_in"], + ) + + +async def _create_url_job( + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + job_metadata: dict, + effective_document_id: str, + effective_namespace: str, +) -> JobResponse: + assert payload.source_url is not None + file_extension = await resolve_file_extension_async(payload.source_url) + if not file_extension: + supported_formats = get_supported_formats() + raise ValidationException( + user_message=f"Unsupported URL file type. Supported formats: {supported_formats}", + violations=[ + { + "field": "source_url", + "description": "URL file type not supported", + } + ], + ) + + source_file_name = _resolve_url_source_file_name( + source_url=payload.source_url, + file_extension=file_extension, + ) + s3_key = f"uploads/{job_id}{file_extension}" + job_metadata.update( + { + "source_file_name": source_file_name, + "source_url": payload.source_url, + "source_type": "url", + } + ) + + job = await _create_waiting_job( + db, + job_id=job_id, + user_id=current_user.user_id, + source_type="url", + webhook_url=payload.webhook.url if payload.webhook else None, + job_metadata=job_metadata, + s3_key=s3_key, + effective_document_id=effective_document_id, + ) + if not job: + raise JobOperationException( + internal_message="Failed to create URL job in database" + ) + + await _cache_job_creation_state( + job_id=job_id, + s3_key=s3_key, + user_id=current_user.user_id, + webhook_enabled=bool(payload.webhook and payload.webhook.url), + source_type="url", + job_metadata=job_metadata, + ) + _schedule_url_upload( + job_id=job_id, + source_url=payload.source_url, + user_id=current_user.user_id, + ) + + return create_job_response( + job_id=job_id, + job=job, + source_type="url", + data_id=payload.data_id, + namespace=effective_namespace, + ) + + +async def create_job_from_request( + db: AsyncSession, + *, + payload: JobCreate, + current_user: CurrentUser, + enforce_capacity, + request: Request, +) -> JobResponse: + try: + job_id = f"job_{uuid.uuid4().hex[:12]}" + await _validate_create_job_payload(payload) + ( + job_metadata, + effective_document_id, + effective_namespace, + ) = await _resolve_job_metadata( + db, + payload=payload, + current_user=current_user, + ) + + await enforce_capacity( + request=request, + db=db, + current_user=current_user, + ) + + if payload.source_type == "file": + return await _create_file_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + job_metadata=job_metadata, + effective_document_id=effective_document_id, + effective_namespace=effective_namespace, + ) + return await _create_url_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + job_metadata=job_metadata, + effective_document_id=effective_document_id, + effective_namespace=effective_namespace, + ) + + except NotFoundException: + raise + except ValidationException: + raise + except ConflictException: + raise + except WebhookConfigException: + raise + except (RateLimitException, UnavailableException): + raise + except JobOperationException: + raise + except Exception as exc: + logger.error(f"Failed to create job: {exc}") + raise JobOperationException( + internal_message=f"Job creation failed: {str(exc)}" + ) diff --git a/apps/api/app/services/job_read_service.py b/apps/api/app/services/job_read_service.py new file mode 100644 index 000000000..4b68d3397 --- /dev/null +++ b/apps/api/app/services/job_read_service.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import math +from datetime import datetime, timedelta, timezone +from typing import Optional + +from app.repositories.job_repository import JobRepository +from app.services.job_response_projection import ( + build_job_result_response, + to_job_status_value, +) +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + JobOperationException, + NotFoundException, + PermissionDeniedException, + ValidationException, +) +from shared.models.schemas.job import JobList, JobResultResponse +from shared.services.redis import RedisServiceFactory +from shared.utils.utc_now import utc_now_naive + + +def check_job_permission(job, user_id: str, job_id: str) -> None: + if not job: + raise NotFoundException( + resource="Job", resource_id=job_id, internal_message="Job not found" + ) + + if str(job.user_id) != user_id: + raise PermissionDeniedException( + user_message="You don't have permission to access this job", + ) + + +def normalize_naive_utc_filter_datetime(dt: Optional[datetime]) -> Optional[datetime]: + if dt is None: + return None + if dt.tzinfo is None or dt.utcoffset() is None: + return dt + return dt.astimezone(timezone.utc).replace(tzinfo=None) + + +async def list_jobs_for_user( + db: AsyncSession, + *, + user_id: str, + page: int, + page_size: int, + job_status: Optional[str], + job_type: Optional[str], + recent_days: Optional[int], + start_time: Optional[datetime], + end_time: Optional[datetime], +) -> JobList: + try: + job_repo = JobRepository() + + if recent_days not in (None, 1, 7, 30): + raise ValidationException( + user_message="recent_days only supports 1, 7, or 30", + violations=[{"field": "recent_days", "description": "Invalid value"}], + ) + + created_after: Optional[datetime] = None + if recent_days: + created_after = utc_now_naive() - timedelta(days=recent_days) + + normalized_start_time = normalize_naive_utc_filter_datetime(start_time) + normalized_end_time = normalize_naive_utc_filter_datetime(end_time) + + if ( + normalized_start_time + and normalized_end_time + and normalized_start_time > normalized_end_time + ): + raise ValidationException( + user_message="start_time cannot be later than end_time", + violations=[ + {"field": "start_time", "description": "Must be before end_time"} + ], + ) + + if normalized_start_time: + created_after = normalized_start_time + created_before = normalized_end_time + + total_count = await job_repo.count_jobs_by_user( + db=db, + user_id=user_id, + created_after=created_after, + created_before=created_before, + job_type=job_type, + job_status=job_status, + ) + jobs = await job_repo.get_jobs_by_user( + db=db, + user_id=user_id, + limit=page_size, + offset=(page - 1) * page_size, + created_after=created_after, + created_before=created_before, + job_type=job_type, + job_status=job_status, + ) + + redis_service = RedisServiceFactory.get_service() + job_responses = [] + for job in jobs: + job_metadata = await job_repo.get_job_metadata( + db, job.job_id, redis_service + ) + job_responses.append( + await build_job_result_response( + job=job, + job_metadata=job_metadata, + progress=None, + ) + ) + + total_pages = math.ceil(total_count / page_size) if total_count > 0 else 0 + return JobList( + jobs=job_responses, + total=total_count, + page=page, + page_size=page_size, + total_pages=total_pages, + ) + + except ValidationException: + raise + except Exception as exc: + logger.error(f"Failed to list jobs: {exc}") + raise JobOperationException( + internal_message=f"Failed to get job list: {str(exc)}" + ) + + +async def get_job_result_for_user( + db: AsyncSession, + *, + job_id: str, + user_id: str, +) -> JobResultResponse: + try: + job_repo = JobRepository() + job = await job_repo.get_job_by_id(db, job_id) + check_job_permission(job, user_id, job_id) + assert job is not None + + progress = None + if to_job_status_value(job.status) == "running": + progress = {"total_pages": 10, "processed_pages": 5} + + redis_service = RedisServiceFactory.get_service() + job_metadata = await job_repo.get_job_metadata(db, job_id, redis_service) + return await build_job_result_response( + job=job, + job_metadata=job_metadata, + progress=progress, + ) + + except NotFoundException: + raise + except PermissionDeniedException: + raise + except Exception as exc: + logger.error(f"Failed to get job result: {exc}") + raise JobOperationException( + internal_message=f"Failed to get job result: {str(exc)}" + ) diff --git a/apps/api/app/services/job_response_projection.py b/apps/api/app/services/job_response_projection.py new file mode 100644 index 000000000..9563206bb --- /dev/null +++ b/apps/api/app/services/job_response_projection.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import os +from datetime import datetime, timedelta, timezone +from typing import Any, Literal, Optional, cast +from urllib.parse import urlparse + +from shared.core.billing import MicroDollar +from shared.core.exceptions.domain_exceptions import JobOperationException +from shared.models.schemas.job import JobResultResponse, StandardErrorObject +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.storage.file_upload_service import FileUploadService +from shared.utils.error_details import normalize_error_details +from shared.utils.utc_now import utc_now_naive + +JobStatusValue = Literal[ + "pending", "waiting-file", "running", "converting", "done", "failed" +] + + +def build_error_response( + job: Any, job_metadata: Optional[dict] = None +) -> Optional[StandardErrorObject]: + if not job.error_message: + return None + + error_details = None + if job_metadata and isinstance(job_metadata, dict): + error_details = normalize_error_details(job_metadata.get("error_details")) + + return StandardErrorObject( + code=job.error_code or "UNKNOWN", + message=job.error_message, + request_id=job.job_id, + details=error_details, + ) + + +def resolve_public_document_id(job: Any) -> Optional[str]: + job_result = getattr(job, "job_result", None) + published_document_id = getattr(job_result, "document_id", None) + if isinstance(published_document_id, str) and published_document_id: + return published_document_id + + return None + + +def ensure_utc(dt: Optional[datetime]) -> Optional[datetime]: + if not dt: + return None + if dt.tzinfo: + return dt.astimezone(timezone.utc) + return dt.replace(tzinfo=timezone.utc) + + +def require_utc(dt: Optional[datetime], *, field_name: str) -> datetime: + normalized_dt = ensure_utc(dt) + if normalized_dt is None: + raise JobOperationException( + internal_message=f"Job is missing required datetime field: {field_name}" + ) + return normalized_dt + + +def to_job_status_value(status: str) -> JobStatusValue: + return cast(JobStatusValue, status) + + +def _resolve_original_request(job_metadata: Optional[dict[str, Any]]) -> dict[str, Any]: + original_request = ( + job_metadata.get("original_request") + if isinstance(job_metadata, dict) + else {} + ) + return original_request if isinstance(original_request, dict) else {} + + +def _resolve_source_file_name(original_request: dict[str, Any]) -> str | None: + source_url = original_request.get("source_url") + file_name = None + if source_url: + parsed_source = urlparse(str(source_url)) + file_name = os.path.basename(parsed_source.path) or None + if not file_name: + file_name = original_request.get("file_name") + return str(file_name) if file_name else None + + +def _resolve_file_extension(file_name: str | None) -> str | None: + if not file_name: + return None + extension = os.path.splitext(file_name)[1] + return extension[1:].upper() if extension else None + + +def _resolve_parsing_params( + job_metadata: Optional[dict[str, Any]], + original_request: dict[str, Any], +) -> dict[str, Any]: + parsing_params = original_request.get("parsing_params") or {} + if not parsing_params and isinstance(job_metadata, dict): + parsing_params = job_metadata.get("parsing_params") or {} + return parsing_params if isinstance(parsing_params, dict) else {} + + +def _resolve_duration_seconds(job: Any) -> float | None: + if job.updated_at and job.created_at: + return (job.updated_at - job.created_at).total_seconds() + return None + + +async def _resolve_result_delivery(job: Any) -> tuple[dict[str, Any] | None, str | None, datetime]: + job_result = job.job_result + result_url = None + result = None + result_url_expires_at = job.created_at + + if job_result and job_result.result_s3_key: + upload_service = FileUploadService() + result_url_info = await upload_service.generate_download_url( + job_result.result_s3_key + ) + result_url = result_url_info["download_url"] + + if job_result.inline_payload: + result = job_result.inline_payload + + if result_url: + expires_in = int(result_url_info.get("expires_in", 3600)) + result_url_expires_at = utc_now_naive() + timedelta(seconds=expires_in) + + return result, result_url, result_url_expires_at + + +async def build_job_result_response( + *, + job: Any, + job_metadata: Optional[dict[str, Any]], + progress: dict[str, Any] | None, +) -> JobResultResponse: + original_request = _resolve_original_request(job_metadata) + file_name = _resolve_source_file_name(original_request) + parsing_params = _resolve_parsing_params(job_metadata, original_request) + result, result_url, result_url_expires_at = await _resolve_result_delivery(job) + + return JobResultResponse( + job_id=job.job_id, + namespace=JobMetadataHelper.get_field(job_metadata, "namespace"), + document_id=resolve_public_document_id(job), + status=to_job_status_value(job.status), + source_type=job.source_type, + data_id=JobMetadataHelper.get_field(job_metadata, "data_id"), + created_at=require_utc(job.created_at, field_name="created_at"), + progress=progress, + error=build_error_response(job, job_metadata), + result=result, + result_url=result_url, + result_url_expires_at=require_utc( + result_url_expires_at, + field_name="result_url_expires_at", + ), + file_name=file_name, + file_extension=_resolve_file_extension(file_name), + model=parsing_params.get("model"), + ocr_enabled=parsing_params.get("ocr_enabled"), + duration_seconds=_resolve_duration_seconds(job), + credits_spent=( + MicroDollar(job.credits_charged).to_credit() + if hasattr(job, "credits_charged") + else 0 + ), + ) diff --git a/apps/api/app/services/job_upload_confirmation_service.py b/apps/api/app/services/job_upload_confirmation_service.py new file mode 100644 index 000000000..421915d58 --- /dev/null +++ b/apps/api/app/services/job_upload_confirmation_service.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from typing import Optional + +from app.repositories.job_repository import JobRepository +from app.services.job_read_service import check_job_permission +from app.services.knowledge.kb_orchestrator import KBOrchestrator +from app.services.state_machine import JobStateMachine +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + JobOperationException, + NotFoundException, + PermissionDeniedException, + ValidationException, +) +from shared.core.state_machine.states import JobStatus +from shared.models.schemas.job import ConfirmUploadRequest +from shared.services.storage.file_upload_service import FileUploadService + + +async def transition_to_uploaded( + db: AsyncSession, + job_id: str, + trigger: str = "manual_upload_completed", +) -> None: + state_machine = JobStateMachine() + await state_machine.transition( + db, job_id, JobStatus.PENDING.value, trigger, None, "system" + ) + + +async def start_workflow_for_job( + db: AsyncSession, + job_id: str, + job_type: str, + source_type: str, + user_id: str, + file_path: Optional[str] = None, + file_url: Optional[str] = None, +) -> None: + if job_type == "kb_management": + orchestrator = KBOrchestrator() + await orchestrator.start_workflow( + db=db, + job_id=job_id, + source_type=source_type, + file_path=file_path, + file_url=file_url, + user_id=user_id, + ) + return + + raise ValidationException( + user_message="Unsupported job type", + violations=[ + { + "field": "job_type", + "description": f"Job type '{job_type}' is not supported", + } + ], + ) + + +async def confirm_job_upload( + db: AsyncSession, + *, + job_id: str, + request: ConfirmUploadRequest | None, + user_id: str, +) -> dict[str, str]: + del request + + try: + job_repo = JobRepository() + job = await job_repo.get_job_by_id(db, job_id) + check_job_permission(job, user_id, job_id) + assert job is not None + + logger.info(f"Confirm upload - Job {job_id} current status: {job.status}") + if job.status not in [JobStatus.PENDING.value, JobStatus.WAITING_FILE.value]: + logger.info(f"Job {job_id} already processed, status: {job.status}") + return {"message": "Job status already updated"} + + if not job.s3_key: + raise ValidationException( + user_message="Job is missing S3 key information", + violations=[ + {"field": "s3_key", "description": "S3 key not set for this job"} + ], + ) + + upload_service = FileUploadService() + file_info = await upload_service.verify_s3_file_exists(job.s3_key) + + if not file_info.get("exists"): + raise ValidationException( + user_message="S3 file does not exist, please upload the file first", + violations=[{"field": "file", "description": "File not found in S3"}], + ) + + await transition_to_uploaded(db, job_id) + await start_workflow_for_job( + db=db, + job_id=job_id, + job_type=job.job_type, + source_type="file", + user_id=user_id, + ) + + return {"message": "File upload confirmed; processing started"} + + except NotFoundException: + raise + except PermissionDeniedException: + raise + except ValidationException: + raise + except Exception as exc: + logger.error(f"Failed to confirm upload: {exc}") + raise JobOperationException( + internal_message=f"Failed to confirm upload: {str(exc)}" + ) diff --git a/apps/api/app/services/qstash_callback_service.py b/apps/api/app/services/qstash_callback_service.py new file mode 100644 index 000000000..bc7e9518c --- /dev/null +++ b/apps/api/app/services/qstash_callback_service.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +import json +from datetime import datetime, timezone +from typing import Any, Optional +from uuid import NAMESPACE_URL, uuid5 + +from fastapi import Response +from loguru import logger +from sqlalchemy import select + +from shared.core.config import app_config +from shared.core.database_sync import get_sync_db_context +from shared.models.database.webhook import WebhookEvent, WebhookEventStatus +from shared.models.database.webhook_log import WebhookLog + + +def get_qstash_verification_url(callback_path: str, request_url: str) -> str: + callback_base_url = app_config.QSTASH_CALLBACK_BASE_URL + if callback_base_url: + return f"{callback_base_url.rstrip('/')}{callback_path}" + return request_url + + +def verify_qstash_signature(raw_body: bytes, signature: str, url: str) -> bool: + current_key = app_config.QSTASH_CURRENT_SIGNING_KEY + next_key = app_config.QSTASH_NEXT_SIGNING_KEY + + if not current_key or not next_key: + logger.error("QStash signing keys not configured — rejecting callback") + return False + + try: + from qstash import Receiver + + receiver = Receiver( + current_signing_key=current_key, + next_signing_key=next_key, + ) + receiver.verify( + body=raw_body.decode("utf-8"), + signature=signature, + url=url, + ) + return True + except Exception as exc: + logger.warning( + "QStash signature verification failed: error_type={error_type}, url={url}", + error_type=type(exc).__name__, + url=url, + ) + return False + + +def handle_qstash_success_callback(raw_body: bytes) -> Response: + data = extract_callback_data(raw_body) + event_id = find_event_id(data) + + if not event_id: + logger.warning("QStash callback: missing event_id, cannot correlate") + return Response(status_code=200, content="OK (no event_id)") + + retried = data.get("retried", 0) + logger.info( + f"QStash callback: event_id={event_id}, status={data.get('status')}, " + f"retried={retried}, qstash_message_id={data.get('sourceMessageId')}" + ) + + return process_qstash_callback( + data, + event_id, + get_callback_event_status(data), + "callback", + ) + + +def handle_qstash_failure_callback(raw_body: bytes) -> Response: + data = extract_callback_data(raw_body) + event_id = find_event_id(data) + + if not event_id: + logger.warning("QStash failure callback: missing event_id, cannot correlate") + return Response(status_code=200, content="OK (no event_id)") + + retried = data.get("retried", 0) + max_retries = data.get("maxRetries", 0) + logger.warning( + f"QStash failure: event_id={event_id}, status={data.get('status')}, " + f"retried={retried}/{max_retries}, qstash_message_id={data.get('sourceMessageId')}" + ) + + return process_qstash_callback( + data, + event_id, + WebhookEventStatus.FAILED, + "failure", + ) + + +def extract_callback_data(body: bytes) -> dict[str, Any]: + try: + return json.loads(body) + except (json.JSONDecodeError, ValueError): + return {"raw": body.decode("utf-8", errors="replace")} + + +def find_event_id(data: dict[str, Any]) -> Optional[str]: + source_header = data.get("sourceHeader", {}) or {} + event_id = normalize_header_value( + source_header.get("X-Knowhere-Event-Id") + or source_header.get("x-knowhere-event-id") + ) + if not event_id: + for key, value in source_header.items(): + if key.lower() == "x-knowhere-event-id": + event_id = normalize_header_value(value) + break + return event_id + + +def normalize_header_value(value: Any) -> Optional[str]: + if isinstance(value, list): + if not value: + return None + first_value = value[0] + return first_value if isinstance(first_value, str) else str(first_value) + + if isinstance(value, str): + return value + + if value is None: + return None + + return str(value) + + +def build_callback_log_idempotency_key( + qstash_message_id: Optional[str], + event_id: str, +) -> str: + if qstash_message_id: + return str(uuid5(NAMESPACE_URL, qstash_message_id)) + return event_id + + +def get_response_status_code(value: Any) -> Optional[int]: + if value is None: + return None + + try: + return int(value) + except (TypeError, ValueError): + return None + + +def is_success_response_status(status_code: Optional[int]) -> bool: + return status_code is not None and 200 <= status_code < 300 + + +def get_callback_event_status(data: dict[str, Any]) -> str: + response_status = get_response_status_code(data.get("status")) + if is_success_response_status(response_status): + return WebhookEventStatus.DELIVERED + return WebhookEventStatus.DELIVERING + + +def resolve_event_status(current_status: str, callback_status: str) -> str: + if current_status in ( + WebhookEventStatus.DELIVERED, + WebhookEventStatus.FAILED, + WebhookEventStatus.CANCELED, + ): + return current_status + return callback_status + + +def process_qstash_callback( + data: dict[str, Any], + event_id: str, + callback_status: str, + log_label: str, +) -> Response: + response_status_code = get_response_status_code(data.get("status")) + response_body = data.get("body", "") + qstash_message_id = data.get("sourceMessageId") + retried = data.get("retried", 0) + is_failed_delivery_attempt = ( + callback_status == WebhookEventStatus.FAILED + or ( + callback_status == WebhookEventStatus.DELIVERING + and not is_success_response_status(response_status_code) + ) + ) + error_message = None + if is_failed_delivery_attempt: + error_message = data.get("error") or response_body + + with get_sync_db_context() as db: + event = db.execute( + select(WebhookEvent).where(WebhookEvent.id == event_id) + ).scalar_one_or_none() + + if not event: + logger.warning(f"QStash {log_label}: event {event_id} not found in DB") + return Response(status_code=200, content="OK (event not found)") + + now = datetime.now(timezone.utc).replace(tzinfo=None) + event_status = resolve_event_status(event.status, callback_status) + attempt_number = retried + 1 + event.status = event_status + event.attempts = max(event.attempts, attempt_number) + event.updated_at = now + + log = WebhookLog( + job_id=event.job_id, + event_id=event.id, + webhook_url=event.target_url, + attempt_number=attempt_number, + request_payload=event.payload, + signature="", + idempotency_key=build_callback_log_idempotency_key( + qstash_message_id, + event.id, + ), + response_status_code=response_status_code, + response_body=response_body[:4096] if response_body else None, + error_message=str(error_message)[:4096] if error_message else None, + duration_ms=0, + qstash_message_id=qstash_message_id, + ) + db.add(log) + db.commit() + + return Response(status_code=200, content="OK") diff --git a/apps/api/app/services/s3_events/__init__.py b/apps/api/app/services/s3_events/__init__.py new file mode 100644 index 000000000..b68b3a02b --- /dev/null +++ b/apps/api/app/services/s3_events/__init__.py @@ -0,0 +1 @@ +"""S3-compatible storage event services.""" diff --git a/apps/api/app/services/s3_events/event_handlers.py b/apps/api/app/services/s3_events/event_handlers.py new file mode 100644 index 000000000..3d146954a --- /dev/null +++ b/apps/api/app/services/s3_events/event_handlers.py @@ -0,0 +1,192 @@ +"""Storage event protocol handlers.""" +from __future__ import annotations + +import base64 +import json +import os +from typing import Any + +from app.services.s3_events.signature_verification import ( + verify_minio_signature, + verify_oss_signature, +) +from app.services.s3_events.subscription_service import confirm_sns_subscription +from app.services.s3_events.upload_event_service import process_upload_events +from loguru import logger + +from shared.models.schemas.oss_event import OSSEvent +from shared.models.schemas.s3_event import S3Event + + +async def handle_sns_event(body: bytes) -> dict[str, str] | None: + try: + sns_message = json.loads(body.decode("utf-8")) + message_type = sns_message.get("Type") + logger.info(f"SNS message type: {message_type}") + + if message_type == "SubscriptionConfirmation": + logger.info("Received an SNS subscription confirmation request") + subscribe_url = sns_message.get("SubscribeURL") + if subscribe_url: + logger.info(f"SNS subscription confirmation URL: {subscribe_url}") + return await confirm_sns_subscription(subscribe_url) + + logger.warning("SNS subscription confirmation did not include SubscribeURL") + return {"message": "SNS subscription confirmation failed"} + + if message_type == "Notification": + logger.info("Received an SNS notification") + logger.info(f"SNS message payload: {sns_message}") + await _handle_sns_notification(sns_message) + return None + + logger.warning(f"Unknown SNS message type: {message_type}") + return {"message": f"Unknown SNS message type: {message_type}"} + + except Exception as exc: + logger.error(f"Failed to handle SNS event: {exc}") + raise + + +async def handle_minio_event(body: bytes, auth_token: str) -> None: + try: + from shared.core.config import settings + + expected_token = getattr(settings, "S3_WEBHOOK_AUTH_TOKEN", "") + if not verify_minio_signature(auth_token, expected_token): + logger.warning("MinIO webhook authentication failed") + return + + s3_event_data = json.loads(body.decode("utf-8")) + await process_upload_events(S3Event(**s3_event_data)) + + except Exception as exc: + logger.error(f"Failed to handle MinIO event: {exc}") + + +async def handle_direct_s3_event(body: bytes) -> None: + try: + s3_event_data = json.loads(body.decode("utf-8")) + await process_upload_events(S3Event(**s3_event_data)) + + except Exception as exc: + logger.error(f"Failed to handle direct S3 event: {exc}") + + +def is_oss_event(headers: dict[str, str]) -> bool: + storage_type = os.getenv("S3_TYPE", "s3").lower() + if storage_type == "oss": + return True + + if "x-oss-pub-key-url" in headers: + return True + + if "x-mns-version" in headers or "x-mns-signing-cert-url" in headers: + return True + user_agent = headers.get("user-agent") or headers.get("User-Agent") + return bool(user_agent and "Aliyun Notification Service Agent" in user_agent) + + +async def handle_oss_event(body: bytes, headers: dict[str, str]) -> None: + try: + if not verify_oss_signature(body, headers): + logger.warning("OSS event signature verification failed") + return + + event_data = json.loads(body.decode("utf-8")) + logger.info(f"OSS event payload: {event_data}") + event_data = _unwrap_mns_message(event_data) + + if "events" in event_data: + oss_event = OSSEvent(**event_data) + elif "Records" in event_data: + oss_event = convert_s3_format_to_oss(event_data) + else: + logger.error(f"Unknown OSS event format: {event_data}") + return + + await process_upload_events(oss_event.to_s3_event()) + + except Exception as exc: + logger.error(f"Failed to handle OSS event: {exc}") + raise + + +def convert_s3_format_to_oss(event_data: dict[str, Any]) -> OSSEvent: + from shared.models.schemas.oss_event import OSSEventRecord + + records = event_data.get("Records", []) + oss_records = [ + OSSEventRecord( + eventName=record.get("eventName", "").replace("s3:", ""), + eventSource="acs:oss", + eventTime=record.get("eventTime", ""), + region=record.get("awsRegion", ""), + oss={ + "bucket": record.get("s3", {}).get("bucket", {}), + "object": record.get("s3", {}).get("object", {}), + }, + ) + for record in records + ] + + return OSSEvent(events=oss_records) + + +async def _handle_sns_notification(sns_message: dict[str, Any]) -> None: + try: + s3_event_data = json.loads(sns_message["Message"]) + logger.info(f"S3 event payload: {s3_event_data}") + + if ( + isinstance(s3_event_data, dict) + and s3_event_data.get("Event") == "s3:TestEvent" + ): + logger.info("Skip S3 test event") + return + + await process_upload_events(S3Event(**s3_event_data)) + except Exception as exc: + logger.error(f"Failed to parse the S3 event payload: {exc}") + logger.error(f"SNS payload: {sns_message}") + try: + await process_upload_events(S3Event(**sns_message)) + except Exception as fallback_exc: + logger.error( + "Fallback parsing of the SNS payload as an S3 event also failed: " + f"{fallback_exc}" + ) + raise + + +def _unwrap_mns_message(event_data: dict[str, Any]) -> dict[str, Any]: + if not isinstance(event_data, dict) or "Message" not in event_data: + return event_data + + inner = event_data.get("Message") + if isinstance(inner, dict): + return inner + if not isinstance(inner, str): + return event_data + + decoded = _decode_mns_inner_json(inner) + if decoded is not None: + logger.info(f"Decoded MNS Message payload: {decoded}") + return decoded + return event_data + + +def _decode_mns_inner_json(inner: str) -> dict[str, Any] | None: + try: + decoded_bytes = base64.b64decode(inner, validate=True) + decoded_str = decoded_bytes.decode("utf-8") + decoded = json.loads(decoded_str) + return decoded if isinstance(decoded, dict) else None + except Exception: + pass + + try: + decoded = json.loads(inner) + return decoded if isinstance(decoded, dict) else None + except Exception: + return None diff --git a/apps/api/app/services/s3_events/service.py b/apps/api/app/services/s3_events/service.py new file mode 100644 index 000000000..e06acedc3 --- /dev/null +++ b/apps/api/app/services/s3_events/service.py @@ -0,0 +1,52 @@ +"""Application service for S3-compatible storage event webhooks.""" +from __future__ import annotations + +from loguru import logger + +from app.services.s3_events.event_handlers import ( + handle_direct_s3_event, + handle_minio_event, + handle_oss_event, + handle_sns_event, + is_oss_event, +) + + +async def handle_s3_event_post( + *, + body: bytes, + headers: dict[str, str], + sns_message_type: str | None, + minio_auth_token: str | None, +) -> dict[str, str]: + if sns_message_type: + result = await handle_sns_event(body) + if result: + return result + elif is_oss_event(headers): + await handle_oss_event(body, headers) + elif minio_auth_token: + await handle_minio_event(body, minio_auth_token) + else: + await handle_direct_s3_event(body) + + return {"message": "Event handled successfully"} + + +async def safely_handle_s3_event_post( + *, + body: bytes, + headers: dict[str, str], + sns_message_type: str | None, + minio_auth_token: str | None, +) -> dict[str, str]: + try: + return await handle_s3_event_post( + body=body, + headers=headers, + sns_message_type=sns_message_type, + minio_auth_token=minio_auth_token, + ) + except Exception as exc: + logger.error(f"Failed to handle S3 event: {exc}") + return {"message": "Event handling completed"} diff --git a/apps/api/app/services/s3_events/signature_verification.py b/apps/api/app/services/s3_events/signature_verification.py new file mode 100644 index 000000000..7665ef07d --- /dev/null +++ b/apps/api/app/services/s3_events/signature_verification.py @@ -0,0 +1,39 @@ +"""Signature and token checks for storage event callbacks.""" +from __future__ import annotations + +from loguru import logger + + +def verify_sns_signature(request_body: bytes, signature: str, message: str) -> bool: + try: + return True + except Exception as exc: + logger.error(f"SNS signature verification failed: {exc}") + return False + + +def verify_minio_signature(auth_token: str, expected_token: str) -> bool: + if not expected_token: + return True + return auth_token == expected_token + + +def verify_oss_signature(request_body: bytes, headers: dict[str, str]) -> bool: + try: + from shared.core.config import settings + + if not getattr(settings, "OSS_EVENT_VERIFY_SIGNATURE", True): + return True + + callback_key = getattr(settings, "OSS_EVENT_CALLBACK_KEY", "") + if not callback_key: + logger.warning( + "OSS_EVENT_CALLBACK_KEY is not configured; skipping signature verification" + ) + return True + + # TODO: Implement OSS callback signature verification. + return True + except Exception as exc: + logger.error(f"OSS signature verification failed: {exc}") + return False diff --git a/apps/api/app/services/s3_events/subscription_service.py b/apps/api/app/services/s3_events/subscription_service.py new file mode 100644 index 000000000..140b2eba3 --- /dev/null +++ b/apps/api/app/services/s3_events/subscription_service.py @@ -0,0 +1,47 @@ +"""SNS subscription confirmation handling.""" +from __future__ import annotations + +from loguru import logger + +from shared.utils.pinned_outbound_http import send_pinned_outbound_request +from shared.utils.url_security import validate_http_url_and_resolve_ip_async + + +SNS_SUBSCRIPTION_TIMEOUT_SECONDS = 10 + + +async def confirm_sns_subscription(subscribe_url: str) -> dict[str, str]: + validation = await validate_http_url_and_resolve_ip_async(subscribe_url) + if not validation.is_valid: + logger.warning( + f"SNS subscription confirmation URL failed validation: {validation.error_message}" + ) + return {"message": "SNS subscription confirmation failed"} + + if not validation.validated_ip: + logger.warning("SNS subscription confirmation URL validation returned no IP") + return {"message": "SNS subscription confirmation failed"} + + try: + response = await send_pinned_outbound_request( + method="GET", + url=subscribe_url, + pinned_ip=validation.validated_ip, + timeout_seconds=SNS_SUBSCRIPTION_TIMEOUT_SECONDS, + ) + if response.status == 200: + logger.info("SNS subscription confirmed successfully") + return {"message": "SNS subscription confirmed"} + + if 300 <= response.status < 400: + logger.warning( + f"SNS subscription confirmation redirect blocked, status={response.status}" + ) + else: + logger.error( + f"SNS subscription confirmation failed, status={response.status}" + ) + return {"message": "SNS subscription confirmation failed"} + except Exception as exc: + logger.error(f"Failed to reach the SNS confirmation URL: {exc}") + return {"message": "SNS subscription confirmation failed"} diff --git a/apps/api/app/services/s3_events/upload_event_service.py b/apps/api/app/services/s3_events/upload_event_service.py new file mode 100644 index 000000000..df124f61f --- /dev/null +++ b/apps/api/app/services/s3_events/upload_event_service.py @@ -0,0 +1,96 @@ +"""Process storage upload-complete events into job workflow handoffs.""" +from __future__ import annotations + +import os + +from app.repositories.job_repository import JobRepository +from app.services.knowledge.kb_orchestrator import KBOrchestrator +from app.services.state_machine import JobStateMachine +from loguru import logger + +from shared.core.database import get_db_context +from shared.core.state_machine.states import JobStatus +from shared.models.schemas.s3_event import S3Event + + +def extract_job_id_from_s3_key(s3_key: str) -> str | None: + if not s3_key.startswith("uploads/"): + return None + + filename = s3_key[8:] + return os.path.splitext(filename)[0] + + +async def process_upload_events(s3_event: S3Event) -> None: + try: + upload_events = s3_event.get_upload_events() + job_repo = JobRepository() + + for event in upload_events: + s3_key = event.object_key or event.s3.get("object", {}).get("key") + if not s3_key: + continue + + job_id = extract_job_id_from_s3_key(s3_key) + if not job_id: + logger.warning(f"Could not extract job_id from S3 key: {s3_key}") + continue + + logger.info(f"Processing S3 upload event: {s3_key} -> job_id={job_id}") + + async with get_db_context() as db: + job = await job_repo.get_job_by_id(db, job_id) + if not job: + logger.warning(f"No job found for upload event: {job_id}") + continue + + if job.status != "waiting-file": + logger.info( + f"Job {job_id} is not in waiting-file status: {job.status}" + ) + continue + + from shared.core.config import settings + from shared.core.state_machine.states import is_job_expired + + if is_job_expired(job.updated_at, settings.JOB_WAITING_EXPIRE_SECONDS): + logger.warning(f"Job {job_id} upload expired, marking failed") + state_machine = JobStateMachine() + await state_machine.mark_failed( + db, + job_id, + "Upload expired: file was not uploaded within the allowed time window", + error_code="UPLOAD_EXPIRED", + ) + continue + + state_machine = JobStateMachine() + await state_machine.transition( + db, + job_id, + JobStatus.PENDING.value, + "s3_upload_completed", + None, + "system", + ) + + if job.job_type == "kb_management": + orchestrator = KBOrchestrator() + await orchestrator.start_workflow( + db=db, + job_id=job_id, + source_type="file", + file_path=None, + file_url=None, + user_id=str(job.user_id), + ) + else: + logger.warning( + f"Unsupported job type for upload event: {job.job_type}, job_id={job_id}" + ) + + logger.info(f"Triggered processing for job {job_id}") + + except Exception as exc: + logger.error(f"Failed to process upload events: {exc}") + raise diff --git a/apps/api/main.py b/apps/api/main.py index ed2c2ae6d..bd7dec810 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -83,7 +83,7 @@ async def lifespan(app: FastAPI): yield try: - from shared.services.retrieval.app_service import ( + from shared.services.retrieval.hit_stats_recorder import ( drain_retrieval_hit_stats_updates, ) diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index 75986d40e..a2bc658cb 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -452,8 +452,14 @@ async def create_checkout_session_for_credits_package( return "https://checkout.stripe.test/session/contract-package" async with developer_api_client_factory() as api_client: - billing_module = importlib.import_module("app.api.v1.routes.billing") - monkeypatch.setattr(billing_module, "StripeService", FakeStripeService) + billing_service_module = importlib.import_module( + "app.services.billing.billing_app_service" + ) + monkeypatch.setattr( + billing_service_module, + "StripeService", + FakeStripeService, + ) response = await api_client.post( "/api/v1/billing/buy-credits-package", json={"price_id": "price_contract_package", "quantity": 2}, @@ -491,8 +497,14 @@ async def create_payment_intent( } async with developer_api_client_factory() as api_client: - billing_module = importlib.import_module("app.api.v1.routes.billing") - monkeypatch.setattr(billing_module, "StripeService", FakeStripeService) + billing_service_module = importlib.import_module( + "app.services.billing.billing_app_service" + ) + monkeypatch.setattr( + billing_service_module, + "StripeService", + FakeStripeService, + ) response = await api_client.post( "/api/v1/billing/buy-credits", json={"credits_amount": 350}, diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index b6cfa1375..8cfa4993c 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -623,9 +623,18 @@ def __init__(self) -> None: def signature(self, task_name: str) -> _FakeCeleryTask: return _FakeCeleryTask(task_name) + def resolve_public_address( + host: str, + port: int | None, + *args: object, + **kwargs: object, + ) -> list[tuple[socket.AddressFamily, socket.SocketKind, int, str, tuple[str, int]]]: + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))] + import shared.core.celery_app as celery_app_module import shared.utils.http_clients as http_clients_module + monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address) monkeypatch.setattr( http_clients_module, "get_async_client", @@ -993,7 +1002,7 @@ async def _fake_start_workflow_for_job( ) async with developer_api_client_factory() as api_client: - import app.api.v1.routes.jobs as jobs_route_module + import app.services.job_upload_confirmation_service as upload_confirmation_service import shared.services.storage.file_upload_service as file_upload_service_module monkeypatch.setattr( @@ -1002,7 +1011,7 @@ async def _fake_start_workflow_for_job( _fake_verify_s3_file_exists, ) monkeypatch.setattr( - jobs_route_module, + upload_confirmation_service, "start_workflow_for_job", _fake_start_workflow_for_job, ) diff --git a/apps/api/tests/contract/test_qstash_callback_contract.py b/apps/api/tests/contract/test_qstash_callback_contract.py index 9b6734d2d..e25c3d581 100644 --- a/apps/api/tests/contract/test_qstash_callback_contract.py +++ b/apps/api/tests/contract/test_qstash_callback_contract.py @@ -48,8 +48,8 @@ async def test_should_return_unauthorized_for_an_invalid_qstash_callback_signatu monkeypatch: MonkeyPatch, ) -> None: async with api_client_factory() as api_client: - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: False) + qstash_module = importlib.import_module("app.services.qstash_callback_service") + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: False) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={"status": 200}, @@ -70,8 +70,8 @@ async def test_should_mark_the_matching_event_delivered_and_persist_a_webhook_lo async with api_client_factory() as api_client: job_id, event_id = await _insert_qstash_event() - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module("app.services.qstash_callback_service") + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={ @@ -133,8 +133,8 @@ async def test_should_keep_the_matching_event_delivering_for_retry_callback_with status="delivering", qstash_message_id="qstash-message-retry", ) - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module("app.services.qstash_callback_service") + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={ @@ -197,8 +197,8 @@ async def test_should_not_downgrade_terminal_event_when_retry_callback_arrives_l attempts=4, qstash_message_id="qstash-message-late-retry", ) - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module("app.services.qstash_callback_service") + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={ @@ -257,8 +257,8 @@ async def test_should_mark_the_matching_event_failed_and_persist_the_error_on_fa async with api_client_factory() as api_client: job_id, event_id = await _insert_qstash_event() - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module("app.services.qstash_callback_service") + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/failure", json={ @@ -318,8 +318,8 @@ async def test_should_return_ok_without_mutating_state_when_the_callback_has_no_ async with api_client_factory() as api_client: _, event_id = await _insert_qstash_event() - qstash_module = importlib.import_module("app.api.v1.routes.qstash_callbacks") - monkeypatch.setattr(qstash_module, "_verify_qstash_signature", lambda *args: True) + qstash_module = importlib.import_module("app.services.qstash_callback_service") + monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", json={ diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index a95ed4d72..cade4fbf7 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -1,12 +1,16 @@ -from collections.abc import Callable -from contextlib import AbstractAsyncContextManager +from collections.abc import AsyncIterator, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from datetime import datetime, timezone from typing import cast from uuid import uuid4 import pytest from httpx import AsyncClient +from pytest import MonkeyPatch from tests.support.contract_database import ContractDatabase +from shared.services.retrieval.agentic.types import AgenticResult +from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, WorkflowResult async def _seed_retrieval_document( @@ -16,12 +20,13 @@ async def _seed_retrieval_document( source_file_name: str, section_path: str, content: str, + chunk_id: str | None = None, ) -> dict[str, str]: document_id = f"doc_{uuid4().hex[:12]}" job_id = f"job_{uuid4().hex[:12]}" job_result_id = str(uuid4()) section_id = f"sec_{uuid4().hex[:12]}" - chunk_id = f"chunk_{uuid4().hex[:12]}" + resolved_chunk_id = chunk_id or f"chunk_{uuid4().hex[:12]}" await ContractDatabase.insert_job( job_id=job_id, @@ -67,7 +72,7 @@ async def _seed_retrieval_document( section_title=section_path.split("/")[-1], ) await ContractDatabase.insert_document_chunk( - chunk_id=chunk_id, + chunk_id=resolved_chunk_id, user_id=user_id, namespace=namespace, document_id=document_id, @@ -83,6 +88,48 @@ async def _seed_retrieval_document( "job_id": job_id, "job_result_id": job_result_id, "section_id": section_id, + "chunk_id": resolved_chunk_id, + "section_path": section_path, + } + + +async def _seed_retrieval_chunk_for_existing_document( + *, + user_id: str, + namespace: str, + document: dict[str, str], + section_path: str, + content: str, + chunk_id: str, +) -> dict[str, str]: + section_id = f"sec_{uuid4().hex[:12]}" + + await ContractDatabase.insert_document_section( + section_id=section_id, + user_id=user_id, + namespace=namespace, + document_id=document["document_id"], + job_result_id=document["job_result_id"], + section_path=section_path, + section_title=section_path.split("/")[-1], + ) + await ContractDatabase.insert_document_chunk( + chunk_id=chunk_id, + user_id=user_id, + namespace=namespace, + document_id=document["document_id"], + job_result_id=document["job_result_id"], + section_id=section_id, + chunk_type="text", + content=content, + section_path=section_path, + ) + + return { + "document_id": document["document_id"], + "job_id": document["job_id"], + "job_result_id": document["job_result_id"], + "section_id": section_id, "chunk_id": chunk_id, "section_path": section_path, } @@ -187,6 +234,622 @@ async def test_should_return_empty_results_for_an_empty_query( } +@pytest.mark.asyncio +async def test_retrieval_hit_stats_should_use_the_current_database_context( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + @asynccontextmanager + async def stale_db_context() -> AsyncIterator[object]: + raise RuntimeError("stale db context should not be used") + yield object() + + async with developer_api_client_factory() as _api_client: + seeded_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-hit-stats", + source_file_name="hit-stats.pdf", + section_path="hit-stats/section", + content="hit stats content", + ) + + from shared.services.retrieval import hit_stats_recorder + + monkeypatch.setattr( + hit_stats_recorder, + "get_db_context", + stale_db_context, + raising=False, + ) + hit_stats_recorder.schedule_retrieval_hit_stats_update( + user_id="local-dev-user", + namespace="contract-hit-stats", + results=[ + { + "document_id": seeded_document["document_id"], + "chunk_id": seeded_document["chunk_id"], + } + ], + ) + await hit_stats_recorder.drain_retrieval_hit_stats_updates() + + hit_stats = await ContractDatabase.fetch_all( + """ + SELECT hit_kind, document_id, chunk_id, hit_count + FROM retrieval_hit_stats + WHERE user_id = :user_id + AND namespace = :namespace + AND document_id = :document_id + ORDER BY hit_kind + """, + { + "user_id": "local-dev-user", + "namespace": "contract-hit-stats", + "document_id": seeded_document["document_id"], + }, + ) + + hit_stats_by_kind = { + cast(str, row["hit_kind"]): row for row in hit_stats + } + + assert set(hit_stats_by_kind) == {"chunk", "document"} + assert hit_stats_by_kind["chunk"]["chunk_id"] == seeded_document["chunk_id"] + assert hit_stats_by_kind["document"]["chunk_id"] is None + assert hit_stats_by_kind["chunk"]["hit_count"] == 1 + assert hit_stats_by_kind["document"]["hit_count"] == 1 + + +@pytest.mark.asyncio +async def test_legacy_retrieval_should_rank_hot_chunk_before_cold_chunk_when_discovery_scores_tie( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + async with developer_api_client_factory() as api_client: + cold_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-hot-ranking", + source_file_name="cold.pdf", + section_path="ranking/cold", + content="same ranking marker cold", + ) + hot_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-hot-ranking", + source_file_name="hot.pdf", + section_path="ranking/hot", + content="same ranking marker hot", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-hot-ranking", + source_file_name="filler.pdf", + section_path="ranking/filler", + content="same ranking marker filler", + ) + + now = datetime.now(timezone.utc).replace(tzinfo=None) + await ContractDatabase.execute( + """ + INSERT INTO retrieval_hit_stats ( + id, + user_id, + namespace, + hit_kind, + document_id, + chunk_id, + hit_count, + last_hit_at, + created_at, + updated_at + ) VALUES ( + :id, + :user_id, + :namespace, + 'chunk', + :document_id, + :chunk_id, + :hit_count, + :now, + :now, + :now + ) + """, + { + "id": f"rhs_{uuid4().hex[:12]}", + "user_id": "local-dev-user", + "namespace": "contract-hot-ranking", + "document_id": hot_document["document_id"], + "chunk_id": hot_document["chunk_id"], + "hit_count": 100, + "now": now, + }, + ) + + def to_channel_row(document: dict[str, str]) -> dict[str, object]: + return { + "document_id": document["document_id"], + "chunk_id": document["chunk_id"], + "section_id": document["section_id"], + "section_path": document["section_path"], + "source_file_name": "cold.pdf" + if document["document_id"] == cold_document["document_id"] + else "hot.pdf", + "chunk_type": "text", + "content": "same ranking marker", + "score": 1.0, + "file_path": None, + "chunk_metadata": {}, + "job_result_id": document["job_result_id"], + "job_id": document["job_id"], + "sort_order": 0, + } + + async def fake_content_channel(*_args: object, **_kwargs: object) -> list[dict[str, object]]: + return [ + to_channel_row(hot_document), + to_channel_row(cold_document), + ] + + async def fake_path_channel(*_args: object, **_kwargs: object) -> list[dict[str, object]]: + return [ + to_channel_row(cold_document), + to_channel_row(hot_document), + ] + + async def fake_graph_routing(*_args: object, **_kwargs: object) -> list[dict[str, object]]: + return [] + + monkeypatch.setattr( + "shared.services.retrieval.app_service.path_channel", + fake_path_channel, + ) + monkeypatch.setattr( + "shared.services.retrieval.app_service.content_channel", + fake_content_channel, + ) + monkeypatch.setattr( + "shared.services.retrieval.app_service.list_graph_routed_chunks", + fake_graph_routing, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-hot-ranking", + "query": "same ranking marker", + "top_k": 1, + "channels": ["path", "content"], + "channel_weights": {"path": 1.0, "content": 1.0}, + "use_agentic": False, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], response_json["results"]) + + assert len(results) == 1 + assert results[0]["source"]["document_id"] == hot_document["document_id"] + + +@pytest.mark.asyncio +async def test_agentic_retrieval_should_reference_root_only_document_content( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("LLM_MOCK_ENABLED", "true") + + async with developer_api_client_factory() as api_client: + rooted_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-root-retrieval", + source_file_name="root-only.pdf", + section_path="Root", + content="root only diluted earnings marker content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-root-retrieval", + source_file_name="filler.pdf", + section_path="filler/section", + content="unrelated filler content", + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-root-retrieval", + "query": "diluted earnings marker", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + assert response_json["router_used"] == "workflow_single_step" + assert { + "chunk_id": rooted_document["chunk_id"], + "document_id": rooted_document["document_id"], + "chunk_type": "text", + "section_path": "root-only.pdf", + "file_path": None, + "job_id": rooted_document["job_id"], + } in referenced_chunks + assert results[0]["content"] == "root only diluted earnings marker content" + assert results[0]["source"] == { + "document_id": rooted_document["document_id"], + "source_file_name": "root-only.pdf", + "section_path": "Root", + } + + +@pytest.mark.asyncio +async def test_agentic_retrieval_should_reference_discovery_content_when_navigation_selects_nothing( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("LLM_MOCK_ENABLED", "true") + + async with developer_api_client_factory() as api_client: + discovered_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-discovery-fallback", + source_file_name="discovery.pdf", + section_path="Findings", + content="discovery fallback EBITDA margin marker content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-discovery-fallback", + source_file_name="filler.pdf", + section_path="filler/section", + content="unrelated filler content", + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-discovery-fallback", + "query": "EBITDA margin marker", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + assert response_json["router_used"] == "workflow_single_step" + assert { + "chunk_id": discovered_document["chunk_id"], + "document_id": discovered_document["document_id"], + "chunk_type": "text", + "section_path": discovered_document["section_path"], + "file_path": None, + "job_id": discovered_document["job_id"], + } in referenced_chunks + assert results[0]["content"] == "discovery fallback EBITDA margin marker content" + assert results[0]["source"] == { + "document_id": discovered_document["document_id"], + "source_file_name": "discovery.pdf", + "section_path": discovered_document["section_path"], + } + + +@pytest.mark.asyncio +async def test_agentic_retrieval_should_not_hydrate_references_outside_request_scope( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + class FakeWorkflowOrchestrator: + async def run(self, *_args: object, **kwargs: object) -> WorkflowResult: + return WorkflowResult( + namespace=str(kwargs["namespace"]), + query=str(kwargs["query"]), + router_used="workflow_single_step", + answer_text="foreign reference answer", + referenced_chunks=[ + { + "chunk_id": foreign_document["chunk_id"], + "document_id": foreign_document["document_id"], + "chunk_type": "text", + "section_path": foreign_document["section_path"], + "file_path": None, + "job_id": foreign_document["job_id"], + } + ], + ) + + async with developer_api_client_factory() as api_client: + request_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-visible-scope", + source_file_name="visible.pdf", + section_path="visible/section", + content="visible scoped content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-visible-scope", + source_file_name="visible-filler.pdf", + section_path="visible/filler", + content="visible scoped filler content", + ) + foreign_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-foreign-scope", + source_file_name="foreign.pdf", + section_path="foreign/section", + content="foreign scoped content should not leak", + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", + FakeWorkflowOrchestrator, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-visible-scope", + "query": "visible", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + assert request_document["document_id"] != foreign_document["document_id"] + assert referenced_chunks == [] + assert results == [] + + +@pytest.mark.asyncio +async def test_agentic_workflow_should_preserve_references_with_the_same_chunk_id_across_documents( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + shared_chunk_id = f"chunk_{uuid4().hex[:12]}" + + async def fake_plan( + self: object, + *, + query: str, + kb_total_docs: int = 0, + kb_total_chunks: int = 0, + ) -> QueryPlan: + return QueryPlan( + original_query=query, + steps=[ + PlannedStep(id="first", sub_query="first shared reference"), + PlannedStep(id="second", sub_query="second shared reference"), + ], + final_strategy="concat_final_parts", + reasoning_summary=( + f"forced two-step contract plan for {kb_total_docs} docs " + f"and {kb_total_chunks} chunks" + ), + ) + + async def fake_retrieval_run( + self: object, + db: object, + **kwargs: object, + ) -> AgenticResult: + query = str(kwargs["query"]) + document = first_document if query == "first shared reference" else second_document + return AgenticResult( + evidence_text=f"evidence for {document['document_id']}", + answer_text=f"answer for {document['document_id']}", + referenced_chunks=[ + { + "chunk_id": shared_chunk_id, + "document_id": document["document_id"], + "chunk_type": "text", + "section_path": document["section_path"], + "file_path": "", + "job_id": document["job_id"], + } + ], + router_used="contract_fake_agent", + ) + + async with developer_api_client_factory() as api_client: + first_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-shared-chunk-id", + source_file_name="first.pdf", + section_path="first/section", + content="shared deterministic content", + chunk_id=shared_chunk_id, + ) + second_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-shared-chunk-id", + source_file_name="second.pdf", + section_path="second/section", + content="shared deterministic content", + chunk_id=shared_chunk_id, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.planner.QueryPlanner.plan", + fake_plan, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator.RetrievalAgent.run", + fake_retrieval_run, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-shared-chunk-id", + "query": "show both shared references", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + referenced_document_ids = { + cast(str, reference["document_id"]) for reference in referenced_chunks + } + result_document_ids = { + cast(str, result["source"]["document_id"]) for result in results + } + + assert referenced_document_ids == { + first_document["document_id"], + second_document["document_id"], + } + assert result_document_ids == { + first_document["document_id"], + second_document["document_id"], + } + + +@pytest.mark.asyncio +async def test_agentic_workflow_should_preserve_references_with_the_same_chunk_id_across_sections( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + shared_chunk_id = f"chunk_{uuid4().hex[:12]}" + + async def fake_plan( + self: object, + *, + query: str, + kb_total_docs: int = 0, + kb_total_chunks: int = 0, + ) -> QueryPlan: + return QueryPlan( + original_query=query, + steps=[ + PlannedStep(id="first", sub_query="first shared section"), + PlannedStep(id="second", sub_query="second shared section"), + ], + final_strategy="concat_final_parts", + reasoning_summary=( + f"forced section identity contract plan for {kb_total_docs} docs " + f"and {kb_total_chunks} chunks" + ), + ) + + async def fake_retrieval_run( + self: object, + db: object, + **kwargs: object, + ) -> AgenticResult: + query = str(kwargs["query"]) + chunk = first_chunk if query == "first shared section" else second_chunk + return AgenticResult( + evidence_text=f"evidence for {chunk['section_path']}", + answer_text=f"answer for {chunk['section_path']}", + referenced_chunks=[ + { + "chunk_id": shared_chunk_id, + "document_id": chunk["document_id"], + "chunk_type": "text", + "section_path": chunk["section_path"], + "file_path": "", + "job_id": chunk["job_id"], + } + ], + router_used="contract_fake_agent", + ) + + async with developer_api_client_factory() as api_client: + first_chunk = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-shared-section-chunk-id", + source_file_name="same-document.pdf", + section_path="first/section", + content="repeated deterministic content", + chunk_id=shared_chunk_id, + ) + second_chunk = await _seed_retrieval_chunk_for_existing_document( + user_id="local-dev-user", + namespace="contract-shared-section-chunk-id", + document=first_chunk, + section_path="second/section", + content="repeated deterministic content", + chunk_id=shared_chunk_id, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.planner.QueryPlanner.plan", + fake_plan, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator.RetrievalAgent.run", + fake_retrieval_run, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-shared-section-chunk-id", + "query": "show both shared section references", + "top_k": 1, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) + results = cast(list[dict[str, object]], response_json["results"]) + + referenced_section_paths = { + cast(str, reference["section_path"]) for reference in referenced_chunks + } + result_section_paths = { + cast(str, result["source"]["section_path"]) for result in results + } + + assert referenced_section_paths == { + first_chunk["section_path"], + second_chunk["section_path"], + } + assert result_section_paths == { + first_chunk["section_path"], + second_chunk["section_path"], + } + + @pytest.mark.asyncio async def test_should_return_request_validation_failure_for_an_invalid_channel( developer_api_client_factory: Callable[ diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index 949722136..ac5bf3a27 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -93,8 +93,14 @@ async def start_workflow( async with api_client_factory() as api_client: user_id, job_id = await _insert_waiting_file_job() - s3_events_module = importlib.import_module("app.api.v1.routes.s3_events") - monkeypatch.setattr(s3_events_module, "KBOrchestrator", FakeKBOrchestrator) + upload_event_service = importlib.import_module( + "app.services.s3_events.upload_event_service" + ) + monkeypatch.setattr( + upload_event_service, + "KBOrchestrator", + FakeKBOrchestrator, + ) response = await api_client.post( "/api/v1/internal/s3-events", json=_build_s3_event_payload(job_id), @@ -139,8 +145,14 @@ async def start_workflow( async with api_client_factory() as api_client: _, job_id = await _insert_waiting_file_job() - s3_events_module = importlib.import_module("app.api.v1.routes.s3_events") - monkeypatch.setattr(s3_events_module, "KBOrchestrator", FakeKBOrchestrator) + upload_event_service = importlib.import_module( + "app.services.s3_events.upload_event_service" + ) + monkeypatch.setattr( + upload_event_service, + "KBOrchestrator", + FakeKBOrchestrator, + ) response = await api_client.post( "/api/v1/internal/s3-events", content=json.dumps( diff --git a/apps/worker/app/core/tasks/kb_tasks.py b/apps/worker/app/core/tasks/kb_tasks.py index 19acc3605..5a4416a90 100644 --- a/apps/worker/app/core/tasks/kb_tasks.py +++ b/apps/worker/app/core/tasks/kb_tasks.py @@ -5,69 +5,21 @@ All I/O operations use sync services that yield cooperatively under gevent. """ -import os -from datetime import datetime, timezone - -import pandas as pd - # Base task class from app.core.tasks.base_task import KBBaseTask -from app.core.tasks.task_utils import ( - cleanup_task_workspace, - create_task_workspace, - download_s3_file_to_temp, -) -from app.services.common.job_start_service import mark_job_running -from app.services.document_parser.stage_profiler import stage_timer - -# Storage operations -from app.services.storage.sync_storage_service import ( - download_file_from_url, - generate_download_url, - upload_to_s3, - verify_s3_file_exists, -) -from app.services.workload.page_estimator import PageEstimator +from app.services.workload.parse_job_service import parse_uploaded_file_job +from app.services.workload.url_upload_service import upload_url_file from loguru import logger -from sqlalchemy import select from shared.core.celery_app import get_celery_app from shared.core.config import settings -from shared.core.database_sync import get_sync_db_context from shared.core.exceptions import RETRYABLE_EXCEPTIONS # Exception handling from shared.core.exceptions.domain_exceptions import ( - InsufficientCreditsException, - NotFoundException, - StorageServiceException, - ValidationException, WorkerHandlingException, ) from shared.core.logging import LogEvent, log_context -from shared.models.database.job import Job - -# Domain services -from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.billing.work_billing_service import WorkBillingService -from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks -from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service -from shared.services.redis.distributed_lock import RedisJobLock - -# Sync services for gevent worker -from shared.services.redis.redis_sync_service import ( - SyncJobInfoRedisService, - SyncJobMetadataService, - SyncRedisServiceFactory, -) -from shared.services.storage.result_storage import get_result_storage -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, -) # Get Celery application celery_app = get_celery_app() @@ -112,143 +64,7 @@ def _upload_url_file( job_id: str, source_url: str, user_id: str | None, job_type: str | None = None ): """Sync URL file download and upload to S3.""" - lifecycle_service = get_sync_job_lifecycle_service() - - # Get job info from Redis - redis_service = SyncRedisServiceFactory.get_service() - job_info_service = SyncJobInfoRedisService(redis_service) - job_info = job_info_service.get_job_info(job_id) - - if not job_info: - metadata_service = SyncJobMetadataService(redis_service) - job_metadata = metadata_service.get_metadata(job_id) - if job_metadata: - s3_key = job_metadata.get("s3_key") - else: - raise NotFoundException( - resource="JobInfo", - resource_id=job_id, - internal_message="Job info not found in Redis or Metadata", - ) - else: - s3_key = job_info.get("s3_key") - - if not s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id="s3_key", - internal_message=f"Missing s3_key in Redis job info for job_id={job_id}", - ) - - # Publish progress: validating file type - lifecycle_service.update_progress( - job_id, progress=3, message="Validating URL file type..." - ) - - # Step 1: Validate URL file type (path first, then Content-Type header) - from shared.utils.url_file_type import resolve_file_extension_sync - - file_extension = resolve_file_extension_sync(source_url) - - if not file_extension: - all_supported_extensions = settings.get_supported_extensions() - supported_formats = ", ".join(sorted(all_supported_extensions)) - raise ValidationException( - user_message="Unsupported file type", - violations=[ - { - "field": "file_extension", - "description": f"Must be one of: {supported_formats}", - } - ], - ) - - # Publish progress: downloading - lifecycle_service.update_progress( - job_id, progress=10, message="Downloading file from URL..." - ) - - # Step 2: Download file to temp directory - try: - temp_file_path = download_file_from_url(source_url) - except Exception as e: - raise ValidationException( - user_message="Failed to download file from URL", - violations=[ - { - "field": "source_url", - "description": "Could not download file from the provided URL", - } - ], - internal_message=f"Failed to download file from URL: {source_url}, error: {e}", - ) - - try: - # Publish progress: validating file size - lifecycle_service.update_progress( - job_id, progress=30, message="Validating file size..." - ) - - # Step 3: Validate file size - file_size = os.path.getsize(temp_file_path) - - if file_size > settings.MAX_FILE_SIZE: - limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) - raise ValidationException( - user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", - violations=[ - { - "field": "file_size", - "description": f"Size {file_size} bytes exceeds limit of {settings.MAX_FILE_SIZE} bytes", - } - ], - ) - - # Publish progress: uploading to S3 - lifecycle_service.update_progress( - job_id, progress=50, message="Uploading file to S3..." - ) - - # Step 4: Upload to S3 - uploads_bucket = settings.S3_BUCKET_NAME - upload_to_s3(temp_file_path, s3_key, uploads_bucket) - logger.info(f"File uploaded to S3: {s3_key}") - - finally: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - logger.debug(f"Temp file cleaned up: {temp_file_path}") - - # Publish progress: verifying upload - lifecycle_service.update_progress( - job_id, progress=80, message="Verifying upload result..." - ) - - # Step 5: Verify S3 file exists - file_info = verify_s3_file_exists(s3_key) - if not file_info.get("exists"): - raise StorageServiceException( - user_message="We failed to verify your file upload", - internal_message=f"S3 file verification failed for {s3_key}", - ) - - # Publish progress: complete - lifecycle_service.update_progress( - job_id, - progress=100, - message="URL file upload complete, waiting for processing...", - ) - - logger.info( - f"URL file upload complete, waiting for S3 webhook: {job_id} -> {s3_key}" - ) - - return { - "status": "success", - "job_id": job_id, - "s3_key": s3_key, - "file_size": file_info.get("size"), - } + return upload_url_file(job_id, source_url, user_id, job_type) @celery_app.task( @@ -283,430 +99,4 @@ def parse_task( def _parse(job_id: str, user_id: str | None): """Sync parse and vectorize (file already uploaded to S3).""" - logger.info(f"Parse started: job_id={job_id}, user_id={user_id}") - lifecycle_service = get_sync_job_lifecycle_service() - - # Get job info from Redis (sync) - redis_service = SyncRedisServiceFactory.get_service() - job_info_service = SyncJobInfoRedisService(redis_service) - job_info = job_info_service.get_job_info(job_id) - - if not job_info: - # Redis JobInfo has expired or been flushed — fall back to the DB, which is - # the durable source of truth for s3_key and user_id written at job creation. - logger.warning( - f"JobInfo not found in Redis for job_id={job_id}; falling back to database" - ) - with get_sync_db_context() as fallback_db: - job_row = fallback_db.execute( - select(Job).where(Job.job_id == job_id) - ).scalar_one_or_none() - - if not job_row or not job_row.s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id=job_id, - internal_message="job info not found in Redis or database", - ) - - s3_key: str = job_row.s3_key - job_user_id: str | None = str(job_row.user_id) if job_row.user_id else user_id - logger.info( - f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}" - ) - else: - raw_s3_key = job_info.get("s3_key") - if not isinstance(raw_s3_key, str) or not raw_s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id="s3_key", - internal_message="Missing s3_key in job_info", - ) - - s3_key = raw_s3_key - raw_job_user_id = job_info.get("user_id") - job_user_id = raw_job_user_id if isinstance(raw_job_user_id, str) else user_id - - # Verify S3 file exists (sync) - file_info = verify_s3_file_exists(s3_key) - if not file_info.get("exists"): - raise NotFoundException( - resource="S3File", - resource_id=s3_key, - internal_message=f"S3 file not found: {s3_key}", - ) - - logger.info(f"S3 file verified: {s3_key}") - - # Validate file size - file_size = file_info.get("size", 0) - file_extension = os.path.splitext(s3_key)[1].lower() - - if file_size > settings.MAX_FILE_SIZE: - limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) - raise ValidationException( - user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", - violations=[ - { - "field": "file_size", - "description": f"Size {file_size} bytes exceeds limit of {settings.MAX_FILE_SIZE} bytes", - } - ], - ) - - # Get job_metadata from Redis - metadata_service = SyncJobMetadataService(redis_service) - job_metadata = metadata_service.get_metadata(job_id) - if not job_metadata: - raise NotFoundException( - resource="JobMetadata", - resource_id=job_id, - internal_message=f"Job metadata not found for job_id={job_id}", - ) - - should_process = mark_job_running(job_id, redis_service) - if not should_process: - logger.warning(f"Skipping parse_task for inactive job: job_id={job_id}") - return { - "status": "skipped", - "job_id": job_id, - "reason": "job_already_terminal", - } - - # Acquire distributed lock to prevent concurrent processing of the same - # job when the broker redelivers a task before the original worker acks. - # If another worker already holds the lock, UnavailableException is raised - # and Celery auto-retries after KB_TASK_RETRY_COUNTDOWN seconds. - with RedisJobLock(redis_service, job_id): - task_workspace_dir = create_task_workspace(job_id) - input_dir = os.path.join(task_workspace_dir, "input") - output_dir = os.path.join(task_workspace_dir, "output") - os.makedirs(input_dir, exist_ok=True) - os.makedirs(output_dir, exist_ok=True) - logger.info( - f"Task workspace ready: job_id={job_id}, workspace={task_workspace_dir}" - ) - - try: - # Publish progress: start parsing - lifecycle_service.update_progress( - job_id, progress=10, message="Parsing document..." - ) - - # Generate download URL and download file (sync) - file_url_response = generate_download_url(s3_key, settings.S3_BUCKET_NAME) - file_url = file_url_response["download_url"] - - filename = JobMetadataHelper.get_field(job_metadata, "source_file_name") - - # Download file to the task workspace - page_count = 1 - - # Derive file extension from s3_key (always has the correct extension) - # rather than filename, which may not have a real extension for URLs - # like arxiv.org/pdf/1706.03762 - file_ext = os.path.splitext(s3_key)[1].lower() if s3_key else "" - local_temp_path = download_s3_file_to_temp(file_url, file_ext, input_dir) - - logger.info( - f"File downloaded: job_id={job_id}, local_path={local_temp_path}" - ) - - from app.services.document_parser.internal_parse_name import ( - prepare_internal_parse_input, - ) - from app.services.document_parser.parse_service import ( - checkerboard_inject_parse, - ) - - prepared_parse_input = prepare_internal_parse_input( - local_temp_path, - filename, - fallback_ext=file_ext, - prefer_fallback_ext=True, - ) - internal_parse_name = prepared_parse_input.internal_filename - local_temp_path = prepared_parse_input.file_path - logger.info( - f"File prepared for parsing: job_id={job_id}, " - f"internal_filename={internal_parse_name}, local_path={local_temp_path}" - ) - - # Estimate workload - page_count = PageEstimator.estimate(local_temp_path) - logger.info( - f"Workload estimation: job_id={job_id}, page_count={page_count}" - ) - - processing_started_at = datetime.now(timezone.utc) - - if not job_user_id: - raise NotFoundException( - resource="JobInfo", - resource_id="user_id", - internal_message=f"Missing user_id in job info for job_id={job_id}", - ) - - billing_service = WorkBillingService() - billing_status = "skipped" - billing_amount_micro_dollars = 0 - billing_credits = 0.0 - with get_sync_db_context() as db: - job_result = db.execute( - select(Job).where(Job.job_id == job_id).with_for_update() - ) - job = job_result.scalar_one_or_none() - - if job and getattr(job, "billing_status", "") == "charged": - logger.info(f"Job already charged: {job_id}") - billing_status = "charged" - billing_amount_micro_dollars = int(job.credits_charged or 0) - billing_credits = billing_amount_micro_dollars / 1_000_000 - else: - try: - billing_result = billing_service.charge_for_pages( - session=db, - user_id=job_user_id, - page_count=page_count, - filename=filename, - ) - except InsufficientCreditsException: - logger.warning( - f"Billing failed: job_id={job_id}, user_id={job_user_id}" - ) - billing_amount = billing_service.estimate_page_charge( - page_count=page_count - ) - if job: - job.page_count = page_count - job.credits_charged = billing_amount.amount_micro_dollars - job.billing_status = "billing_failed" - db.commit() - - raise InsufficientCreditsException( - user_message=( - "Insufficient credits to process this document " - f"({page_count} pages required, cost: " - f"{billing_amount.credits})." - ), - required_credits=billing_amount.credits, - internal_message=( - f"job_id={job_id}, user_id={job_user_id}, " - f"page_count={page_count}" - ), - ) - - billing_status = billing_result.billing_status - billing_amount_micro_dollars = billing_result.amount_micro_dollars - billing_credits = billing_result.credits - if job: - job.page_count = page_count - job.credits_charged = billing_amount_micro_dollars - job.billing_status = billing_status - - # Store billing info in Redis - metadata_updates = { - "page_count": page_count, - "billing_status": billing_status, - "billing_amount_micro_dollars": billing_amount_micro_dollars, - "billing_credits": billing_credits, - "processing_started_at": processing_started_at.isoformat(), - } - metadata_service.update_metadata(job_id, metadata_updates) - job_metadata.update(metadata_updates) - - doc_type = JobMetadataHelper.get_parsing_param( - job_metadata, "doc_type", "auto" - ) - logger.info( - f"Start parse: job_id={job_id}, filename={filename}, " - f"internal_filename={internal_parse_name}, type={doc_type}" - ) - - with stage_timer( - "worker.parse.document", - job_id=job_id, - filename=filename, - doc_type=doc_type, - ): - add_dir, add_contents_df = checkerboard_inject_parse( - file_full_path=local_temp_path, - filename=filename, - output_dir=output_dir, - job_id=job_id, - internal_output_filename=internal_parse_name, - kb_dir=JobMetadataHelper.get_parsing_param( - job_metadata, "kb_dir", "Default_Root" - ), - doc_type=doc_type, - smart_title_parse=JobMetadataHelper.get_parsing_param( - job_metadata, "smart_title_parse", True - ), - summary_image=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_image", True - ), - summary_table=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_table", True - ), - summary_txt=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_txt", True - ), - add_frag_desc=JobMetadataHelper.get_parsing_param( - job_metadata, "add_frag_desc", "" - ), - s3_key=s3_key, - ) - parsed_contents_df: pd.DataFrame | None = add_contents_df - - logger.info( - f"File parsing completed: job_id={job_id}, add_dir={add_dir}, chunks={len(parsed_contents_df) if parsed_contents_df is not None else 0}" - ) - - if parsed_contents_df is None: - raise WorkerHandlingException( - user_message="We could not extract content from your file", - internal_message="File parsing failed, no content returned from parser", - ) - - if parsed_contents_df.empty: - logger.warning( - f"No content returned from file parsing: job_id={job_id}, filename={filename}" - ) - - lifecycle_service.update_progress( - job_id, progress=30, message="Parse completed, preparing chunks..." - ) - - chunks = dataframe_to_chunks(parsed_contents_df) - - lifecycle_service.update_progress( - job_id, progress=70, message="Chunks ready, generating zip..." - ) - logger.info(f"Chunks prepared: job_id={job_id}, count={len(chunks)}") - - # Get source file name - source_file_name = JobMetadataHelper.get_field( - job_metadata, "source_file_name" - ) or JobMetadataHelper.get_field(job_metadata, "source_url") - if isinstance(source_file_name, str) and "/" in source_file_name: - source_file_name = os.path.basename(source_file_name) - - document_top_summary = "" - section_summaries: dict[str, str] = {} - if add_dir and source_file_name: - if add_contents_df is not None and "path" in add_contents_df.columns: - ensure_doc_nav_json( - str(add_dir), - chunks, - source_file_name=str(source_file_name), - ) - # Enrich non-leaf section summaries (bottom-up aggregation) - try: - kb_dir_for_enrich = os.path.dirname(str(add_dir)) - summary_use_llm = JobMetadataHelper.get_parsing_param( - job_metadata, "summary_use_llm", False - ) - enrich_doc_nav_summaries( - kb_dir_for_enrich, - source_file=str(source_file_name), - use_llm=summary_use_llm, - ) - section_summaries = build_section_summary_lookup(str(add_dir)) - except Exception as _e: - logger.warning(f"doc_nav enrichment failed (non-fatal): {_e}") - document_top_summary = load_nav_top_summary(str(add_dir), str(source_file_name)) - if document_top_summary: - for chunk in chunks: - metadata = chunk.get("metadata") - if not isinstance(metadata, dict): - metadata = {} - chunk["metadata"] = metadata - metadata["document_top_summary"] = document_top_summary - - data_id = JobMetadataHelper.get_field(job_metadata, "data_id") - - lifecycle_service.update_progress( - job_id, progress=80, message="Generating ZIP package..." - ) - processing_completed_at = datetime.now(timezone.utc) - processing_timing_updates = { - "processing_completed_at": processing_completed_at.isoformat(), - "processing_duration_ms": max( - 0, - int( - ( - processing_completed_at - processing_started_at - ).total_seconds() - * 1000 - ), - ), - } - metadata_service.update_metadata(job_id, processing_timing_updates) - job_metadata.update(processing_timing_updates) - - # Generate ZIP package - zip_service = ZipResultService() - zip_file_path, checksum, statistics, zip_size = ( - zip_service.generate_zip_package( - job_id=job_id, - chunks=chunks, - add_dir=str(add_dir) if add_dir else "", - source_file_name=source_file_name, - data_id=data_id, - job_metadata=job_metadata, - parsed_df=parsed_contents_df, - temp_dir=task_workspace_dir, - ) - ) - - checksum_value = ( - checksum.get("value", "") - if isinstance(checksum, dict) - else (checksum or "") - ) - - lifecycle_service.update_progress( - job_id, progress=90, message="Uploading results to S3..." - ) - - result_bundle = get_result_storage().upload( - job_id=job_id, - result_dir=str(add_dir) if add_dir else "", - zip_file_path=zip_file_path, - ) - result_s3_key = result_bundle.zip_key - - stored_count = 0 - - lifecycle_service.update_progress( - job_id, progress=100, message="Task complete!" - ) - - # Finalize job success directly to the database - lifecycle_service.finalize_job_success( - job_id=job_id, - chunks=chunks, - result_s3_key=result_s3_key, - checksum=checksum_value, - zip_size=zip_size, - stored_count=stored_count, - delivery_mode="url", - section_summaries=section_summaries, - ) - - logger.info( - f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" - ) - - return { - "status": "success", - "job_id": job_id, - "add_dir": None, - "vectors_count": 0, - "contents_count": len(parsed_contents_df), - "stored_count": stored_count, - "delivery_mode": "url", - "result_s3_key": result_s3_key, - } - finally: - cleanup_task_workspace(task_workspace_dir) + return parse_uploaded_file_job(job_id, user_id) diff --git a/apps/worker/app/services/workload/parse_job_service.py b/apps/worker/app/services/workload/parse_job_service.py new file mode 100644 index 000000000..f8b6d453f --- /dev/null +++ b/apps/worker/app/services/workload/parse_job_service.py @@ -0,0 +1,467 @@ +from __future__ import annotations + +import os +from datetime import datetime, timezone + +import pandas as pd +from app.core.tasks.task_utils import ( + cleanup_task_workspace, + create_task_workspace, + download_s3_file_to_temp, +) +from app.services.common.job_start_service import mark_job_running +from app.services.connect_builder.summary_builder import ( + build_section_summary_lookup, + enrich_doc_nav_summaries, + ensure_doc_nav_json, + load_nav_top_summary, +) +from app.services.document_parser.stage_profiler import stage_timer +from app.services.storage.sync_storage_service import ( + generate_download_url, + verify_s3_file_exists, +) +from app.services.workload.page_estimator import PageEstimator +from loguru import logger +from sqlalchemy import select + +from shared.core.config import settings +from shared.core.database_sync import get_sync_db_context +from shared.core.exceptions.domain_exceptions import ( + InsufficientCreditsException, + NotFoundException, + ValidationException, + WorkerHandlingException, +) +from shared.models.database.job import Job +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.billing.work_billing_service import WorkBillingService +from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks +from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service +from shared.services.redis.distributed_lock import RedisJobLock +from shared.services.redis.redis_sync_service import ( + SyncJobInfoRedisService, + SyncJobMetadataService, + SyncRedisServiceFactory, +) +from shared.services.storage.result_storage import get_result_storage +from shared.services.storage.zip_result_service import ZipResultService + + +def parse_uploaded_file_job(job_id: str, user_id: str | None) -> dict[str, object]: + logger.info(f"Parse started: job_id={job_id}, user_id={user_id}") + lifecycle_service = get_sync_job_lifecycle_service() + + redis_service = SyncRedisServiceFactory.get_service() + job_info_service = SyncJobInfoRedisService(redis_service) + job_info = job_info_service.get_job_info(job_id) + + if not job_info: + logger.warning( + f"JobInfo not found in Redis for job_id={job_id}; falling back to database" + ) + with get_sync_db_context() as fallback_db: + job_row = fallback_db.execute( + select(Job).where(Job.job_id == job_id) + ).scalar_one_or_none() + + if not job_row or not job_row.s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id=job_id, + internal_message="job info not found in Redis or database", + ) + + s3_key: str = job_row.s3_key + job_user_id: str | None = str(job_row.user_id) if job_row.user_id else user_id + logger.info( + f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}" + ) + else: + raw_s3_key = job_info.get("s3_key") + if not isinstance(raw_s3_key, str) or not raw_s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id="s3_key", + internal_message="Missing s3_key in job_info", + ) + + s3_key = raw_s3_key + raw_job_user_id = job_info.get("user_id") + job_user_id = raw_job_user_id if isinstance(raw_job_user_id, str) else user_id + + file_info = verify_s3_file_exists(s3_key) + if not file_info.get("exists"): + raise NotFoundException( + resource="S3File", + resource_id=s3_key, + internal_message=f"S3 file not found: {s3_key}", + ) + + logger.info(f"S3 file verified: {s3_key}") + + file_size = file_info.get("size", 0) + file_extension = os.path.splitext(s3_key)[1].lower() + + if file_size > settings.MAX_FILE_SIZE: + limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) + raise ValidationException( + user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", + violations=[ + { + "field": "file_size", + "description": ( + f"Size {file_size} bytes exceeds limit of " + f"{settings.MAX_FILE_SIZE} bytes" + ), + } + ], + ) + + metadata_service = SyncJobMetadataService(redis_service) + job_metadata = metadata_service.get_metadata(job_id) + if not job_metadata: + raise NotFoundException( + resource="JobMetadata", + resource_id=job_id, + internal_message=f"Job metadata not found for job_id={job_id}", + ) + + should_process = mark_job_running(job_id, redis_service) + if not should_process: + logger.warning(f"Skipping parse_task for inactive job: job_id={job_id}") + return { + "status": "skipped", + "job_id": job_id, + "reason": "job_already_terminal", + } + + with RedisJobLock(redis_service, job_id): + task_workspace_dir = create_task_workspace(job_id) + input_dir = os.path.join(task_workspace_dir, "input") + output_dir = os.path.join(task_workspace_dir, "output") + os.makedirs(input_dir, exist_ok=True) + os.makedirs(output_dir, exist_ok=True) + logger.info( + f"Task workspace ready: job_id={job_id}, workspace={task_workspace_dir}" + ) + + try: + lifecycle_service.update_progress( + job_id, progress=10, message="Parsing document..." + ) + + file_url_response = generate_download_url(s3_key, settings.S3_BUCKET_NAME) + file_url = file_url_response["download_url"] + + filename = JobMetadataHelper.get_field(job_metadata, "source_file_name") + + file_ext = os.path.splitext(s3_key)[1].lower() if s3_key else "" + local_temp_path = download_s3_file_to_temp(file_url, file_ext, input_dir) + + logger.info( + f"File downloaded: job_id={job_id}, local_path={local_temp_path}" + ) + + from app.services.document_parser.internal_parse_name import ( + prepare_internal_parse_input, + ) + from app.services.document_parser.parse_service import ( + checkerboard_inject_parse, + ) + + prepared_parse_input = prepare_internal_parse_input( + local_temp_path, + filename, + fallback_ext=file_ext, + prefer_fallback_ext=True, + ) + internal_parse_name = prepared_parse_input.internal_filename + local_temp_path = prepared_parse_input.file_path + logger.info( + f"File prepared for parsing: job_id={job_id}, " + f"internal_filename={internal_parse_name}, local_path={local_temp_path}" + ) + + page_count = PageEstimator.estimate(local_temp_path) + logger.info( + f"Workload estimation: job_id={job_id}, page_count={page_count}" + ) + + processing_started_at = datetime.now(timezone.utc) + + if not job_user_id: + raise NotFoundException( + resource="JobInfo", + resource_id="user_id", + internal_message=f"Missing user_id in job info for job_id={job_id}", + ) + + billing_service = WorkBillingService() + billing_status = "skipped" + billing_amount_micro_dollars = 0 + billing_credits = 0.0 + with get_sync_db_context() as db: + job_result = db.execute( + select(Job).where(Job.job_id == job_id).with_for_update() + ) + job = job_result.scalar_one_or_none() + + if job and getattr(job, "billing_status", "") == "charged": + logger.info(f"Job already charged: {job_id}") + billing_status = "charged" + billing_amount_micro_dollars = int(job.credits_charged or 0) + billing_credits = billing_amount_micro_dollars / 1_000_000 + else: + try: + billing_result = billing_service.charge_for_pages( + session=db, + user_id=job_user_id, + page_count=page_count, + filename=filename, + ) + except InsufficientCreditsException: + logger.warning( + f"Billing failed: job_id={job_id}, user_id={job_user_id}" + ) + billing_amount = billing_service.estimate_page_charge( + page_count=page_count + ) + if job: + job.page_count = page_count + job.credits_charged = billing_amount.amount_micro_dollars + job.billing_status = "billing_failed" + db.commit() + + raise InsufficientCreditsException( + user_message=( + "Insufficient credits to process this document " + f"({page_count} pages required, cost: " + f"{billing_amount.credits})." + ), + required_credits=billing_amount.credits, + internal_message=( + f"job_id={job_id}, user_id={job_user_id}, " + f"page_count={page_count}" + ), + ) + + billing_status = billing_result.billing_status + billing_amount_micro_dollars = billing_result.amount_micro_dollars + billing_credits = billing_result.credits + if job: + job.page_count = page_count + job.credits_charged = billing_amount_micro_dollars + job.billing_status = billing_status + + metadata_updates = { + "page_count": page_count, + "billing_status": billing_status, + "billing_amount_micro_dollars": billing_amount_micro_dollars, + "billing_credits": billing_credits, + "processing_started_at": processing_started_at.isoformat(), + } + metadata_service.update_metadata(job_id, metadata_updates) + job_metadata.update(metadata_updates) + + doc_type = JobMetadataHelper.get_parsing_param( + job_metadata, "doc_type", "auto" + ) + logger.info( + f"Start parse: job_id={job_id}, filename={filename}, " + f"internal_filename={internal_parse_name}, type={doc_type}" + ) + + with stage_timer( + "worker.parse.document", + job_id=job_id, + filename=filename, + doc_type=doc_type, + ): + add_dir, add_contents_df = checkerboard_inject_parse( + file_full_path=local_temp_path, + filename=filename, + output_dir=output_dir, + job_id=job_id, + internal_output_filename=internal_parse_name, + kb_dir=JobMetadataHelper.get_parsing_param( + job_metadata, "kb_dir", "Default_Root" + ), + doc_type=doc_type, + smart_title_parse=JobMetadataHelper.get_parsing_param( + job_metadata, "smart_title_parse", True + ), + summary_image=JobMetadataHelper.get_parsing_param( + job_metadata, "summary_image", True + ), + summary_table=JobMetadataHelper.get_parsing_param( + job_metadata, "summary_table", True + ), + summary_txt=JobMetadataHelper.get_parsing_param( + job_metadata, "summary_txt", True + ), + add_frag_desc=JobMetadataHelper.get_parsing_param( + job_metadata, "add_frag_desc", "" + ), + s3_key=s3_key, + ) + parsed_contents_df: pd.DataFrame | None = add_contents_df + + logger.info( + "File parsing completed: " + f"job_id={job_id}, add_dir={add_dir}, " + f"chunks={len(parsed_contents_df) if parsed_contents_df is not None else 0}" + ) + + if parsed_contents_df is None: + raise WorkerHandlingException( + user_message="We could not extract content from your file", + internal_message="File parsing failed, no content returned from parser", + ) + + if parsed_contents_df.empty: + logger.warning( + f"No content returned from file parsing: job_id={job_id}, filename={filename}" + ) + + lifecycle_service.update_progress( + job_id, progress=30, message="Parse completed, preparing chunks..." + ) + + chunks = dataframe_to_chunks(parsed_contents_df) + + lifecycle_service.update_progress( + job_id, progress=70, message="Chunks ready, generating zip..." + ) + logger.info(f"Chunks prepared: job_id={job_id}, count={len(chunks)}") + + source_file_name = JobMetadataHelper.get_field( + job_metadata, "source_file_name" + ) or JobMetadataHelper.get_field(job_metadata, "source_url") + if isinstance(source_file_name, str) and "/" in source_file_name: + source_file_name = os.path.basename(source_file_name) + + document_top_summary = "" + section_summaries: dict[str, str] = {} + if add_dir and source_file_name: + if add_contents_df is not None and "path" in add_contents_df.columns: + ensure_doc_nav_json( + str(add_dir), + chunks, + source_file_name=str(source_file_name), + ) + try: + kb_dir_for_enrich = os.path.dirname(str(add_dir)) + summary_use_llm = JobMetadataHelper.get_parsing_param( + job_metadata, "summary_use_llm", False + ) + enrich_doc_nav_summaries( + kb_dir_for_enrich, + source_file=str(source_file_name), + use_llm=summary_use_llm, + ) + section_summaries = build_section_summary_lookup(str(add_dir)) + except Exception as exc: + logger.warning(f"doc_nav enrichment failed (non-fatal): {exc}") + document_top_summary = load_nav_top_summary( + str(add_dir), str(source_file_name) + ) + if document_top_summary: + for chunk in chunks: + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + chunk["metadata"] = metadata + metadata["document_top_summary"] = document_top_summary + + data_id = JobMetadataHelper.get_field(job_metadata, "data_id") + + lifecycle_service.update_progress( + job_id, progress=80, message="Generating ZIP package..." + ) + processing_completed_at = datetime.now(timezone.utc) + processing_timing_updates = { + "processing_completed_at": processing_completed_at.isoformat(), + "processing_duration_ms": max( + 0, + int( + ( + processing_completed_at - processing_started_at + ).total_seconds() + * 1000 + ), + ), + } + metadata_service.update_metadata(job_id, processing_timing_updates) + job_metadata.update(processing_timing_updates) + + zip_service = ZipResultService() + zip_file_path, checksum, statistics, zip_size = ( + zip_service.generate_zip_package( + job_id=job_id, + chunks=chunks, + add_dir=str(add_dir) if add_dir else "", + source_file_name=source_file_name, + data_id=data_id, + job_metadata=job_metadata, + parsed_df=parsed_contents_df, + temp_dir=task_workspace_dir, + ) + ) + del statistics + + checksum_value = ( + checksum.get("value", "") + if isinstance(checksum, dict) + else (checksum or "") + ) + + lifecycle_service.update_progress( + job_id, progress=90, message="Uploading results to S3..." + ) + + result_bundle = get_result_storage().upload( + job_id=job_id, + result_dir=str(add_dir) if add_dir else "", + zip_file_path=zip_file_path, + ) + result_s3_key = result_bundle.zip_key + + stored_count = 0 + + lifecycle_service.update_progress( + job_id, progress=100, message="Task complete!" + ) + + lifecycle_service.finalize_job_success( + job_id=job_id, + chunks=chunks, + result_s3_key=result_s3_key, + checksum=checksum_value, + zip_size=zip_size, + stored_count=stored_count, + delivery_mode="url", + section_summaries=section_summaries, + ) + + logger.info( + f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" + ) + + return { + "status": "success", + "job_id": job_id, + "add_dir": None, + "vectors_count": 0, + "contents_count": len(parsed_contents_df), + "stored_count": stored_count, + "delivery_mode": "url", + "result_s3_key": result_s3_key, + } + finally: + cleanup_task_workspace(task_workspace_dir) + + raise WorkerHandlingException( + user_message="We could not complete document processing", + internal_message=f"Parse workflow exited without a result for job_id={job_id}", + ) diff --git a/apps/worker/app/services/workload/url_upload_service.py b/apps/worker/app/services/workload/url_upload_service.py new file mode 100644 index 000000000..3618a31fd --- /dev/null +++ b/apps/worker/app/services/workload/url_upload_service.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import os +from typing import Any + +from app.services.storage.sync_storage_service import ( + download_file_from_url, + upload_to_s3, + verify_s3_file_exists, +) +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + NotFoundException, + StorageServiceException, + ValidationException, +) +from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service +from shared.services.redis.redis_sync_service import ( + SyncJobInfoRedisService, + SyncJobMetadataService, + SyncRedisServiceFactory, +) +from shared.utils.url_file_type import resolve_file_extension_sync + + +def upload_url_file( + job_id: str, + source_url: str, + user_id: str | None, + job_type: str | None = None, +) -> dict[str, Any]: + del user_id, job_type + + lifecycle_service = get_sync_job_lifecycle_service() + redis_service = SyncRedisServiceFactory.get_service() + job_info_service = SyncJobInfoRedisService(redis_service) + job_info = job_info_service.get_job_info(job_id) + + if not job_info: + metadata_service = SyncJobMetadataService(redis_service) + job_metadata = metadata_service.get_metadata(job_id) + if job_metadata: + s3_key = job_metadata.get("s3_key") + else: + raise NotFoundException( + resource="JobInfo", + resource_id=job_id, + internal_message="Job info not found in Redis or Metadata", + ) + else: + s3_key = job_info.get("s3_key") + + if not s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id="s3_key", + internal_message=f"Missing s3_key in Redis job info for job_id={job_id}", + ) + + lifecycle_service.update_progress( + job_id, progress=3, message="Validating URL file type..." + ) + file_extension = resolve_file_extension_sync(source_url) + if not file_extension: + supported_formats = ", ".join(sorted(settings.get_supported_extensions())) + raise ValidationException( + user_message="Unsupported file type", + violations=[ + { + "field": "file_extension", + "description": f"Must be one of: {supported_formats}", + } + ], + ) + + lifecycle_service.update_progress( + job_id, progress=10, message="Downloading file from URL..." + ) + try: + temp_file_path = download_file_from_url(source_url) + except Exception as exc: + raise ValidationException( + user_message="Failed to download file from URL", + violations=[ + { + "field": "source_url", + "description": "Could not download file from the provided URL", + } + ], + internal_message=( + f"Failed to download file from URL: {source_url}, error: {exc}" + ), + ) + + try: + lifecycle_service.update_progress( + job_id, progress=30, message="Validating file size..." + ) + file_size = os.path.getsize(temp_file_path) + if file_size > settings.MAX_FILE_SIZE: + limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) + raise ValidationException( + user_message=( + f"File size exceeds limit (max {limit_mb}MB for {file_extension})" + ), + violations=[ + { + "field": "file_size", + "description": ( + f"Size {file_size} bytes exceeds limit of " + f"{settings.MAX_FILE_SIZE} bytes" + ), + } + ], + ) + + lifecycle_service.update_progress( + job_id, progress=50, message="Uploading file to S3..." + ) + upload_to_s3(temp_file_path, str(s3_key), settings.S3_BUCKET_NAME) + logger.info(f"File uploaded to S3: {s3_key}") + + finally: + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + logger.debug(f"Temp file cleaned up: {temp_file_path}") + + lifecycle_service.update_progress( + job_id, progress=80, message="Verifying upload result..." + ) + file_info = verify_s3_file_exists(str(s3_key)) + if not file_info.get("exists"): + raise StorageServiceException( + user_message="We failed to verify your file upload", + internal_message=f"S3 file verification failed for {s3_key}", + ) + + lifecycle_service.update_progress( + job_id, + progress=100, + message="URL file upload complete, waiting for processing...", + ) + logger.info( + f"URL file upload complete, waiting for S3 webhook: {job_id} -> {s3_key}" + ) + + return { + "status": "success", + "job_id": job_id, + "s3_key": s3_key, + "file_size": file_info.get("size"), + } diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 12f672267..8612bb376 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -36,7 +36,7 @@ def _build_pending_file_job_metadata(source_file_name: str) -> dict[str, Any]: def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: import app.core.tasks.kb_tasks as kb_tasks import app.services.document_parser.parse_service as parse_service - import app.services.storage.sync_storage_service as sync_storage_service + import app.services.workload.parse_job_service as parse_job_service from shared.core.database_sync import get_sync_engine from shared.services.redis.redis_sync_service import ( SyncJobInfoRedisService, @@ -47,7 +47,7 @@ def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: return ( kb_tasks, parse_service, - sync_storage_service, + parse_job_service, get_sync_engine(), SyncJobInfoRedisService, SyncJobMetadataService, @@ -124,7 +124,7 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state( ( kb_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, @@ -166,8 +166,8 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state( ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", billing_enabled) + monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(parse_job_service.settings, "BILLING_ENABLED", billing_enabled) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { @@ -178,15 +178,13 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: return {"download_url": f"https://example.test/{storage_key}"} - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "generate_download_url", fake_generate_download_url, ) @@ -309,9 +307,9 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: raw_files={}, ) - monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) - monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage()) + monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") @@ -338,8 +336,8 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: }, }, ] - expected_credits_charged = 3 * int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE) - expected_initial_balance = int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + expected_credits_charged = 3 * int(parse_job_service.settings.MICRO_DOLLARS_PER_PAGE) + expected_initial_balance = int(parse_job_service.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 assert result == { "status": "success", @@ -515,7 +513,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: SELECT transition_reason, to_state FROM job_state_audit_logs WHERE job_id = :job_id - ORDER BY created_at ASC + ORDER BY id ASC """ ), {"job_id": job_id}, @@ -577,7 +575,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks ( kb_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, @@ -765,8 +763,8 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", False) + monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(parse_job_service.settings, "BILLING_ENABLED", False) def fake_cleanup_task_workspace(workspace_dir: str | None) -> bool: captured_artifacts["workspace_dir"] = workspace_dir @@ -867,22 +865,20 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: raw_files={}, ) - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "generate_download_url", fake_generate_download_url, ) - monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) - monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage()) - monkeypatch.setattr(kb_tasks, "cleanup_task_workspace", fake_cleanup_task_workspace) + monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) + monkeypatch.setattr(parse_job_service, "cleanup_task_workspace", fake_cleanup_task_workspace) result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") @@ -965,7 +961,7 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks( ( kb_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, @@ -1011,8 +1007,8 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks( ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", True) + monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(parse_job_service.settings, "BILLING_ENABLED", True) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { @@ -1072,22 +1068,20 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: raw_files={}, ) - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "generate_download_url", fake_generate_download_url, ) - monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) - monkeypatch.setattr(kb_tasks.PageEstimator, "estimate", fake_estimate_page_count) + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) + monkeypatch.setattr(parse_job_service.PageEstimator, "estimate", fake_estimate_page_count) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) - monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage()) + monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) def run_parse_task(job_id: str) -> dict[str, Any]: return dict(kb_tasks.parse_task.run(job_id, user_id, "kb_management")) @@ -1095,9 +1089,9 @@ def run_parse_task(job_id: str) -> dict[str, Any]: with ThreadPoolExecutor(max_workers=len(job_ids)) as executor: results = list(executor.map(run_parse_task, job_ids)) - expected_credits_charged = int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE) + expected_credits_charged = int(parse_job_service.settings.MICRO_DOLLARS_PER_PAGE) expected_initial_balance = ( - int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + int(parse_job_service.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 ) with engine.begin() as connection: @@ -1215,7 +1209,7 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( ( kb_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, @@ -1253,18 +1247,17 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return {"exists": storage_key == s3_key, "size": 1024} - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "verify_s3_file_exists", fake_verify_s3_file_exists, ) monkeypatch.setattr( - kb_tasks, + parse_job_service, "generate_download_url", lambda *_args, **_kwargs: (_ for _ in ()).throw( AssertionError("terminal parse task should not request a download URL") @@ -1316,7 +1309,7 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi ( kb_tasks, parse_service, - sync_storage_service, + parse_job_service, engine, sync_job_info_service_cls, sync_job_metadata_service_cls, @@ -1354,7 +1347,7 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { "exists": storage_key == s3_key, @@ -1364,15 +1357,13 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: return {"download_url": f"https://example.test/{storage_key}"} - monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url) monkeypatch.setattr( - sync_storage_service, + parse_job_service, "generate_download_url", fake_generate_download_url, ) @@ -1384,14 +1375,14 @@ def fake_download_s3_file_to_temp( shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) return str(downloaded_path) - monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp) + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr( parse_service, "checkerboard_inject_parse", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("parse failed")), ) monkeypatch.setattr( - kb_tasks, + parse_job_service, "get_result_storage", lambda: (_ for _ in ()).throw( AssertionError("result storage should not run after parser failure") @@ -1406,8 +1397,8 @@ def fake_download_s3_file_to_temp( assert result.status == "FAILURE" assert _find_task_workspaces(tmp_path, job_id) == [] - expected_credits_charged = 3 * int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE) - expected_initial_balance = int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + expected_credits_charged = 3 * int(parse_job_service.settings.MICRO_DOLLARS_PER_PAGE) + expected_initial_balance = int(parse_job_service.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 with engine.begin() as connection: job_row = ( @@ -1460,7 +1451,7 @@ def fake_download_s3_file_to_temp( SELECT transition_reason, to_state FROM job_state_audit_logs WHERE job_id = :job_id - ORDER BY created_at ASC + ORDER BY id ASC """ ), {"job_id": job_id}, diff --git a/apps/worker/tests/contract/test_url_upload_contract.py b/apps/worker/tests/contract/test_url_upload_contract.py index 8c3ade508..0c3966620 100644 --- a/apps/worker/tests/contract/test_url_upload_contract.py +++ b/apps/worker/tests/contract/test_url_upload_contract.py @@ -13,15 +13,22 @@ from support.contract_database import insert_contract_job, insert_contract_user -def _load_upload_task_modules() -> tuple[Any, Engine, Any, Any]: +def _load_upload_task_modules() -> tuple[Any, Any, Engine, Any, Any]: import app.core.tasks.kb_tasks as kb_tasks + import app.services.workload.url_upload_service as url_upload_service from shared.core.database_sync import get_sync_engine from shared.services.redis.redis_sync_service import ( SyncJobInfoRedisService, SyncRedisServiceFactory, ) - return kb_tasks, get_sync_engine(), SyncJobInfoRedisService, SyncRedisServiceFactory + return ( + kb_tasks, + url_upload_service, + get_sync_engine(), + SyncJobInfoRedisService, + SyncRedisServiceFactory, + ) def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progress( @@ -29,9 +36,13 @@ def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progres monkeypatch: MonkeyPatch, tmp_path: Path, ) -> None: - kb_tasks, engine, sync_job_info_service_cls, sync_redis_service_factory = ( - _load_upload_task_modules() - ) + ( + kb_tasks, + url_upload_service, + engine, + sync_job_info_service_cls, + sync_redis_service_factory, + ) = _load_upload_task_modules() user_id = f"worker-user-{uuid4().hex[:12]}" job_id = f"job_url_upload_{uuid4().hex[:12]}" @@ -47,19 +58,19 @@ def resolve_public_address( return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))] monkeypatch.setattr( - kb_tasks, + url_upload_service, "download_file_from_url", lambda _source_url: str(downloaded_path), ) monkeypatch.setattr( - kb_tasks, + url_upload_service, "upload_to_s3", lambda local_path, storage_key, bucket: uploaded_calls.append( (local_path, storage_key, bucket) ), ) monkeypatch.setattr( - kb_tasks, + url_upload_service, "verify_s3_file_exists", lambda storage_key: {"exists": storage_key == s3_key, "size": 3}, ) @@ -112,7 +123,7 @@ def resolve_public_address( "file_size": 3, } assert uploaded_calls == [ - (str(downloaded_path), s3_key, kb_tasks.settings.S3_BUCKET_NAME), + (str(downloaded_path), s3_key, url_upload_service.settings.S3_BUCKET_NAME), ] assert os.path.exists(downloaded_path) is False @@ -140,4 +151,3 @@ def resolve_public_address( assert job_row["status"] == "waiting-file" assert job_row["source_type"] == "url" assert job_row["s3_key"] == s3_key - diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py deleted file mode 100644 index d5b1bed78..000000000 --- a/packages/shared-python/shared/services/retrieval/agent_navigate.py +++ /dev/null @@ -1,1137 +0,0 @@ -"""Shared helpers for agentic KG document routing and scope navigation.""" -from __future__ import annotations - -import json -import re -from typing import Any, Sequence, TYPE_CHECKING - -if TYPE_CHECKING: - from shared.services.retrieval.agentic.types import DocTreeNode - -from loguru import logger -from sqlalchemy import func, select, or_ -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import Document, DocumentChunk, DocumentSection, GraphNode, GraphEdge -from shared.services.retrieval.lexical_text import normalize_section_path, split_section_path -from shared.utils.text_utils import tokenize_for_retrieval - -_MAX_OVERVIEW_FILES = 50 - -_FILE_SELECT_PROMPT = """\ -You are a document routing assistant. - -{budget_block} -Below is a knowledge base overview showing all available documents, -their navigation summaries, chunk counts, and media counts. - -=== Knowledge Base Overview === -{overview} -=== End Overview === - -User query: {query} -{revision_context} -Based on the query, select documents that may contain relevant information. -If NO document in the knowledge base is relevant to the query, return an EMPTY array []. -Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"] -Do not include any explanation. -""" - - - -_DISCOVERY_SELECT_PROMPT = """\ -You are a document navigation assistant. - -Document: "{doc_name}" - -{budget_block} -After navigating the document's section tree, the following section paths -were additionally discovered via keyword and semantic search. -They may contain relevant evidence not found through hierarchical navigation. - -=== Discovery Candidates === -{items} -=== End Discovery Candidates === - -User query: {query} -{revision_context} -Select section paths whose content is needed to answer the query. -If none are relevant, return an EMPTY list []. - -Return ONLY a JSON object: -{{"selections": [{{"path": "...", "confidence": }}, ...]}} -Do not include any explanation. -""" - - -_ACTION_PROMPT = """\ -You are a document navigation agent. - -Document: "{doc_name}" (id: {doc_id}) - -{budget_block} -{scope_header} -Below is the document's section tree. -Sections tagged [SELECT] are within the current scope and may be selected. -Other sections are shown as structural context only (not selectable). -Nodes marked [Leaf] have no further sub-sections. - -=== Section Tree === -{items_overview} -=== End Section Tree === - -User query: {query} - -=== Available Actions === - -Choose ONE action: - -NAVIGATE — Drill into selected sections for detailed content. - Consider this when the query targets specific topics and you need deeper text evidence. - Select one or more [SELECT] sections. - -STOP — Current scope evidence is sufficient. No further drill-down. - Consider this when: - - The query asks for an outline, overview, or summary - - The query is broad/global, the tree section can fulfill it without drilling into individual sections. - - You have already collected enough evidence at this level. - -{tools_block} - -When action is NAVIGATE, provide selections: -- You may ONLY select sections marked with [SELECT]. - -When action is STOP, selections must be empty. - -Return ONLY a JSON object: -{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}} -or -{{"action": "STOP", "tools": [...], "selections": []}} -Do not include any explanation. -""" - - -def _parse_action_response(text: str) -> dict: - """Parse the unified action response from LLM. - - Returns dict with keys: - action: 'NAVIGATE' | 'STOP' - tools: list[str] (subset of FIND_IMAGES, FIND_TABLES) - selections: list[dict] (each has 'path' and optional 'confidence') - - When action is STOP, selections are forced to empty. - """ - import json as _json - import re as _re - - text = text.strip() - _ASSET_TOOLS = {'FIND_IMAGES', 'FIND_TABLES'} - default = {'action': 'NAVIGATE', 'tools': [], 'selections': []} - - def _extract(data: dict) -> dict: - action = str(data.get('action', 'NAVIGATE')).strip().upper() - if action not in ('NAVIGATE', 'STOP'): - action = 'NAVIGATE' - - tools_val = data.get('tools') or [] - if isinstance(tools_val, list): - tools = [str(t).strip().upper() for t in tools_val if str(t).strip().upper() in _ASSET_TOOLS] - else: - tools = [] - - # STOP → no selections allowed - if action == 'STOP': - return {'action': action, 'tools': tools, 'selections': []} - - selections_val = data.get('selections') or [] - selections = [] - if isinstance(selections_val, list): - for s in selections_val: - if isinstance(s, dict) and s.get('path'): - conf = _normalize_confidence(s.get('confidence', 0.7)) - selections.append({'path': str(s['path']), 'confidence': conf or 0.7}) - - return {'action': action, 'tools': tools, 'selections': selections} - - # Try JSON parse - try: - data = _json.loads(text) - if isinstance(data, dict): - return _extract(data) - except (ValueError, _json.JSONDecodeError): - pass - - # Try extracting JSON from markdown fences - fence_match = _re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, _re.DOTALL) - if fence_match: - try: - data = _json.loads(fence_match.group(1).strip()) - if isinstance(data, dict): - return _extract(data) - except (ValueError, _json.JSONDecodeError): - pass - - # Try finding a JSON object anywhere - brace_match = _re.search(r'\{.*\}', text, _re.DOTALL) - if brace_match: - try: - data = _json.loads(brace_match.group()) - if isinstance(data, dict): - return _extract(data) - except (ValueError, _json.JSONDecodeError): - pass - - return default - - -def _format_budget_block(snapshot: dict | None) -> str: - if not snapshot: - return "" - planning = snapshot.get("planning") or {} - context = snapshot.get("context") or {} - return ( - "=== Resource Status ===\n" - f"Planning Budget: {planning.get('status', 'HEALTHY')} " - f"({planning.get('used_pct', 0)}% used)\n" - f"Context Budget: {context.get('status', 'HEALTHY')} " - f"({context.get('used_pct', 0)}% used)\n" - f"KG Coverage: {snapshot.get('explored_chunks', 0)}/" - f"{snapshot.get('total_chunks', 0)} chunks explored\n" - f"Docs Explored: {snapshot.get('explored_docs', 0)}/" - f"{snapshot.get('total_docs', 0)}\n" - "When budget is TIGHT, prefer fewer high-confidence selections over broad exploration. " - "When CRITICAL, be very selective — only pick paths with strong relevance. Return empty if evidence suffices.\n" - "=== End Resource Status ===\n" - ) - - -def _extract_json_array_payload(text: str) -> list[Any]: - """Best-effort extraction of a JSON array payload from LLM response text.""" - text = text.strip() - try: - result = json.loads(text) - if isinstance(result, list): - return result - except (json.JSONDecodeError, ValueError): - pass - match = re.search(r'\[.*?\]', text, re.DOTALL) - if match: - try: - result = json.loads(match.group()) - if isinstance(result, list): - return result - except (json.JSONDecodeError, ValueError): - pass - return [] - - -def _parse_json_array(text: str) -> list[str]: - """Best-effort extraction of a JSON array of strings from LLM response text.""" - result = _extract_json_array_payload(text) - return [str(x) for x in result] - - -def _normalize_confidence(value: Any) -> float | None: - if value is None: - return None - if isinstance(value, str): - value = value.strip().rstrip('%') - try: - parsed = float(value) - except (TypeError, ValueError): - return None - if parsed > 1.0: - parsed = parsed / 100.0 - return max(0.0, min(parsed, 1.0)) - - -async def _build_knowledge_map_overview( - db: AsyncSession, - *, - user_id: str, - namespace: str, -) -> tuple[str, dict[str, str]]: - """Build a file-level knowledge map overview for LLM file selection. - - Returns (overview_text, doc_id_to_name) where doc_id_to_name maps - document_id -> source_file_name for validation after LLM response. - """ - doc_stmt = ( - select(Document) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(Document.current_job_result_id.is_not(None)) - .order_by(Document.updated_at.desc()) - .limit(_MAX_OVERVIEW_FILES) - ) - doc_result = await db.execute(doc_stmt) - documents = list(doc_result.scalars()) - - if not documents: - return '(empty)', {} - - doc_ids = [d.document_id for d in documents] - doc_id_to_name: dict[str, str] = { - d.document_id: (d.source_file_name or d.document_id) - for d in documents - } - - chunk_stats_stmt = ( - select( - DocumentChunk.document_id, - func.count(DocumentChunk.id).label('chunk_count'), - func.count(func.nullif(DocumentChunk.chunk_type, 'text')).label('media_count'), - ) - .join(Document, (Document.document_id == DocumentChunk.document_id) & (Document.current_job_result_id == DocumentChunk.job_result_id)) - .where(DocumentChunk.document_id.in_(doc_ids)) - .group_by(DocumentChunk.document_id) - ) - chunk_stats_result = await db.execute(chunk_stats_stmt) - chunk_stats: dict[str, dict[str, int]] = {} - for row in chunk_stats_result.all(): - chunk_stats[row[0]] = {'total': row[1], 'media': row[2]} - - graph_summary_stmt = ( - select(GraphNode.owner_document_id, GraphNode.properties) - .where(GraphNode.owner_document_id.in_(doc_ids)) - .where(GraphNode.node_kind == 'document') - ) - graph_summary_result = await db.execute(graph_summary_stmt) - doc_top_summaries: dict[str, str] = {} - for did, properties in graph_summary_result.all(): - if not isinstance(properties, dict): - continue - top_summary = str(properties.get('top_summary') or '').strip() - if top_summary: - doc_top_summaries[did] = top_summary - - lines: list[str] = [] - for doc in documents: - did = doc.document_id - name = doc_id_to_name[did] - stats = chunk_stats.get(did, {'total': 0, 'media': 0}) - top_summary = doc_top_summaries.get(did, '') - - line = f'- [{did}] {name} chunks={stats["total"]}' - if stats['media'] > 0: - line += f' media={stats["media"]}' - if top_summary: - line += f'\n top_summary:\n{_indent_block(top_summary, 4)}' - lines.append(line) - - return '\n'.join(lines), doc_id_to_name - - -def _indent_block(text: str, spaces: int) -> str: - prefix = ' ' * spaces - return '\n'.join(f'{prefix}{line}' for line in str(text or '').splitlines()) - - -def _format_items_for_llm( - items: list[dict], - max_chars: int = 20000, -) -> tuple[str, bool]: - """Format items with ▸ └ [Leaf] hierarchy for scope navigation. - - Supports arbitrary depth levels via absolute ``level`` field. - Items with ``show_summary=False`` render title only (structural context). - ``[LN]`` tags indicate the absolute document depth of each section. - ``[Leaf]`` tags indicate bottom-level sections with no further children. - Summaries are included when within budget, dropped on overflow. - - Returns (text, overflowed). - """ - from shared.utils.text_utils import truncate_content_preview - - if not items: - return '(no items available)', False - - SUMMARY_HEAD_TOKENS = 80 - - def _render_item(item: dict, include_summary: bool) -> str: - level = item.get('level', 1) - show = item.get('show_summary', True) - is_leaf = item.get('is_leaf', False) - leaf_tag = ' [Leaf]' if is_leaf else '' - path = item.get('path', '') - summary = item.get('summary') or '' - - # Build chunk count tags (only for current-scope items) - counts_str = '' - if show: - count_parts: list[str] = [] - chunk_count = item.get('chunk_count', 0) - if chunk_count > 0: - count_parts.append(f'text={chunk_count}') - image_count = item.get('image_count', 0) - if image_count > 0: - count_parts.append(f'image={image_count}') - table_count = item.get('table_count', 0) - if table_count > 0: - count_parts.append(f'table={table_count}') - counts_str = f' [{" ".join(count_parts)}]' if count_parts else '' - - indent = " " * (level - 1) - prefix = '▸' if level == 1 else '└' - level_tag = f'[L{level}]' - select_tag = '[SELECT] ' if item.get('selectable', False) else '' - - lines: list[str] = [] - lines.append(f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}') - - if include_summary and show and summary: - sub_indent = " " * level - clipped = truncate_content_preview(summary, head=SUMMARY_HEAD_TOKENS, tail=0) - lines.append(f'{sub_indent}{clipped}') - - return '\n'.join(lines) - - # Try full render (with summaries for show_summary=True items) - full_lines = [_render_item(item, include_summary=True) for item in items] - full_text = '\n'.join(full_lines) - if len(full_text) <= max_chars: - return full_text, False - - # Overflow: render without summaries - slim_lines = [_render_item(item, include_summary=False) for item in items] - slim_text = '\n'.join(slim_lines) - return slim_text[:max_chars], True - - -# ------------------------------------------------------------------ -# GREP document discovery (aligned with KB do_discover_files) -# ------------------------------------------------------------------ - -async def _grep_discover_document_ids( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - exclude_document_ids: Sequence[str] = (), - limit: int = 10, -) -> list[str]: - """GREP discovery: search term_search_text for query terms, return parent document_ids. - - Aligned with KB's do_discover_files(): if a chunk's term_search_text - contains query terms, its parent document is included in the KG scope. - """ - units = tokenize_for_retrieval(query, dedupe=True) - logger.info(f' GREP tokenized units (cap 8): {units[:8]} (total={len(units)})') - if not units: - return [] - - # Build OR conditions for ILIKE matching - conditions = [] - params: dict[str, str] = { - 'user_id': user_id, - 'namespace': namespace, - } - for i, unit in enumerate(units[:8]): # cap at 8 terms to avoid huge queries - param_name = f'unit_{i}' - params[param_name] = f'%{unit}%' - conditions.append(DocumentChunk.term_search_text.ilike(f'%{unit}%')) - - if not conditions: - return [] - - stmt = ( - select(Document.document_id) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(DocumentChunk.term_search_text.is_not(None)) - .where(or_(*conditions)) - .distinct() - .limit(limit) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - - result = await db.execute(stmt) - return [row[0] for row in result.all()] - - -# ------------------------------------------------------------------ -# Edge expansion (aligned with KB KGIndex.neighbors) -# ------------------------------------------------------------------ - -async def _expand_by_edges( - db: AsyncSession, - *, - document_ids: list[str], - user_id: str, - namespace: str, - hops: int = 1, -) -> list[str]: - """Expand document set by following GraphEdge relationships. - - Aligned with KB's KGIndex.neighbors(): traverse edges to include - related documents. Only queries document-level nodes (no section nodes). - No weight filtering — edges already passed threshold during publication. - """ - if not document_ids: - return document_ids - - current = set(document_ids) - - for hop_idx in range(hops): - # Find document-level graph nodes for current document set - doc_node_ids = [f"doc:{did}" for did in current] - node_stmt = ( - select(GraphNode.node_id, GraphNode.owner_document_id) - .where(GraphNode.user_id == user_id) - .where(GraphNode.namespace == namespace) - .where(GraphNode.node_kind == 'document') - .where(GraphNode.node_id.in_(doc_node_ids)) - ) - node_result = await db.execute(node_stmt) - node_rows = node_result.all() - logger.info(f' edge_expand hop={hop_idx}: doc_nodes_found={len(node_rows)} (of {len(doc_node_ids)} requested)') - - if not node_rows: - break - - node_ids = {row[0] for row in node_rows} - - # Follow edges from/to these document nodes - edge_stmt = ( - select(GraphEdge.source_node_id, GraphEdge.target_node_id) - .where(GraphEdge.user_id == user_id) - .where(GraphEdge.namespace == namespace) - .where(or_( - GraphEdge.source_node_id.in_(list(node_ids)), - GraphEdge.target_node_id.in_(list(node_ids)), - )) - ) - edge_result = await db.execute(edge_stmt) - edge_rows = edge_result.all() - - neighbor_node_ids: set[str] = set() - for src, tgt in edge_rows: - if src in node_ids: - neighbor_node_ids.add(tgt) - if tgt in node_ids: - neighbor_node_ids.add(src) - logger.info(f' edge_expand hop={hop_idx}: edges_traversed={len(edge_rows)} neighbor_nodes={len(neighbor_node_ids)}') - - if not neighbor_node_ids: - break - - # Resolve neighbor nodes to document_ids - neighbor_doc_stmt = ( - select(GraphNode.owner_document_id) - .where(GraphNode.node_id.in_(list(neighbor_node_ids))) - .where(GraphNode.node_kind == 'document') - ) - neighbor_doc_result = await db.execute(neighbor_doc_stmt) - for (doc_id,) in neighbor_doc_result.all(): - current.add(doc_id) - - # Preserve original order, append new ones at end - ordered = list(document_ids) - for doc_id in current: - if doc_id not in document_ids: - ordered.append(doc_id) - return ordered - - -# --------------------------------------------------------------------------- -# Unified scope navigation: load child sections (2-level) -# --------------------------------------------------------------------------- - -async def _load_child_sections( - db: AsyncSession, - document_id: str, - job_result_id: str, - scope_path: str | list[str] | None = None, - exclude_paths: set[str] | None = None, -) -> list[dict]: - """Load the Continuous Context Tree for *scope_path*. - - Returns a flat list sorted by document order, each item: - {path, title, summary, chunk_count, image_count, table_count, - level, show_summary, is_leaf} - - scope_path can be: - - None: root scope, all items are selectable (2 depth bands). - - str: single scope, descendants are selectable. - - list[str]: multi-scope, descendants of ALL paths are selectable - simultaneously — used when the LLM selected multiple drill-down - paths in the previous step. - - - level: absolute depth in the document (1-based) - - show_summary: controls whether _format_items_for_llm renders summary - - exclude_paths: paths already hydrated; skipped from selectable items - """ - # ── Fetch all sections for this document revision ──────────────────── - stmt = ( - select( - DocumentSection.section_id, - DocumentSection.section_title, - DocumentSection.section_path, - DocumentSection.summary, - DocumentSection.sort_order, - ) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .order_by(DocumentSection.sort_order) - ) - section_rows = (await db.execute(stmt)).all() - if not section_rows: - return [] - - # ── Normalize scope(s) ─────────────────────────────────────────────── - # Multi-scope: list of paths to expand simultaneously - if isinstance(scope_path, list): - scope_list = [normalize_section_path(p) for p in scope_path] - elif scope_path: - scope_list = [normalize_section_path(scope_path)] - else: - scope_list = [] # root - - # For logging, derive representative scope info - scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0 - - logger.debug( - f' _load_child_sections: scopes={scope_list or ["root"]} ' - f'scope_depth={scope_depth} exclude_paths={_excl if (_excl := exclude_paths or set()) else "none"} ' - f'total_sections={len(section_rows)}' - ) - - # Build full section metadata index - all_sections: dict[str, dict] = {} # path → {title, summary, sort_order, section_id, parts, depth} - for section_id, title, path, summary, sort_order in section_rows: - if not path: - continue - path = normalize_section_path(path) - parts = split_section_path(path) - all_sections[path] = { - 'title': title or parts[-1] if parts else path, - 'summary': summary or '', - 'sort_order': int(sort_order or 0), - 'section_id': section_id, - 'parts': parts, - 'depth': len(parts), - } - - # ── Build the set of ancestor prefixes for pruning ──────────────────── - # For multi-scope, union all ancestor prefixes from all scope paths - ancestor_prefixes: set[str] = set() - for sp in scope_list: - sp_parts = split_section_path(sp) - for i in range(1, len(sp_parts) + 1): - ancestor_prefixes.add(' / '.join(sp_parts[:i])) - - # ── Classify each section ──────────────────────────────────────────── - _excl = exclude_paths or set() - items_by_path: dict[str, dict] = {} - # Per-scope depth bands: track child depths separately per scope - per_scope_child_depths: dict[str, set[int]] = {sp: set() for sp in scope_list} if scope_list else {} - root_child_depths: set[int] = set() # used when scope_list is empty (root) - - def _make_item(path: str, meta: dict, show_summary: bool) -> dict: - return { - 'path': path, - 'title': meta['title'], - 'summary': meta['summary'], - 'level': meta['depth'], - 'sort_order': meta['sort_order'], - 'chunk_count': 0, - 'image_count': 0, - 'table_count': 0, - 'section_id': meta['section_id'], - 'show_summary': show_summary, - } - - def _is_excluded(path: str) -> bool: - return bool(_excl and any( - path == ep or path.startswith(ep + ' / ') for ep in _excl - )) - - for path, meta in all_sections.items(): - parts = meta['parts'] - depth = meta['depth'] - - if not scope_list: - # Root scope: everything is a potential child - if depth < 1 or _is_excluded(path): - continue - root_child_depths.add(depth) - items_by_path[path] = _make_item(path, meta, show_summary=True) - continue - - # --- Non-root scope(s) --- - # Check if this path is a descendant of ANY scope in scope_list - matched_scope: str | None = None - for sp in scope_list: - sp_parts = split_section_path(sp) - sp_depth = len(sp_parts) - if depth > sp_depth and parts[:sp_depth] == sp_parts: - matched_scope = sp - break - - if matched_scope: - # Category 2: descendant of a scope path → selectable - if _is_excluded(path): - continue - per_scope_child_depths[matched_scope].add(depth) - items_by_path[path] = _make_item(path, meta, show_summary=True) - continue - - # Category 1: structural context (ancestors of scope paths only) - # Only show nodes that are on the ancestor chain of a scope path. - # Non-scope siblings (e.g. 法律声明, 前言 when navigating into - # chapters 一~六) are pruned to reduce token waste and prevent - # summary overflow in _format_items_for_llm. - max_scope_depth = max(len(split_section_path(sp)) for sp in scope_list) - if depth <= max_scope_depth: - if depth == 1: - if path in ancestor_prefixes: - items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) - else: - parent_prefix = ' / '.join(parts[:-1]) - if parent_prefix in ancestor_prefixes: - items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) - continue - - # Category 3: pruned - - if not items_by_path: - return [] - - # ── Limit children to 2 depth bands (relative to each scope) ──────── - allowed_set: set[int] = set() - if scope_list: - for sp, depths in per_scope_child_depths.items(): - if depths: - allowed_set.update(sorted(depths)[:2]) - else: - if root_child_depths: - allowed_set.update(sorted(root_child_depths)[:2]) - - if allowed_set: - to_remove = [ - path for path, item in items_by_path.items() - if item['show_summary'] and item['level'] not in allowed_set - ] - for path in to_remove: - del items_by_path[path] - - if not items_by_path: - return [] - - # ── Count chunks per section (text / image / table) ────────────────── - # Only count for show_summary=True items (current scope children) - scope_item_sids = {item['section_id'] for item in items_by_path.values() if item['show_summary']} - # Also need all section_ids for upward aggregation - all_section_ids = [meta['section_id'] for meta in all_sections.values()] - if all_section_ids and scope_item_sids: - from sqlalchemy import case, literal_column - chunk_stmt = ( - select( - DocumentChunk.section_id, - func.count( - case( - (DocumentChunk.chunk_type.notin_(['image', 'table']), literal_column('1')), - ) - ).label('text_count'), - func.count( - case( - (DocumentChunk.chunk_type == 'image', literal_column('1')), - ) - ).label('image_count'), - func.count( - case( - (DocumentChunk.chunk_type == 'table', literal_column('1')), - ) - ).label('table_count'), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_section_ids)) - .group_by(DocumentChunk.section_id) - ) - chunk_rows = (await db.execute(chunk_stmt)).all() - section_id_counts: dict[str, tuple[int, int, int]] = { - sid: (int(tc), int(ic), int(tbc)) for sid, tc, ic, tbc in chunk_rows - } - else: - section_id_counts = {} - - # Build section_id → path mapping for aggregation - sid_to_path = {meta['section_id']: path for path, meta in all_sections.items()} - - # Aggregate chunk counts upward: each show_summary item gets counts from itself + descendants - # Phase 1: Direct section assignment — counts from chunks directly under each section - for sid, (text_c, img_c, tbl_c) in section_id_counts.items(): - chunk_path = sid_to_path.get(sid, '') - if not chunk_path: - continue - - for item_path, item in items_by_path.items(): - if not item['show_summary']: - continue - if chunk_path == item_path or chunk_path.startswith(item_path + ' / '): - item['chunk_count'] += text_c - item['image_count'] += img_c - item['table_count'] += tbl_c - - # Phase 2: connect_to reference tracing — Root-level standalone assets - # Images/tables often live in the Root section but are referenced via connect_to - # from text chunks in deeper sections. Trace these references to attribute - # assets to the sections that actually use them. - # - # Algorithm: for each show_summary item, find all text chunks under its subtree, - # collect their connect_to targets, and count how many are image/table chunks. - scope_items_with_zero_assets = [ - item for item in items_by_path.values() - if item['show_summary'] and item['image_count'] == 0 and item['table_count'] == 0 - ] - if scope_items_with_zero_assets: - # Load connect_to metadata for text chunks under all scope sections - scope_section_ids = {item['section_id'] for item in items_by_path.values() if item.get('section_id')} - if scope_section_ids: - from sqlalchemy import literal_column - connect_stmt = ( - select( - DocumentChunk.section_id, - DocumentChunk.chunk_metadata, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(scope_section_ids))) - .where(DocumentChunk.chunk_type == 'text') - ) - connect_result = (await db.execute(connect_stmt)).all() - - # Map section_id → set of connected target chunk_ids - section_target_ids: dict[str, set[str]] = {} - for sec_id, metadata in connect_result: - if not isinstance(metadata, dict): - continue - for conn in metadata.get('connect_to') or []: - target_id = conn.get('target', '') - if target_id: - section_target_ids.setdefault(sec_id, set()).add(target_id) - - if section_target_ids: - # Collect all target chunk_ids and look up their types - all_target_ids = set() - for tids in section_target_ids.values(): - all_target_ids.update(tids) - - target_type_stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_id.in_(list(all_target_ids))) - .where(DocumentChunk.chunk_type.in_(['image', 'table'])) - ) - target_type_result = (await db.execute(target_type_stmt)).all() - target_types: dict[str, str] = {cid: ctype for cid, ctype in target_type_result} - - # Aggregate connected asset counts per section path → upward to items - for sec_id, target_ids in section_target_ids.items(): - ref_path = sid_to_path.get(sec_id, '') - if not ref_path: - continue - ref_img = sum(1 for tid in target_ids if target_types.get(tid) == 'image') - ref_tbl = sum(1 for tid in target_ids if target_types.get(tid) == 'table') - if ref_img == 0 and ref_tbl == 0: - continue - for item_path, item in items_by_path.items(): - if not item['show_summary']: - continue - if ref_path == item_path or ref_path.startswith(item_path + ' / '): - item['image_count'] += ref_img - item['table_count'] += ref_tbl - - # ── Sort by native document order ───────────────────────────────────── - sorted_items = sorted(items_by_path.values(), key=lambda x: x['sort_order']) - # Clean up internal fields - for item in sorted_items: - item.pop('sort_order', None) - item.pop('section_id', None) - - # ── Detect leaf status ──────────────────────────────────────────────── - # A section is a leaf if no other section in the database for this - # document has a path that descends from it. - all_section_paths = set(all_sections.keys()) - for item in sorted_items: - item_path = item['path'] - has_descendants = any( - p != item_path and p.startswith(item_path + ' / ') - for p in all_section_paths - ) - item['is_leaf'] = not has_descendants - - # ── Assign selectability ────────────────────────────────────────────── - # Rule: in the 2-band window, only the DEEPER band is selectable. - # Leaf nodes at the shallower band are still selectable (no children - # to drill into). Structural context (show_summary=False) is never - # selectable. - if allowed_set: - shallowest_band = min(allowed_set) - for item in sorted_items: - if not item.get('show_summary', True): - # Structural context → never selectable - item['selectable'] = False - elif item['level'] == shallowest_band and not item.get('is_leaf', False): - # Shallowest band, non-leaf → grouping header, not selectable - item['selectable'] = False - else: - item['selectable'] = True - else: - for item in sorted_items: - item['selectable'] = item.get('show_summary', True) - - return sorted_items - - -# --------------------------------------------------------------------------- - - -# --------------------------------------------------------------------------- -# Unified document tree rendering (DocTreeNode → single coherent hierarchy) -# --------------------------------------------------------------------------- - -def _render_leaf_chunks( - parts: list[str], - chunks: list[dict[str, Any]], - indent: str, - asset_lookup: dict[str, str] | None = None, -) -> None: - """Render hydrated leaf chunks inline with table/image inlining and dedup. - - Uses ``connect_to`` metadata to resolve asset references — the same - pattern as ``assemble_retrieval_results``: - - **Tables**: inline HTML content at the ``ref`` placeholder - - **Images**: inline the ``file_path`` (S3-compatible URL) at the - placeholder for multimodal LLMs - - Connected target chunks (images/tables) are expected to already be - present in ``chunks`` via ``hydrate_connected_target_rows``. - - Phase 2: After rendering all text chunks, standalone image/table - chunks that were NOT inlined via connect_to are rendered separately. - This handles cases where assets exist at root/section level without - a parent text chunk referencing them. - """ - chunk_by_id: dict[str, dict] = { - c.get('chunk_id', ''): c for c in chunks if c.get('chunk_id') - } - rendered_ids: set[str] = set() - - # Phase 1: Render text chunks with inline asset resolution - for chunk in chunks: - cid = chunk.get('chunk_id', '') - if cid and cid in rendered_ids: - continue - - chunk_type = (chunk.get('chunk_type') or chunk.get('type') or 'text').strip().lower() - - # Skip standalone image/table chunks — they'll be rendered in Phase 2 - # if not inlined via connect_to from a parent text chunk. - # NOTE: do NOT add to rendered_ids here — Phase 2 needs to see them. - if chunk_type in ('image', 'table'): - continue - - if cid: - rendered_ids.add(cid) - - content = str(chunk.get('content', '')).strip() - - # Resolve connected assets via connect_to metadata - for conn in (chunk.get('chunk_metadata') or {}).get('connect_to') or []: - target = chunk_by_id.get(conn.get('target', '')) - if not target: - continue - target_cid = target.get('chunk_id', '') - target_type = (target.get('chunk_type') or target.get('type') or '').strip().lower() - ref_str = conn.get('ref', '') - if not ref_str or ref_str not in content: - continue - - if target_cid: - rendered_ids.add(target_cid) - - if target_type == 'table': - table_html = str(target.get('content', '')).strip() - content = content.replace(ref_str, f'\n[表格内容]\n{table_html}\n') - elif target_type == 'image': - file_path = target.get('file_path') or '' - img_desc = str(target.get('content', '')).strip() - # Strip self-reference from image description - if ref_str in img_desc: - img_desc = img_desc.replace(ref_str, '').strip() - # Use pre-generated asset URL if available, fall back to file_path - asset_url = (asset_lookup or {}).get(target_cid, '') if target_cid else '' - display_ref = asset_url or file_path - if display_ref: - content = content.replace(ref_str, f'\n[图片: {display_ref}]\n{img_desc}\n') - elif img_desc: - content = content.replace(ref_str, f'\n[图片描述]\n{img_desc}\n') - - for line in content.split('\n'): - if line.strip(): - parts.append(f'{indent}┈ {line}') - - # Phase 2: Render standalone image/table chunks not inlined via connect_to - for chunk in chunks: - cid = chunk.get('chunk_id', '') - if cid and cid in rendered_ids: - continue - if cid: - rendered_ids.add(cid) - - chunk_type = (chunk.get('chunk_type') or chunk.get('type') or '').strip().lower() - if chunk_type == 'image': - file_path = chunk.get('file_path') or '' - img_desc = str(chunk.get('content', '')).strip() - asset_url = (asset_lookup or {}).get(cid, '') if cid else '' - display_ref = asset_url or file_path - if display_ref: - parts.append(f'{indent}┈ [图片: {display_ref}]') - if img_desc: - for line in img_desc.split('\n'): - if line.strip(): - parts.append(f'{indent}┈ {line}') - elif chunk_type == 'table': - table_html = str(chunk.get('content', '')).strip() - parts.append(f'{indent}┈ [表格内容]') - if table_html: - for line in table_html.split('\n'): - if line.strip(): - parts.append(f'{indent}┈ {line}') - - -def render_unified_doc_tree( - node: DocTreeNode, - doc_name: str, - depth: int = 0, - asset_lookup: dict[str, str] | None = None, -) -> str: - """Render a DocTreeNode as a single coherent hierarchy. - - Summaries are navigation-only aids and NEVER appear in evidence. - The rendered output contains: - 1. Structural titles for ALL sections (positioning context) - 2. Hydrated chunk content (┈ lines) ONLY for selected leaf paths - - Asset references (tables/images) are resolved via ``connect_to`` - metadata in hydrated chunks — no separate lookup needed. - """ - - parts: list[str] = [] - indent = ' ' * depth - - if depth == 0: - parts.append(f'【文档】{doc_name}\n') - - # Collect children keys for path-hierarchy dedup: - child_prefixes = set(node.children.keys()) - - # Helper: min sort_order of a leaf_content entry - def _min_sort(path: str) -> float: - chunks = node.leaf_content.get(path, []) - return min((c.get('sort_order') or float('inf') for c in chunks), default=float('inf')) - - # ── Build a unified render queue ── - # Each entry: (sort_key, render_type, data) - # render_type: 'outline' | 'orphan_leaf' | 'orphan_child' - render_queue: list[tuple[float, str, dict | str]] = [] - - outline_paths: set[str] = set() - # Position counter for outline-only items (no leaf content) to preserve - # their relative ordering among themselves. - outline_position = 0.0 - - for item in node.outline_items: - path = item.get('path', '') - # Skip items belonging to a drilled-into child's subtree - if any(path.startswith(cp + ' / ') for cp in child_prefixes): - continue - outline_paths.add(path) - - # Determine sort_key: use chunk sort_order if content exists, - # else use a synthetic position to maintain outline ordering. - if path in node.leaf_content or path in node.children: - sort_key = _min_sort(path) if path in node.leaf_content else outline_position - else: - sort_key = outline_position - outline_position = max(outline_position, sort_key) + 0.001 - - render_queue.append((sort_key, 'outline', item)) - - # Add orphan leaf_content paths (not covered by outline_items) - for path in node.leaf_content: - if path not in outline_paths: - render_queue.append((_min_sort(path), 'orphan_leaf', path)) - - # Add orphan children (not covered by outline_items) - for path in node.children: - if path not in outline_paths: - render_queue.append((float('inf'), 'orphan_child', path)) - - # Sort by sort_key (stable sort preserves insertion order for ties) - render_queue.sort(key=lambda x: x[0]) - - from typing import cast - - # ── Render the unified queue ── - for _sort_key, rtype, data in render_queue: - if rtype == 'outline': - item = cast(dict, data) - path = item.get('path', '') - title = item.get('title', '') - is_leaf = item.get('is_leaf', False) - level = item.get('level', 1) - leaf_tag = ' [Leaf]' if is_leaf else '' - - level_tag = f'[L{level}] ' if level else '' - if level <= 1: - parts.append(f'{indent}▸ {level_tag}{title}{leaf_tag}') - else: - parts.append(f'{indent}└ {level_tag}{title}{leaf_tag}') - - sub_indent = indent + ' ' - - # Case 1: drilled-into child → render child tree - if path in node.children: - child = node.children[path] - if path in node.leaf_content: - _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup) - if child_text.strip(): - parts.append(child_text) - - # Case 2: hydrated leaf → show chunk content - elif path in node.leaf_content: - _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - - # Case 3: unselected → title only (already rendered above) - - elif rtype == 'orphan_leaf': - path = cast(str, data) - title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path - parts.append(f'{indent}▸ [Leaf] {title}') - sub_indent = indent + ' ' - _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - - elif rtype == 'orphan_child': - path = cast(str, data) - title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path - parts.append(f'{indent}▸ {title} [DrillDown]') - child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) - if child_text.strip(): - parts.append(child_text) - - return '\n'.join(parts) - diff --git a/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py b/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py new file mode 100644 index 000000000..395120e9d --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py @@ -0,0 +1,266 @@ +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult + + +def build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, str]: + owner_map: dict[str, str] = {} + for chunk in text_chunks: + if (chunk.get("chunk_type") or "text") != "text": + continue + section_path = chunk.get("section_path") or "" + if not section_path: + continue + metadata = chunk.get("chunk_metadata") or {} + if not isinstance(metadata, dict): + continue + for conn in metadata.get("connect_to") or []: + if not isinstance(conn, dict): + continue + target_id = str(conn.get("target") or "").strip() + if target_id and target_id not in owner_map: + owner_map[target_id] = section_path + return owner_map + + +async def resolve_root_asset_owners( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + chunks: list[dict[str, Any]], +) -> dict[str, str]: + root_asset_ids = [ + str(chunk.get("chunk_id") or "") + for chunk in chunks + if not chunk.get("owner_section_path") + and (chunk.get("section_path") or "") == "Root" + and (chunk.get("chunk_type") or "").lower() in ("image", "table") + and chunk.get("chunk_id") + ] + if not root_asset_ids: + return {} + + root_asset_set = set(root_asset_ids) + text_stmt = ( + select( + DocumentChunk.chunk_metadata, + DocumentSection.section_path, + ) + .outerjoin( + DocumentSection, + DocumentSection.section_id == DocumentChunk.section_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_type == "text") + ) + result = await db.execute(text_stmt) + + owner_map: dict[str, str] = {} + for metadata, section_path in result.all(): + if not isinstance(metadata, dict) or not section_path: + continue + for conn in metadata.get("connect_to") or []: + if not isinstance(conn, dict): + continue + target_id = str(conn.get("target") or "").strip() + if target_id in root_asset_set and target_id not in owner_map: + owner_map[target_id] = section_path + + if owner_map: + logger.info( + f" resolve_root_asset_owners: resolved {len(owner_map)}/{len(root_asset_ids)} " + f"Root assets to their owner sections" + ) + return owner_map + + +async def asset_filter_step( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + scope_path: str | list[str] | None, + asset_type: str, +) -> list[dict[str, Any]]: + t0 = time.monotonic() + try: + scope_list = ( + scope_path + if isinstance(scope_path, list) + else [scope_path] + if scope_path + else [] + ) + + section_stmt = ( + select(DocumentSection.section_id, DocumentSection.section_path) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + ) + if scope_list: + from sqlalchemy import or_ + + scope_filters = [] + for scope in scope_list: + scope_filters.append(DocumentSection.section_path == scope) + scope_filters.append(DocumentSection.section_path.like(f"{scope} / %")) + section_stmt = section_stmt.where(or_(*scope_filters)) + section_rows = (await db.execute(section_stmt)).all() + section_ids = {row[0] for row in section_rows} + + if not section_ids: + logger.info(f" asset_filter_step: no sections found under scope={scope_path}") + return [] + + section_path_by_id = { + section_id: section_path for section_id, section_path in section_rows + } + asset_rows = ( + await db.execute( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.file_path, + DocumentChunk.section_id, + DocumentChunk.source_chunk_path, + DocumentChunk.chunk_metadata, + DocumentChunk.sort_order, + DocumentChunk.job_result_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(section_ids))) + .where(DocumentChunk.chunk_type == asset_type) + .order_by(DocumentChunk.sort_order) + ) + ).all() + + text_rows = ( + await db.execute( + select( + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.chunk_metadata, + DocumentChunk.source_chunk_path, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(section_ids))) + .where(DocumentChunk.chunk_type == "text") + ) + ).all() + text_row_dicts = [ + { + "chunk_type": chunk_type, + "chunk_metadata": metadata or {}, + "section_id": section_id, + "section_path": section_path_by_id.get(section_id, ""), + "source_chunk_path": source_chunk_path, + } + for section_id, chunk_type, metadata, source_chunk_path in text_rows + ] + owner_by_target_id = build_connected_owner_map(text_row_dicts) + + if any(value == "Root" for value in owner_by_target_id.values()): + doc_stmt = select(Document.source_file_name).where( + Document.document_id == document_id + ) + doc_file_name = (await db.execute(doc_stmt)).scalar() or "" + if doc_file_name: + for target_id in list(owner_by_target_id): + if owner_by_target_id[target_id] == "Root": + owner_by_target_id[target_id] = doc_file_name + + connected_target_ids: set[str] = set(owner_by_target_id.keys()) + if connected_target_ids: + connected_rows = ( + await db.execute( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.file_path, + DocumentChunk.section_id, + DocumentChunk.source_chunk_path, + DocumentChunk.chunk_metadata, + DocumentChunk.sort_order, + DocumentChunk.job_result_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_id.in_(list(connected_target_ids))) + .where(DocumentChunk.chunk_type == asset_type) + .order_by(DocumentChunk.sort_order) + ) + ).all() + else: + connected_rows = [] + + job_id = ( + await db.execute(select(JobResult.job_id).where(JobResult.id == job_result_id)) + ).scalar() or "" + seen_ids: set[str] = set() + chunks: list[dict[str, Any]] = [] + for row in list(asset_rows) + list(connected_rows): + chunk_id = row[0] + if chunk_id in seen_ids: + continue + seen_ids.add(chunk_id) + + owner_section_path = owner_by_target_id.get(chunk_id) + if not owner_section_path: + own_section_path = section_path_by_id.get(row[4]) + if own_section_path and own_section_path == "Root": + logger.warning( + " asset_filter_step: rejecting root-level owner fallback " + f"chunk_id={chunk_id} section_path={own_section_path}" + ) + own_section_path = None + owner_section_path = own_section_path + + if not owner_section_path: + logger.warning( + f" asset_filter_step unresolved owner: chunk_id={chunk_id} " + f"file_path={row[3]} scope={scope_path or 'root'}" + ) + continue + + chunks.append( + { + "document_id": document_id, + "chunk_id": chunk_id, + "chunk_type": row[1], + "content": row[2], + "file_path": row[3], + "section_id": row[4], + "section_path": owner_section_path, + "owner_section_path": owner_section_path, + "source_chunk_path": row[5], + "chunk_metadata": row[6] or {}, + "sort_order": row[7], + "job_result_id": job_result_id, + "job_id": job_id, + } + ) + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" asset_filter_step scope={scope_path or 'root'} " + f"type={asset_type}: {len(chunks)} chunks found, {latency}ms" + ) + return chunks + + except Exception as exc: + logger.error(f" asset_filter_step failed: {exc}") + return [] diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py new file mode 100644 index 000000000..5f32e36cb --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py @@ -0,0 +1,294 @@ +"""Agentic retrieval discovery tools. + +This Module owns phase-1 retrieval: lexical bottom discovery and document +selection from the document-level knowledge map. The public tool adapter stays +in ``tools.py`` so orchestrator call sites keep a stable interface. +""" +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document +from shared.services.retrieval.agentic.budget import BudgetExceeded +from shared.services.retrieval.agentic.knowledge_map import build_knowledge_map_overview +from shared.services.retrieval.agentic.prompts import ( + FILE_SELECT_PROMPT, + format_budget_block, + parse_json_array, +) +from shared.services.retrieval.agentic.types import ToolResult +from shared.services.retrieval.channels import content_channel, path_channel, term_channel +from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.scoring import ( + merge_channels_rrf, + merge_same_section_rows, + normalize_row_scores, +) +from shared.services.retrieval.settings import ( + CHANNEL_WEIGHT_CONTENT, + CHANNEL_WEIGHT_PATH, + CHANNEL_WEIGHT_TERM, + INTERNAL_RECALL_K_MULTIPLIER, + resolve_allowed_chunk_types, +) + + +async def bottom_discovery( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int = 1, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + channels: list[str] | None = None, + channel_weights: dict[str, float] | None = None, + internal_recall_k: int | None = None, + **_kwargs: Any, +) -> ToolResult: + """Run 3-channel BM25 discovery plus RRF fusion.""" + t0 = time.monotonic() + try: + allowed_chunk_types = resolve_allowed_chunk_types(data_type) + effective_recall_k = ( + internal_recall_k + if internal_recall_k is not None + else top_k * INTERNAL_RECALL_K_MULTIPLIER + ) + active_channels = set(channels) if channels else {"path", "content", "term"} + + path_rows: list[dict[str, Any]] = [] + content_rows: list[dict[str, Any]] = [] + term_rows: list[dict[str, Any]] = [] + + if "path" in active_channels: + path_rows = await path_channel( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=effective_recall_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, + filter_mode=filter_mode, + ) + + if "content" in active_channels: + content_rows = await content_channel( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=effective_recall_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, + filter_mode=filter_mode, + ) + + if "term" in active_channels: + term_rows = await term_channel( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=effective_recall_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, + filter_mode=filter_mode, + ) + + default_weights = { + "path": CHANNEL_WEIGHT_PATH, + "content": CHANNEL_WEIGHT_CONTENT, + "term": CHANNEL_WEIGHT_TERM, + } + effective_weights = {**default_weights, **(channel_weights or {})} + + channel_lists: list[list[dict[str, Any]]] = [] + weight_list: list[float] = [] + if path_rows: + channel_lists.append(path_rows) + weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) + if content_rows: + channel_lists.append(content_rows) + weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) + if term_rows: + channel_lists.append(term_rows) + weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) + + fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else [] + fused_rows = merge_same_section_rows(fused_rows) + + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + + doc_id_counts: dict[str, int] = {} + for row in fused_rows: + did = row.get("document_id", "") + if did: + doc_id_counts[did] = doc_id_counts.get(did, 0) + 1 + top_doc_ids = sorted( + doc_id_counts, + key=lambda document_id: doc_id_counts[document_id], + reverse=True, + )[:5] + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" agentic.bottom_discovery: {len(fused_rows)} fused rows, " + f"top_doc_ids={top_doc_ids}, {latency}ms" + ) + return ToolResult( + status="discovery_done", + payload={ + "fused_rows": fused_rows, + "top_doc_ids": top_doc_ids, + "channel_counts": { + "path": len(path_rows), + "content": len(content_rows), + "term": len(term_rows), + }, + }, + latency_ms=latency, + ) + except Exception as exc: + latency = int((time.monotonic() - t0) * 1000) + logger.error(f" agentic.bottom_discovery failed: {exc}") + return ToolResult(status="error", error=str(exc), latency_ms=latency) + + +async def kg_document_select( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + llm_fn: LLMFn | None, + exclude_document_ids: list[str], + revision_hint: str | None = None, + **_kwargs: Any, +) -> ToolResult: + """Select candidate documents from document-level KG.""" + t0 = time.monotonic() + try: + overview_text, doc_id_to_name = await build_knowledge_map_overview( + db, + user_id=user_id, + namespace=namespace, + ) + if overview_text == "(empty)": + latency = int((time.monotonic() - t0) * 1000) + return ToolResult( + status="no_confident_doc", + payload={"reason": "no active documents in namespace"}, + latency_ms=latency, + ) + + if llm_fn is None: + latency = int((time.monotonic() - t0) * 1000) + return ToolResult( + status="no_confident_doc", + payload={"reason": "LLM not available"}, + latency_ms=latency, + ) + + revision_context = "" + if revision_hint: + revision_context = ( + "\nIMPORTANT: This is a REVISION round. " + "The previous search attempt failed because:\n" + f'"{revision_hint}"\n' + "Adjust your document selection accordingly. " + "If no document can address this, return an EMPTY array [].\n" + ) + + file_prompt = FILE_SELECT_PROMPT.format( + overview=overview_text, + query=query, + revision_context=revision_context, + budget_block=format_budget_block(_kwargs.get("budget_snapshot")), + ) + file_response = await llm_fn(file_prompt) + selected_ids = parse_json_array(file_response) + + exclude_set = set(exclude_document_ids) + valid_ids = [ + document_id + for document_id in selected_ids + if document_id in doc_id_to_name and document_id not in exclude_set + ] + + if not valid_ids: + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" agentic.kg_document_select: LLM returned no valid docs, {latency}ms" + ) + return ToolResult( + status="no_confident_doc", + payload={ + "reason": "LLM returned no valid document IDs", + "raw_ids": selected_ids, + }, + latency_ms=latency, + ) + + doc_job_map: dict[str, str] = {} + doc_stmt = ( + select(Document.document_id, Document.current_job_result_id) + .where(Document.document_id.in_(valid_ids)) + ) + doc_result = await db.execute(doc_stmt) + for document_id, job_result_id in doc_result.all(): + if job_result_id: + doc_job_map[document_id] = job_result_id + + candidate_docs = [ + { + "document_id": document_id, + "source_file_name": doc_id_to_name.get(document_id, ""), + "confidence": 1.0, + "reason": "LLM selected from KG overview", + "source": "kg_llm_select", + } + for document_id in valid_ids + ] + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" agentic.kg_document_select: {len(candidate_docs)} docs selected, {latency}ms" + ) + return ToolResult( + status="selected_docs", + payload={ + "candidate_docs": candidate_docs, + "doc_id_to_name": doc_id_to_name, + "doc_job_map": doc_job_map, + }, + latency_ms=latency, + ) + except BudgetExceeded: + raise + except Exception as exc: + latency = int((time.monotonic() - t0) * 1000) + logger.error(f" agentic.kg_document_select failed: {exc}") + return ToolResult(status="error", error=str(exc), latency_ms=latency) diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence.py b/packages/shared-python/shared/services/retrieval/agentic/evidence.py new file mode 100644 index 000000000..1d88efef7 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence.py @@ -0,0 +1,345 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import RetrievalHitStat +from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.types import DocTreeNode +from shared.services.retrieval.assets import ( + generate_retrieval_asset_url, + is_client_result_artifact_ref, +) +from shared.services.retrieval.hit_stats_service import compute_importance_score +from shared.utils.token_estimate import estimate_tokens + + +def with_context_prompt_projection( + snapshot: dict[str, object], + *, + prompt_tokens: int, +) -> dict[str, object]: + projected: dict[str, object] = dict(snapshot) + context_raw = projected.get("context") or {} + if not isinstance(context_raw, dict): + return projected + + context = dict(context_raw) + used = int(context.get("used", 0) or 0) + reserved = int(context.get("reserved", 0) or 0) + capacity = int(context.get("capacity", 0) or 0) + projected_used = min(capacity, used + max(int(prompt_tokens), 0)) + projected_remaining = max(capacity - projected_used - reserved, 0) + context.update( + { + "used_projected_before_answer": projected_used, + "answer_prompt_estimate": max(int(prompt_tokens), 0), + "remaining": projected_remaining, + "used_pct": 100 + if capacity <= 0 + else min(100, int(round((projected_used + reserved) * 100 / capacity))), + } + ) + if projected_remaining <= 0: + context["status"] = "EXHAUSTED" + elif context["used_pct"] >= 80: + context["status"] = "CRITICAL" + elif context["used_pct"] >= 50: + context["status"] = "TIGHT" + else: + context["status"] = "HEALTHY" + projected["context"] = context + return projected + + +def collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]: + media: list[dict[str, Any]] = [] + for chunks in node.leaf_content.values(): + for chunk in chunks: + chunk_type = ( + chunk.get("chunk_type") or chunk.get("type") or "" + ).strip().lower() + if chunk_type in ("image", "table"): + media.append(chunk) + for child in node.children.values(): + media.extend(collect_media_chunks(child)) + return media + + +def collect_media_chunks_all( + doc_trees: dict[str, DocTreeNode], +) -> list[dict[str, Any]]: + media: list[dict[str, Any]] = [] + for tree in doc_trees.values(): + media.extend(collect_media_chunks(tree)) + return media + + +async def build_asset_url_map( + media_chunks: list[dict[str, Any]], +) -> dict[str, str]: + url_map: dict[str, str] = {} + for chunk in media_chunks: + chunk_id = str(chunk.get("chunk_id") or "").strip() + file_path = chunk.get("file_path") or "" + job_id = chunk.get("job_id") or "" + if not chunk_id or not file_path or not job_id: + continue + if not is_client_result_artifact_ref(file_path): + continue + try: + url = await generate_retrieval_asset_url( + job_id=str(job_id), + artifact_ref=str(file_path), + ) + if url: + url_map[chunk_id] = url + except Exception as exc: + logger.warning( + f"Failed to generate asset URL for {chunk_id} (ignored): {exc}" + ) + return url_map + + +def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]: + paths = set(node.leaf_content.keys()) + for child in node.children.values(): + paths.update(_collect_all_leaf_paths(child)) + return paths + + +def _collect_visible_paths(node: DocTreeNode) -> set[str]: + paths = {item["path"] for item in node.outline_items if item.get("path")} + for child in node.children.values(): + paths.update(_collect_visible_paths(child)) + return paths + + +def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None: + parts = path.split(" / ") + for i in range(len(parts) - 1, 0, -1): + ancestor = " / ".join(parts[:i]) + if ancestor in target_paths: + return ancestor + return None + + +def reconcile_deferred_assets( + tree: DocTreeNode, + pending_assets: list[dict], +) -> None: + final_paths = _collect_all_leaf_paths(tree) + visible_paths = _collect_visible_paths(tree) + all_target_paths = final_paths | visible_paths + + if not all_target_paths: + return + + existing_ids = { + str(row.get("chunk_id") or "") + for row in tree.flatten_chunk_rows() + if row.get("chunk_id") + } + + placed = 0 + ancestor_placed = 0 + for asset in pending_assets: + chunk_id = str(asset.get("chunk_id") or "") + if chunk_id and chunk_id in existing_ids: + continue + + owner_path = asset.get("owner_section_path") or asset.get("section_path") + if not owner_path: + continue + + target_path = owner_path if owner_path in all_target_paths else None + if target_path is None: + target_path = _find_closest_ancestor(owner_path, all_target_paths) + if target_path: + ancestor_placed += 1 + + if target_path is None: + continue + + tree.add_leaf_chunks(target_path, [asset]) + if chunk_id: + existing_ids.add(chunk_id) + placed += 1 + + if placed: + tree.reparent_leaf_content() + logger.info( + f" deferred asset reconcile: {placed}/{len(pending_assets)} " + f"assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths " + f"(ancestor_fallback={ancestor_placed})" + ) + + +async def render_evidence( + db: AsyncSession, + doc_trees: dict[str, DocTreeNode], + doc_id_to_name: dict[str, str], +) -> str: + del db + + from shared.services.retrieval.agentic.evidence_renderer import render_unified_doc_tree + + asset_url_map = await build_asset_url_map(collect_media_chunks_all(doc_trees)) + + evidence_parts: list[str] = [] + for doc_id, doc_tree in doc_trees.items(): + if doc_tree.has_content(): + doc_name = doc_id_to_name.get(doc_id, doc_id) + rendered = render_unified_doc_tree( + doc_tree, + doc_name, + asset_lookup=asset_url_map, + ) + if rendered.strip(): + evidence_parts.append(rendered) + + return "\n\n".join(evidence_parts) if evidence_parts else "(no evidence collected)" + + +def _iter_leaf_content(node: DocTreeNode): + for path, chunks in node.leaf_content.items(): + yield path, chunks + for child in node.children.values(): + yield from _iter_leaf_content(child) + + +def _collect_confidences(node: DocTreeNode) -> dict[str, float]: + values = dict(node.confidence) + for child in node.children.values(): + for path, score in _collect_confidences(child).items(): + values[path] = max(values.get(path, 0.0), score) + return values + + +def _pop_leaf_path(node: DocTreeNode, path: str) -> bool: + if path in node.leaf_content: + node.leaf_content.pop(path) + return True + for child in node.children.values(): + if _pop_leaf_path(child, path): + return True + return False + + +def _estimate_chunks_tokens(chunks: list[dict[str, Any]]) -> int: + text = "\n".join(str(chunk.get("content") or "") for chunk in chunks) + return estimate_tokens(text) + + +async def _fetch_importance_norm_scores( + db: AsyncSession, + *, + user_id: str, + namespace: str, + chunk_ids: list[str], +) -> dict[str, float]: + if not chunk_ids: + return {} + stmt = ( + select( + RetrievalHitStat.chunk_id, + RetrievalHitStat.hit_count, + RetrievalHitStat.last_hit_at, + RetrievalHitStat.created_at, + ) + .where(RetrievalHitStat.user_id == user_id) + .where(RetrievalHitStat.namespace == namespace) + .where(RetrievalHitStat.hit_kind == "chunk") + .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) + ) + result = await db.execute(stmt) + scores: dict[str, float] = {} + for chunk_id, hit_count, last_hit_at, created_at in result.all(): + if chunk_id and last_hit_at and created_at: + scores[str(chunk_id)] = compute_importance_score( + hit_count, + last_hit_at, + created_at, + ) + return scores + + +async def trim_evidence_to_budget( + db: AsyncSession, + *, + doc_trees: dict[str, DocTreeNode], + doc_id_to_name: dict[str, str], + context_remaining: int, + user_id: str, + namespace: str, + ledger: BudgetLedger | None, + safety_margin: float = 0.9, +) -> str: + full_text = await render_evidence(db, doc_trees, doc_id_to_name) + target = int(max(context_remaining, 0) * safety_margin) + if estimate_tokens(full_text) <= target: + return full_text + + candidates: list[tuple[str, str, tuple[float, float, float], int]] = [] + for doc_id, tree in doc_trees.items(): + confidence = _collect_confidences(tree) + for path, chunks in _iter_leaf_content(tree): + chunk_ids = [ + str(chunk.get("chunk_id")) + for chunk in chunks + if chunk.get("chunk_id") + ] + importance = 0.0 + importance_scores = await _fetch_importance_norm_scores( + db, + user_id=user_id, + namespace=namespace, + chunk_ids=chunk_ids, + ) + if importance_scores: + importance = max(importance_scores.values()) + discovery_score = ( + float(chunks[0].get("discovery_score", 0.0) or 0.0) + if chunks + else 0.0 + ) + score = ( + float(confidence.get(path, 0.0) or 0.0), + discovery_score, + importance, + ) + candidates.append((doc_id, path, score, _estimate_chunks_tokens(chunks))) + + current_estimate = estimate_tokens(full_text) + removed: list[dict[str, Any]] = [] + for doc_id, path, score, token_estimate in sorted( + candidates, + key=lambda item: (item[2], -item[3]), + ): + if current_estimate <= target: + break + if _pop_leaf_path(doc_trees[doc_id], path): + confidence_score, discovery_score, importance_score = score + removed.append( + { + "document_id": doc_id, + "document_name": doc_id_to_name.get(doc_id, doc_id), + "path": path, + "confidence_score": round(confidence_score, 4), + "discovery_score": round(discovery_score, 4), + "importance_score": round(importance_score, 4), + "token_estimate": token_estimate, + } + ) + current_estimate = max(current_estimate - token_estimate, 0) + + if ledger is not None: + ledger.trimmed_paths.extend(removed) + logger.info( + f" agentic.trim_evidence: removed={len(removed)} " + f"est_tokens={current_estimate} target={target}" + ) + return await render_evidence(db, doc_trees, doc_id_to_name) diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py b/packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py new file mode 100644 index 000000000..22326d484 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py @@ -0,0 +1,182 @@ +"""Render agentic document trees into evidence text.""" +from __future__ import annotations + +from typing import Any, cast + +from shared.services.retrieval.agentic.types import DocTreeNode + + +def render_unified_doc_tree( + node: DocTreeNode, + doc_name: str, + depth: int = 0, + asset_lookup: dict[str, str] | None = None, +) -> str: + """Render a DocTreeNode as one coherent hierarchy.""" + parts: list[str] = [] + indent = " " * depth + + if depth == 0: + parts.append(f"【文档】{doc_name}\n") + + child_prefixes = set(node.children.keys()) + + def min_sort(path: str) -> float: + chunks = node.leaf_content.get(path, []) + return min((chunk.get("sort_order") or float("inf") for chunk in chunks), default=float("inf")) + + render_queue: list[tuple[float, str, dict | str]] = [] + outline_paths: set[str] = set() + outline_position = 0.0 + + for item in node.outline_items: + path = item.get("path", "") + if any(path.startswith(child_prefix + " / ") for child_prefix in child_prefixes): + continue + outline_paths.add(path) + + if path in node.leaf_content or path in node.children: + sort_key = min_sort(path) if path in node.leaf_content else outline_position + else: + sort_key = outline_position + outline_position = max(outline_position, sort_key) + 0.001 + + render_queue.append((sort_key, "outline", item)) + + for path in node.leaf_content: + if path not in outline_paths: + render_queue.append((min_sort(path), "orphan_leaf", path)) + + for path in node.children: + if path not in outline_paths: + render_queue.append((float("inf"), "orphan_child", path)) + + render_queue.sort(key=lambda item: item[0]) + + for _sort_key, render_type, data in render_queue: + if render_type == "outline": + item = cast(dict, data) + path = item.get("path", "") + title = item.get("title", "") + is_leaf = item.get("is_leaf", False) + level = item.get("level", 1) + leaf_tag = " [Leaf]" if is_leaf else "" + + level_tag = f"[L{level}] " if level else "" + if level <= 1: + parts.append(f"{indent}▸ {level_tag}{title}{leaf_tag}") + else: + parts.append(f"{indent}└ {level_tag}{title}{leaf_tag}") + + sub_indent = indent + " " + if path in node.children: + child = node.children[path] + if path in node.leaf_content: + render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) + child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup) + if child_text.strip(): + parts.append(child_text) + elif path in node.leaf_content: + render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) + + elif render_type == "orphan_leaf": + path = cast(str, data) + title = path.rsplit(" / ", 1)[-1] if " / " in path else path + parts.append(f"{indent}▸ [Leaf] {title}") + render_leaf_chunks(parts, node.leaf_content[path], indent + " ", asset_lookup=asset_lookup) + + elif render_type == "orphan_child": + path = cast(str, data) + title = path.rsplit(" / ", 1)[-1] if " / " in path else path + parts.append(f"{indent}▸ {title} [DrillDown]") + child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) + if child_text.strip(): + parts.append(child_text) + + return "\n".join(parts) + + +def render_leaf_chunks( + parts: list[str], + chunks: list[dict[str, Any]], + indent: str, + asset_lookup: dict[str, str] | None = None, +) -> None: + chunk_by_id = { + chunk.get("chunk_id", ""): chunk + for chunk in chunks + if chunk.get("chunk_id") + } + rendered_ids: set[str] = set() + + for chunk in chunks: + chunk_id = chunk.get("chunk_id", "") + if chunk_id and chunk_id in rendered_ids: + continue + + chunk_type = (chunk.get("chunk_type") or chunk.get("type") or "text").strip().lower() + if chunk_type in ("image", "table"): + continue + + if chunk_id: + rendered_ids.add(chunk_id) + + content = str(chunk.get("content", "")).strip() + for connection in (chunk.get("chunk_metadata") or {}).get("connect_to") or []: + target = chunk_by_id.get(connection.get("target", "")) + if not target: + continue + target_id = target.get("chunk_id", "") + target_type = (target.get("chunk_type") or target.get("type") or "").strip().lower() + ref_str = connection.get("ref", "") + if not ref_str or ref_str not in content: + continue + + if target_id: + rendered_ids.add(target_id) + + if target_type == "table": + table_html = str(target.get("content", "")).strip() + content = content.replace(ref_str, f"\n[表格内容]\n{table_html}\n") + elif target_type == "image": + file_path = target.get("file_path") or "" + image_description = str(target.get("content", "")).strip() + if ref_str in image_description: + image_description = image_description.replace(ref_str, "").strip() + asset_url = (asset_lookup or {}).get(target_id, "") if target_id else "" + display_ref = asset_url or file_path + if display_ref: + content = content.replace(ref_str, f"\n[图片: {display_ref}]\n{image_description}\n") + elif image_description: + content = content.replace(ref_str, f"\n[图片描述]\n{image_description}\n") + + for line in content.split("\n"): + if line.strip(): + parts.append(f"{indent}┈ {line}") + + for chunk in chunks: + chunk_id = chunk.get("chunk_id", "") + if chunk_id and chunk_id in rendered_ids: + continue + if chunk_id: + rendered_ids.add(chunk_id) + + chunk_type = (chunk.get("chunk_type") or chunk.get("type") or "").strip().lower() + if chunk_type == "image": + file_path = chunk.get("file_path") or "" + image_description = str(chunk.get("content", "")).strip() + asset_url = (asset_lookup or {}).get(chunk_id, "") if chunk_id else "" + display_ref = asset_url or file_path + if display_ref: + parts.append(f"{indent}┈ [图片: {display_ref}]") + if image_description: + for line in image_description.split("\n"): + if line.strip(): + parts.append(f"{indent}┈ {line}") + elif chunk_type == "table": + table_html = str(chunk.get("content", "")).strip() + parts.append(f"{indent}┈ [表格内容]") + if table_html: + for line in table_html.split("\n"): + if line.strip(): + parts.append(f"{indent}┈ {line}") diff --git a/packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py b/packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py new file mode 100644 index 000000000..c0196be66 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py @@ -0,0 +1,93 @@ +"""Knowledge-map overview for agentic document selection.""" +from __future__ import annotations + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, GraphNode + + +_MAX_OVERVIEW_FILES = 50 + + +async def build_knowledge_map_overview( + db: AsyncSession, + *, + user_id: str, + namespace: str, +) -> tuple[str, dict[str, str]]: + """Build a file-level knowledge map overview for LLM file selection.""" + doc_stmt = ( + select(Document) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + .order_by(Document.updated_at.desc()) + .limit(_MAX_OVERVIEW_FILES) + ) + doc_result = await db.execute(doc_stmt) + documents = list(doc_result.scalars()) + + if not documents: + return "(empty)", {} + + doc_ids = [document.document_id for document in documents] + doc_id_to_name = { + document.document_id: (document.source_file_name or document.document_id) + for document in documents + } + + chunk_stats_stmt = ( + select( + DocumentChunk.document_id, + func.count(DocumentChunk.id).label("chunk_count"), + func.count(func.nullif(DocumentChunk.chunk_type, "text")).label("media_count"), + ) + .join( + Document, + (Document.document_id == DocumentChunk.document_id) + & (Document.current_job_result_id == DocumentChunk.job_result_id), + ) + .where(DocumentChunk.document_id.in_(doc_ids)) + .group_by(DocumentChunk.document_id) + ) + chunk_stats_result = await db.execute(chunk_stats_stmt) + chunk_stats: dict[str, dict[str, int]] = {} + for document_id, chunk_count, media_count in chunk_stats_result.all(): + chunk_stats[document_id] = {"total": chunk_count, "media": media_count} + + graph_summary_stmt = ( + select(GraphNode.owner_document_id, GraphNode.properties) + .where(GraphNode.owner_document_id.in_(doc_ids)) + .where(GraphNode.node_kind == "document") + ) + graph_summary_result = await db.execute(graph_summary_stmt) + doc_top_summaries: dict[str, str] = {} + for document_id, properties in graph_summary_result.all(): + if not isinstance(properties, dict): + continue + top_summary = str(properties.get("top_summary") or "").strip() + if top_summary: + doc_top_summaries[document_id] = top_summary + + lines: list[str] = [] + for document in documents: + document_id = document.document_id + name = doc_id_to_name[document_id] + stats = chunk_stats.get(document_id, {"total": 0, "media": 0}) + top_summary = doc_top_summaries.get(document_id, "") + + line = f'- [{document_id}] {name} chunks={stats["total"]}' + if stats["media"] > 0: + line += f' media={stats["media"]}' + if top_summary: + line += f"\n top_summary:\n{indent_block(top_summary, 4)}" + lines.append(line) + + return "\n".join(lines), doc_id_to_name + + +def indent_block(text: str, spaces: int) -> str: + prefix = " " * spaces + return "\n".join(f"{prefix}{line}" for line in str(text or "").splitlines()) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py new file mode 100644 index 000000000..8d489f8be --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py @@ -0,0 +1,564 @@ +"""Agentic retrieval navigation tools. + +This Module owns document-scope navigation and post-navigation discovery +selection. It keeps the LLM prompt, section traversal, hydration, and asset +owner reconciliation local to the navigation seam. +""" +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy import func as sa_func +from sqlalchemy import or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import DocumentChunk, DocumentSection +from shared.services.retrieval.agentic import asset_tools +from shared.services.retrieval.agentic.budget import BudgetExceeded +from shared.services.retrieval.agentic.prompts import ( + ACTION_PROMPT, + DISCOVERY_SELECT_PROMPT, + format_budget_block, + parse_action_response, +) +from shared.services.retrieval.agentic.section_tree import ( + format_items_for_llm, + load_child_sections, +) +from shared.services.retrieval.agentic.types import DocTreeNode +from shared.services.retrieval.hydration import ( + hydrate_connected_target_rows, + hydrate_paths_to_rows, +) +from shared.services.retrieval.lexical_text import normalize_section_path +from shared.services.retrieval.llm_adapter import LLMFn + + +_MAX_DISCOVERY_PER_DOC = 3 + + +async def navigate_step( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + query: str, + llm_fn: LLMFn, + user_id: str, + namespace: str, + doc_name: str = "", + scope_path: str | list[str] | None = None, + exclude_paths: set[str] | None = None, + revision_hint: str | None = None, + budget_snapshot: dict | None = None, +) -> tuple[str, list[str], DocTreeNode, list[dict]]: + """Navigate one document scope and hydrate selected sections.""" + scope_paths = ( + scope_path if isinstance(scope_path, list) + else [scope_path] if scope_path + else [] + ) + scope_path_set = set(scope_paths) + + empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None) + + try: + items = await load_child_sections( + db, + document_id, + job_result_id, + scope_path, + exclude_paths=exclude_paths, + ) + if not items: + return "STOP", [], empty, [] + + selectable = { + item["path"]: item for item in items if item.get("selectable", False) + } + total_images, total_tables = await _count_assets_under_scope( + db, + document_id=document_id, + job_result_id=job_result_id, + scope_paths=scope_paths, + ) + tools_block = _build_tools_block(total_images, total_tables) + + items_text, overflowed = format_items_for_llm(items) + prompt = _build_navigation_prompt( + document_id=document_id, + doc_name=doc_name, + query=query, + scope_paths=scope_paths, + budget_snapshot=budget_snapshot, + items_text=items_text, + tools_block=tools_block, + revision_hint=revision_hint, + ) + + response = await llm_fn(prompt) + parsed = parse_action_response(response) + action = parsed["action"] + selected_tools = parsed["tools"] + selections = parsed["selections"] + + scope_label = ", ".join(scope_paths) if scope_paths else "root" + logger.info( + f" navigate_step scope={scope_label}: " + f"action={action} tools={selected_tools} " + f"selections={len(selections)} selectable={len(selectable)} " + f"overflowed={overflowed}" + ) + + node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None) + node.outline_items = [item for item in items if item.get("show_summary", True)] + + valid_selections = [ + selection + for selection in selections + if selection["path"] in selectable and selection["path"] not in scope_path_set + ] + + pending: list[dict] = [] + path_selections: list[dict[str, Any]] = [] + for selection in valid_selections: + path = selection["path"] + confidence = selection.get("confidence", 0.7) + item = selectable[path] + node.confidence[path] = confidence + + if item.get("is_leaf"): + path_selections.append({ + "path": path, + "confidence": confidence, + "hydrate_mode": "chunks", + }) + else: + pending.append({"path": path, "confidence": confidence}) + path_selections.append({ + "path": path, + "confidence": confidence, + "hydrate_mode": "self_only", + }) + + await _hydrate_selections_into_node( + db, + node=node, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + ) + + return action, selected_tools, node, pending + + except BudgetExceeded: + raise + except Exception as exc: + logger.error(f" navigate_step failed for doc={document_id}: {exc}") + return "STOP", [], empty, [] + + +async def discovery_select_step( + db: AsyncSession, + *, + document_id: str, + query: str, + llm_fn: LLMFn, + user_id: str, + namespace: str, + doc_name: str = "", + discovery_hints: list[dict[str, Any]], + exclude_paths: set[str] | None = None, + revision_hint: str | None = None, + budget_snapshot: dict | None = None, +) -> DocTreeNode: + """Select and hydrate discovery-found sections after BFS navigation.""" + node = DocTreeNode(scope_path=None) + if not discovery_hints: + return node + + hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC] + + t0 = time.monotonic() + try: + hint_lines, hint_by_path, root_path_selections = _project_discovery_hints( + hints, + exclude_paths=exclude_paths, + ) + if not hint_lines and not root_path_selections: + return node + + selections: list[dict[str, Any]] = [] + if hint_lines: + prompt = _build_discovery_selection_prompt( + document_id=document_id, + doc_name=doc_name, + query=query, + hint_lines=hint_lines, + revision_hint=revision_hint, + budget_snapshot=budget_snapshot, + ) + response = await llm_fn(prompt) + parsed = parse_action_response(response) + selections = parsed.get("selections", []) + + logger.info( + f' discovery_select_step doc="{doc_name}": ' + f"hints={len(hints)} selections={len(selections)} " + f"root_selections={len(root_path_selections)}" + ) + + path_selections = _build_discovery_path_selections( + selections=selections, + hint_by_path=hint_by_path, + root_path_selections=root_path_selections, + node=node, + ) + await _hydrate_discovery_selections_into_node( + db, + node=node, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" discovery_select_step done: hydrated={len(node.leaf_content)} " + f"latency={latency}ms" + ) + return node + + except BudgetExceeded: + raise + except Exception as exc: + logger.error(f" discovery_select_step failed for doc={document_id}: {exc}") + return node + + +async def _count_assets_under_scope( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + scope_paths: list[str], +) -> tuple[int, int]: + scope_section_stmt = ( + select(DocumentSection.section_id) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + ) + if scope_paths: + scope_filters = [] + for scope in scope_paths: + scope_filters.append(DocumentSection.section_path == scope) + scope_filters.append(DocumentSection.section_path.like(f"{scope} / %")) + scope_section_stmt = scope_section_stmt.where(or_(*scope_filters)) + scope_section_ids = await db.execute(scope_section_stmt) + all_section_ids = [row[0] for row in scope_section_ids.all()] + + if not all_section_ids: + return 0, 0 + + count_stmt = ( + select( + DocumentChunk.chunk_type, + sa_func.count(DocumentChunk.id), + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(all_section_ids)) + .where(DocumentChunk.chunk_type.in_(["image", "table"])) + .group_by(DocumentChunk.chunk_type) + ) + count_result = await db.execute(count_stmt) + + total_images = 0 + total_tables = 0 + for chunk_type, count in count_result.all(): + if chunk_type == "image": + total_images = count + elif chunk_type == "table": + total_tables = count + return total_images, total_tables + + +def _build_tools_block(total_images: int, total_tables: int) -> str: + if total_images <= 0 and total_tables <= 0: + return "" + + tools_lines = ["\nOptional asset tools (usable with NAVIGATE or STOP):\n"] + if total_images > 0: + tools_lines.append( + f" FIND_IMAGES — Extract image/chart assets under the current scope ({total_images} available).\n" + ) + if total_tables > 0: + tools_lines.append( + f" FIND_TABLES — Extract table/data assets under the current scope ({total_tables} available).\n" + ) + tools_lines.append( + " Note: with NAVIGATE selections, asset tools are limited to the selected sections; " + "with STOP or no selections, they use the current scope.\n" + ) + return "".join(tools_lines) + + +def _build_navigation_prompt( + *, + document_id: str, + doc_name: str, + query: str, + scope_paths: list[str], + budget_snapshot: dict | None, + items_text: str, + tools_block: str, + revision_hint: str | None, +) -> str: + if not scope_paths: + scope_header = "Current scope: root (document top level)" + elif len(scope_paths) == 1: + scope_header = f'Current scope: navigating into "{scope_paths[0]}"' + else: + scope_header = f"Current scope: navigating into {len(scope_paths)} sections" + + prompt = ACTION_PROMPT.format( + doc_name=doc_name or document_id, + doc_id=document_id, + scope_header=scope_header, + budget_block=format_budget_block(budget_snapshot), + items_overview=items_text, + query=query, + tools_block=tools_block, + ) + if revision_hint: + prompt += ( + "\n\nIMPORTANT: Previous round feedback: " + f'"{revision_hint}". Adjust your selections accordingly.' + ) + return prompt + + +async def _hydrate_selections_into_node( + db: AsyncSession, + *, + node: DocTreeNode, + path_selections: list[dict[str, Any]], + user_id: str, + namespace: str, + document_id: str, + job_result_id: str, +) -> None: + chunks = await hydrate_paths_to_rows( + db, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + if not chunks: + return + + connected = await hydrate_connected_target_rows( + db=db, + rows=chunks, + exclude_document_ids=[], + exclude_sections=[], + ) + if connected: + owner_map = asset_tools.build_connected_owner_map(chunks) + for chunk in connected: + if not chunk.get("owner_section_path"): + chunk["owner_section_path"] = owner_map.get(str(chunk.get("chunk_id") or "")) + chunks = chunks + connected + + root_map = await asset_tools.resolve_root_asset_owners( + db, + document_id=document_id, + job_result_id=job_result_id, + chunks=chunks, + ) + if root_map: + for chunk in chunks: + if chunk.get("owner_section_path"): + continue + chunk_id = str(chunk.get("chunk_id") or "") + if chunk_id in root_map: + chunk["owner_section_path"] = root_map[chunk_id] + + _add_chunks_to_node(node, chunks) + + +def _project_discovery_hints( + hints: list[dict[str, Any]], + *, + exclude_paths: set[str] | None, +) -> tuple[list[str], dict[str, dict], list[dict[str, Any]]]: + exclude_set = { + normalize_section_path(path) + for path in (exclude_paths or set()) + if path + } + hint_lines: list[str] = [] + hint_by_path: dict[str, dict] = {} + root_path_selections: list[dict[str, Any]] = [] + for hint in hints: + section_path = normalize_section_path(hint.get("section_path", "")) + if not section_path: + continue + if section_path in exclude_set: + continue + if section_path in hint_by_path: + continue + + hint_by_path[section_path] = hint + if section_path == "Root": + root_path_selections.append({ + "path": section_path, + "confidence": float( + hint.get("discovery_score") or hint.get("score") or 0.7 + ), + "hydrate_mode": "self_only", + }) + continue + + summary = hint.get("summary", "") or "" + hint_lines.append(f'▸ path="{section_path}"') + if summary: + hint_lines.append(f" {summary[:300]}") + + return hint_lines, hint_by_path, root_path_selections + + +def _build_discovery_selection_prompt( + *, + document_id: str, + doc_name: str, + query: str, + hint_lines: list[str], + revision_hint: str | None, + budget_snapshot: dict | None, +) -> str: + revision_context = "" + if revision_hint: + revision_context = ( + "\nIMPORTANT: This is a REVISION round. " + "The previous search attempt failed because:\n" + f'"{revision_hint}"\n' + "Adjust your selection accordingly. " + "If no candidate is relevant, return an EMPTY list [].\n" + ) + + return DISCOVERY_SELECT_PROMPT.format( + doc_name=doc_name or document_id, + budget_block=format_budget_block(budget_snapshot), + items="\n".join(hint_lines), + query=query, + revision_context=revision_context, + ) + + +def _build_discovery_path_selections( + *, + selections: list[dict[str, Any]], + hint_by_path: dict[str, dict], + root_path_selections: list[dict[str, Any]], + node: DocTreeNode, +) -> list[dict[str, Any]]: + valid_selections = [ + selection for selection in selections if selection["path"] in hint_by_path + ] + path_selections = list(root_path_selections) + for selection in valid_selections: + path = selection["path"] + confidence = selection.get("confidence", 0.7) + node.confidence[path] = confidence + path_selections.append({"path": path, "confidence": confidence}) + + if not path_selections and hint_by_path: + fallback_path, fallback_hint = next(iter(hint_by_path.items())) + fallback_confidence = float( + fallback_hint.get("discovery_score") + or fallback_hint.get("score") + or 0.5 + ) + node.confidence[fallback_path] = fallback_confidence + path_selections.append({ + "path": fallback_path, + "confidence": fallback_confidence, + "hydrate_mode": "self_only", + }) + + return path_selections + + +async def _hydrate_discovery_selections_into_node( + db: AsyncSession, + *, + node: DocTreeNode, + path_selections: list[dict[str, Any]], + user_id: str, + namespace: str, + document_id: str, +) -> None: + chunks = await hydrate_paths_to_rows( + db, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + if not chunks: + return + + connected = await hydrate_connected_target_rows( + db=db, + rows=chunks, + exclude_document_ids=[], + exclude_sections=[], + ) + if connected: + owner_map = asset_tools.build_connected_owner_map(chunks) + for chunk in connected: + if not chunk.get("owner_section_path"): + chunk["owner_section_path"] = owner_map.get(str(chunk.get("chunk_id") or "")) + chunks = chunks + connected + + job_result_id = next( + (str(chunk["job_result_id"]) for chunk in chunks if chunk.get("job_result_id")), + None, + ) + root_map = ( + await asset_tools.resolve_root_asset_owners( + db, + document_id=document_id, + job_result_id=job_result_id, + chunks=chunks, + ) + if job_result_id + else {} + ) + if root_map: + for chunk in chunks: + if chunk.get("owner_section_path"): + continue + chunk_id = str(chunk.get("chunk_id") or "") + if chunk_id in root_map: + chunk["owner_section_path"] = root_map[chunk_id] + + _add_chunks_to_node(node, chunks) + + +def _add_chunks_to_node(node: DocTreeNode, chunks: list[dict[str, Any]]) -> None: + for chunk in chunks: + real_path = ( + chunk.get("owner_section_path") + or chunk.get("section_path") + or chunk.get("source_chunk_path") + ) + if real_path: + node.add_leaf_chunks(str(real_path), [chunk]) diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index 8aa58964d..f5bbd4889 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -25,9 +25,17 @@ from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document, DocumentChunk, RetrievalHitStat +from shared.models.database.document import Document, DocumentChunk from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger, BudgetPoolName +from shared.services.retrieval.agentic.evidence import ( + build_asset_url_map as _build_asset_url_map, + collect_media_chunks_all as _collect_media_chunks_all, + reconcile_deferred_assets as _reconcile_deferred_assets, + render_evidence as _render_evidence, + trim_evidence_to_budget as _trim_evidence_to_budget, + with_context_prompt_projection as _with_context_prompt_projection, +) from shared.services.retrieval.agentic.trace import TraceRecorder from shared.services.retrieval.agentic.types import ( AgentRunConfig, @@ -37,217 +45,12 @@ DocTreeNode, ToolResult, ) -from shared.services.retrieval.app_service import ( - generate_retrieval_asset_url, - _is_client_result_artifact_ref, -) from shared.services.retrieval.llm_adapter import LLMFn from shared.services.retrieval.llm_adapter import current_llm_usage -from shared.services.retrieval.hit_stats_service import compute_importance_score from shared.utils.token_estimate import estimate_tokens -def _with_context_prompt_projection( - snapshot: dict[str, object], - *, - prompt_tokens: int, -) -> dict[str, object]: - """Return a display snapshot that includes the upcoming answer prompt.""" - projected: dict[str, object] = dict(snapshot) - context_raw = projected.get('context') or {} - if not isinstance(context_raw, dict): - return projected - context = dict(context_raw) - used = int(context.get('used', 0) or 0) - reserved = int(context.get('reserved', 0) or 0) - capacity = int(context.get('capacity', 0) or 0) - projected_used = min(capacity, used + max(int(prompt_tokens), 0)) - projected_remaining = max(capacity - projected_used - reserved, 0) - context.update({ - 'used_projected_before_answer': projected_used, - 'answer_prompt_estimate': max(int(prompt_tokens), 0), - 'remaining': projected_remaining, - 'used_pct': 100 if capacity <= 0 else min( - 100, - int(round((projected_used + reserved) * 100 / capacity)), - ), - }) - if projected_remaining <= 0: - context['status'] = 'EXHAUSTED' - elif context['used_pct'] >= 80: - context['status'] = 'CRITICAL' - elif context['used_pct'] >= 50: - context['status'] = 'TIGHT' - else: - context['status'] = 'HEALTHY' - projected['context'] = context - return projected - - - - - -def _collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]: - """Recursively collect image/table chunks from a doc tree's leaf_content.""" - media: list[dict[str, Any]] = [] - for chunks in node.leaf_content.values(): - for chunk in chunks: - ct = (chunk.get('chunk_type') or chunk.get('type') or '').strip().lower() - if ct in ('image', 'table'): - media.append(chunk) - for child in node.children.values(): - media.extend(_collect_media_chunks(child)) - return media - - -def _collect_media_chunks_all(doc_trees: dict[str, DocTreeNode]) -> list[dict[str, Any]]: - """Collect media chunks from all doc trees.""" - media: list[dict[str, Any]] = [] - for tree in doc_trees.values(): - media.extend(_collect_media_chunks(tree)) - return media - - -async def _build_asset_url_map( - media_chunks: list[dict[str, Any]], -) -> dict[str, str]: - """Generate presigned asset URLs for media chunks. - - Uses the same ``generate_retrieval_asset_url`` as ``_to_public_response`` - in ``app_service.py`` — no separate logic. - """ - url_map: dict[str, str] = {} - for chunk in media_chunks: - chunk_id = str(chunk.get('chunk_id') or '').strip() - file_path = chunk.get('file_path') or '' - job_id = chunk.get('job_id') or '' - if not chunk_id or not file_path or not job_id: - continue - if not _is_client_result_artifact_ref(file_path): - continue - try: - url = await generate_retrieval_asset_url( - job_id=str(job_id), - artifact_ref=str(file_path), - ) - if url: - url_map[chunk_id] = url - except Exception as e: - logger.warning(f'Failed to generate asset URL for {chunk_id} (ignored): {e}') - return url_map - - -def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]: - """Recursively collect all leaf_content keys across the entire tree.""" - paths = set(node.leaf_content.keys()) - for child in node.children.values(): - paths.update(_collect_all_leaf_paths(child)) - return paths - - -def _collect_visible_paths(node: DocTreeNode) -> set[str]: - """Collect all outline_items paths across the entire tree. - - These are sections that are "visible" in the rendered tree (shown to the - LLM during navigation) even if no chunks have been hydrated into them yet. - Used as fallback targets for asset reconciliation. - """ - paths = {item['path'] for item in node.outline_items if item.get('path')} - for child in node.children.values(): - paths.update(_collect_visible_paths(child)) - return paths - - -def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None: - """Walk up a section path to find the closest ancestor in target_paths. - - Example: path="kb/file/Ch1/S1.1/S1.1.1", target_paths={"kb/file/Ch1/S1.1"} - → returns "kb/file/Ch1/S1.1" - - Uses the ' / ' separator convention from the section path format. - """ - parts = path.split(' / ') - # Walk from most specific to least specific (skip the full path itself) - for i in range(len(parts) - 1, 0, -1): - ancestor = ' / '.join(parts[:i]) - if ancestor in target_paths: - return ancestor - return None - - -def _reconcile_deferred_assets( - tree: DocTreeNode, - pending_assets: list[dict], -) -> None: - """Place collected assets into the tree based on final navigated paths. - - Called ONCE after the entire BFS + discovery merge completes for a - document. Asset placement uses a two-tier strategy: - - 1. **Exact match**: If the asset's ``owner_section_path`` matches a - leaf_content key, place directly (existing behavior). - 2. **Closest visible ancestor**: If exact match fails, walk up the - owner_section_path hierarchy to find the nearest ancestor that - appears in either leaf_content or outline_items. This handles - the case where the LLM stopped navigation early (e.g. at root) - but still requested images/tables — assets at L3 get attributed - to the visible L2 section on their path. - """ - final_paths = _collect_all_leaf_paths(tree) - visible_paths = _collect_visible_paths(tree) - all_target_paths = final_paths | visible_paths - - if not all_target_paths: - return - - # Collect existing chunk_ids to avoid duplicates - existing_ids = { - str(row.get('chunk_id') or '') - for row in tree.flatten_chunk_rows() - if row.get('chunk_id') - } - - placed = 0 - ancestor_placed = 0 - for asset in pending_assets: - chunk_id = str(asset.get('chunk_id') or '') - if chunk_id and chunk_id in existing_ids: - continue # already in tree via hydrate_connected_target_rows - - owner_path = ( - asset.get('owner_section_path') - or asset.get('section_path') - ) - if not owner_path: - continue - - # Tier 1: exact match in leaf_content or visible outline - target_path = owner_path if owner_path in all_target_paths else None - - # Tier 2: closest visible ancestor fallback - if target_path is None: - target_path = _find_closest_ancestor(owner_path, all_target_paths) - if target_path: - ancestor_placed += 1 - - if target_path is None: - continue # no visible ancestor → discard - - # Place into root; reparent_leaf_content will move to correct child - tree.add_leaf_chunks(target_path, [asset]) - if chunk_id: - existing_ids.add(chunk_id) - placed += 1 - - if placed: - tree.reparent_leaf_content() - logger.info( - f' deferred asset reconcile: {placed}/{len(pending_assets)} ' - f'assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths ' - f'(ancestor_fallback={ancestor_placed})' - ) - def _build_config_from_env() -> AgentRunConfig: """Read agent config from environment, with sensible defaults.""" @@ -272,36 +75,6 @@ def _stringify_llm_input(prompt: Any) -> str: return str(prompt) -async def _render_evidence( - db: AsyncSession, - doc_trees: dict[str, DocTreeNode], - doc_id_to_name: dict[str, str], -) -> str: - """Render unified evidence text from doc trees. - - Discovery paths are now handled by ``discovery_select_step`` in Phase 2 - and merged into doc_trees — no separate fallback needed. - """ - from shared.services.retrieval.agent_navigate import render_unified_doc_tree - - # Build asset URL map for all media chunks (images/tables) - # — same pattern as _to_public_response in app_service.py - all_media_chunks: list[dict[str, Any]] = [] - for doc_tree in doc_trees.values(): - all_media_chunks.extend(_collect_media_chunks(doc_tree)) - asset_url_map = await _build_asset_url_map(all_media_chunks) - - # Render unified evidence from doc trees - evidence_parts: list[str] = [] - for doc_id, doc_tree in doc_trees.items(): - if doc_tree.has_content(): - doc_name = doc_id_to_name.get(doc_id, doc_id) - rendered = render_unified_doc_tree(doc_tree, doc_name, asset_lookup=asset_url_map) - if rendered.strip(): - evidence_parts.append(rendered) - - return '\n\n'.join(evidence_parts) if evidence_parts else '(no evidence collected)' - async def _load_budget_inventory( db: AsyncSession, @@ -330,135 +103,6 @@ async def _load_budget_inventory( return sum(doc_chunks.values()), len(doc_chunks), doc_chunks -def _iter_leaf_content(node: DocTreeNode): - for path, chunks in node.leaf_content.items(): - yield path, chunks - for child in node.children.values(): - yield from _iter_leaf_content(child) - - -def _collect_confidences(node: DocTreeNode) -> dict[str, float]: - values = dict(node.confidence) - for child in node.children.values(): - for path, score in _collect_confidences(child).items(): - values[path] = max(values.get(path, 0.0), score) - return values - - -def _pop_leaf_path(node: DocTreeNode, path: str) -> bool: - if path in node.leaf_content: - node.leaf_content.pop(path) - return True - for child in node.children.values(): - if _pop_leaf_path(child, path): - return True - return False - - -def _estimate_chunks_tokens(chunks: list[dict[str, Any]]) -> int: - text = '\n'.join(str(chunk.get('content') or '') for chunk in chunks) - return estimate_tokens(text) - - -async def _fetch_importance_norm_scores( - db: AsyncSession, - *, - user_id: str, - namespace: str, - chunk_ids: list[str], -) -> dict[str, float]: - if not chunk_ids: - return {} - stmt = ( - select( - RetrievalHitStat.chunk_id, - RetrievalHitStat.hit_count, - RetrievalHitStat.last_hit_at, - RetrievalHitStat.created_at, - ) - .where(RetrievalHitStat.user_id == user_id) - .where(RetrievalHitStat.namespace == namespace) - .where(RetrievalHitStat.hit_kind == 'chunk') - .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) - ) - result = await db.execute(stmt) - scores: dict[str, float] = {} - for chunk_id, hit_count, last_hit_at, created_at in result.all(): - if chunk_id and last_hit_at and created_at: - scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) - return scores - - -async def _trim_evidence_to_budget( - db: AsyncSession, - *, - doc_trees: dict[str, DocTreeNode], - doc_id_to_name: dict[str, str], - context_remaining: int, - user_id: str, - namespace: str, - ledger: BudgetLedger | None, - safety_margin: float = 0.9, -) -> str: - full_text = await _render_evidence(db, doc_trees, doc_id_to_name) - target = int(max(context_remaining, 0) * safety_margin) - if estimate_tokens(full_text) <= target: - return full_text - - candidates: list[tuple[str, str, tuple[float, float, float], int]] = [] - for doc_id, tree in doc_trees.items(): - confidence = _collect_confidences(tree) - for path, chunks in _iter_leaf_content(tree): - chunk_ids = [ - str(chunk.get('chunk_id')) - for chunk in chunks - if chunk.get('chunk_id') - ] - importance = 0.0 - importance_scores = await _fetch_importance_norm_scores( - db, - user_id=user_id, - namespace=namespace, - chunk_ids=chunk_ids, - ) - if importance_scores: - importance = max(importance_scores.values()) - discovery_score = ( - float(chunks[0].get('discovery_score', 0.0) or 0.0) - if chunks else 0.0 - ) - score = (float(confidence.get(path, 0.0) or 0.0), discovery_score, importance) - candidates.append((doc_id, path, score, _estimate_chunks_tokens(chunks))) - - current_estimate = estimate_tokens(full_text) - removed: list[dict[str, Any]] = [] - for doc_id, path, _score, token_estimate in sorted( - candidates, - key=lambda item: (item[2], -item[3]), - ): - if current_estimate <= target: - break - if _pop_leaf_path(doc_trees[doc_id], path): - confidence_score, discovery_score, importance_score = _score - removed.append({ - 'document_id': doc_id, - 'document_name': doc_id_to_name.get(doc_id, doc_id), - 'path': path, - 'confidence_score': round(confidence_score, 4), - 'discovery_score': round(discovery_score, 4), - 'importance_score': round(importance_score, 4), - 'token_estimate': token_estimate, - }) - current_estimate = max(current_estimate - token_estimate, 0) - - if ledger is not None: - ledger.trimmed_paths.extend(removed) - logger.info( - f' agentic.trim_evidence: removed={len(removed)} ' - f'est_tokens={current_estimate} target={target}' - ) - return await _render_evidence(db, doc_trees, doc_id_to_name) - class RetrievalAgent: """Agentic retrieval orchestrator — navigate-then-answer loop. diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py new file mode 100644 index 000000000..7c6295913 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/prompts.py @@ -0,0 +1,223 @@ +"""Prompt templates and response parsers for agentic retrieval.""" +from __future__ import annotations + +import json +import re +from typing import Any + + +FILE_SELECT_PROMPT = """\ +You are a document routing assistant. + +{budget_block} +Below is a knowledge base overview showing all available documents, +their navigation summaries, chunk counts, and media counts. + +=== Knowledge Base Overview === +{overview} +=== End Overview === + +User query: {query} +{revision_context} +Based on the query, select documents that may contain relevant information. +If NO document in the knowledge base is relevant to the query, return an EMPTY array []. +Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"] +Do not include any explanation. +""" + + +DISCOVERY_SELECT_PROMPT = """\ +You are a document navigation assistant. + +Document: "{doc_name}" + +{budget_block} +After navigating the document's section tree, the following section paths +were additionally discovered via keyword and semantic search. +They may contain relevant evidence not found through hierarchical navigation. + +=== Discovery Candidates === +{items} +=== End Discovery Candidates === + +User query: {query} +{revision_context} +Select section paths whose content is needed to answer the query. +If none are relevant, return an EMPTY list []. + +Return ONLY a JSON object: +{{"selections": [{{"path": "...", "confidence": }}, ...]}} +Do not include any explanation. +""" + + +ACTION_PROMPT = """\ +You are a document navigation agent. + +Document: "{doc_name}" (id: {doc_id}) + +{budget_block} +{scope_header} +Below is the document's section tree. +Sections tagged [SELECT] are within the current scope and may be selected. +Other sections are shown as structural context only (not selectable). +Nodes marked [Leaf] have no further sub-sections. + +=== Section Tree === +{items_overview} +=== End Section Tree === + +User query: {query} + +=== Available Actions === + +Choose ONE action: + +NAVIGATE — Drill into selected sections for detailed content. + Consider this when the query targets specific topics and you need deeper text evidence. + Select one or more [SELECT] sections. + +STOP — Current scope evidence is sufficient. No further drill-down. + Consider this when: + - The query asks for an outline, overview, or summary + - The query is broad/global, the tree section can fulfill it without drilling into individual sections. + - You have already collected enough evidence at this level. + +{tools_block} + +When action is NAVIGATE, provide selections: +- You may ONLY select sections marked with [SELECT]. + +When action is STOP, selections must be empty. + +Return ONLY a JSON object: +{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}} +or +{{"action": "STOP", "tools": [...], "selections": []}} +Do not include any explanation. +""" + + +def parse_action_response(text: str) -> dict: + """Parse the unified navigation response from an LLM.""" + text = text.strip() + asset_tools = {"FIND_IMAGES", "FIND_TABLES"} + default = {"action": "NAVIGATE", "tools": [], "selections": []} + + def extract(data: dict) -> dict: + action = str(data.get("action", "NAVIGATE")).strip().upper() + if action not in ("NAVIGATE", "STOP"): + action = "NAVIGATE" + + tools_val = data.get("tools") or [] + if isinstance(tools_val, list): + tools = [ + str(tool).strip().upper() + for tool in tools_val + if str(tool).strip().upper() in asset_tools + ] + else: + tools = [] + + if action == "STOP": + return {"action": action, "tools": tools, "selections": []} + + selections_val = data.get("selections") or [] + selections = [] + if isinstance(selections_val, list): + for selection in selections_val: + if isinstance(selection, dict) and selection.get("path"): + confidence = normalize_confidence(selection.get("confidence", 0.7)) + selections.append({ + "path": str(selection["path"]), + "confidence": confidence or 0.7, + }) + + return {"action": action, "tools": tools, "selections": selections} + + try: + data = json.loads(text) + if isinstance(data, dict): + return extract(data) + except (ValueError, json.JSONDecodeError): + pass + + fence_match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", text, re.DOTALL) + if fence_match: + try: + data = json.loads(fence_match.group(1).strip()) + if isinstance(data, dict): + return extract(data) + except (ValueError, json.JSONDecodeError): + pass + + brace_match = re.search(r"\{.*\}", text, re.DOTALL) + if brace_match: + try: + data = json.loads(brace_match.group()) + if isinstance(data, dict): + return extract(data) + except (ValueError, json.JSONDecodeError): + pass + + return default + + +def format_budget_block(snapshot: dict | None) -> str: + if not snapshot: + return "" + planning = snapshot.get("planning") or {} + context = snapshot.get("context") or {} + return ( + "=== Resource Status ===\n" + f"Planning Budget: {planning.get('status', 'HEALTHY')} " + f"({planning.get('used_pct', 0)}% used)\n" + f"Context Budget: {context.get('status', 'HEALTHY')} " + f"({context.get('used_pct', 0)}% used)\n" + f"KG Coverage: {snapshot.get('explored_chunks', 0)}/" + f"{snapshot.get('total_chunks', 0)} chunks explored\n" + f"Docs Explored: {snapshot.get('explored_docs', 0)}/" + f"{snapshot.get('total_docs', 0)}\n" + "When budget is TIGHT, prefer fewer high-confidence selections over broad exploration. " + "When CRITICAL, be very selective — only pick paths with strong relevance. Return empty if evidence suffices.\n" + "=== End Resource Status ===\n" + ) + + +def parse_json_array(text: str) -> list[str]: + """Best-effort extraction of a JSON array of strings from LLM response text.""" + result = extract_json_array_payload(text) + return [str(item) for item in result] + + +def extract_json_array_payload(text: str) -> list[Any]: + text = text.strip() + try: + result = json.loads(text) + if isinstance(result, list): + return result + except (json.JSONDecodeError, ValueError): + pass + match = re.search(r"\[.*?\]", text, re.DOTALL) + if match: + try: + result = json.loads(match.group()) + if isinstance(result, list): + return result + except (json.JSONDecodeError, ValueError): + pass + return [] + + +def normalize_confidence(value: Any) -> float | None: + if value is None: + return None + if isinstance(value, str): + value = value.strip().rstrip("%") + try: + parsed = float(value) + except (TypeError, ValueError): + return None + if parsed > 1.0: + parsed = parsed / 100.0 + return max(0.0, min(parsed, 1.0)) diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py new file mode 100644 index 000000000..8c395c07c --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py @@ -0,0 +1,454 @@ +"""Section-tree loading and prompt projection for agentic navigation.""" +from __future__ import annotations + +from loguru import logger +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import DocumentChunk, DocumentSection +from shared.services.retrieval.lexical_text import normalize_section_path, split_section_path +from shared.utils.text_utils import truncate_content_preview + + +async def load_child_sections( + db: AsyncSession, + document_id: str, + job_result_id: str, + scope_path: str | list[str] | None = None, + exclude_paths: set[str] | None = None, +) -> list[dict]: + """Load the continuous context tree for a navigation scope.""" + stmt = ( + select( + DocumentSection.section_id, + DocumentSection.section_title, + DocumentSection.section_path, + DocumentSection.summary, + DocumentSection.sort_order, + ) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + .order_by(DocumentSection.sort_order) + ) + section_rows = (await db.execute(stmt)).all() + if not section_rows: + return [] + + if isinstance(scope_path, list): + scope_list = [normalize_section_path(path) for path in scope_path] + elif scope_path: + scope_list = [normalize_section_path(scope_path)] + else: + scope_list = [] + + scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0 + excluded_paths = exclude_paths or set() + + logger.debug( + f" load_child_sections: scopes={scope_list or ['root']} " + f"scope_depth={scope_depth} exclude_paths={excluded_paths if excluded_paths else 'none'} " + f"total_sections={len(section_rows)}" + ) + + all_sections: dict[str, dict] = {} + for section_id, title, path, summary, sort_order in section_rows: + if not path: + continue + normalized_path = normalize_section_path(path) + parts = split_section_path(normalized_path) + all_sections[normalized_path] = { + "title": title or parts[-1] if parts else normalized_path, + "summary": summary or "", + "sort_order": int(sort_order or 0), + "section_id": section_id, + "parts": parts, + "depth": len(parts), + } + + ancestor_prefixes: set[str] = set() + for scope in scope_list: + scope_parts = split_section_path(scope) + for index in range(1, len(scope_parts) + 1): + ancestor_prefixes.add(" / ".join(scope_parts[:index])) + + items_by_path = _select_scope_items( + all_sections, + scope_list=scope_list, + ancestor_prefixes=ancestor_prefixes, + exclude_paths=excluded_paths, + ) + if not items_by_path: + return [] + + allowed_set = _resolve_allowed_depths(items_by_path, scope_list) + if allowed_set: + to_remove = [ + path + for path, item in items_by_path.items() + if item["show_summary"] and item["level"] not in allowed_set + ] + for path in to_remove: + del items_by_path[path] + + if not items_by_path: + return [] + + await _attach_chunk_counts( + db, + document_id=document_id, + job_result_id=job_result_id, + all_sections=all_sections, + items_by_path=items_by_path, + ) + + sorted_items = sorted(items_by_path.values(), key=lambda item: item["sort_order"]) + for item in sorted_items: + item.pop("sort_order", None) + item.pop("section_id", None) + + _mark_leaf_and_selectable(sorted_items, all_section_paths=set(all_sections.keys()), allowed_set=allowed_set) + return sorted_items + + +def format_items_for_llm( + items: list[dict], + max_chars: int = 20000, +) -> tuple[str, bool]: + """Format section items with hierarchy, selectability, counts, and summaries.""" + if not items: + return "(no items available)", False + + full_text = "\n".join(_render_item(item, include_summary=True) for item in items) + if len(full_text) <= max_chars: + return full_text, False + + slim_text = "\n".join(_render_item(item, include_summary=False) for item in items) + return slim_text[:max_chars], True + + +def _select_scope_items( + all_sections: dict[str, dict], + *, + scope_list: list[str], + ancestor_prefixes: set[str], + exclude_paths: set[str], +) -> dict[str, dict]: + items_by_path: dict[str, dict] = {} + + def is_excluded(path: str) -> bool: + return bool( + exclude_paths + and any(path == excluded or path.startswith(excluded + " / ") for excluded in exclude_paths) + ) + + for path, meta in all_sections.items(): + parts = meta["parts"] + depth = meta["depth"] + + if not scope_list: + if depth < 1 or is_excluded(path): + continue + items_by_path[path] = _make_item(path, meta, show_summary=True) + continue + + matched_scope = _find_matched_scope(parts, depth=depth, scope_list=scope_list) + if matched_scope: + if is_excluded(path): + continue + items_by_path[path] = _make_item(path, meta, show_summary=True) + continue + + max_scope_depth = max(len(split_section_path(scope)) for scope in scope_list) + if depth <= max_scope_depth: + if depth == 1 and path in ancestor_prefixes: + items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) + elif depth > 1: + parent_prefix = " / ".join(parts[:-1]) + if parent_prefix in ancestor_prefixes: + items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) + + return items_by_path + + +def _make_item(path: str, meta: dict, show_summary: bool) -> dict: + return { + "path": path, + "title": meta["title"], + "summary": meta["summary"], + "level": meta["depth"], + "sort_order": meta["sort_order"], + "chunk_count": 0, + "image_count": 0, + "table_count": 0, + "section_id": meta["section_id"], + "show_summary": show_summary, + } + + +def _find_matched_scope(parts: list[str], *, depth: int, scope_list: list[str]) -> str | None: + for scope in scope_list: + scope_parts = split_section_path(scope) + scope_depth = len(scope_parts) + if depth > scope_depth and parts[:scope_depth] == scope_parts: + return scope + return None + + +def _resolve_allowed_depths(items_by_path: dict[str, dict], scope_list: list[str]) -> set[int]: + if not scope_list: + depths = { + item["level"] + for item in items_by_path.values() + if item.get("show_summary", True) + } + return set(sorted(depths)[:2]) + + allowed_set: set[int] = set() + for scope in scope_list: + scope_parts = split_section_path(scope) + scope_depth = len(scope_parts) + child_depths = { + item["level"] + for item in items_by_path.values() + if item.get("show_summary", True) + and item["level"] > scope_depth + and split_section_path(item["path"])[:scope_depth] == scope_parts + } + if child_depths: + allowed_set.update(sorted(child_depths)[:2]) + return allowed_set + + +async def _attach_chunk_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + all_sections: dict[str, dict], + items_by_path: dict[str, dict], +) -> None: + scope_item_sids = { + item["section_id"] + for item in items_by_path.values() + if item["show_summary"] + } + all_section_ids = [meta["section_id"] for meta in all_sections.values()] + if not all_section_ids or not scope_item_sids: + return + + section_id_counts = await _load_direct_chunk_counts( + db, + document_id=document_id, + job_result_id=job_result_id, + all_section_ids=all_section_ids, + ) + + sid_to_path = {meta["section_id"]: path for path, meta in all_sections.items()} + for section_id, (text_count, image_count, table_count) in section_id_counts.items(): + chunk_path = sid_to_path.get(section_id, "") + if not chunk_path: + continue + + for item_path, item in items_by_path.items(): + if not item["show_summary"]: + continue + if chunk_path == item_path or chunk_path.startswith(item_path + " / "): + item["chunk_count"] += text_count + item["image_count"] += image_count + item["table_count"] += table_count + + await _attach_connected_asset_counts( + db, + document_id=document_id, + job_result_id=job_result_id, + items_by_path=items_by_path, + sid_to_path=sid_to_path, + ) + + +async def _load_direct_chunk_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + all_section_ids: list[str], +) -> dict[str, tuple[int, int, int]]: + from sqlalchemy import case, literal_column + + chunk_stmt = ( + select( + DocumentChunk.section_id, + func.count( + case( + (DocumentChunk.chunk_type.notin_(["image", "table"]), literal_column("1")), + ) + ).label("text_count"), + func.count( + case( + (DocumentChunk.chunk_type == "image", literal_column("1")), + ) + ).label("image_count"), + func.count( + case( + (DocumentChunk.chunk_type == "table", literal_column("1")), + ) + ).label("table_count"), + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(all_section_ids)) + .group_by(DocumentChunk.section_id) + ) + chunk_rows = (await db.execute(chunk_stmt)).all() + return { + section_id: (int(text_count), int(image_count), int(table_count)) + for section_id, text_count, image_count, table_count in chunk_rows + } + + +async def _attach_connected_asset_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + items_by_path: dict[str, dict], + sid_to_path: dict[str, str], +) -> None: + scope_items_with_zero_assets = [ + item + for item in items_by_path.values() + if item["show_summary"] and item["image_count"] == 0 and item["table_count"] == 0 + ] + if not scope_items_with_zero_assets: + return + + scope_section_ids = { + item["section_id"] + for item in items_by_path.values() + if item.get("section_id") + } + if not scope_section_ids: + return + + connect_stmt = ( + select( + DocumentChunk.section_id, + DocumentChunk.chunk_metadata, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(scope_section_ids))) + .where(DocumentChunk.chunk_type == "text") + ) + connect_result = (await db.execute(connect_stmt)).all() + + section_target_ids: dict[str, set[str]] = {} + for section_id, metadata in connect_result: + if not isinstance(metadata, dict): + continue + for connection in metadata.get("connect_to") or []: + target_id = connection.get("target", "") + if target_id: + section_target_ids.setdefault(section_id, set()).add(target_id) + + if not section_target_ids: + return + + all_target_ids: set[str] = set() + for target_ids in section_target_ids.values(): + all_target_ids.update(target_ids) + + target_type_stmt = ( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_id.in_(list(all_target_ids))) + .where(DocumentChunk.chunk_type.in_(["image", "table"])) + ) + target_type_result = (await db.execute(target_type_stmt)).all() + target_types = {chunk_id: chunk_type for chunk_id, chunk_type in target_type_result} + + for section_id, target_ids in section_target_ids.items(): + ref_path = sid_to_path.get(section_id, "") + if not ref_path: + continue + referenced_images = sum(1 for target_id in target_ids if target_types.get(target_id) == "image") + referenced_tables = sum(1 for target_id in target_ids if target_types.get(target_id) == "table") + if referenced_images == 0 and referenced_tables == 0: + continue + for item_path, item in items_by_path.items(): + if not item["show_summary"]: + continue + if ref_path == item_path or ref_path.startswith(item_path + " / "): + item["image_count"] += referenced_images + item["table_count"] += referenced_tables + + +def _mark_leaf_and_selectable( + sorted_items: list[dict], + *, + all_section_paths: set[str], + allowed_set: set[int], +) -> None: + for item in sorted_items: + item_path = item["path"] + has_descendants = any( + path != item_path and path.startswith(item_path + " / ") + for path in all_section_paths + ) + item["is_leaf"] = not has_descendants + + if allowed_set: + shallowest_band = min(allowed_set) + for item in sorted_items: + if not item.get("show_summary", True): + item["selectable"] = False + elif item["level"] == shallowest_band and not item.get("is_leaf", False): + item["selectable"] = False + else: + item["selectable"] = True + else: + for item in sorted_items: + item["selectable"] = item.get("show_summary", True) + + +def _render_item(item: dict, include_summary: bool) -> str: + level = item.get("level", 1) + show_summary = item.get("show_summary", True) + is_leaf = item.get("is_leaf", False) + leaf_tag = " [Leaf]" if is_leaf else "" + path = item.get("path", "") + summary = item.get("summary") or "" + + counts_str = "" + if show_summary: + count_parts: list[str] = [] + chunk_count = item.get("chunk_count", 0) + if chunk_count > 0: + count_parts.append(f"text={chunk_count}") + image_count = item.get("image_count", 0) + if image_count > 0: + count_parts.append(f"image={image_count}") + table_count = item.get("table_count", 0) + if table_count > 0: + count_parts.append(f"table={table_count}") + counts_str = f' [{" ".join(count_parts)}]' if count_parts else "" + + indent = " " * (level - 1) + prefix = "▸" if level == 1 else "└" + level_tag = f"[L{level}]" + select_tag = "[SELECT] " if item.get("selectable", False) else "" + + lines = [ + f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}' + ] + + if include_summary and show_summary and summary: + sub_indent = " " * level + clipped = truncate_content_preview(summary, head=80, tail=0) + lines.append(f"{sub_indent}{clipped}") + + return "\n".join(lines) diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 10ab54757..29c0c4e73 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -1,151 +1,19 @@ -"""Agentic retrieval tools — thin wrappers around existing retrieval components. +"""Agentic retrieval tool adapters. -Each tool: - 1. Calls existing functions from channels.py, agent_navigate.py, app_service.py - 2. Returns a unified ToolResult - 3. Never raises — errors are captured in ToolResult.error +Concrete tool implementations live in focused Modules. This file is the stable +adapter seam used by the workflow orchestrator and contract tests. """ from __future__ import annotations -import time from typing import Any -from loguru import logger -from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document -from shared.services.retrieval.agentic.budget import BudgetExceeded +from shared.services.retrieval.agentic import asset_tools, discovery_tools, navigation_tools from shared.services.retrieval.agentic.types import DocTreeNode, ToolResult -from shared.services.retrieval.agent_navigate import ( - _build_knowledge_map_overview, - _format_items_for_llm, - _load_child_sections, - _parse_json_array, - _parse_action_response, - _ACTION_PROMPT, - _DISCOVERY_SELECT_PROMPT, - _FILE_SELECT_PROMPT, - _format_budget_block, -) -from shared.services.retrieval.app_service import ( - _CHANNEL_WEIGHT_CONTENT, - _CHANNEL_WEIGHT_PATH, - _CHANNEL_WEIGHT_TERM, - _INTERNAL_RECALL_K_MULTIPLIER, - _merge_same_section_rows, - _normalize_row_scores, - _resolve_allowed_chunk_types, - hydrate_connected_target_rows, - merge_channels_rrf, -) -from shared.services.retrieval.channels import content_channel, path_channel, term_channel -from shared.services.retrieval.lexical_text import normalize_section_path from shared.services.retrieval.llm_adapter import LLMFn -# --------------------------------------------------------------------------- -# Helper: resolve connected asset → owner text chunk section_path -# --------------------------------------------------------------------------- - -def _build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, str]: - """Build target_chunk_id → owner text chunk section_path mapping. - - When text chunks reference images/tables via connect_to metadata, - the referenced assets live in Root section. This map lets us attribute - those assets back to the text chunk's section for correct tree placement. - """ - owner_map: dict[str, str] = {} - for chunk in text_chunks: - if (chunk.get('chunk_type') or 'text') != 'text': - continue - section_path = chunk.get('section_path') or '' - if not section_path: - continue - metadata = chunk.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - continue - for conn in metadata.get('connect_to') or []: - if not isinstance(conn, dict): - continue - target_id = str(conn.get('target') or '').strip() - if target_id and target_id not in owner_map: - owner_map[target_id] = section_path - return owner_map - - -async def _resolve_root_asset_owners( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - chunks: list[dict[str, Any]], -) -> dict[str, str]: - """Resolve owner section_path for Root-stranded image/table chunks. - - When Root is hydrated directly (e.g. via discovery selection), the - batch contains standalone image/table chunks whose section_path is - 'Root'. ``_build_connected_owner_map`` cannot help because the - referencing text chunks live in other sections outside the batch. - - This function queries the *entire document* for text chunks with - connect_to metadata, using the same logic as - ``_build_connected_owner_map``, to resolve the true owner. - - Returns target_chunk_id → owner_section_path for Root assets only. - Returns empty dict when there are no Root assets (zero DB overhead). - """ - from shared.models.database.document import DocumentChunk, DocumentSection - - root_asset_ids = [ - str(c.get('chunk_id') or '') - for c in chunks - if not c.get('owner_section_path') # skip if already resolved by batch-level owner map - and (c.get('section_path') or '') == 'Root' - and (c.get('chunk_type') or '').lower() in ('image', 'table') - and c.get('chunk_id') - ] - if not root_asset_ids: - return {} - - root_asset_set = set(root_asset_ids) - - # Query all text chunks in this document for connect_to metadata - text_stmt = ( - select( - DocumentChunk.chunk_metadata, - DocumentSection.section_path, - ) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_type == 'text') - ) - result = await db.execute(text_stmt) - - owner_map: dict[str, str] = {} - for metadata, section_path in result.all(): - if not isinstance(metadata, dict) or not section_path: - continue - for conn in metadata.get('connect_to') or []: - if not isinstance(conn, dict): - continue - target_id = str(conn.get('target') or '').strip() - if target_id in root_asset_set and target_id not in owner_map: - owner_map[target_id] = section_path - - if owner_map: - logger.info( - f' _resolve_root_asset_owners: resolved {len(owner_map)}/{len(root_asset_ids)} ' - f'Root assets to their owner sections' - ) - return owner_map - - -# --------------------------------------------------------------------------- -# Tool: bottom_discovery -# --------------------------------------------------------------------------- - async def bottom_discovery( db: AsyncSession, *, @@ -157,108 +25,29 @@ async def bottom_discovery( exclude_sections: list[dict[str, str]], data_type: int = 1, signal_paths: list[str] | None = None, - filter_mode: str = 'delete', + filter_mode: str = "delete", channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, internal_recall_k: int | None = None, - **_kwargs: Any, + **kwargs: Any, ) -> ToolResult: - """Run 3-channel BM25 discovery + RRF fusion.""" - t0 = time.monotonic() - try: - allowed_chunk_types = _resolve_allowed_chunk_types(data_type) - effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER - active_channels = set(channels) if channels else {'path', 'content', 'term'} - - path_rows: list[dict[str, Any]] = [] - content_rows: list[dict[str, Any]] = [] - term_rows: list[dict[str, Any]] = [] - - if 'path' in active_channels: - path_rows = await path_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - - if 'content' in active_channels: - content_rows = await content_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - - if 'term' in active_channels: - term_rows = await term_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - - # RRF fusion - default_weights = { - 'path': _CHANNEL_WEIGHT_PATH, - 'content': _CHANNEL_WEIGHT_CONTENT, - 'term': _CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM)) - - fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else [] - fused_rows = _merge_same_section_rows(fused_rows) - - if fused_rows: - _normalize_row_scores(fused_rows, source_field='score', target_field='discovery_score', default=0.5) - - # Extract top document IDs as hints for KG selection - doc_id_counts: dict[str, int] = {} - for row in fused_rows: - did = row.get('document_id', '') - if did: - doc_id_counts[did] = doc_id_counts.get(did, 0) + 1 - top_doc_ids = sorted(doc_id_counts, key=lambda d: doc_id_counts[d], reverse=True)[:5] - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f' agentic.bottom_discovery: {len(fused_rows)} fused rows, ' - f'top_doc_ids={top_doc_ids}, {latency}ms' - ) - return ToolResult( - status='discovery_done', - payload={ - 'fused_rows': fused_rows, - 'top_doc_ids': top_doc_ids, - 'channel_counts': { - 'path': len(path_rows), - 'content': len(content_rows), - 'term': len(term_rows), - }, - }, - latency_ms=latency, - ) - except Exception as e: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.bottom_discovery failed: {e}') - return ToolResult(status='error', error=str(e), latency_ms=latency) - + return await discovery_tools.bottom_discovery( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + internal_recall_k=internal_recall_k, + **kwargs, + ) -# --------------------------------------------------------------------------- -# Tool: kg_document_select -# --------------------------------------------------------------------------- async def kg_document_select( db: AsyncSession, @@ -269,104 +58,19 @@ async def kg_document_select( llm_fn: LLMFn | None, exclude_document_ids: list[str], revision_hint: str | None = None, - **_kwargs: Any, + **kwargs: Any, ) -> ToolResult: - """Select candidate documents from document-level KG.""" - t0 = time.monotonic() - try: - overview_text, doc_id_to_name = await _build_knowledge_map_overview( - db, user_id=user_id, namespace=namespace, - ) - if overview_text == '(empty)': - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_confident_doc', - payload={'reason': 'no active documents in namespace'}, - latency_ms=latency, - ) - - if llm_fn is None: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_confident_doc', - payload={'reason': 'LLM not available'}, - latency_ms=latency, - ) - - revision_context = '' - if revision_hint: - revision_context = ( - f'\nIMPORTANT: This is a REVISION round. ' - f'The previous search attempt failed because:\n' - f'"{revision_hint}"\n' - f'Adjust your document selection accordingly. ' - f'If no document can address this, return an EMPTY array [].\n' - ) - - file_prompt = _FILE_SELECT_PROMPT.format( - overview=overview_text, query=query, - revision_context=revision_context, - budget_block=_format_budget_block(_kwargs.get('budget_snapshot')), - ) - file_response = await llm_fn(file_prompt) - selected_ids = _parse_json_array(file_response) - - exclude_set = set(exclude_document_ids) - valid_ids = [did for did in selected_ids if did in doc_id_to_name and did not in exclude_set] - - if not valid_ids: - latency = int((time.monotonic() - t0) * 1000) - logger.info(f' agentic.kg_document_select: LLM returned no valid docs, {latency}ms') - return ToolResult( - status='no_confident_doc', - payload={'reason': 'LLM returned no valid document IDs', 'raw_ids': selected_ids}, - latency_ms=latency, - ) - - # Load job_result_ids for selected documents - doc_job_map: dict[str, str] = {} - doc_stmt = ( - select(Document.document_id, Document.current_job_result_id) - .where(Document.document_id.in_(valid_ids)) - ) - doc_result = await db.execute(doc_stmt) - for did, jrid in doc_result.all(): - if jrid: - doc_job_map[did] = jrid - - candidate_docs = [] - for did in valid_ids: - candidate_docs.append({ - 'document_id': did, - 'source_file_name': doc_id_to_name.get(did, ''), - 'confidence': 1.0, - 'reason': 'LLM selected from KG overview', - 'source': 'kg_llm_select', - }) - - latency = int((time.monotonic() - t0) * 1000) - logger.info(f' agentic.kg_document_select: {len(candidate_docs)} docs selected, {latency}ms') - return ToolResult( - status='selected_docs', - payload={ - 'candidate_docs': candidate_docs, - 'doc_id_to_name': doc_id_to_name, - 'doc_job_map': doc_job_map, - }, - latency_ms=latency, - ) - except BudgetExceeded: - raise - except Exception as e: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.kg_document_select failed: {e}') - return ToolResult(status='error', error=str(e), latency_ms=latency) - - + return await discovery_tools.kg_document_select( + db, + user_id=user_id, + namespace=namespace, + query=query, + llm_fn=llm_fn, + exclude_document_ids=exclude_document_ids, + revision_hint=revision_hint, + **kwargs, + ) -# --------------------------------------------------------------------------- -# Tool: asset_filter_step (programmatic asset extraction) -# --------------------------------------------------------------------------- async def asset_filter_step( db: AsyncSession, @@ -374,218 +78,16 @@ async def asset_filter_step( document_id: str, job_result_id: str, scope_path: str | list[str] | None, - asset_type: str, # 'image' | 'table' + asset_type: str, ) -> list[dict[str, Any]]: - """Extract assets from all descendants under scope_path. - - Terminal action — no LLM involved. - Algorithm: load all text chunks under scope → parse connect_to metadata → - batch-load target image/table chunks → return directly. - - Also collects standalone asset chunks (image/table) that exist directly - under the scope but are not referenced via connect_to. - - scope_path can be: - - None: root scope (entire document) - - str: single scope path - - list[str]: multiple scope paths (queried simultaneously) - """ - from shared.models.database.document import DocumentChunk, DocumentSection - - t0 = time.monotonic() - try: - # 1. Find all section_ids under scope_path(s) - # Normalize scope to list for uniform handling - scope_list = ( - scope_path if isinstance(scope_path, list) - else [scope_path] if scope_path - else [] - ) - - section_stmt = ( - select(DocumentSection.section_id, DocumentSection.section_path) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - ) - if scope_list: - from sqlalchemy import or_ - scope_filters = [] - for sp in scope_list: - scope_filters.append(DocumentSection.section_path == sp) - scope_filters.append(DocumentSection.section_path.like(f'{sp} / %')) - section_stmt = section_stmt.where(or_(*scope_filters)) - section_result = await db.execute(section_stmt) - section_rows = section_result.all() - section_ids = {row[0] for row in section_rows} - - if not section_ids: - logger.info(f' asset_filter_step: no sections found under scope={scope_path}') - return [] - - # 2. Load target asset chunks directly (standalone assets in the scope) - asset_stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.file_path, - DocumentChunk.section_id, - DocumentChunk.source_chunk_path, - DocumentChunk.chunk_metadata, - DocumentChunk.sort_order, - DocumentChunk.job_result_id, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(section_ids))) - .where(DocumentChunk.chunk_type == asset_type) - .order_by(DocumentChunk.sort_order) - ) - asset_result = await db.execute(asset_stmt) - asset_rows = asset_result.all() - - section_path_by_id = {section_id: section_path for section_id, section_path in section_rows} - - # 3. Resolve media → owner text section via connect_to tracing - text_stmt = ( - select( - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.chunk_metadata, - DocumentChunk.source_chunk_path, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(section_ids))) - .where(DocumentChunk.chunk_type == 'text') - ) - text_result = await db.execute(text_stmt) - text_row_dicts = [ - { - 'chunk_type': chunk_type, - 'chunk_metadata': metadata or {}, - 'section_id': sid, - 'section_path': section_path_by_id.get(sid, ''), - 'source_chunk_path': scp, - } - for sid, chunk_type, metadata, scp in text_result.all() - ] - owner_by_target_id = _build_connected_owner_map(text_row_dicts) - - # Replace synthetic "Root" owner with the document's source_file_name. - # Root is a hybrid node whose real path is the file name (e.g. - # "32_安全大模型技术与市场研究报告_1.docx"); the DB stores the - # synthetic label "Root" which cannot match any outline node. - if any(v == 'Root' for v in owner_by_target_id.values()): - doc_stmt = select(Document.source_file_name).where( - Document.document_id == document_id - ) - doc_file_name = (await db.execute(doc_stmt)).scalar() or '' - if doc_file_name: - for tid in list(owner_by_target_id): - if owner_by_target_id[tid] == 'Root': - owner_by_target_id[tid] = doc_file_name - - # Collect connected target IDs for batch-loading - connected_target_ids: set[str] = set(owner_by_target_id.keys()) - - # Load connected targets that match asset_type - if connected_target_ids: - connected_stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.file_path, - DocumentChunk.section_id, - DocumentChunk.source_chunk_path, - DocumentChunk.chunk_metadata, - DocumentChunk.sort_order, - DocumentChunk.job_result_id, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_id.in_(list(connected_target_ids))) - .where(DocumentChunk.chunk_type == asset_type) - .order_by(DocumentChunk.sort_order) - ) - connected_result = await db.execute(connected_stmt) - connected_rows = connected_result.all() - else: - connected_rows = [] - - # 4. Merge and deduplicate - seen_ids: set[str] = set() - chunks: list[dict[str, Any]] = [] - - # Helper to look up job_id from job_result - from shared.models.database.job_result import JobResult - job_stmt = ( - select(JobResult.job_id) - .where(JobResult.id == job_result_id) - ) - job_result_row = await db.execute(job_stmt) - job_id = job_result_row.scalar() or '' - - for row in list(asset_rows) + list(connected_rows): - chunk_id = row[0] - if chunk_id in seen_ids: - continue - seen_ids.add(chunk_id) - - # Owner resolution: prefer connect_to-based owner - owner_section_path = owner_by_target_id.get(chunk_id) - - # Fallback: media's own section_id path, but guard against - # Root / top-level aggregation sections - if not owner_section_path: - own_section_path = section_path_by_id.get(row[4]) - if own_section_path and own_section_path == 'Root': - # Reject only the synthetic Root aggregation label; - # legitimate L1 sections (e.g. "前言") are valid owners. - logger.warning( - f' asset_filter_step: rejecting root-level owner fallback ' - f'chunk_id={chunk_id} section_path={own_section_path}' - ) - own_section_path = None - owner_section_path = own_section_path - - if not owner_section_path: - logger.warning( - f' asset_filter_step unresolved owner: chunk_id={chunk_id} ' - f'file_path={row[3]} scope={scope_path or "root"}' - ) - continue - chunks.append({ - 'document_id': document_id, - 'chunk_id': chunk_id, - 'chunk_type': row[1], - 'content': row[2], - 'file_path': row[3], - 'section_id': row[4], - 'section_path': owner_section_path, - 'owner_section_path': owner_section_path, - 'source_chunk_path': row[5], - 'chunk_metadata': row[6] or {}, - 'sort_order': row[7], - 'job_result_id': job_result_id, - 'job_id': job_id, - }) - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f' asset_filter_step scope={scope_path or "root"} ' - f'type={asset_type}: {len(chunks)} chunks found, {latency}ms' - ) - return chunks - - except Exception as e: - logger.error(f' asset_filter_step failed: {e}') - return [] + return await asset_tools.asset_filter_step( + db, + document_id=document_id, + job_result_id=job_result_id, + scope_path=scope_path, + asset_type=asset_type, + ) -# --------------------------------------------------------------------------- -# Tool: navigate_step (unified action — merges tool_select + scope_navigate) -# --------------------------------------------------------------------------- async def navigate_step( db: AsyncSession, @@ -596,226 +98,26 @@ async def navigate_step( llm_fn: LLMFn, user_id: str, namespace: str, - doc_name: str = '', + doc_name: str = "", scope_path: str | list[str] | None = None, exclude_paths: set[str] | None = None, revision_hint: str | None = None, budget_snapshot: dict | None = None, ) -> tuple[str, list[str], DocTreeNode, list[dict]]: - """Unified navigation step — one LLM call for action + tools + selections. - - scope_path can be: - - None: root scope - - str: single scope to drill into - - list[str]: multiple scopes to expand simultaneously - - Returns: - - action: 'STOP' | 'NAVIGATE' - - asset_tools: list of asset tools to run (FIND_IMAGES, FIND_TABLES) - - node: DocTreeNode with outline_items and leaf_content - - pending: list of {path, confidence} for non-leaf drill-downs (empty when STOP) - """ - from shared.services.retrieval.app_service import _hydrate_paths_to_rows - - # Normalize scope for internal use - scope_paths: list[str] = ( - scope_path if isinstance(scope_path, list) - else [scope_path] if scope_path - else [] + return await navigation_tools.navigate_step( + db, + document_id=document_id, + job_result_id=job_result_id, + query=query, + llm_fn=llm_fn, + user_id=user_id, + namespace=namespace, + doc_name=doc_name, + scope_path=scope_path, + exclude_paths=exclude_paths, + revision_hint=revision_hint, + budget_snapshot=budget_snapshot, ) - # Set of scope path strings (for filtering selections) - scope_path_set = set(scope_paths) - - empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None) - - try: - # 1. Load continuous context tree (supports multi-scope) - items = await _load_child_sections( - db, document_id, job_result_id, scope_path, - exclude_paths=exclude_paths, - ) - if not items: - return 'STOP', [], empty, [] - - # 2. Build selectable index - selectable = {item['path']: item for item in items if item.get('selectable', False)} - - # 3. Count ALL image/table chunks under the scope subtree(s) - from shared.models.database.document import DocumentChunk, DocumentSection - from sqlalchemy import func as sa_func - - scope_section_stmt = ( - select(DocumentSection.section_id) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - ) - if scope_paths: - from sqlalchemy import or_ - scope_filters = [] - for sp in scope_paths: - scope_filters.append(DocumentSection.section_path == sp) - scope_filters.append(DocumentSection.section_path.like(f'{sp} / %')) - scope_section_stmt = scope_section_stmt.where(or_(*scope_filters)) - scope_section_ids = await db.execute(scope_section_stmt) - all_section_ids = [r[0] for r in scope_section_ids.all()] - - total_images = 0 - total_tables = 0 - if all_section_ids: - count_stmt = ( - select( - DocumentChunk.chunk_type, - sa_func.count(DocumentChunk.id), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_section_ids)) - .where(DocumentChunk.chunk_type.in_(['image', 'table'])) - .group_by(DocumentChunk.chunk_type) - ) - count_result = await db.execute(count_stmt) - for chunk_type, cnt in count_result.all(): - if chunk_type == 'image': - total_images = cnt - elif chunk_type == 'table': - total_tables = cnt - - tools_block = '' - if total_images > 0 or total_tables > 0: - tools_lines = ['\nOptional asset tools (usable with NAVIGATE or STOP):\n'] - if total_images > 0: - tools_lines.append( - f' FIND_IMAGES — Extract image/chart assets under the current scope ({total_images} available).\n' - ) - if total_tables > 0: - tools_lines.append( - f' FIND_TABLES — Extract table/data assets under the current scope ({total_tables} available).\n' - ) - tools_lines.append( - ' Note: with NAVIGATE selections, asset tools are limited to the selected sections; ' - 'with STOP or no selections, they use the current scope.\n' - ) - tools_block = ''.join(tools_lines) - - # 4. Format tree and build prompt - text, overflowed = _format_items_for_llm(items) - if not scope_paths: - scope_header = 'Current scope: root (document top level)' - elif len(scope_paths) == 1: - scope_header = f'Current scope: navigating into "{scope_paths[0]}"' - else: - scope_header = f'Current scope: navigating into {len(scope_paths)} sections' - prompt = _ACTION_PROMPT.format( - doc_name=doc_name or document_id, - doc_id=document_id, - scope_header=scope_header, - budget_block=_format_budget_block(budget_snapshot), - items_overview=text, - query=query, - tools_block=tools_block, - ) - if revision_hint: - prompt += ( - f'\n\nIMPORTANT: Previous round feedback: ' - f'"{revision_hint}". Adjust your selections accordingly.' - ) - - # 5. Single LLM call - response = await llm_fn(prompt) - parsed = _parse_action_response(response) - action = parsed['action'] - asset_tools = parsed['tools'] - selections = parsed['selections'] - - scope_label = ', '.join(scope_paths) if scope_paths else 'root' - logger.info( - f' navigate_step scope={scope_label}: ' - f'action={action} tools={asset_tools} ' - f'selections={len(selections)} selectable={len(selectable)} ' - f'overflowed={overflowed}' - ) - - # 6. Build node with LOCAL items only - node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None) - local_items = [item for item in items if item.get('show_summary', True)] - node.outline_items = local_items - - # 7. Dispatch selections (only present when action == NAVIGATE) - valid_selections = [ - s for s in selections - if s['path'] in selectable and s['path'] not in scope_path_set - ] - - pending: list[dict] = [] - path_selections = [] - for sel in valid_selections: - path = sel['path'] - conf = sel.get('confidence', 0.7) - item = selectable[path] - node.confidence[path] = conf - - if item.get('is_leaf'): - path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'chunks'}) - else: - # Non-leaf → will be batched into a single next call - pending.append({'path': path, 'confidence': conf}) - path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'self_only'}) - - if path_selections: - chunks = await _hydrate_paths_to_rows( - db, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - if chunks: - connected = await hydrate_connected_target_rows( - db=db, - rows=chunks, - exclude_document_ids=[], - exclude_sections=[], - ) - if connected: - _owner_map = _build_connected_owner_map(chunks) - for c in connected: - if not c.get('owner_section_path'): - c['owner_section_path'] = _owner_map.get(str(c.get('chunk_id') or '')) - chunks = chunks + connected - - _root_map = await _resolve_root_asset_owners( - db, - document_id=document_id, - job_result_id=job_result_id, - chunks=chunks, - ) - if _root_map: - for c in chunks: - if c.get('owner_section_path'): - continue - cid = str(c.get('chunk_id') or '') - if cid in _root_map: - c['owner_section_path'] = _root_map[cid] - - for chunk in chunks: - real_path = chunk.get('owner_section_path') or chunk.get('section_path') or chunk.get('source_chunk_path') - if real_path: - node.add_leaf_chunks(str(real_path), [chunk]) - - return action, asset_tools, node, pending - - except BudgetExceeded: - raise - except Exception as e: - logger.error(f' navigate_step failed for doc={document_id}: {e}') - return 'STOP', [], empty, [] - - -# --------------------------------------------------------------------------- -# Tool: discovery_select_step (post-navigation discovery selection) -# --------------------------------------------------------------------------- - -_MAX_DISCOVERY_PER_DOC = 3 async def discovery_select_step( @@ -826,152 +128,22 @@ async def discovery_select_step( llm_fn: LLMFn, user_id: str, namespace: str, - doc_name: str = '', + doc_name: str = "", discovery_hints: list[dict[str, Any]], exclude_paths: set[str] | None = None, revision_hint: str | None = None, budget_snapshot: dict | None = None, ) -> DocTreeNode: - """Post-navigation discovery selection step. - - After BFS navigation exhausts for a document, present discovery-found - section paths (from bottom_discovery BM25) to the LLM for selection. - Selected paths are hydrated as leaf content. - - For B-class documents (discovery-only, not KG-selected), this is the - only navigation step — no prior BFS. - """ - from shared.services.retrieval.app_service import _hydrate_paths_to_rows - - node = DocTreeNode(scope_path=None) - if not discovery_hints: - return node - - # Limit hints per document - hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC] - - t0 = time.monotonic() - try: - # 1. Format hints for LLM (deduplicate by section_path) - exclude_set = { - normalize_section_path(path) - for path in (exclude_paths or set()) - if path - } - hint_lines: list[str] = [] - hint_by_path: dict[str, dict] = {} - for h in hints: - sp = normalize_section_path(h.get('section_path', '')) - if not sp or sp == 'Root': - continue - if sp in exclude_set: - continue - if sp in hint_by_path: - continue # skip duplicate section_path - summary = h.get('summary', '') or '' - hint_lines.append(f'▸ path="{sp}"') - if summary: - clipped = summary[:300] - hint_lines.append(f' {clipped}') - hint_by_path[sp] = h - - if not hint_lines: - return node - - items_text = '\n'.join(hint_lines) - - revision_context = '' - if revision_hint: - revision_context = ( - f'\nIMPORTANT: This is a REVISION round. ' - f'The previous search attempt failed because:\n' - f'"{revision_hint}"\n' - f'Adjust your selection accordingly. ' - f'If no candidate is relevant, return an EMPTY list [].\n' - ) - - prompt = _DISCOVERY_SELECT_PROMPT.format( - doc_name=doc_name or document_id, - budget_block=_format_budget_block(budget_snapshot), - items=items_text, - query=query, - revision_context=revision_context, - ) - response = await llm_fn(prompt) - # Parse {"selections": [...]} response — reuse action parser's extraction - parsed = _parse_action_response(response) - selections = parsed.get('selections', []) - - logger.info( - f' discovery_select_step doc="{doc_name}": ' - f'hints={len(hints)} selections={len(selections)}' - ) - - # 2. Hydrate selected paths - valid_selections = [s for s in selections if s['path'] in hint_by_path] - path_selections = [] - for sel in valid_selections: - path = sel['path'] - conf = sel.get('confidence', 0.7) - node.confidence[path] = conf - path_selections.append({'path': path, 'confidence': conf}) - - if path_selections: - chunks = await _hydrate_paths_to_rows( - db, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - if chunks: - connected = await hydrate_connected_target_rows( - db=db, - rows=chunks, - exclude_document_ids=[], - exclude_sections=[], - ) - if connected: - _owner_map = _build_connected_owner_map(chunks) - for c in connected: - if not c.get('owner_section_path'): - c['owner_section_path'] = _owner_map.get(str(c.get('chunk_id') or '')) - chunks = chunks + connected - - # Resolve Root-stranded assets to their true owner sections - _disc_job_result_id = next( - (str(c['job_result_id']) for c in chunks if c.get('job_result_id')), - None, - ) - _root_map = await _resolve_root_asset_owners( - db, - document_id=document_id, - job_result_id=_disc_job_result_id, - chunks=chunks, - ) if _disc_job_result_id else {} - if _root_map: - for c in chunks: - if c.get('owner_section_path'): - continue # already resolved by batch-level owner map - cid = str(c.get('chunk_id') or '') - if cid in _root_map: - c['owner_section_path'] = _root_map[cid] - - for chunk in chunks: - # Distribute chunk to its real path or fallback to the selection path - real_path = chunk.get('owner_section_path') or chunk.get('section_path') or chunk.get('source_chunk_path') - if real_path: - node.add_leaf_chunks(str(real_path), [chunk]) - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f' discovery_select_step done: hydrated={len(node.leaf_content)} ' - f'latency={latency}ms' - ) - return node - - except BudgetExceeded: - raise - except Exception as e: - logger.error(f' discovery_select_step failed for doc={document_id}: {e}') - return node + return await navigation_tools.discovery_select_step( + db, + document_id=document_id, + query=query, + llm_fn=llm_fn, + user_id=user_id, + namespace=namespace, + doc_name=doc_name, + discovery_hints=discovery_hints, + exclude_paths=exclude_paths, + revision_hint=revision_hint, + budget_snapshot=budget_snapshot, + ) diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 96c125e49..76b70c8c5 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -1,315 +1,40 @@ from __future__ import annotations -import asyncio import os -import re import time from typing import Any from loguru import logger -from sqlalchemy import and_, func, or_, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.database import get_db_context -from shared.models.database.document import Document, DocumentChunk, DocumentSection, RetrievalHitStat -from shared.services.retrieval.graph_service import GraphQueryService, is_excluded_section -from shared.services.retrieval.lexical_text import normalize_section_path from shared.services.retrieval.cache_service import get_cached_retrieval_query_result, set_cached_retrieval_query_result -from shared.services.retrieval.hit_stats_service import compute_importance_score, record_retrieval_hits from shared.services.retrieval.channels import path_channel, content_channel, term_channel -from shared.services.storage.result_storage import get_result_storage -from shared.models.database.job_result import JobResult - - -_MEDIA_CHUNK_TYPES = {'image', 'table'} - -_RRF_K = 60 -_CHANNEL_WEIGHT_PATH = 1.0 -_CHANNEL_WEIGHT_CONTENT = 2.0 -_CHANNEL_WEIGHT_TERM = 1.5 -_INTERNAL_RECALL_K_MULTIPLIER = 2 -_pending_retrieval_hit_stat_tasks: set[asyncio.Task[None]] = set() - -_DATA_TYPE_ALLOWED_CHUNK_TYPES: dict[int, set[str] | None] = { - 1: None, - 2: {'text'}, - 3: {'image'}, - 4: {'table'}, - 5: {'text', 'image'}, - 6: {'text', 'table'}, -} - - -def _resolve_allowed_chunk_types(data_type: int) -> set[str] | None: - return _DATA_TYPE_ALLOWED_CHUNK_TYPES.get(data_type) - - -_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]') - - -def _clean_content(content: str) -> str: - return _PATH_REF_RE.sub('', content).strip() - -_PUBLIC_RESULT_FIELDS = { - 'chunk_type', 'content', 'score', 'asset_url', -} - -_PUBLIC_SOURCE_FIELDS = { - 'document_id', 'source_file_name', 'section_path', -} - - -def _normalize_chunk_type(raw: str | None) -> str: - return str(raw or '').strip().split('\n', 1)[0].lower() - - -def _filter_excluded_rows( - rows: list[dict[str, Any]], - *, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - filtered: list[dict[str, Any]] = [] - excluded_documents = set(exclude_document_ids) - for row in rows: - document_id = row.get('document_id') - if document_id in excluded_documents: - continue - if is_excluded_section( - document_id=document_id, - section_path=row.get('section_path'), - exclude_sections=exclude_sections, - ): - continue - filtered.append(row) - return filtered - - -def _iter_connected_target_ids(row: dict[str, Any]) -> list[str]: - metadata = row.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - return [] - - target_ids: list[str] = [] - for item in metadata.get('connect_to') or []: - if not isinstance(item, dict): - continue - target_id = str(item.get('target') or '').strip() - if target_id: - target_ids.append(target_id) - return target_ids - - -async def hydrate_connected_target_rows( - *, - db: AsyncSession | None, - rows: list[dict[str, Any]], - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - if db is None: - return [] - - existing_chunk_ids = { - str(row.get('chunk_id') or '').strip() - for row in rows - if row.get('chunk_id') - } - target_ids_by_revision: dict[tuple[str, str], set[str]] = {} - for row in rows: - if _normalize_chunk_type(row.get('chunk_type')) != 'text': - continue - document_id = str(row.get('document_id') or '').strip() - job_result_id = str(row.get('job_result_id') or '').strip() - if not document_id or not job_result_id: - continue - for target_id in _iter_connected_target_ids(row): - if target_id in existing_chunk_ids: - continue - target_ids_by_revision.setdefault((document_id, job_result_id), set()).add(target_id) - - if not target_ids_by_revision: - return [] - - revision_filters = [ - and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == job_result_id, - DocumentChunk.chunk_id.in_(sorted(target_ids)), - ) - for (document_id, job_result_id), target_ids in target_ids_by_revision.items() - if target_ids - ] - if not revision_filters: - return [] - - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(or_(*revision_filters)) - .order_by(DocumentChunk.sort_order) - ) - result = await db.execute(stmt) - - hydrated_rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): - section_path = section.section_path if section else None - hydrated_rows.append( - { - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': 0.0, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'sort_order': chunk.sort_order, - } - ) - - return _filter_excluded_rows( - hydrated_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - - -async def assemble_retrieval_results( - *, - db: AsyncSession | None = None, - rows: list[dict[str, Any]], - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None = None, -) -> list[dict[str, Any]]: - filtered_rows = _filter_excluded_rows( - rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if allowed_chunk_types is not None: - filtered_rows = [ - row for row in filtered_rows - if _normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types - ] - hydrated_rows = await hydrate_connected_target_rows( - db=db, - rows=filtered_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - rows_by_chunk_id = { - str(row.get('chunk_id') or ''): row - for row in [*filtered_rows, *hydrated_rows] - if row.get('chunk_id') - } - - embedded_targets: set[str] = set() - for row in filtered_rows: - for target_id in _iter_connected_target_ids(row): - if target_id in rows_by_chunk_id: - embedded_targets.add(target_id) - - assembled: list[dict[str, Any]] = [] - for row in filtered_rows: - if row.get('chunk_id') in embedded_targets: - continue - metadata = row.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - metadata = {} - assembled_row = dict(row) - base_content = str(row.get('content') or '') - if _normalize_chunk_type(row.get('chunk_type')) == 'text': - connected_targets: list[tuple[int, str]] = [] - for target_id in _iter_connected_target_ids(row): - target_row = rows_by_chunk_id.get(target_id) - if not target_row: - continue - if _normalize_chunk_type(target_row.get('chunk_type')) != 'table': - continue - target_content = str(target_row.get('content') or '').strip() - if target_content: - sort_key = int(target_row.get('sort_order', 0) or 0) - connected_targets.append((sort_key, target_content)) - connected_targets.sort(key=lambda x: x[0]) - related_parts = [content for _, content in connected_targets] - if base_content and related_parts: - assembled_row['content'] = '\n\n'.join([base_content, *related_parts]) - else: - assembled_row['content'] = base_content - else: - assembled_row['content'] = base_content - assembled_row['content'] = _clean_content(assembled_row['content']) - assembled.append(assembled_row) - return assembled - - - - - -def _merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: - if not rows: - return rows - groups: dict[str, list[dict[str, Any]]] = {} - order: list[str] = [] - for row in rows: - sp = row.get('section_path') - if sp: - key = f"{row.get('document_id', '')}::{sp}" - else: - key = row.get('chunk_id', '') - if key not in groups: - groups[key] = [] - order.append(key) - groups[key].append(row) - - merged: list[dict[str, Any]] = [] - for key in order: - group = groups[key] - if len(group) == 1: - merged.append(group[0]) - continue - base = dict(group[0]) - base['content'] = '\n'.join(str(r.get('content', '')) for r in group) - base['score'] = max(r.get('score', 0.0) for r in group) - merged.append(base) - return merged - - -def merge_channels_rrf( - channels: list[list[dict[str, Any]]], - weights: list[float], - top_k: int, - k: int = _RRF_K, -) -> list[dict[str, Any]]: - """Reciprocal Rank Fusion across multiple retrieval channels.""" - score_dict: dict[str, float] = {} - row_by_chunk_id: dict[str, dict[str, Any]] = {} - - for channel_idx, channel_rows in enumerate(channels): - w = weights[channel_idx] if channel_idx < len(weights) else 1.0 - for rank, row in enumerate(channel_rows): - chunk_id = str(row.get('chunk_id') or '') - if not chunk_id: - continue - rrf_score = w / (k + rank + 1) - score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score - if chunk_id not in row_by_chunk_id: - row_by_chunk_id[chunk_id] = row - - ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) - results: list[dict[str, Any]] = [] - for chunk_id, fused_score in ranked[:top_k]: - row = row_by_chunk_id[chunk_id] - results.append(dict(row, score=round(fused_score, 6))) - return results +from shared.services.retrieval.graph_service import GraphQueryService +from shared.services.retrieval.hit_stats_recorder import schedule_retrieval_hit_stats_update +from shared.services.retrieval.hydration import ( + assemble_retrieval_results, + hydrate_referenced_chunk_rows, +) +from shared.services.retrieval.response_projection import ( + attach_citation, + enrich_referenced_chunks_with_asset_urls, + project_public_retrieval_response, +) +from shared.services.retrieval.scoring import ( + get_row_path, + merge_channels_rrf, + merge_same_section_rows, + normalize_row_scores, +) +from shared.services.retrieval.ranking import rank_retrieval_candidates +from shared.services.retrieval.scoped_corpus import count_scoped_chunks, load_all_scoped_chunks +from shared.services.retrieval.settings import ( + CHANNEL_WEIGHT_CONTENT as _CHANNEL_WEIGHT_CONTENT, + CHANNEL_WEIGHT_PATH as _CHANNEL_WEIGHT_PATH, + CHANNEL_WEIGHT_TERM as _CHANNEL_WEIGHT_TERM, + INTERNAL_RECALL_K_MULTIPLIER as _INTERNAL_RECALL_K_MULTIPLIER, + resolve_allowed_chunk_types as _resolve_allowed_chunk_types, +) async def list_graph_routed_chunks( @@ -342,627 +67,6 @@ async def list_graph_routed_chunks( ) -def _finalize_retrieval_hit_stats_task(task: asyncio.Task[None]) -> None: - _pending_retrieval_hit_stat_tasks.discard(task) - - try: - task.result() - except asyncio.CancelledError: - pass - except Exception as e: - logger.warning(f'Failed to record retrieval hit stats (ignored): {e}') - - -def schedule_retrieval_hit_stats_update(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None: - try: - task = asyncio.create_task( - _record_retrieval_hit_stats_best_effort( - user_id=user_id, - namespace=namespace, - results=results, - ), - name=f'retrieval_hit_stats:{user_id}:{namespace}', - ) - _pending_retrieval_hit_stat_tasks.add(task) - task.add_done_callback(_finalize_retrieval_hit_stats_task) - except Exception as e: - logger.warning(f'Failed to schedule retrieval hit stats update (ignored): {e}') - - -async def drain_retrieval_hit_stats_updates(timeout_seconds: float = 2.0) -> None: - if not _pending_retrieval_hit_stat_tasks: - return - - pending_tasks = tuple(_pending_retrieval_hit_stat_tasks) - - try: - await asyncio.wait_for( - asyncio.gather(*pending_tasks, return_exceptions=True), - timeout=timeout_seconds, - ) - except asyncio.TimeoutError: - for task in pending_tasks: - if not task.done(): - task.cancel() - - await asyncio.gather(*pending_tasks, return_exceptions=True) - - -async def _record_retrieval_hit_stats_best_effort(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None: - try: - async with get_db_context() as db: - await record_retrieval_hits(db, user_id=user_id, namespace=namespace, results=results) - await db.commit() - except Exception as e: - logger.warning(f'Failed to record retrieval hit stats (ignored): {e}') - - -def _with_citation(row: dict[str, Any]) -> dict[str, Any]: - citation = { - 'document_id': row.get('document_id'), - 'chunk_id': row.get('chunk_id'), - 'source_file_name': row.get('source_file_name'), - 'section_path': row.get('section_path'), - } - return {**row, 'citation': citation} - - -def _to_public_source(row: dict[str, Any]) -> dict[str, Any]: - return {field: row.get(field) for field in _PUBLIC_SOURCE_FIELDS} - - -def _is_media_chunk(row: dict[str, Any]) -> bool: - return _normalize_chunk_type(row.get('chunk_type')) in _MEDIA_CHUNK_TYPES - - -async def generate_retrieval_asset_url(*, job_id: str, artifact_ref: str) -> str | None: - return get_result_storage().generate_artifact_url(job_id=job_id, artifact_ref=artifact_ref) - - -def _is_client_result_artifact_ref(asset_ref: str | None) -> bool: - return get_result_storage().normalize_artifact_ref(asset_ref) is not None - - -async def _to_public_response(response: dict[str, Any]) -> dict[str, Any]: - public_response = { - 'namespace': response.get('namespace'), - 'query': response.get('query'), - 'router_used': response.get('router_used'), - 'results': [], - } - - # Forward agentic evidence fields when present - if response.get('answer_text') is not None: - public_response['answer_text'] = response['answer_text'] - if response.get('referenced_chunks') is not None: - public_response['referenced_chunks'] = response['referenced_chunks'] - - public_results: list[dict[str, Any]] = [] - for row in response.get('results', []): - artifact_ref = row.get('file_path') - asset_url = None - if _is_media_chunk(row) and _is_client_result_artifact_ref(artifact_ref) and row.get('job_id'): - try: - asset_url = await generate_retrieval_asset_url( - job_id=str(row['job_id']), - artifact_ref=str(artifact_ref), - ) - except Exception as e: - logger.warning(f'Failed to generate retrieval asset URL (ignored): {e}') - - public_row: dict[str, Any] = {} - for field in _PUBLIC_RESULT_FIELDS: - if field == 'asset_url': - if asset_url: - public_row['asset_url'] = asset_url - elif field in row: - public_row[field] = row[field] - if 'source' in row: - public_row['source'] = row['source'] - else: - public_row['source'] = _to_public_source(row) - public_results.append(public_row) - - public_response['results'] = public_results - return public_response - - -def _get_row_path(row: dict[str, Any]) -> str: - """Extract the canonical path from a row for deduplication.""" - return str(row.get('section_path') or row.get('source_chunk_path') or '') - - -def _get_candidate_key(row: dict[str, Any]) -> str: - path = _get_row_path(row) - if path: - return f'path:{path}' - chunk_id = str(row.get('chunk_id') or '').strip() - return f'chunk:{chunk_id}' if chunk_id else '' - - -def _normalize_row_scores( - rows: list[dict[str, Any]], - *, - source_field: str, - target_field: str, - default: float, -) -> None: - if not rows: - return - values = [float(row.get(source_field, 0.0) or 0.0) for row in rows] - min_score = min(values) - max_score = max(values) - if max_score <= 0.0 and min_score <= 0.0: - for row in rows: - row[target_field] = 0.0 - return - if max_score == min_score: - for row in rows: - row[target_field] = default - return - denominator = max_score - min_score - for row in rows: - raw_score = float(row.get(source_field, 0.0) or 0.0) - row[target_field] = round((raw_score - min_score) / denominator, 6) - - -def _importance_multiplier( - rows: list[dict[str, Any]], - *, - raw_field: str = 'importance_raw_score', - low: float = 0.1, - high: float = 2.0, -) -> None: - """Apply adaptive sigmoid-based importance boost to agent/discovery scores. - - Uses median of ``raw_field`` as center and IQR as spread so the curve - adapts to any KB size without hard-coded thresholds. When all values - are identical (IQR ≈ 0) the multiplier is 1.0 (neutral). - - Output range ``[low, high]`` — default [0.1, 2.0] — is the only - configured constant: max 2× boost, min 10%. The function modifies - ``agent_score`` and ``discovery_score`` **in place**. - """ - import math - - if not rows: - return - - values = sorted(float(r.get(raw_field, 0.0) or 0.0) for r in rows) - n = len(values) - median = values[n // 2] if n % 2 else (values[n // 2 - 1] + values[n // 2]) / 2 - q1 = values[n // 4] if n >= 4 else values[0] - q3 = values[3 * n // 4] if n >= 4 else values[-1] - iqr = q3 - q1 - - for row in rows: - raw = float(row.get(raw_field, 0.0) or 0.0) - if iqr <= 1e-9: - mult = 1.0 - else: - z = (raw - median) / iqr - s = 1.0 / (1.0 + math.exp(-z)) - mult = low + (high - low) * s - row['importance_multiplier'] = round(mult, 4) - row['agent_score'] = round( - float(row.get('agent_score', 0.0) or 0.0) * mult, 6, - ) - row['discovery_score'] = round( - float(row.get('discovery_score', 0.0) or 0.0) * mult, 6, - ) - - -async def _load_chunk_importance_scores( - db: AsyncSession, - *, - user_id: str, - namespace: str, - rows: list[dict[str, Any]], -) -> dict[str, float]: - chunk_ids = sorted({ - str(row.get('chunk_id') or '').strip() - for row in rows - if row.get('chunk_id') - }) - if not chunk_ids: - return {} - stmt = ( - select( - RetrievalHitStat.chunk_id, - RetrievalHitStat.hit_count, - RetrievalHitStat.last_hit_at, - RetrievalHitStat.created_at, - ) - .where(RetrievalHitStat.user_id == user_id) - .where(RetrievalHitStat.namespace == namespace) - .where(RetrievalHitStat.hit_kind == 'chunk') - .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) - ) - result = await db.execute(stmt) - importance_scores: dict[str, float] = {} - for chunk_id, hit_count, last_hit_at, created_at in result.all(): - if not chunk_id: - continue - importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) - return importance_scores - - -def _rank_candidates_by_path( - discovery_rows: list[dict[str, Any]], - routed_rows: list[dict[str, Any]], - top_k: int, -) -> list[dict[str, Any]]: - """Rank discovery and routed candidates in one comparable path space.""" - merged: dict[str, dict[str, Any]] = {} - insertion_order: dict[str, int] = {} - counter = 0 - - for row in discovery_rows: - key = _get_candidate_key(row) - if not key: - continue - candidate = dict(row) - candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) - candidate['agent_score'] = 0.0 - candidate.setdefault('hydrate_mode', 'chunks') - merged[key] = candidate - insertion_order[key] = counter - counter += 1 - - for row in routed_rows: - key = _get_candidate_key(row) - if not key: - continue - routed_agent_score = float(row.get('agent_score', 0.0) or 0.0) - if key not in merged: - candidate = dict(row) - candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) - candidate['agent_score'] = routed_agent_score - merged[key] = candidate - insertion_order[key] = counter - counter += 1 - continue - candidate = merged[key] - candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score) - if not candidate.get('source_chunk_path') and row.get('source_chunk_path'): - candidate['source_chunk_path'] = row.get('source_chunk_path') - if not candidate.get('section_path') and row.get('section_path'): - candidate['section_path'] = row.get('section_path') - - # ── Dual-priority ranking ──────────────────────────────────────────── - # When the agent produced results (routed_rows non-empty), rows with - # agent_score=0 are demoted to a fallback pool. Primary sort is by - # agent_score (includes importance boost from _importance_multiplier), - # with discovery_score as tiebreaker. - has_agent_results = len(routed_rows) > 0 - - primary_rows: list[dict[str, Any]] = [] - fallback_rows: list[dict[str, Any]] = [] - - for key, row in merged.items(): - agent_score = float(row.get('agent_score', 0.0) or 0.0) - discovery_score = float(row.get('discovery_score', 0.0) or 0.0) - row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) - row['score'] = row['evidence_score'] - row['_candidate_order'] = insertion_order[key] - - if has_agent_results and agent_score <= 0.0: - fallback_rows.append(row) - else: - primary_rows.append(row) - - def _sort_key(row): - return ( - float(row.get('agent_score', 0.0) or 0.0), - float(row.get('discovery_score', 0.0) or 0.0), - -int(row.get('_candidate_order', 0) or 0), - ) - - primary_rows.sort(key=_sort_key, reverse=True) - ranked_rows = primary_rows[:top_k] - - # Back-fill from fallback if primary results are insufficient - if len(ranked_rows) < top_k and fallback_rows: - fallback_rows.sort(key=_sort_key, reverse=True) - ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) - - for row in ranked_rows: - row.pop('_candidate_order', None) - return ranked_rows - - -async def _count_scoped_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], - allowed_chunk_types: set[str] | None, -) -> int: - stmt = ( - select(func.count(DocumentChunk.id)) - .join(Document, (Document.document_id == DocumentChunk.document_id) & (Document.current_job_result_id == DocumentChunk.job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - if allowed_chunk_types is not None: - stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) - result = await db.execute(stmt) - return result.scalar() or 0 - - -async def _load_all_scoped_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None, - signal_paths: list[str], - filter_mode: str, -) -> list[dict[str, Any]]: - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .order_by(DocumentChunk.sort_order) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - if allowed_chunk_types is not None: - stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) - - result = await db.execute(stmt) - rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): - section_path = section.section_path if section else None - if is_excluded_section(document_id=document.document_id, section_path=section_path, exclude_sections=exclude_sections): - continue - if signal_paths and section_path: - path_lower = section_path.lower() - matches_any = any(kw.lower() in path_lower for kw in signal_paths) - if filter_mode == 'keep' and not matches_any: - continue - if filter_mode == 'delete' and matches_any: - continue - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': 1.0, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'sort_order': chunk.sort_order, - }) - return rows - - -async def _hydrate_paths_to_rows( - db: AsyncSession, - *, - path_selections: list[dict[str, Any]], - user_id: str, - namespace: str, - document_id: str | None = None, -) -> list[dict[str, Any]]: - """Load full chunk rows by section_path or source_chunk_path. - - When *document_id* is provided the query is scoped to that single - document, preventing cross-document collisions on generic paths - such as ``Root``. - - Supports hydrate_mode branching: - - 'chunks' (default): all chunk types under the section subtree - - 'outline': synthetic row from section metadata, no real chunks - - 'assets_only': only image + table chunks - - 'image_only': only image chunks - - 'table_only': only table chunks - """ - if not path_selections: - return [] - - # Group selections by hydrate_mode - confidence_by_path: dict[str, float] = {} - mode_by_path: dict[str, str] = {} - ordered_paths: list[str] = [] - for item in path_selections: - raw_path = str(item.get('path') or '').strip() - path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path - if not path: - continue - confidence = float(item.get('confidence', 0.0) or 0.0) - hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower() - if path not in confidence_by_path: - ordered_paths.append(path) - confidence_by_path[path] = confidence - mode_by_path[path] = hydrate_mode - else: - confidence_by_path[path] = max(confidence_by_path[path], confidence) - if not ordered_paths: - return [] - - # Separate outline paths from chunk-loading paths - outline_paths = [p for p in ordered_paths if mode_by_path.get(p) == 'outline'] - chunk_paths = [p for p in ordered_paths if mode_by_path.get(p) != 'outline'] - - rows: list[dict[str, Any]] = [] - - # ── Outline mode: synthesize rows from section metadata ────────────── - if outline_paths: - outline_section_filters = [] - for path in outline_paths: - outline_section_filters.append(DocumentSection.section_path == path) - - outline_stmt = ( - select(Document, DocumentSection) - .join(DocumentSection, (DocumentSection.document_id == Document.document_id) - & (DocumentSection.job_result_id == Document.current_job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(or_(*outline_section_filters)) - ) - if document_id: - outline_stmt = outline_stmt.where(Document.document_id == document_id) - outline_result = await db.execute(outline_stmt) - for document, section in outline_result.all(): - agent_score = confidence_by_path.get(section.section_path, 0.0) - summary_text = (section.summary or '').strip() - title_text = (section.section_title or '').strip() - content = f'[Outline] {title_text}' - if summary_text: - content += f'\n{summary_text}' - rows.append({ - 'document_id': document.document_id, - 'chunk_id': f'outline_{section.section_id}', - 'section_id': section.section_id, - 'section_path': section.section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': 'outline', - 'content': content, - 'score': agent_score, - 'agent_score': agent_score, - 'file_path': None, - 'chunk_metadata': {}, - 'job_result_id': section.job_result_id, - 'job_id': None, - 'source_chunk_path': None, - 'sort_order': section.sort_order, - 'hydrate_mode': 'outline', - }) - - # ── Chunk modes: load real chunks with optional type filters ───────── - if chunk_paths: - section_path_filters = [] - # Separate self_only paths (exact match only, no descendant LIKE) - # from regular chunk paths (exact + descendant subtree match) - self_only_paths = {p for p in chunk_paths if mode_by_path.get(p) == 'self_only'} - for path in chunk_paths: - section_path_filters.append(DocumentSection.section_path == path) - if path not in self_only_paths: - section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) - - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where( - or_( - *section_path_filters, - DocumentChunk.source_chunk_path.in_(chunk_paths), - ) - ) - ) - if document_id: - stmt = stmt.where(Document.document_id == document_id) - result = await db.execute(stmt) - - # Build a map of path → allowed chunk_types based on hydrate_mode - _MODE_ALLOWED_TYPES: dict[str, set[str] | None] = { - 'chunks': None, # all types - 'self_only': None, # all types, but without descendant filtering - 'assets_only': {'image', 'table'}, - 'image_only': {'image'}, - 'table_only': {'table'}, - } - - seen_paths: set[str] = set() - for document, chunk, section, job_result in result.all(): - row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' - if row_path in seen_paths: - continue - - # Find which ordered path this row belongs to - matched_path = row_path - if section and section.section_path not in confidence_by_path: - matched_path = next( - ( - path for path in chunk_paths - if section.section_path == path or section.section_path.startswith(f'{path} / ') - ), - row_path, - ) - - # Check chunk_type filter based on hydrate_mode - path_mode = mode_by_path.get(matched_path, 'chunks') - allowed_types = _MODE_ALLOWED_TYPES.get(path_mode) - if allowed_types is not None: - chunk_type_lower = (chunk.chunk_type or '').strip().lower() - if chunk_type_lower not in allowed_types: - continue - - seen_paths.add(row_path) - agent_score = confidence_by_path.get(matched_path, 0.0) - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section.section_path if section else None, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': agent_score, - 'agent_score': agent_score, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'source_chunk_path': chunk.source_chunk_path, - 'sort_order': chunk.sort_order, - 'hydrate_mode': path_mode, - }) - - # ── Sort by agent-selected order ───────────────────────────────────── - path_order = {p: idx for idx, p in enumerate(ordered_paths)} - - def _row_sort_key(row: dict[str, Any]) -> int: - row_path = _get_row_path(row) - if row_path in path_order: - return path_order[row_path] - for path, idx in path_order.items(): - if row_path.startswith(f'{path} / '): - return idx - return 10**9 - - rows.sort(key=_row_sort_key) - hydrated_paths = {_get_row_path(r) for r in rows} - resolved_inputs = { - path for path in ordered_paths - if path in hydrated_paths or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths) - } - # Outline paths are always resolved (synthesized) - resolved_inputs |= set(outline_paths) - missed = len(ordered_paths) - len(resolved_inputs) - if missed > 0: - missing_paths = [p for p in ordered_paths if p not in resolved_inputs] - logger.warning( - f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); ' - f'missing[:5]={missing_paths[:5]}' - ) - else: - logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved') - return rows - - async def run_retrieval_query( *, db: AsyncSession, @@ -1039,7 +143,7 @@ async def run_retrieval_query( ) except Exception as e: logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - return await _to_public_response(cached) + return await project_public_retrieval_response(cached) except Exception as e: logger.warning(f"Failed to read retrieval cache (ignored): {e}") @@ -1047,7 +151,7 @@ async def run_retrieval_query( # ── Small KB optimization ── try: - total_chunk_count = await _count_scoped_chunks( + total_chunk_count = await count_scoped_chunks( db, user_id=user_id, namespace=namespace, exclude_document_ids=exclude_document_ids, allowed_chunk_types=allowed_chunk_types, @@ -1058,7 +162,7 @@ async def run_retrieval_query( logger.info(f'\n 📊 Total chunks in scope: {total_chunk_count}') if total_chunk_count <= top_k: logger.info(f' ⚡ Small KB optimization: {total_chunk_count} chunks <= top_k={top_k}, returning all') - all_rows = await _load_all_scoped_chunks( + all_rows = await load_all_scoped_chunks( db, user_id=user_id, namespace=namespace, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, @@ -1073,7 +177,7 @@ async def run_retrieval_query( exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, ) - results = [_with_citation(row) for row in assembled_rows] + results = [attach_citation(row) for row in assembled_rows] response = { "namespace": namespace, "query": query, "router_used": "small_kb_all", "results": results, @@ -1095,7 +199,7 @@ async def run_retrieval_query( logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") elapsed_total = round((time.monotonic() - t_start) * 1000) logger.info(f' ✅ Small KB: {len(results)} results in {elapsed_total}ms') - return await _to_public_response(response) + return await project_public_retrieval_response(response) # ══ Route: agentic (unified workflow) vs legacy ══ if use_agentic is not None: @@ -1125,27 +229,41 @@ async def run_retrieval_query( channel_weights=channel_weights, ) - # Enrich referenced_chunks with asset URLs (images/tables) - enriched_refs: list[dict[str, Any]] = [] - for ref in workflow_result.referenced_chunks: - enriched = dict(ref) - chunk_type = _normalize_chunk_type(ref.get('chunk_type')) - artifact_ref = ref.get('file_path', '') - job_id = ref.get('job_id', '') - if chunk_type in _MEDIA_CHUNK_TYPES and job_id and _is_client_result_artifact_ref(artifact_ref): - try: - asset_url = await generate_retrieval_asset_url( - job_id=str(job_id), artifact_ref=str(artifact_ref), - ) - if asset_url: - enriched['asset_url'] = asset_url - except Exception as e: - logger.warning(f'Failed to generate agentic asset URL (ignored): {e}') - enriched_refs.append(enriched) + enriched_refs = await enrich_referenced_chunks_with_asset_urls( + workflow_result.referenced_chunks, + ) + workflow_result_rows = await hydrate_referenced_chunk_rows( + db=db, + user_id=user_id, + namespace=namespace, + refs=enriched_refs, + ) + scoped_reference_keys = { + ( + str(row.get('document_id') or '').strip(), + str(row.get('chunk_id') or '').strip(), + ) + for row in workflow_result_rows + } + enriched_refs = [ + ref for ref in enriched_refs + if ( + str(ref.get('document_id') or '').strip(), + str(ref.get('chunk_id') or '').strip(), + ) in scoped_reference_keys + ] + assembled_workflow_rows = await assemble_retrieval_results( + db=db, + rows=workflow_result_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + ) response = workflow_result.to_api_response() # Override referenced_chunks with enriched versions response['referenced_chunks'] = enriched_refs + response['results'] = [attach_citation(row) for row in assembled_workflow_rows] if cache_version is not None: try: @@ -1177,7 +295,7 @@ async def run_retrieval_query( f'{"█" * 70}' ) - return await _to_public_response(response) + return await project_public_retrieval_response(response) else: @@ -1258,7 +376,7 @@ async def run_retrieval_query( channel_lists.append(term_rows) weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM)) - fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else [] + fused_rows = merge_channels_rrf(channel_lists, weight_list, effective_recall_k) if channel_lists else [] logger.info(f'\n 🔀 RRF Fusion: {len(fused_rows)} rows from {len(channel_lists)} channels (weights={dict(zip(["path","content","term"][:len(weight_list)], weight_list))})') for i, r in enumerate(fused_rows[:5]): logger.info(f' [{i}] rrf_score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")}') @@ -1267,7 +385,7 @@ async def run_retrieval_query( # ── Section merging ── pre_merge = len(fused_rows) - fused_rows = _merge_same_section_rows(fused_rows) + fused_rows = merge_same_section_rows(fused_rows) if len(fused_rows) != pre_merge: logger.info(f'retrieval: section_merge={pre_merge}->{len(fused_rows)}') @@ -1278,7 +396,7 @@ async def run_retrieval_query( logger.info(f'retrieval: threshold_filter={pre_count}->{len(fused_rows)} (threshold={threshold})') if fused_rows: - _normalize_row_scores( + normalize_row_scores( fused_rows, source_field='score', target_field='discovery_score', @@ -1304,29 +422,21 @@ async def run_retrieval_query( agent_rows = [] if agent_rows: - _normalize_row_scores( + normalize_row_scores( agent_rows, source_field='score', target_field='agent_score', default=0.5, ) - combined_rows = [*fused_rows, *agent_rows] - if combined_rows: - try: - chunk_importance_scores = await _load_chunk_importance_scores( - db, - user_id=user_id, - namespace=namespace, - rows=combined_rows, - ) - except Exception as exc: - logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}') - chunk_importance_scores = {} - for row in combined_rows: - row['importance_raw_score'] = float(chunk_importance_scores.get(str(row.get('chunk_id') or ''), 0.0) or 0.0) - - ranked_rows = _rank_candidates_by_path(fused_rows, agent_rows, top_k) + ranked_rows = await rank_retrieval_candidates( + db, + user_id=user_id, + namespace=namespace, + discovery_rows=fused_rows, + routed_rows=agent_rows, + top_k=top_k, + ) if ranked_rows: logger.info(f'\n 🧮 Unified candidate ranking: {len(ranked_rows)} rows') for i, row in enumerate(ranked_rows[:10]): @@ -1335,7 +445,7 @@ async def run_retrieval_query( f'[{i}] evidence={row.get("evidence_score", 0.0):.4f} ' f'discovery={row.get("discovery_score", 0.0):.4f} ' f'agent={row.get("agent_score", 0.0):.4f} ' - f'path={_get_row_path(row)}' + f'path={get_row_path(row)}' ) assembled_rows = await assemble_retrieval_results( @@ -1345,7 +455,7 @@ async def run_retrieval_query( exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, ) - results = [_with_citation(row) for row in assembled_rows] + results = [attach_citation(row) for row in assembled_rows] response = { "namespace": namespace, @@ -1393,4 +503,4 @@ async def run_retrieval_query( logger.info(f' ... and {len(results) - 10} more') logger.info(f'{"█" * 70}') - return await _to_public_response(response) + return await project_public_retrieval_response(response) diff --git a/packages/shared-python/shared/services/retrieval/assets.py b/packages/shared-python/shared/services/retrieval/assets.py new file mode 100644 index 000000000..39badabea --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/assets.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from shared.services.storage.result_storage import get_result_storage + + +async def generate_retrieval_asset_url(*, job_id: str, artifact_ref: str) -> str | None: + return get_result_storage().generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + ) + + +def is_client_result_artifact_ref(asset_ref: str | None) -> bool: + return get_result_storage().normalize_artifact_ref(asset_ref) is not None diff --git a/packages/shared-python/shared/services/retrieval/hit_stats_recorder.py b/packages/shared-python/shared/services/retrieval/hit_stats_recorder.py new file mode 100644 index 000000000..b101292fa --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hit_stats_recorder.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +from loguru import logger + +from shared.services.retrieval.hit_stats_service import record_retrieval_hits + + +_pending_retrieval_hit_stat_tasks: set[asyncio.Task[None]] = set() + + +def _finalize_retrieval_hit_stats_task(task: asyncio.Task[None]) -> None: + _pending_retrieval_hit_stat_tasks.discard(task) + + try: + task.result() + except asyncio.CancelledError: + pass + except Exception as exc: + logger.warning(f'Failed to record retrieval hit stats (ignored): {exc}') + + +def schedule_retrieval_hit_stats_update(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None: + try: + task = asyncio.create_task( + _record_retrieval_hit_stats_best_effort( + user_id=user_id, + namespace=namespace, + results=results, + ), + name=f'retrieval_hit_stats:{user_id}:{namespace}', + ) + _pending_retrieval_hit_stat_tasks.add(task) + task.add_done_callback(_finalize_retrieval_hit_stats_task) + except Exception as exc: + logger.warning(f'Failed to schedule retrieval hit stats update (ignored): {exc}') + + +async def drain_retrieval_hit_stats_updates(timeout_seconds: float = 2.0) -> None: + if not _pending_retrieval_hit_stat_tasks: + return + + pending_tasks = tuple(_pending_retrieval_hit_stat_tasks) + + try: + await asyncio.wait_for( + asyncio.gather(*pending_tasks, return_exceptions=True), + timeout=timeout_seconds, + ) + except asyncio.TimeoutError: + for task in pending_tasks: + if not task.done(): + task.cancel() + + await asyncio.gather(*pending_tasks, return_exceptions=True) + + +async def _record_retrieval_hit_stats_best_effort(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None: + try: + from shared.core.database import get_db_context + + async with get_db_context() as db: + await record_retrieval_hits(db, user_id=user_id, namespace=namespace, results=results) + await db.commit() + except Exception as exc: + logger.warning(f'Failed to record retrieval hit stats (ignored): {exc}') diff --git a/packages/shared-python/shared/services/retrieval/hydration.py b/packages/shared-python/shared/services/retrieval/hydration.py new file mode 100644 index 000000000..2b6821f96 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration.py @@ -0,0 +1,549 @@ +from __future__ import annotations + +import re +from typing import Any + +from loguru import logger +from sqlalchemy import and_, or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.graph_service import is_excluded_section +from shared.services.retrieval.lexical_text import normalize_section_path +from shared.services.retrieval.scoring import get_row_path + +MEDIA_CHUNK_TYPES = {'image', 'table'} +PUBLIC_RESULT_FIELDS = { + 'chunk_type', 'content', 'score', 'asset_url', +} +PUBLIC_SOURCE_FIELDS = { + 'document_id', 'source_file_name', 'section_path', +} + +ReferenceLookupKey = tuple[str, str, str, str] + +_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]') + + +def clean_content(content: str) -> str: + return _PATH_REF_RE.sub('', content).strip() + + +def normalize_chunk_type(raw: object) -> str: + return str(raw or '').strip().split('\n', 1)[0].lower() + + +def is_media_chunk(row: dict[str, Any]) -> bool: + return normalize_chunk_type(row.get('chunk_type')) in MEDIA_CHUNK_TYPES + + +def build_reference_lookup_key( + *, + document_id: object, + chunk_id: object, + section_path: object = '', + file_path: object = '', +) -> ReferenceLookupKey: + return ( + str(document_id or '').strip(), + str(chunk_id or '').strip(), + str(section_path or '').strip(), + str(file_path or '').strip(), + ) + + +def filter_excluded_rows( + rows: list[dict[str, Any]], + *, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + filtered: list[dict[str, Any]] = [] + excluded_documents = set(exclude_document_ids) + for row in rows: + document_id = row.get('document_id') + if document_id in excluded_documents: + continue + if is_excluded_section( + document_id=document_id, + section_path=row.get('section_path'), + exclude_sections=exclude_sections, + ): + continue + filtered.append(row) + return filtered + + +def iter_connected_target_ids(row: dict[str, Any]) -> list[str]: + metadata = row.get('chunk_metadata') or {} + if not isinstance(metadata, dict): + return [] + + target_ids: list[str] = [] + for item in metadata.get('connect_to') or []: + if not isinstance(item, dict): + continue + target_id = str(item.get('target') or '').strip() + if target_id: + target_ids.append(target_id) + return target_ids + + +async def hydrate_connected_target_rows( + *, + db: AsyncSession | None, + rows: list[dict[str, Any]], + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + if db is None: + return [] + + existing_chunk_ids = { + str(row.get('chunk_id') or '').strip() + for row in rows + if row.get('chunk_id') + } + target_ids_by_revision: dict[tuple[str, str], set[str]] = {} + for row in rows: + if normalize_chunk_type(row.get('chunk_type')) != 'text': + continue + document_id = str(row.get('document_id') or '').strip() + job_result_id = str(row.get('job_result_id') or '').strip() + if not document_id or not job_result_id: + continue + for target_id in iter_connected_target_ids(row): + if target_id in existing_chunk_ids: + continue + target_ids_by_revision.setdefault((document_id, job_result_id), set()).add(target_id) + + if not target_ids_by_revision: + return [] + + revision_filters = [ + and_( + DocumentChunk.document_id == document_id, + DocumentChunk.job_result_id == job_result_id, + DocumentChunk.chunk_id.in_(sorted(target_ids)), + ) + for (document_id, job_result_id), target_ids in target_ids_by_revision.items() + if target_ids + ] + if not revision_filters: + return [] + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(or_(*revision_filters)) + .order_by(DocumentChunk.sort_order) + ) + result = await db.execute(stmt) + + hydrated_rows: list[dict[str, Any]] = [] + for document, chunk, section, job_result in result.all(): + section_path = section.section_path if section else None + hydrated_rows.append( + { + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 0.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'sort_order': chunk.sort_order, + } + ) + + return filter_excluded_rows( + hydrated_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + + +async def hydrate_referenced_chunk_rows( + *, + db: AsyncSession | None, + user_id: str, + namespace: str, + refs: list[dict[str, Any]], +) -> list[dict[str, Any]]: + if db is None or not refs: + return [] + + ref_keys = [ + build_reference_lookup_key( + document_id=ref.get('document_id'), + chunk_id=ref.get('chunk_id'), + section_path=ref.get('section_path'), + file_path=ref.get('file_path'), + ) + for ref in refs + ] + ref_keys = [key for key in ref_keys if key[0] and key[1]] + if not ref_keys: + return [] + + document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) + chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(Document.document_id.in_(document_ids)) + .where(DocumentChunk.chunk_id.in_(chunk_ids)) + .order_by(DocumentChunk.sort_order) + ) + result = await db.execute(stmt) + + rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} + rows_by_base_key: dict[tuple[str, str], list[dict[str, Any]]] = {} + for document, chunk, section, job_result in result.all(): + row = { + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 1.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, + } + key = build_reference_lookup_key( + document_id=row['document_id'], + chunk_id=row['chunk_id'], + section_path=row['section_path'], + file_path=row['file_path'], + ) + rows_by_key[key] = row + rows_by_base_key.setdefault((key[0], key[1]), []).append(row) + + rows: list[dict[str, Any]] = [] + seen_keys: set[ReferenceLookupKey] = set() + for key in ref_keys: + row = rows_by_key.get(key) + if row is None: + candidates = rows_by_base_key.get((key[0], key[1]), []) + row = next( + ( + candidate for candidate in candidates + if key[2] and str(candidate.get('section_path') or '').strip() == key[2] + ), + None, + ) + if row is None: + row = next( + ( + candidate for candidate in candidates + if build_reference_lookup_key( + document_id=candidate.get('document_id'), + chunk_id=candidate.get('chunk_id'), + section_path=candidate.get('section_path'), + file_path=candidate.get('file_path'), + ) + not in seen_keys + ), + None, + ) + if row is not None: + row_key = build_reference_lookup_key( + document_id=row.get('document_id'), + chunk_id=row.get('chunk_id'), + section_path=row.get('section_path'), + file_path=row.get('file_path'), + ) + if row_key in seen_keys: + continue + seen_keys.add(row_key) + rows.append(row) + return rows + + +async def assemble_retrieval_results( + *, + db: AsyncSession | None = None, + rows: list[dict[str, Any]], + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + allowed_chunk_types: set[str] | None = None, +) -> list[dict[str, Any]]: + filtered_rows = filter_excluded_rows( + rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + if allowed_chunk_types is not None: + filtered_rows = [ + row for row in filtered_rows + if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types + ] + hydrated_rows = await hydrate_connected_target_rows( + db=db, + rows=filtered_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + rows_by_chunk_id = { + str(row.get('chunk_id') or ''): row + for row in [*filtered_rows, *hydrated_rows] + if row.get('chunk_id') + } + + embedded_targets: set[str] = set() + for row in filtered_rows: + for target_id in iter_connected_target_ids(row): + if target_id in rows_by_chunk_id: + embedded_targets.add(target_id) + + assembled: list[dict[str, Any]] = [] + for row in filtered_rows: + if row.get('chunk_id') in embedded_targets: + continue + metadata = row.get('chunk_metadata') or {} + if not isinstance(metadata, dict): + metadata = {} + assembled_row = dict(row) + base_content = str(row.get('content') or '') + if normalize_chunk_type(row.get('chunk_type')) == 'text': + connected_targets: list[tuple[int, str]] = [] + for target_id in iter_connected_target_ids(row): + target_row = rows_by_chunk_id.get(target_id) + if not target_row: + continue + if normalize_chunk_type(target_row.get('chunk_type')) != 'table': + continue + target_content = str(target_row.get('content') or '').strip() + if target_content: + sort_key = int(target_row.get('sort_order', 0) or 0) + connected_targets.append((sort_key, target_content)) + connected_targets.sort(key=lambda x: x[0]) + related_parts = [content for _, content in connected_targets] + if base_content and related_parts: + assembled_row['content'] = '\n\n'.join([base_content, *related_parts]) + else: + assembled_row['content'] = base_content + else: + assembled_row['content'] = base_content + assembled_row['content'] = clean_content(assembled_row['content']) + assembled.append(assembled_row) + return assembled + + +async def hydrate_paths_to_rows( + db: AsyncSession, + *, + path_selections: list[dict[str, Any]], + user_id: str, + namespace: str, + document_id: str | None = None, +) -> list[dict[str, Any]]: + """Load full chunk rows by section_path or source_chunk_path.""" + if not path_selections: + return [] + + confidence_by_path: dict[str, float] = {} + mode_by_path: dict[str, str] = {} + ordered_paths: list[str] = [] + for item in path_selections: + raw_path = str(item.get('path') or '').strip() + path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path + if not path: + continue + confidence = float(item.get('confidence', 0.0) or 0.0) + hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower() + if path not in confidence_by_path: + ordered_paths.append(path) + confidence_by_path[path] = confidence + mode_by_path[path] = hydrate_mode + else: + confidence_by_path[path] = max(confidence_by_path[path], confidence) + if not ordered_paths: + return [] + + outline_paths = [p for p in ordered_paths if mode_by_path.get(p) == 'outline'] + chunk_paths = [p for p in ordered_paths if mode_by_path.get(p) != 'outline'] + + rows: list[dict[str, Any]] = [] + + if outline_paths: + outline_section_filters = [ + DocumentSection.section_path == path + for path in outline_paths + ] + outline_stmt = ( + select(Document, DocumentSection) + .join( + DocumentSection, + (DocumentSection.document_id == Document.document_id) + & (DocumentSection.job_result_id == Document.current_job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(or_(*outline_section_filters)) + ) + if document_id: + outline_stmt = outline_stmt.where(Document.document_id == document_id) + outline_result = await db.execute(outline_stmt) + for document, section in outline_result.all(): + agent_score = confidence_by_path.get(section.section_path, 0.0) + summary_text = (section.summary or '').strip() + title_text = (section.section_title or '').strip() + content = f'[Outline] {title_text}' + if summary_text: + content += f'\n{summary_text}' + rows.append({ + 'document_id': document.document_id, + 'chunk_id': f'outline_{section.section_id}', + 'section_id': section.section_id, + 'section_path': section.section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': 'outline', + 'content': content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': None, + 'chunk_metadata': {}, + 'job_result_id': section.job_result_id, + 'job_id': None, + 'source_chunk_path': None, + 'sort_order': section.sort_order, + 'hydrate_mode': 'outline', + }) + + if chunk_paths: + section_path_filters = [] + self_only_paths = {p for p in chunk_paths if mode_by_path.get(p) == 'self_only'} + for path in chunk_paths: + section_path_filters.append(DocumentSection.section_path == path) + if path not in self_only_paths: + section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where( + or_( + *section_path_filters, + DocumentChunk.source_chunk_path.in_(chunk_paths), + ) + ) + ) + if document_id: + stmt = stmt.where(Document.document_id == document_id) + result = await db.execute(stmt) + + mode_allowed_types: dict[str, set[str] | None] = { + 'chunks': None, + 'self_only': None, + 'assets_only': {'image', 'table'}, + 'image_only': {'image'}, + 'table_only': {'table'}, + } + + seen_paths: set[str] = set() + for document, chunk, section, job_result in result.all(): + row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' + if row_path in seen_paths: + continue + + matched_path = row_path + if section and section.section_path not in confidence_by_path: + matched_path = next( + ( + path for path in chunk_paths + if section.section_path == path or section.section_path.startswith(f'{path} / ') + ), + row_path, + ) + + path_mode = mode_by_path.get(matched_path, 'chunks') + allowed_types = mode_allowed_types.get(path_mode) + if allowed_types is not None: + chunk_type_lower = (chunk.chunk_type or '').strip().lower() + if chunk_type_lower not in allowed_types: + continue + + seen_paths.add(row_path) + agent_score = confidence_by_path.get(matched_path, 0.0) + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, + 'hydrate_mode': path_mode, + }) + + path_order = {path: index for index, path in enumerate(ordered_paths)} + + def _row_sort_key(row: dict[str, Any]) -> int: + row_path = get_row_path(row) + if row_path in path_order: + return path_order[row_path] + for path, index in path_order.items(): + if row_path.startswith(f'{path} / '): + return index + return 10**9 + + rows.sort(key=_row_sort_key) + hydrated_paths = {get_row_path(row) for row in rows} + resolved_inputs = { + path for path in ordered_paths + if path in hydrated_paths + or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths) + } + resolved_inputs |= set(outline_paths) + missed = len(ordered_paths) - len(resolved_inputs) + if missed > 0: + missing_paths = [path for path in ordered_paths if path not in resolved_inputs] + logger.warning( + f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); ' + f'missing[:5]={missing_paths[:5]}' + ) + else: + logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved') + return rows diff --git a/packages/shared-python/shared/services/retrieval/ranking.py b/packages/shared-python/shared/services/retrieval/ranking.py new file mode 100644 index 000000000..bd1cfd9fd --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/ranking.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import math +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import RetrievalHitStat +from shared.services.retrieval.hit_stats_service import compute_importance_score +from shared.services.retrieval.scoring import get_row_path + + +def get_candidate_key(row: dict[str, Any]) -> str: + path = get_row_path(row) + if path: + return f'path:{path}' + chunk_id = str(row.get('chunk_id') or '').strip() + return f'chunk:{chunk_id}' if chunk_id else '' + + +async def load_chunk_importance_scores( + db: AsyncSession, + *, + user_id: str, + namespace: str, + rows: list[dict[str, Any]], +) -> dict[str, float]: + chunk_ids = sorted({ + str(row.get('chunk_id') or '').strip() + for row in rows + if row.get('chunk_id') + }) + if not chunk_ids: + return {} + stmt = ( + select( + RetrievalHitStat.chunk_id, + RetrievalHitStat.hit_count, + RetrievalHitStat.last_hit_at, + RetrievalHitStat.created_at, + ) + .where(RetrievalHitStat.user_id == user_id) + .where(RetrievalHitStat.namespace == namespace) + .where(RetrievalHitStat.hit_kind == 'chunk') + .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) + ) + result = await db.execute(stmt) + importance_scores: dict[str, float] = {} + for chunk_id, hit_count, last_hit_at, created_at in result.all(): + if not chunk_id: + continue + importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) + return importance_scores + + +def apply_importance_multiplier( + rows: list[dict[str, Any]], + *, + raw_field: str = 'importance_raw_score', + low: float = 0.1, + high: float = 2.0, +) -> None: + if not rows: + return + + values = sorted(float(row.get(raw_field, 0.0) or 0.0) for row in rows) + item_count = len(values) + median = values[item_count // 2] if item_count % 2 else (values[item_count // 2 - 1] + values[item_count // 2]) / 2 + q1 = values[item_count // 4] if item_count >= 4 else values[0] + q3 = values[3 * item_count // 4] if item_count >= 4 else values[-1] + iqr = q3 - q1 + + for row in rows: + raw_score = float(row.get(raw_field, 0.0) or 0.0) + if iqr <= 1e-9: + multiplier = 1.0 + else: + z_score = (raw_score - median) / iqr + sigmoid_score = 1.0 / (1.0 + math.exp(-z_score)) + multiplier = low + (high - low) * sigmoid_score + row['importance_multiplier'] = round(multiplier, 4) + row['agent_score'] = round( + float(row.get('agent_score', 0.0) or 0.0) * multiplier, + 6, + ) + row['discovery_score'] = round( + float(row.get('discovery_score', 0.0) or 0.0) * multiplier, + 6, + ) + + +def rank_candidates_by_path( + discovery_rows: list[dict[str, Any]], + routed_rows: list[dict[str, Any]], + top_k: int, + *, + importance_scores: dict[str, float] | None = None, +) -> list[dict[str, Any]]: + merged: dict[str, dict[str, Any]] = {} + insertion_order: dict[str, int] = {} + counter = 0 + + for row in discovery_rows: + key = get_candidate_key(row) + if not key: + continue + candidate = dict(row) + candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) + candidate['agent_score'] = 0.0 + candidate.setdefault('hydrate_mode', 'chunks') + merged[key] = candidate + insertion_order[key] = counter + counter += 1 + + for row in routed_rows: + key = get_candidate_key(row) + if not key: + continue + routed_agent_score = float(row.get('agent_score', 0.0) or 0.0) + if key not in merged: + candidate = dict(row) + candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0) + candidate['agent_score'] = routed_agent_score + merged[key] = candidate + insertion_order[key] = counter + counter += 1 + continue + candidate = merged[key] + candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score) + if not candidate.get('source_chunk_path') and row.get('source_chunk_path'): + candidate['source_chunk_path'] = row.get('source_chunk_path') + if not candidate.get('section_path') and row.get('section_path'): + candidate['section_path'] = row.get('section_path') + + for row in merged.values(): + row['importance_raw_score'] = float( + (importance_scores or {}).get(str(row.get('chunk_id') or ''), 0.0) or 0.0 + ) + apply_importance_multiplier(list(merged.values())) + + has_agent_results = len(routed_rows) > 0 + primary_rows: list[dict[str, Any]] = [] + fallback_rows: list[dict[str, Any]] = [] + + for key, row in merged.items(): + agent_score = float(row.get('agent_score', 0.0) or 0.0) + discovery_score = float(row.get('discovery_score', 0.0) or 0.0) + row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) + row['score'] = row['evidence_score'] + row['_candidate_order'] = insertion_order[key] + + if has_agent_results and agent_score <= 0.0: + fallback_rows.append(row) + else: + primary_rows.append(row) + + def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]: + return ( + float(row.get('agent_score', 0.0) or 0.0), + float(row.get('discovery_score', 0.0) or 0.0), + -int(row.get('_candidate_order', 0) or 0), + ) + + primary_rows.sort(key=get_sort_key, reverse=True) + ranked_rows = primary_rows[:top_k] + + if len(ranked_rows) < top_k and fallback_rows: + fallback_rows.sort(key=get_sort_key, reverse=True) + ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) + + for row in ranked_rows: + row.pop('_candidate_order', None) + return ranked_rows + + +async def rank_retrieval_candidates( + db: AsyncSession, + *, + user_id: str, + namespace: str, + discovery_rows: list[dict[str, Any]], + routed_rows: list[dict[str, Any]], + top_k: int, +) -> list[dict[str, Any]]: + try: + importance_scores = await load_chunk_importance_scores( + db, + user_id=user_id, + namespace=namespace, + rows=[*discovery_rows, *routed_rows], + ) + except Exception as exc: + logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}') + importance_scores = {} + return rank_candidates_by_path( + discovery_rows, + routed_rows, + top_k, + importance_scores=importance_scores, + ) diff --git a/packages/shared-python/shared/services/retrieval/response_projection.py b/packages/shared-python/shared/services/retrieval/response_projection.py new file mode 100644 index 000000000..36e01dfbc --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/response_projection.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from shared.services.retrieval.assets import generate_retrieval_asset_url, is_client_result_artifact_ref +from shared.services.retrieval.hydration import ( + MEDIA_CHUNK_TYPES, + PUBLIC_RESULT_FIELDS, + PUBLIC_SOURCE_FIELDS, + is_media_chunk, + normalize_chunk_type, +) + + +def attach_citation(row: dict[str, Any]) -> dict[str, Any]: + citation = { + 'document_id': row.get('document_id'), + 'chunk_id': row.get('chunk_id'), + 'source_file_name': row.get('source_file_name'), + 'section_path': row.get('section_path'), + } + return {**row, 'citation': citation} + + +def to_public_source(row: dict[str, Any]) -> dict[str, Any]: + return {field: row.get(field) for field in PUBLIC_SOURCE_FIELDS} + + +async def enrich_referenced_chunks_with_asset_urls(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + enriched_refs: list[dict[str, Any]] = [] + for ref in refs: + enriched = dict(ref) + chunk_type = normalize_chunk_type(ref.get('chunk_type')) + artifact_ref = ref.get('file_path', '') + job_id = ref.get('job_id', '') + if chunk_type in MEDIA_CHUNK_TYPES and job_id and is_client_result_artifact_ref(artifact_ref): + try: + asset_url = await generate_retrieval_asset_url( + job_id=str(job_id), + artifact_ref=str(artifact_ref), + ) + if asset_url: + enriched['asset_url'] = asset_url + except Exception as exc: + logger.warning(f'Failed to generate agentic asset URL (ignored): {exc}') + enriched_refs.append(enriched) + return enriched_refs + + +async def project_public_retrieval_response(response: dict[str, Any]) -> dict[str, Any]: + public_response = { + 'namespace': response.get('namespace'), + 'query': response.get('query'), + 'router_used': response.get('router_used'), + 'results': [], + } + + if response.get('answer_text') is not None: + public_response['answer_text'] = response['answer_text'] + if response.get('referenced_chunks') is not None: + public_response['referenced_chunks'] = response['referenced_chunks'] + + public_results: list[dict[str, Any]] = [] + for row in response.get('results', []): + artifact_ref = row.get('file_path') + asset_url = None + if is_media_chunk(row) and is_client_result_artifact_ref(artifact_ref) and row.get('job_id'): + try: + asset_url = await generate_retrieval_asset_url( + job_id=str(row['job_id']), + artifact_ref=str(artifact_ref), + ) + except Exception as exc: + logger.warning(f'Failed to generate retrieval asset URL (ignored): {exc}') + + public_row: dict[str, Any] = {} + for field in PUBLIC_RESULT_FIELDS: + if field == 'asset_url': + if asset_url: + public_row['asset_url'] = asset_url + elif field in row: + public_row[field] = row[field] + if 'source' in row: + public_row['source'] = row['source'] + else: + public_row['source'] = to_public_source(row) + public_results.append(public_row) + + public_response['results'] = public_results + return public_response diff --git a/packages/shared-python/shared/services/retrieval/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/scoped_corpus.py new file mode 100644 index 000000000..604400253 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/scoped_corpus.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.graph_service import is_excluded_section + + +async def count_scoped_chunks( + db: AsyncSession, + *, + user_id: str, + namespace: str, + exclude_document_ids: list[str], + allowed_chunk_types: set[str] | None, +) -> int: + stmt = ( + select(func.count(DocumentChunk.id)) + .join( + Document, + (Document.document_id == DocumentChunk.document_id) + & (Document.current_job_result_id == DocumentChunk.job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + if allowed_chunk_types is not None: + stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) + result = await db.execute(stmt) + return result.scalar() or 0 + + +async def load_all_scoped_chunks( + db: AsyncSession, + *, + user_id: str, + namespace: str, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + allowed_chunk_types: set[str] | None, + signal_paths: list[str], + filter_mode: str, +) -> list[dict[str, Any]]: + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .order_by(DocumentChunk.sort_order) + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + if allowed_chunk_types is not None: + stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) + + result = await db.execute(stmt) + rows: list[dict[str, Any]] = [] + for document, chunk, section, job_result in result.all(): + section_path = section.section_path if section else None + if is_excluded_section( + document_id=document.document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + if signal_paths and section_path: + path_lower = section_path.lower() + matches_any = any(keyword.lower() in path_lower for keyword in signal_paths) + if filter_mode == 'keep' and not matches_any: + continue + if filter_mode == 'delete' and matches_any: + continue + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 1.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'sort_order': chunk.sort_order, + }) + return rows diff --git a/packages/shared-python/shared/services/retrieval/scoring.py b/packages/shared-python/shared/services/retrieval/scoring.py new file mode 100644 index 000000000..848a3adac --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/scoring.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from typing import Any + +from shared.services.retrieval.settings import RRF_K + + +def get_row_path(row: dict[str, Any]) -> str: + """Extract the canonical path from a row for deduplication.""" + return str(row.get('section_path') or row.get('source_chunk_path') or '') + + +def merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + if not rows: + return rows + groups: dict[str, list[dict[str, Any]]] = {} + order: list[str] = [] + for row in rows: + section_path = row.get('section_path') + if section_path: + key = f"{row.get('document_id', '')}::{section_path}" + else: + key = row.get('chunk_id', '') + if key not in groups: + groups[key] = [] + order.append(key) + groups[key].append(row) + + merged: list[dict[str, Any]] = [] + for key in order: + group = groups[key] + if len(group) == 1: + merged.append(group[0]) + continue + base = dict(group[0]) + base['content'] = '\n'.join(str(row.get('content', '')) for row in group) + base['score'] = max(row.get('score', 0.0) for row in group) + merged.append(base) + return merged + + +def merge_channels_rrf( + channels: list[list[dict[str, Any]]], + weights: list[float], + top_k: int, + k: int = RRF_K, +) -> list[dict[str, Any]]: + """Reciprocal Rank Fusion across multiple retrieval channels.""" + score_dict: dict[str, float] = {} + row_by_chunk_id: dict[str, dict[str, Any]] = {} + + for channel_idx, channel_rows in enumerate(channels): + weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 + for rank, row in enumerate(channel_rows): + chunk_id = str(row.get('chunk_id') or '') + if not chunk_id: + continue + rrf_score = weight / (k + rank + 1) + score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score + if chunk_id not in row_by_chunk_id: + row_by_chunk_id[chunk_id] = row + + ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) + results: list[dict[str, Any]] = [] + for chunk_id, fused_score in ranked[:top_k]: + row = row_by_chunk_id[chunk_id] + results.append(dict(row, score=round(fused_score, 6))) + return results + + +def normalize_row_scores( + rows: list[dict[str, Any]], + *, + source_field: str, + target_field: str, + default: float, +) -> None: + if not rows: + return + values = [float(row.get(source_field, 0.0) or 0.0) for row in rows] + min_score = min(values) + max_score = max(values) + if max_score <= 0.0 and min_score <= 0.0: + for row in rows: + row[target_field] = 0.0 + return + if max_score == min_score: + for row in rows: + row[target_field] = default + return + denominator = max_score - min_score + for row in rows: + raw_score = float(row.get(source_field, 0.0) or 0.0) + row[target_field] = round((raw_score - min_score) / denominator, 6) diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py new file mode 100644 index 000000000..49f2461b5 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/settings.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +CHANNEL_WEIGHT_PATH = 1.0 +CHANNEL_WEIGHT_CONTENT = 2.0 +CHANNEL_WEIGHT_TERM = 1.5 +INTERNAL_RECALL_K_MULTIPLIER = 2 +RRF_K = 60 + +DATA_TYPE_ALLOWED_CHUNK_TYPES: dict[int, set[str] | None] = { + 1: None, + 2: {'text'}, + 3: {'image'}, + 4: {'table'}, + 5: {'text', 'image'}, + 6: {'text', 'table'}, +} + + +def resolve_allowed_chunk_types(data_type: int) -> set[str] | None: + return DATA_TYPE_ALLOWED_CHUNK_TYPES.get(data_type) diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index abce68b36..d63735be5 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -4,13 +4,14 @@ import asyncio import os import time +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager from typing import Any from uuid import uuid4 from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.database import get_db_context from shared.services.retrieval.agentic.budget import BudgetLedger from shared.services.retrieval.agentic.orchestrator import RetrievalAgent, _load_budget_inventory from shared.services.retrieval.agentic.types import AgenticResult @@ -27,12 +28,23 @@ from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult, WorkflowResult from shared.services.retrieval.workflow.wallet import BudgetWallet +DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] + class WorkflowOrchestrator: """Plan and execute a query workflow DAG.""" - def __init__(self) -> None: + def __init__(self, db_factory: DbSessionFactory | None = None) -> None: self.parent_run_id = f'wret_{uuid4().hex[:12]}' + self._db_factory = db_factory + + def _get_db_factory(self) -> DbSessionFactory: + if self._db_factory is not None: + return self._db_factory + + from shared.core.database import get_db_context + + return get_db_context async def run( self, @@ -100,7 +112,6 @@ async def run( await asyncio.gather( *[ self._run_step( - db, step=step, ledger=ledgers[step.id], results_by_id=results_by_id, @@ -198,7 +209,6 @@ async def _load_or_plan( async def _run_step( self, - db: AsyncSession, *, step: PlannedStep, ledger: BudgetLedger, @@ -221,7 +231,6 @@ async def _run_step( await self._run_synthesize_step(step, ledger, results_by_id, llm_fn) return await self._run_retrieve_step( - db, step=step, ledger=ledger, results_by_id=results_by_id, @@ -240,7 +249,6 @@ async def _run_step( async def _run_retrieve_step( self, - db: AsyncSession, *, step: PlannedStep, ledger: BudgetLedger, @@ -261,7 +269,8 @@ async def _run_retrieve_step( # AsyncSession is not safe for concurrent use. Workflow steps may # run in the same topological batch, so each retrieve step opens an # isolated session and leaves the parent session untouched. - async with get_db_context() as step_db: + db_factory = self._get_db_factory() + async with db_factory() as step_db: agentic_result = await RetrievalAgent().run( step_db, user_id=user_id, @@ -375,8 +384,15 @@ def _dedupe_references(refs) -> list[dict[str, Any]]: seen: set[str] = set() out: list[dict[str, Any]] = [] for ref in refs: - chunk_id = str(ref.get('chunk_id') or '') - key = chunk_id or str(ref) + document_id = str(ref.get('document_id') or '').strip() + chunk_id = str(ref.get('chunk_id') or '').strip() + section_path = str(ref.get('section_path') or '').strip() + file_path = str(ref.get('file_path') or '').strip() + key = ( + f'{document_id}:{chunk_id}:{section_path}:{file_path}' + if document_id and chunk_id + else str(ref) + ) if key in seen: continue seen.add(key) From 4a5cfd3515d60b172a8541712cba988a0921ee29 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 02:29:31 +0800 Subject: [PATCH 02/40] refactor: deepen apps/api workflow seams --- CONTEXT.md | 190 ++++++ apps/api/app/api/v1/routes/billing.py | 34 +- apps/api/app/api/v1/routes/jobs.py | 12 +- .../services/billing/billing_app_service.py | 364 ----------- .../billing/billing_workflow_service.py | 378 +++++++++++ .../api/app/services/demo_document_service.py | 499 +-------------- apps/api/app/services/demo_source_catalog.py | 480 ++++++++++++++ .../services/document_ingestion_service.py | 605 ++++++++++++++++++ apps/api/app/services/job_creation_service.py | 463 -------------- .../job_upload_confirmation_service.py | 124 ---- .../app/services/rate_limit/dependencies.py | 409 +----------- .../rate_limit/job_admission_service.py | 356 +++++++++++ .../tests/contract/test_billing_contract.py | 4 +- .../contract/test_job_creation_contract.py | 6 +- 14 files changed, 2080 insertions(+), 1844 deletions(-) create mode 100644 CONTEXT.md delete mode 100644 apps/api/app/services/billing/billing_app_service.py create mode 100644 apps/api/app/services/billing/billing_workflow_service.py create mode 100644 apps/api/app/services/demo_source_catalog.py create mode 100644 apps/api/app/services/document_ingestion_service.py delete mode 100644 apps/api/app/services/job_creation_service.py delete mode 100644 apps/api/app/services/job_upload_confirmation_service.py create mode 100644 apps/api/app/services/rate_limit/job_admission_service.py diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 000000000..0c956ac22 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,190 @@ +# CONTEXT + +## Purpose + +Knowhere API turns authenticated requests into document ingestion, document +lifecycle, retrieval, billing, and webhook workflows. + +Within this repository, `apps/api` is the coordination layer between HTTP +adapters and the shared implementations in `packages/shared-python/shared`. + +## Core Terms + +### User + +The authenticated owner of jobs, documents, credits, API keys, and webhooks. + +### Namespace + +The isolation scope for retrieval-visible data. The default namespace is +`default`. + +### Job + +The API-side intake and execution handle for a workflow such as file parsing, +URL ingestion, or demo source materialization. + +### Job Result + +The terminal artifact record attached to a Job. It stores delivery metadata, +result bundle references, and the revision that publication uses. + +### Document + +The retrieval-visible knowledge object produced from a Job Result after +publication. + +### Document Section + +The hierarchical navigation node derived from parsed headings and section paths. + +### Document Chunk + +The retrieval-visible text, image, or table row attached to a Document Section. + +### Document Ingestion + +The workflow that creates a Job, accepts a file or URL source, confirms upload +state, and starts parsing work. + +### Job Admission + +The policy checks that must pass before a new Job is created: authentication, +guest scope, system limits, billing RPM, concurrent job limits, and daily +quota. + +### Publication + +The shared workflow that turns parsed chunks into Documents, Document Sections, +Document Chunks, and document graph state. + +### Retrieval + +The query workflow that returns cited evidence from published documents. + +### Demo Source + +An API-owned canonical document shipped with the repository for demo and guest +flows. + +### Demo Source Materialization + +The workflow that copies a Demo Source into a user's Namespace as normal Job, +Job Result, Document, and Document Chunk records. + +### Billing Workflow + +The credits purchase, checkout, webhook handling, refund reconciliation, and +tier refresh flows. + +### Guest API Key + +A guest-tier API key with a restricted route surface. + +### Webhook Management + +The user-facing workflow for storing outbound webhook configuration and reading +delivery logs. + +### QStash Callback + +The verified async callback used to continue background work after external +delivery. + +## apps/api Module Map + +### HTTP Adapters + +`apps/api/app/api/v1/routes/*` + +These modules translate HTTP requests into application workflow calls. + +### Application Workflows + +`apps/api/app/services/*` + +These modules coordinate Job Admission, Document Ingestion, document lifecycle, +Billing Workflow, Demo Source Materialization, webhook handling, and internal +callbacks. + +### Persistence Adapters + +`apps/api/app/repositories/*` + +These modules own database reads and writes for API-side workflows. + +### Shared Implementations + +`packages/shared-python/shared/*` + +These modules own the lower-level implementations for publication, retrieval, +state machines, storage, Redis-backed metadata, billing primitives, and core +exceptions. + +## apps/api Workflow Ownership + +### Document Ingestion + +- `app/api/v1/routes/jobs.py` +- `app/services/job_creation_service.py` +- `app/services/job_upload_confirmation_service.py` +- `app/services/job_document_scope_service.py` +- `app/repositories/job_repository.py` + +### Job Admission + +- `app/services/rate_limit/*` +- `app/core/dependencies.py` + +### Document Lifecycle + +- `app/api/v1/routes/documents.py` +- `app/services/document_service.py` +- `app/repositories/document_repository.py` + +### Retrieval + +- `app/api/v1/routes/retrieval.py` +- shared retrieval modules in `packages/shared-python/shared/services/retrieval/*` + +### Demo Source Materialization + +- `app/api/v1/routes/demo.py` +- `app/services/demo_document_service.py` + +### Billing Workflow + +- `app/api/v1/routes/billing.py` +- `app/services/billing/*` +- `app/repositories/payment_record_repository.py` +- shared billing modules in `packages/shared-python/shared/services/billing/*` + +### Webhook Management + +- `app/api/v1/routes/webhook.py` +- `app/api/v1/routes/webhook_secrets.py` +- `app/services/webhook_service.py` +- `app/repositories/webhook_repository.py` + +### Internal Storage Events + +- `app/api/v1/routes/s3_events.py` +- `app/services/s3_events/*` + +### Async Callbacks + +- `app/api/v1/routes/qstash_callbacks.py` +- `app/services/qstash_callback_service.py` + +## Invariants + +- `apps/api` coordinates workflows. Parsing, publication, retrieval internals, + storage mechanics, and state-machine implementation mostly live outside the + route modules. +- A Job and a Document are not the same thing. Jobs track intake and processing; + Documents track retrieval-visible knowledge state. +- `current_job_result_id` selects the active revision of a Document. +- Namespace is part of the retrieval contract, not a UI-only label. +- Demo Sources should behave like normal Documents after materialization. +- Billing Workflow and Job Admission shape whether work is allowed to start; + they are not worker-only concerns. diff --git a/apps/api/app/api/v1/routes/billing.py b/apps/api/app/api/v1/routes/billing.py index f9c2a2dd7..34e1334f7 100644 --- a/apps/api/app/api/v1/routes/billing.py +++ b/apps/api/app/api/v1/routes/billing.py @@ -2,16 +2,9 @@ from typing import Optional -from app.services.billing.billing_app_service import ( +from app.services.billing.billing_workflow_service import ( + BillingWorkflowService, ParseUsageResponse, - buy_credits_for_user, - buy_credits_package_for_user, - get_credits_balance_for_user, - get_parse_usage_overview_for_user, - get_price_configs_payload, - get_transaction_history_for_user, - get_usage_stats_for_user, - handle_stripe_webhook, ) from app.services.rate_limit.dependencies import CurrentUser, with_current_user from fastapi import APIRouter, Depends, Query, Request @@ -28,6 +21,7 @@ ) router = APIRouter(tags=["Billing"]) +_billing_workflow_service = BillingWorkflowService() @router.post("/buy-credits", summary="Buy Credits", response_model=PaymentIntentResponse) @@ -36,7 +30,7 @@ async def buy_credits( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ) -> PaymentIntentResponse: - return await buy_credits_for_user( + return await _billing_workflow_service.buy_credits( request=request, user_id=current_user.user_id, ) @@ -51,7 +45,10 @@ async def get_credits_balance( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ) -> CreditsBalanceResponse: - return await get_credits_balance_for_user(db, user_id=current_user.user_id) + return await _billing_workflow_service.get_credits_balance( + db, + user_id=current_user.user_id, + ) @router.get( @@ -64,7 +61,7 @@ async def get_usage_stats( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ) -> UsageStatsResponse: - return await get_usage_stats_for_user( + return await _billing_workflow_service.get_usage_stats( db, user_id=current_user.user_id, period=period, @@ -80,7 +77,7 @@ async def parse_usage_overview( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ) -> ParseUsageResponse: - return await get_parse_usage_overview_for_user( + return await _billing_workflow_service.get_parse_usage_overview( db, user_id=current_user.user_id, ) @@ -92,7 +89,7 @@ async def get_transaction_history( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ): - return await get_transaction_history_for_user( + return await _billing_workflow_service.get_transaction_history( db, user_id=current_user.user_id, limit=limit, @@ -107,7 +104,10 @@ async def get_price_configs( ), db: AsyncSession = Depends(get_db), ) -> dict[str, list[dict]]: - return await get_price_configs_payload(db, product_type=product_type) + return await _billing_workflow_service.get_price_configs( + db, + product_type=product_type, + ) @router.post( @@ -120,7 +120,7 @@ async def buy_credits_package( current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ) -> CheckoutSessionResponse: - return await buy_credits_package_for_user( + return await _billing_workflow_service.buy_credits_package( db, request=request, user_id=current_user.user_id, @@ -129,7 +129,7 @@ async def buy_credits_package( @router.post("/webhook", summary="Stripe Webhook") async def stripe_webhook(request: Request, db: AsyncSession = Depends(get_db)): - return await handle_stripe_webhook( + return await _billing_workflow_service.handle_stripe_webhook( db, payload=await request.body(), stripe_signature=request.headers.get("stripe-signature"), diff --git a/apps/api/app/api/v1/routes/jobs.py b/apps/api/app/api/v1/routes/jobs.py index b6afb84f6..befada76c 100644 --- a/apps/api/app/api/v1/routes/jobs.py +++ b/apps/api/app/api/v1/routes/jobs.py @@ -7,18 +7,16 @@ from datetime import datetime from typing import Optional -from app.services.job_creation_service import create_job_from_request +from app.services.document_ingestion_service import DocumentIngestionService from app.services.job_read_service import ( get_job_result_for_user, list_jobs_for_user, ) from app.services.rate_limit.dependencies import ( CurrentUser, - enforce_job_creation_capacity, require_billing_limits, with_current_user, ) -from app.services.job_upload_confirmation_service import confirm_job_upload from fastapi import APIRouter, Depends, Query, Request from sqlalchemy.ext.asyncio import AsyncSession @@ -32,6 +30,7 @@ ) router = APIRouter(tags=["Jobs"]) +_document_ingestion_service = DocumentIngestionService() # ==================== Shared Helpers ==================== @@ -48,11 +47,10 @@ async def create_job( # pyright: ignore[reportGeneralTypeIssues] """ Create a parsing job. """ - return await create_job_from_request( + return await _document_ingestion_service.create_job( db, payload=payload, current_user=current_user, - enforce_capacity=enforce_job_creation_capacity, request=http_request, ) @@ -122,9 +120,9 @@ async def confirm_upload( """ Confirm a completed file upload as a fallback path. """ - return await confirm_job_upload( + return await _document_ingestion_service.confirm_upload( db, job_id=job_id, - request=request, + request_payload=request, user_id=current_user.user_id, ) diff --git a/apps/api/app/services/billing/billing_app_service.py b/apps/api/app/services/billing/billing_app_service.py deleted file mode 100644 index c7c3e6b5b..000000000 --- a/apps/api/app/services/billing/billing_app_service.py +++ /dev/null @@ -1,364 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -from app.services.billing.price_config_service import PriceConfigService -from app.services.billing.stripe_service import StripeService -from pydantic import BaseModel -from sqlalchemy import func, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.billing import MicroDollar -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import StripeServiceException -from shared.models.database.credits_transaction import CreditsTransaction -from shared.models.database.job import Job -from shared.models.database.stripe_price_config import StripePriceConfig -from shared.models.database.user import User -from shared.models.schemas.billing import ( - BuyCreditsPackageRequest, - BuyCreditsRequest, - CheckoutSessionResponse, - CreditsBalanceResponse, - PaymentIntentResponse, - TransactionHistoryResponse, - UsageStatsResponse, -) -from shared.services.billing import CreditsService - - -class ParseUsageResponse(BaseModel): - request_total: int - mom_growth: float - credits_used: float - estimated_amount: Optional[float] - success_rate: float - avg_processing_time: float - - -async def buy_credits_for_user( - *, - request: BuyCreditsRequest, - user_id: str, -) -> PaymentIntentResponse: - stripe_service = StripeService() - try: - amount_cny = request.credits_amount * 0.02 - amount_cents = int(amount_cny * 100) - payment_intent = await stripe_service.create_payment_intent( - user_id=user_id, - amount=amount_cents, - credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, - currency="cny", - ) - - return PaymentIntentResponse( - client_secret=payment_intent["client_secret"], - payment_intent_id=payment_intent["payment_intent_id"], - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to buy credits: {str(exc)}" - ) - - -async def get_credits_balance_for_user( - db: AsyncSession, - *, - user_id: str, -) -> CreditsBalanceResponse: - credits_service = CreditsService() - try: - await credits_service.ensure_user_initialized(db, user_id) - await db.commit() - - balance_micro_dollar = await credits_service.get_balance(db, user_id) - return CreditsBalanceResponse( - credits_balance=MicroDollar(balance_micro_dollar).to_credit() - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get credits balance: {str(exc)}" - ) - - -async def get_usage_stats_for_user( - db: AsyncSession, - *, - user_id: str, - period: str, -) -> UsageStatsResponse: - credits_service = CreditsService() - try: - stats = await credits_service.get_usage_stats(db, user_id, period) - return UsageStatsResponse( - period=stats["period"], - total_credits_used=MicroDollar(stats["total_used"]).to_credit(), - api_calls_count=stats["transaction_count"], - success_rate=95.0, - average_response_time=stats.get("avg_response_time", 0), - top_endpoints=[], - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get usage statistics: {str(exc)}" - ) - - -async def get_parse_usage_overview_for_user( - db: AsyncSession, - *, - user_id: str, -) -> ParseUsageResponse: - try: - total_micro_credits_used = await _load_total_parse_micro_credits_used( - db, - user_id=user_id, - ) - success_rate, avg_processing_time = await _load_parse_job_usage_stats( - db, - user_id=user_id, - ) - estimated_amount = await _estimate_parse_usage_amount( - db, - total_micro_credits_used=total_micro_credits_used, - ) - - return ParseUsageResponse( - request_total=0, - mom_growth=0.0, - credits_used=MicroDollar(total_micro_credits_used).to_credit() or 0, - estimated_amount=estimated_amount, - success_rate=round(success_rate, 2), - avg_processing_time=avg_processing_time, - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get parse usage overview: {str(exc)}" - ) - - -async def get_transaction_history_for_user( - db: AsyncSession, - *, - user_id: str, - limit: int, -) -> list[TransactionHistoryResponse]: - credits_service = CreditsService() - try: - transactions = await credits_service.get_transaction_history( - db, - user_id, - limit, - ) - return [ - TransactionHistoryResponse( - id=transaction.id, - credits_amount=MicroDollar(transaction.credits_amount).to_credit(), - transaction_type=transaction.transaction_type, - description=transaction.description, - created_at=transaction.created_at, - ) - for transaction in transactions - ] - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get transaction history: {str(exc)}" - ) - - -async def get_price_configs_payload( - db: AsyncSession, - *, - product_type: str | None, -) -> dict[str, list[dict]]: - try: - price_config_service = PriceConfigService() - - if product_type == "subscription": - configs = await price_config_service.repository.get_all_active(db) - return { - "subscriptions": [ - _subscription_config_payload(config) - for config in configs - if config.product_type == "subscription" - ], - "credits_packages": [], - } - - if product_type == "credits_package": - credits_configs = await price_config_service.get_all_credits_packages(db) - return { - "subscriptions": [], - "credits_packages": [ - _credits_package_config_payload(config) - for config in credits_configs - ], - } - - configs = await price_config_service.repository.get_all_active(db) - return { - "subscriptions": [ - _subscription_config_payload(config) - for config in configs - if config.product_type == "subscription" - ], - "credits_packages": [ - _credits_package_config_payload(config) - for config in configs - if config.product_type == "credits_package" - ], - } - - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get price configurations: {str(exc)}" - ) - - -async def buy_credits_package_for_user( - db: AsyncSession, - *, - request: BuyCreditsPackageRequest, - user_id: str, -) -> CheckoutSessionResponse: - stripe_service = StripeService() - try: - result = await db.execute(select(User.email).where(User.id == user_id)) - user_email = result.scalar_one_or_none() - - frontend_url = settings.FRONTEND_URL - success_url = f"{frontend_url}/billing?success=true&type=credits_package" - cancel_url = f"{frontend_url}/billing?canceled=true" - - checkout_url = await stripe_service.create_checkout_session_for_credits_package( - db=db, - user_id=user_id, - price_id=request.price_id, - success_url=success_url, - cancel_url=cancel_url, - quantity=request.quantity, - email=user_email, - ) - - return CheckoutSessionResponse(checkout_url=checkout_url, session_id="") - - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to create credits package purchase: {str(exc)}" - ) - - -async def handle_stripe_webhook( - db: AsyncSession, - *, - payload: bytes, - stripe_signature: str | None, -): - stripe_service = StripeService() - try: - if not stripe_signature: - raise StripeServiceException( - internal_message="Missing stripe-signature header" - ) - return await stripe_service.handle_webhook(db, payload, stripe_signature) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to handle webhook: {str(exc)}" - ) - - -async def _load_total_parse_micro_credits_used( - db: AsyncSession, - *, - user_id: str, -) -> int: - credits_row = await db.execute( - select(func.coalesce(func.sum(CreditsTransaction.credits_amount), 0)) - .where(CreditsTransaction.user_id == user_id) - .where(CreditsTransaction.transaction_type.in_(["usage", "refund"])) - ) - return int(abs(credits_row.scalar_one() or 0)) - - -async def _load_parse_job_usage_stats( - db: AsyncSession, - *, - user_id: str, -) -> tuple[float, float]: - job_row = await db.execute( - select( - func.count().filter(Job.status == "done").label("done_cnt"), - func.count() - .filter(Job.status.in_(["done", "failed"])) - .label("terminal_cnt"), - func.avg(func.extract("epoch", Job.updated_at - Job.created_at)) - .filter(Job.status.in_(["done", "failed"])) - .label("avg_secs"), - ).where(Job.user_id == user_id) - ) - job_stats = job_row.first() or (0, 0, 0.0) - done_count = getattr(job_stats, "done_cnt", 0) or 0 - terminal_count = getattr(job_stats, "terminal_cnt", 0) or 0 - success_rate = (done_count / terminal_count * 100) if terminal_count > 0 else 0.0 - avg_processing_time = round( - float(getattr(job_stats, "avg_secs", 0.0) or 0.0), - 2, - ) - return success_rate, avg_processing_time - - -async def _estimate_parse_usage_amount( - db: AsyncSession, - *, - total_micro_credits_used: int, -) -> float | None: - price_row = await db.execute( - select(StripePriceConfig) - .where(StripePriceConfig.product_type == "credits_package") - .where(StripePriceConfig.is_active.is_(True)) - .order_by(StripePriceConfig.created_at) - .limit(1) - ) - price_cfg = price_row.scalar_one_or_none() - if not price_cfg or not price_cfg.credits_amount or price_cfg.credits_amount <= 0: - return None - - return round( - price_cfg.amount_cents - * total_micro_credits_used - / (100 * price_cfg.credits_amount), - 4, - ) - - -def _subscription_config_payload(config) -> dict: - metadata = config.extra_metadata or {} - return { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": metadata.get("display_name", config.plan_id.upper()), - "description": metadata.get("description", ""), - "features": metadata.get("features", []), - "popular": metadata.get("frontend_config", {}).get("popular", False), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": metadata, - } - - -def _credits_package_config_payload(config) -> dict: - metadata = config.extra_metadata or {} - credit_amount = MicroDollar(config.credits_amount).to_credit() - return { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": metadata.get("display_name", f"{credit_amount} Credits"), - "description": metadata.get("description", ""), - "credits_amount": credit_amount, - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": metadata, - } diff --git a/apps/api/app/services/billing/billing_workflow_service.py b/apps/api/app/services/billing/billing_workflow_service.py new file mode 100644 index 000000000..51dbd44bb --- /dev/null +++ b/apps/api/app/services/billing/billing_workflow_service.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +from typing import Optional + +from app.services.billing.price_config_service import PriceConfigService +from app.services.billing.stripe_service import StripeService +from pydantic import BaseModel +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.billing import MicroDollar +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import StripeServiceException +from shared.models.database.credits_transaction import CreditsTransaction +from shared.models.database.job import Job +from shared.models.database.stripe_price_config import StripePriceConfig +from shared.models.database.user import User +from shared.models.schemas.billing import ( + BuyCreditsPackageRequest, + BuyCreditsRequest, + CheckoutSessionResponse, + CreditsBalanceResponse, + PaymentIntentResponse, + TransactionHistoryResponse, + UsageStatsResponse, +) +from shared.services.billing import CreditsService + + +class ParseUsageResponse(BaseModel): + request_total: int + mom_growth: float + credits_used: float + estimated_amount: Optional[float] + success_rate: float + avg_processing_time: float + + +class BillingWorkflowService: + def __init__( + self, + *, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + + async def buy_credits( + self, + *, + request: BuyCreditsRequest, + user_id: str, + ) -> PaymentIntentResponse: + stripe_service = self._create_stripe_service() + try: + amount_cny = request.credits_amount * 0.02 + amount_cents = int(amount_cny * 100) + payment_intent = await stripe_service.create_payment_intent( + user_id=user_id, + amount=amount_cents, + credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, + currency="cny", + ) + + return PaymentIntentResponse( + client_secret=payment_intent["client_secret"], + payment_intent_id=payment_intent["payment_intent_id"], + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to buy credits: {str(exc)}" + ) + + async def get_credits_balance( + self, + db: AsyncSession, + *, + user_id: str, + ) -> CreditsBalanceResponse: + try: + await self._credits_service.ensure_user_initialized(db, user_id) + await db.commit() + + balance_micro_dollar = await self._credits_service.get_balance(db, user_id) + return CreditsBalanceResponse( + credits_balance=MicroDollar(balance_micro_dollar).to_credit() + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get credits balance: {str(exc)}" + ) + + async def get_usage_stats( + self, + db: AsyncSession, + *, + user_id: str, + period: str, + ) -> UsageStatsResponse: + try: + stats = await self._credits_service.get_usage_stats(db, user_id, period) + return UsageStatsResponse( + period=stats["period"], + total_credits_used=MicroDollar(stats["total_used"]).to_credit(), + api_calls_count=stats["transaction_count"], + success_rate=95.0, + average_response_time=stats.get("avg_response_time", 0), + top_endpoints=[], + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get usage statistics: {str(exc)}" + ) + + async def get_parse_usage_overview( + self, + db: AsyncSession, + *, + user_id: str, + ) -> ParseUsageResponse: + try: + total_micro_credits_used = await self._load_total_parse_micro_credits_used( + db, + user_id=user_id, + ) + success_rate, avg_processing_time = await self._load_parse_job_usage_stats( + db, + user_id=user_id, + ) + estimated_amount = await self._estimate_parse_usage_amount( + db, + total_micro_credits_used=total_micro_credits_used, + ) + + return ParseUsageResponse( + request_total=0, + mom_growth=0.0, + credits_used=MicroDollar(total_micro_credits_used).to_credit() or 0, + estimated_amount=estimated_amount, + success_rate=round(success_rate, 2), + avg_processing_time=avg_processing_time, + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get parse usage overview: {str(exc)}" + ) + + async def get_transaction_history( + self, + db: AsyncSession, + *, + user_id: str, + limit: int, + ) -> list[TransactionHistoryResponse]: + try: + transactions = await self._credits_service.get_transaction_history( + db, + user_id, + limit, + ) + return [ + TransactionHistoryResponse( + id=transaction.id, + credits_amount=MicroDollar(transaction.credits_amount).to_credit(), + transaction_type=transaction.transaction_type, + description=transaction.description, + created_at=transaction.created_at, + ) + for transaction in transactions + ] + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get transaction history: {str(exc)}" + ) + + async def get_price_configs( + self, + db: AsyncSession, + *, + product_type: str | None, + ) -> dict[str, list[dict[str, object]]]: + try: + if product_type == "subscription": + configs = await self._price_config_service.repository.get_all_active(db) + return { + "subscriptions": [ + _subscription_config_payload(config) + for config in configs + if config.product_type == "subscription" + ], + "credits_packages": [], + } + + if product_type == "credits_package": + credits_configs = await self._price_config_service.get_all_credits_packages( + db + ) + return { + "subscriptions": [], + "credits_packages": [ + _credits_package_config_payload(config) + for config in credits_configs + ], + } + + configs = await self._price_config_service.repository.get_all_active(db) + return { + "subscriptions": [ + _subscription_config_payload(config) + for config in configs + if config.product_type == "subscription" + ], + "credits_packages": [ + _credits_package_config_payload(config) + for config in configs + if config.product_type == "credits_package" + ], + } + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get price configurations: {str(exc)}" + ) + + async def buy_credits_package( + self, + db: AsyncSession, + *, + request: BuyCreditsPackageRequest, + user_id: str, + ) -> CheckoutSessionResponse: + stripe_service = self._create_stripe_service() + try: + result = await db.execute(select(User.email).where(User.id == user_id)) + user_email = result.scalar_one_or_none() + + frontend_url = settings.FRONTEND_URL + success_url = f"{frontend_url}/billing?success=true&type=credits_package" + cancel_url = f"{frontend_url}/billing?canceled=true" + + checkout_url = await stripe_service.create_checkout_session_for_credits_package( + db=db, + user_id=user_id, + price_id=request.price_id, + success_url=success_url, + cancel_url=cancel_url, + quantity=request.quantity, + email=user_email, + ) + + return CheckoutSessionResponse(checkout_url=checkout_url, session_id="") + except Exception as exc: + raise StripeServiceException( + internal_message=( + "Failed to create credits package purchase: " + f"{str(exc)}" + ) + ) + + async def handle_stripe_webhook( + self, + db: AsyncSession, + *, + payload: bytes, + stripe_signature: str | None, + ) -> dict[str, object]: + stripe_service = self._create_stripe_service() + try: + if not stripe_signature: + raise StripeServiceException( + internal_message="Missing stripe-signature header" + ) + return await stripe_service.handle_webhook(db, payload, stripe_signature) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to handle webhook: {str(exc)}" + ) + + def _create_stripe_service(self) -> StripeService: + return StripeService() + + async def _load_total_parse_micro_credits_used( + self, + db: AsyncSession, + *, + user_id: str, + ) -> int: + credits_row = await db.execute( + select(func.coalesce(func.sum(CreditsTransaction.credits_amount), 0)) + .where(CreditsTransaction.user_id == user_id) + .where(CreditsTransaction.transaction_type.in_(["usage", "refund"])) + ) + return int(abs(credits_row.scalar_one() or 0)) + + async def _load_parse_job_usage_stats( + self, + db: AsyncSession, + *, + user_id: str, + ) -> tuple[float, float]: + job_row = await db.execute( + select( + func.count().filter(Job.status == "done").label("done_cnt"), + func.count() + .filter(Job.status.in_(["done", "failed"])) + .label("terminal_cnt"), + func.avg(func.extract("epoch", Job.updated_at - Job.created_at)) + .filter(Job.status.in_(["done", "failed"])) + .label("avg_secs"), + ).where(Job.user_id == user_id) + ) + job_stats = job_row.first() or (0, 0, 0.0) + done_count = getattr(job_stats, "done_cnt", 0) or 0 + terminal_count = getattr(job_stats, "terminal_cnt", 0) or 0 + success_rate = ( + done_count / terminal_count * 100 if terminal_count > 0 else 0.0 + ) + avg_processing_time = round( + float(getattr(job_stats, "avg_secs", 0.0) or 0.0), + 2, + ) + return success_rate, avg_processing_time + + async def _estimate_parse_usage_amount( + self, + db: AsyncSession, + *, + total_micro_credits_used: int, + ) -> float | None: + price_row = await db.execute( + select(StripePriceConfig) + .where(StripePriceConfig.product_type == "credits_package") + .where(StripePriceConfig.is_active.is_(True)) + .order_by(StripePriceConfig.created_at) + .limit(1) + ) + price_cfg = price_row.scalar_one_or_none() + if not price_cfg or not price_cfg.credits_amount or price_cfg.credits_amount <= 0: + return None + + return round( + price_cfg.amount_cents + * total_micro_credits_used + / (100 * price_cfg.credits_amount), + 4, + ) + + +def _subscription_config_payload(config: StripePriceConfig) -> dict[str, object]: + metadata = config.extra_metadata or {} + return { + "id": config.plan_id, + "plan_id": config.plan_id, + "price_id": config.price_id, + "name": metadata.get("display_name", config.plan_id.upper()), + "description": metadata.get("description", ""), + "features": metadata.get("features", []), + "popular": metadata.get("frontend_config", {}).get("popular", False), + "amount_cents": config.amount_cents, + "currency": config.currency, + "metadata": metadata, + } + + +def _credits_package_config_payload(config: StripePriceConfig) -> dict[str, object]: + metadata = config.extra_metadata or {} + credit_amount = MicroDollar(config.credits_amount).to_credit() + return { + "id": config.plan_id, + "plan_id": config.plan_id, + "price_id": config.price_id, + "name": metadata.get("display_name", f"{credit_amount} Credits"), + "description": metadata.get("description", ""), + "credits_amount": credit_amount, + "amount_cents": config.amount_cents, + "currency": config.currency, + "metadata": metadata, + } diff --git a/apps/api/app/services/demo_document_service.py b/apps/api/app/services/demo_document_service.py index ea50a9b1d..d1bbd779c 100644 --- a/apps/api/app/services/demo_document_service.py +++ b/apps/api/app/services/demo_document_service.py @@ -1,23 +1,21 @@ -"""API-owned canonical demo document catalog and materialization.""" +"""Demo Source materialization workflow.""" from __future__ import annotations -import json -import math import shutil import tempfile from dataclasses import dataclass from datetime import datetime, timezone -from functools import lru_cache from hashlib import blake2b from pathlib import Path from typing import Any -from urllib.parse import quote from uuid import uuid4 from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession +from app.services.demo_source_catalog import DemoSourceCatalog, DemoSourceDefinition + from shared.core.exceptions.domain_exceptions import ValidationException from shared.models.database.demo_materialization import DemoMaterialization from shared.models.database.document import Document @@ -28,39 +26,6 @@ from shared.services.storage.result_storage import get_result_storage -@dataclass(frozen=True) -class DemoCitationDefinition: - """Curated answer citation that resolves to a canonical demo chunk.""" - - section_path: str - description: str - content: str - - -@dataclass(frozen=True) -class DemoExampleDefinition: - """Curated user-facing demo question and answer.""" - - id: str - question: str - answer: str - citations: tuple[DemoCitationDefinition, ...] - - -@dataclass(frozen=True) -class DemoSourceDefinition: - """Canonical demo source metadata and local asset pointers.""" - - demo_source_id: str - canonical_document_id: str - title: str - mime_type: str - size_bytes: int - asset_directory: str - chunk_count: int - examples: tuple[DemoExampleDefinition, ...] - - @dataclass(frozen=True) class MaterializedDemoSource: """User-owned copy of one canonical demo source.""" @@ -74,114 +39,20 @@ class MaterializedDemoSource: chunk_count: int -_DATA_ROOT = Path(__file__).resolve().parents[1] / "data" / "demo_documents" -_ASSET_DIRECTORY_NAMES = frozenset({"images", "tables"}) -_DEMO_SOURCE_DEFINITIONS: tuple[DemoSourceDefinition, ...] = ( - DemoSourceDefinition( - demo_source_id="demo-tsla-q4-2025", - canonical_document_id="demo-doc-tsla-q4-2025", - title="TSLA-Q4-2025-Update.pdf", - mime_type="application/pdf", - size_bytes=5_648_867, - asset_directory="tsla-q4-2025", - chunk_count=70, - examples=( - DemoExampleDefinition( - id="demo-tsla-q4-2025-xai", - question="What does the document say about Tesla's xAI investment?", - answer=( - "Tesla entered an agreement on January 16, 2026 to invest " - "approximately $2 billion in xAI Series E Preferred Stock.\n\n" - "The document also says Tesla and xAI entered a framework " - "agreement to evaluate AI collaboration, with the investment " - "expected to close in Q1 2026 subject to customary regulatory " - "conditions." - ), - citations=( - DemoCitationDefinition( - section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->OTHER UPDATES" - ), - description="xAI investment", - content=( - "On January 16, 2026, Tesla entered into an agreement " - "to invest approximately" - ), - ), - ), - ), - DemoExampleDefinition( - id="demo-tsla-q4-2025-energy-storage", - question="What does the document say about energy storage?", - answer=( - "Tesla achieved its highest quarterly energy storage " - "deployments, driven by record Megapack deployments.\n\n" - "Energy gross profit reached a record $1.1 billion, marking " - "the fifth consecutive record quarter.\n\n" - "Tesla also plans to begin Megapack 3 and Megablock " - "production at Megafactory Houston in 2026." - ), - citations=( - DemoCitationDefinition( - section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->" - "Energy generation and storage" - ), - description="Storage deployment growth", - content=( - "We achieved our highest quarterly energy storage " - "deployments, driven by record Megapack deployments." - ), - ), - ), - ), - DemoExampleDefinition( - id="demo-tsla-q4-2025-production-plans", - question="What production plans does Tesla mention for 2026?", - answer=( - "Tesla says Cybercab, Tesla Semi, and Megapack 3 are on " - "schedule for volume production starting in 2026.\n\n" - "The same product update also notes that first-generation " - "Optimus production lines are being installed before volume " - "production." - ), - citations=( - DemoCitationDefinition( - section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->" - "Product" - ), - description="2026 production plans", - content=( - "Cybercab, Tesla Semi and Megapack 3 are on schedule " - "for volume production starting in 2026." - ), - ), - ), - ), - ), - ), -) - - class DemoDocumentService: """Serves canonical demo data and copies it into user namespaces.""" def __init__( self, *, + catalog: DemoSourceCatalog | None = None, publication_service: RetrievalPublicationService | None = None, ) -> None: + self._catalog = catalog or DemoSourceCatalog() self._publication_service = publication_service or RetrievalPublicationService() def get_catalog(self) -> dict[str, Any]: - """Return the cacheable canonical demo source catalog.""" - return { - "sources": [ - self._source_catalog_payload(source) - for source in _DEMO_SOURCE_DEFINITIONS - ], - } + return self._catalog.get_catalog() def list_chunks( self, @@ -190,30 +61,11 @@ def list_chunks( page: int, page_size: int, ) -> dict[str, Any] | None: - """Return paginated canonical demo chunks.""" - source = _get_source_definition(demo_source_id) - if source is None: - return None - - chunks = _load_source_chunks(source) - start = (page - 1) * page_size - page_chunks = chunks[start : start + page_size] - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "title": source.title, - "mime_type": source.mime_type, - "chunks": [ - _chunk_payload(source=source, chunk=chunk) - for chunk in page_chunks - ], - "pagination": { - "page": page, - "page_size": page_size, - "total": len(chunks), - "total_pages": math.ceil(len(chunks) / page_size) if chunks else 0, - }, - } + return self._catalog.list_chunks( + demo_source_id=demo_source_id, + page=page, + page_size=page_size, + ) def get_chunk( self, @@ -221,29 +73,13 @@ def get_chunk( demo_source_id: str, demo_chunk_id: str, ) -> dict[str, Any] | None: - """Return one canonical demo chunk by canonical row id or parser chunk id.""" - source = _get_source_definition(demo_source_id) - if source is None: - return None - - for chunk in _load_source_chunks(source): - if demo_chunk_id in {_canonical_chunk_id(source, chunk), chunk["chunk_id"]}: - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "chunk": _chunk_payload(source=source, chunk=chunk), - } - - return None + return self._catalog.get_chunk( + demo_source_id=demo_source_id, + demo_chunk_id=demo_chunk_id, + ) def get_original_file_path(self, *, demo_source_id: str) -> Path | None: - """Return the canonical original file path for a demo source.""" - source = _get_source_definition(demo_source_id) - if source is None: - return None - - file_path = _source_directory(source) / "original.pdf" - return file_path if file_path.is_file() else None + return self._catalog.get_original_file_path(demo_source_id=demo_source_id) def get_asset_file_path( self, @@ -251,21 +87,10 @@ def get_asset_file_path( demo_source_id: str, asset_path: str, ) -> Path | None: - """Return a canonical parsed media/table asset path.""" - source = _get_source_definition(demo_source_id) - if source is None: - return None - - source_directory = _source_directory(source).resolve() - normalized_asset_path = _normalize_asset_path(asset_path) - if normalized_asset_path is None: - return None - - candidate = (source_directory / normalized_asset_path).resolve() - if not candidate.is_relative_to(source_directory): - return None - - return candidate if candidate.is_file() else None + return self._catalog.get_asset_file_path( + demo_source_id=demo_source_id, + asset_path=asset_path, + ) async def materialize_sources( self, @@ -275,7 +100,6 @@ async def materialize_sources( namespace: str, demo_source_ids: list[str], ) -> list[MaterializedDemoSource]: - """Copy selected canonical demo sources into a user namespace.""" selected_demo_source_ids = _deduplicate_source_ids(demo_source_ids) if not selected_demo_source_ids: raise ValidationException( @@ -289,7 +113,7 @@ async def materialize_sources( ) selected_sources = [ - _require_source_definition(demo_source_id) + self._catalog.require_source(demo_source_id) for demo_source_id in selected_demo_source_ids ] results: list[MaterializedDemoSource] = [] @@ -343,7 +167,10 @@ async def _materialize_source( job_id = f"job_demo_{uuid4().hex[:12]}" job_result_id = str(uuid4()) timestamp = _utc_now() - result_bundle = _upload_demo_result_bundle(job_id=job_id, source=source) + result_bundle = _upload_demo_result_bundle( + job_id=job_id, + source_directory=self._catalog.source_directory(source), + ) db.add( Job( @@ -384,7 +211,7 @@ async def _materialize_source( ) ) await db.flush() - chunks = _publication_chunks(source) + chunks = self._catalog.publication_chunks(source) await db.run_sync( lambda sync_db: self._publication_service.publish_document_state( sync_db, @@ -455,43 +282,6 @@ async def _is_active_document( ) return result.scalar_one_or_none() is not None - def _source_catalog_payload(self, source: DemoSourceDefinition) -> dict[str, Any]: - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "title": source.title, - "mime_type": source.mime_type, - "size_bytes": source.size_bytes, - "status": "ready", - "chunk_count": source.chunk_count, - "original_file": { - "url": f"/api/v1/demo/sources/{source.demo_source_id}/original", - "mime_type": source.mime_type, - "size_bytes": source.size_bytes, - "can_download": False, - }, - "examples": [ - self._example_payload(source=source, example=example) - for example in source.examples - ], - } - - def _example_payload( - self, - *, - source: DemoSourceDefinition, - example: DemoExampleDefinition, - ) -> dict[str, Any]: - return { - "id": example.id, - "question": example.question, - "answer": example.answer, - "citations": [ - _citation_payload(source=source, citation=citation) - for citation in example.citations - ], - } - def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: selected: list[str] = [] @@ -551,10 +341,8 @@ def _materialized_source_payload( def _upload_demo_result_bundle( *, job_id: str, - source: DemoSourceDefinition, + source_directory: Path, ) -> dict[str, int | str]: - """Upload canonical demo result files so copied media URLs resolve.""" - source_directory = _source_directory(source) with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temp_directory: zip_base_path = Path(temp_directory) / job_id zip_file_path = Path( @@ -577,238 +365,5 @@ def _upload_demo_result_bundle( } -def _publication_chunks(source: DemoSourceDefinition) -> list[dict[str, Any]]: - return [ - _publication_chunk(source=source, chunk=chunk) - for chunk in _load_source_chunks(source) - ] - - -def _publication_chunk( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> dict[str, Any]: - materialized_chunk = dict(chunk) - metadata = _metadata(materialized_chunk) - raw_path = _first_string(metadata.get("path"), materialized_chunk.get("path")) - publication_path = _publication_path(source=source, raw_path=raw_path) - file_path = _first_string( - metadata.get("file_path"), - metadata.get("filePath"), - materialized_chunk.get("file_path"), - materialized_chunk.get("path") if _is_media_chunk(materialized_chunk) else None, - ) - - metadata["path"] = publication_path - if file_path: - metadata["file_path"] = file_path - materialized_chunk["file_path"] = file_path - materialized_chunk["path"] = publication_path - materialized_chunk["metadata"] = metadata - return materialized_chunk - - -def _publication_path( - *, - source: DemoSourceDefinition, - raw_path: str | None, -) -> str: - prefix = f"Default_Root/{source.title}" - raw = str(raw_path or "").strip() - if not raw: - return prefix - - if "-->" in raw: - sections = [ - part.strip() - for part in raw.split("-->")[1:] - if part.strip() - ] - return "/".join([prefix, *sections]) if sections else prefix - - if raw.startswith("images/") or raw.startswith("tables/"): - return f"{prefix}/Assets/{raw}" - - parts = [part.strip() for part in raw.split("/") if part.strip()] - if len(parts) >= 2 and parts[0] == "Default_Root": - return raw - return prefix - - -def _normalize_asset_path(asset_path: str) -> Path | None: - normalized = str(asset_path or "").strip().replace("\\", "/").lstrip("/") - parts = [part for part in normalized.split("/") if part and part != "."] - if not parts or parts[0] not in _ASSET_DIRECTORY_NAMES: - return None - if any(part == ".." or part.startswith(".") for part in parts): - return None - return Path(*parts) - - -def _citation_payload( - *, - source: DemoSourceDefinition, - citation: DemoCitationDefinition, -) -> dict[str, Any]: - chunk = _resolve_citation_chunk(source=source, citation=citation) - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "canonical_chunk_id": _canonical_chunk_id(source, chunk), - "chunk_id": chunk["chunk_id"], - "chunk_type": _normalize_chunk_type(chunk.get("type")), - "content": citation.content, - "description": citation.description, - "source": { - "document_id": source.canonical_document_id, - "source_file_name": source.title, - "section_path": citation.section_path, - }, - } - - -def _resolve_citation_chunk( - *, - source: DemoSourceDefinition, - citation: DemoCitationDefinition, -) -> dict[str, Any]: - chunks = _load_source_chunks(source) - normalized_content = _normalize_text(citation.content) - if normalized_content: - for chunk in chunks: - if normalized_content in _normalize_text(str(chunk.get("content") or "")): - return chunk - - for chunk in chunks: - if str(chunk.get("path") or "") == citation.section_path: - return chunk - - raise ValueError( - "Demo citation does not resolve to a canonical chunk: " - f"demo_source_id={source.demo_source_id}, section_path={citation.section_path}" - ) - - -def _chunk_payload( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> dict[str, Any]: - metadata = _metadata(chunk) - file_path = _first_string( - metadata.get("file_path"), - metadata.get("filePath"), - chunk.get("file_path"), - chunk.get("path") if _is_media_chunk(chunk) else None, - ) - return { - "id": _canonical_chunk_id(source, chunk), - "chunk_id": chunk["chunk_id"], - "chunk_type": _normalize_chunk_type(chunk.get("type")), - "content": str(chunk.get("content") or ""), - "section_path": str(chunk.get("path") or "") or None, - "source_chunk_path": str(chunk.get("path") or "") or None, - "file_path": file_path, - "sort_order": _sort_order(source=source, chunk=chunk), - "metadata": metadata, - "asset_url": _asset_url(source=source, file_path=file_path), - "created_at": None, - } - - -def _sort_order( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> int: - try: - return _load_source_chunks(source).index(chunk) - except ValueError: - return 0 - - -def _metadata(chunk: dict[str, Any]) -> dict[str, Any]: - metadata = chunk.get("metadata") - return dict(metadata) if isinstance(metadata, dict) else {} - - -def _is_media_chunk(chunk: dict[str, Any]) -> bool: - return _normalize_chunk_type(chunk.get("type")) in {"image", "table"} - - -def _canonical_chunk_id( - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> str: - return f"{source.demo_source_id}:{chunk['chunk_id']}" - - -def _asset_url( - *, - source: DemoSourceDefinition, - file_path: str | None, -) -> str | None: - if not file_path: - return None - encoded_path = quote(file_path, safe="/") - return f"/api/v1/demo/sources/{source.demo_source_id}/assets/{encoded_path}" - - -def _normalize_chunk_type(value: object) -> str: - raw = str(value or "").strip().split("\n", 1)[0].lower() - return raw if raw in {"text", "image", "table"} else "text" - - -def _first_string(*values: object) -> str | None: - for value in values: - if isinstance(value, str) and value.strip(): - return value.strip() - return None - - -def _normalize_text(value: str) -> str: - return " ".join(value.lower().split()) - - -def _get_source_definition(demo_source_id: str) -> DemoSourceDefinition | None: - return next( - ( - source - for source in _DEMO_SOURCE_DEFINITIONS - if source.demo_source_id == demo_source_id - ), - None, - ) - - -def _require_source_definition(demo_source_id: str) -> DemoSourceDefinition: - source = _get_source_definition(demo_source_id) - if source is None: - raise KeyError(demo_source_id) - return source - - -def _source_directory(source: DemoSourceDefinition) -> Path: - return _DATA_ROOT / source.asset_directory - - -@lru_cache(maxsize=8) -def _load_source_chunks(source: DemoSourceDefinition) -> tuple[dict[str, Any], ...]: - chunks_path = _source_directory(source) / "chunks.json" - with chunks_path.open("r", encoding="utf-8") as file: - payload = json.load(file) - - chunks = payload.get("chunks") if isinstance(payload, dict) else None - if not isinstance(chunks, list): - return () - - return tuple( - dict(chunk) - for chunk in chunks - if isinstance(chunk, dict) and isinstance(chunk.get("chunk_id"), str) - ) - - def _utc_now() -> datetime: return datetime.now(timezone.utc).replace(tzinfo=None) diff --git a/apps/api/app/services/demo_source_catalog.py b/apps/api/app/services/demo_source_catalog.py new file mode 100644 index 000000000..9a7ecef40 --- /dev/null +++ b/apps/api/app/services/demo_source_catalog.py @@ -0,0 +1,480 @@ +"""Canonical Demo Source catalog and payload shaping.""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import Any +from urllib.parse import quote + + +@dataclass(frozen=True) +class DemoCitationDefinition: + section_path: str + description: str + content: str + + +@dataclass(frozen=True) +class DemoExampleDefinition: + id: str + question: str + answer: str + citations: tuple[DemoCitationDefinition, ...] + + +@dataclass(frozen=True) +class DemoSourceDefinition: + demo_source_id: str + canonical_document_id: str + title: str + mime_type: str + size_bytes: int + asset_directory: str + chunk_count: int + examples: tuple[DemoExampleDefinition, ...] + + +_DATA_ROOT = Path(__file__).resolve().parents[1] / "data" / "demo_documents" +_ASSET_DIRECTORY_NAMES = frozenset({"images", "tables"}) +_DEMO_SOURCE_DEFINITIONS: tuple[DemoSourceDefinition, ...] = ( + DemoSourceDefinition( + demo_source_id="demo-tsla-q4-2025", + canonical_document_id="demo-doc-tsla-q4-2025", + title="TSLA-Q4-2025-Update.pdf", + mime_type="application/pdf", + size_bytes=5_648_867, + asset_directory="tsla-q4-2025", + chunk_count=70, + examples=( + DemoExampleDefinition( + id="demo-tsla-q4-2025-xai", + question="What does the document say about Tesla's xAI investment?", + answer=( + "Tesla entered an agreement on January 16, 2026 to invest " + "approximately $2 billion in xAI Series E Preferred Stock.\n\n" + "The document also says Tesla and xAI entered a framework " + "agreement to evaluate AI collaboration, with the investment " + "expected to close in Q1 2026 subject to customary regulatory " + "conditions." + ), + citations=( + DemoCitationDefinition( + section_path=( + "Default_Root/TSLA-Q4-2025-Update.pdf-->OTHER UPDATES" + ), + description="xAI investment", + content=( + "On January 16, 2026, Tesla entered into an agreement " + "to invest approximately" + ), + ), + ), + ), + DemoExampleDefinition( + id="demo-tsla-q4-2025-energy-storage", + question="What does the document say about energy storage?", + answer=( + "Tesla achieved its highest quarterly energy storage " + "deployments, driven by record Megapack deployments.\n\n" + "Energy gross profit reached a record $1.1 billion, marking " + "the fifth consecutive record quarter.\n\n" + "Tesla also plans to begin Megapack 3 and Megablock " + "production at Megafactory Houston in 2026." + ), + citations=( + DemoCitationDefinition( + section_path=( + "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->" + "Energy generation and storage" + ), + description="Storage deployment growth", + content=( + "We achieved our highest quarterly energy storage " + "deployments, driven by record Megapack deployments." + ), + ), + ), + ), + DemoExampleDefinition( + id="demo-tsla-q4-2025-production-plans", + question="What production plans does Tesla mention for 2026?", + answer=( + "Tesla says Cybercab, Tesla Semi, and Megapack 3 are on " + "schedule for volume production starting in 2026.\n\n" + "The same product update also notes that first-generation " + "Optimus production lines are being installed before volume " + "production." + ), + citations=( + DemoCitationDefinition( + section_path=( + "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->" + "Product" + ), + description="2026 production plans", + content=( + "Cybercab, Tesla Semi and Megapack 3 are on schedule " + "for volume production starting in 2026." + ), + ), + ), + ), + ), + ), +) + + +class DemoSourceCatalog: + def get_catalog(self) -> dict[str, Any]: + return { + "sources": [ + self._source_catalog_payload(source) + for source in _DEMO_SOURCE_DEFINITIONS + ], + } + + def list_chunks( + self, + *, + demo_source_id: str, + page: int, + page_size: int, + ) -> dict[str, Any] | None: + source = self.get_source(demo_source_id) + if source is None: + return None + + chunks = _load_source_chunks(source) + start = (page - 1) * page_size + page_chunks = chunks[start : start + page_size] + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "title": source.title, + "mime_type": source.mime_type, + "chunks": [ + _chunk_payload(source=source, chunk=chunk) + for chunk in page_chunks + ], + "pagination": { + "page": page, + "page_size": page_size, + "total": len(chunks), + "total_pages": math.ceil(len(chunks) / page_size) if chunks else 0, + }, + } + + def get_chunk( + self, + *, + demo_source_id: str, + demo_chunk_id: str, + ) -> dict[str, Any] | None: + source = self.get_source(demo_source_id) + if source is None: + return None + + for chunk in _load_source_chunks(source): + if demo_chunk_id in {_canonical_chunk_id(source, chunk), chunk["chunk_id"]}: + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "chunk": _chunk_payload(source=source, chunk=chunk), + } + + return None + + def get_original_file_path(self, *, demo_source_id: str) -> Path | None: + source = self.get_source(demo_source_id) + if source is None: + return None + + file_path = self.source_directory(source) / "original.pdf" + return file_path if file_path.is_file() else None + + def get_asset_file_path( + self, + *, + demo_source_id: str, + asset_path: str, + ) -> Path | None: + source = self.get_source(demo_source_id) + if source is None: + return None + + source_directory = self.source_directory(source).resolve() + normalized_asset_path = _normalize_asset_path(asset_path) + if normalized_asset_path is None: + return None + + candidate = (source_directory / normalized_asset_path).resolve() + if not candidate.is_relative_to(source_directory): + return None + + return candidate if candidate.is_file() else None + + def require_source(self, demo_source_id: str) -> DemoSourceDefinition: + source = self.get_source(demo_source_id) + if source is None: + raise KeyError(demo_source_id) + return source + + def get_source(self, demo_source_id: str) -> DemoSourceDefinition | None: + return next( + ( + source + for source in _DEMO_SOURCE_DEFINITIONS + if source.demo_source_id == demo_source_id + ), + None, + ) + + def source_directory(self, source: DemoSourceDefinition) -> Path: + return _DATA_ROOT / source.asset_directory + + def publication_chunks(self, source: DemoSourceDefinition) -> list[dict[str, Any]]: + return [ + _publication_chunk(source=source, chunk=chunk) + for chunk in _load_source_chunks(source) + ] + + def _source_catalog_payload(self, source: DemoSourceDefinition) -> dict[str, Any]: + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "title": source.title, + "mime_type": source.mime_type, + "size_bytes": source.size_bytes, + "status": "ready", + "chunk_count": source.chunk_count, + "original_file": { + "url": f"/api/v1/demo/sources/{source.demo_source_id}/original", + "mime_type": source.mime_type, + "size_bytes": source.size_bytes, + "can_download": False, + }, + "examples": [ + self._example_payload(source=source, example=example) + for example in source.examples + ], + } + + def _example_payload( + self, + *, + source: DemoSourceDefinition, + example: DemoExampleDefinition, + ) -> dict[str, Any]: + return { + "id": example.id, + "question": example.question, + "answer": example.answer, + "citations": [ + _citation_payload(source=source, citation=citation) + for citation in example.citations + ], + } + + +def _publication_chunk( + *, + source: DemoSourceDefinition, + chunk: dict[str, Any], +) -> dict[str, Any]: + materialized_chunk = dict(chunk) + metadata = _metadata(materialized_chunk) + raw_path = _first_string(metadata.get("path"), materialized_chunk.get("path")) + publication_path = _publication_path(source=source, raw_path=raw_path) + file_path = _first_string( + metadata.get("file_path"), + metadata.get("filePath"), + materialized_chunk.get("file_path"), + materialized_chunk.get("path") if _is_media_chunk(materialized_chunk) else None, + ) + + metadata["path"] = publication_path + if file_path: + metadata["file_path"] = file_path + materialized_chunk["file_path"] = file_path + materialized_chunk["path"] = publication_path + materialized_chunk["metadata"] = metadata + return materialized_chunk + + +def _publication_path( + *, + source: DemoSourceDefinition, + raw_path: str | None, +) -> str: + prefix = f"Default_Root/{source.title}" + raw = str(raw_path or "").strip() + if not raw: + return prefix + + if "-->" in raw: + sections = [part.strip() for part in raw.split("-->")[1:] if part.strip()] + return "/".join([prefix, *sections]) if sections else prefix + + if raw.startswith("images/") or raw.startswith("tables/"): + return f"{prefix}/Assets/{raw}" + + parts = [part.strip() for part in raw.split("/") if part.strip()] + if len(parts) >= 2 and parts[0] == "Default_Root": + return raw + return prefix + + +def _normalize_asset_path(asset_path: str) -> Path | None: + normalized = str(asset_path or "").strip().replace("\\", "/").lstrip("/") + parts = [part for part in normalized.split("/") if part and part != "."] + if not parts or parts[0] not in _ASSET_DIRECTORY_NAMES: + return None + if any(part == ".." or part.startswith(".") for part in parts): + return None + return Path(*parts) + + +def _citation_payload( + *, + source: DemoSourceDefinition, + citation: DemoCitationDefinition, +) -> dict[str, Any]: + chunk = _resolve_citation_chunk(source=source, citation=citation) + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "canonical_chunk_id": _canonical_chunk_id(source, chunk), + "chunk_id": chunk["chunk_id"], + "chunk_type": _normalize_chunk_type(chunk.get("type")), + "content": citation.content, + "description": citation.description, + "source": { + "document_id": source.canonical_document_id, + "source_file_name": source.title, + "section_path": citation.section_path, + }, + } + + +def _resolve_citation_chunk( + *, + source: DemoSourceDefinition, + citation: DemoCitationDefinition, +) -> dict[str, Any]: + chunks = _load_source_chunks(source) + normalized_content = _normalize_text(citation.content) + if normalized_content: + for chunk in chunks: + if normalized_content in _normalize_text(str(chunk.get("content") or "")): + return chunk + + for chunk in chunks: + if str(chunk.get("path") or "") == citation.section_path: + return chunk + + raise ValueError( + "Demo citation does not resolve to a canonical chunk: " + f"demo_source_id={source.demo_source_id}, section_path={citation.section_path}" + ) + + +def _chunk_payload( + *, + source: DemoSourceDefinition, + chunk: dict[str, Any], +) -> dict[str, Any]: + metadata = _metadata(chunk) + file_path = _first_string( + metadata.get("file_path"), + metadata.get("filePath"), + chunk.get("file_path"), + chunk.get("path") if _is_media_chunk(chunk) else None, + ) + return { + "id": _canonical_chunk_id(source, chunk), + "chunk_id": chunk["chunk_id"], + "chunk_type": _normalize_chunk_type(chunk.get("type")), + "content": str(chunk.get("content") or ""), + "section_path": str(chunk.get("path") or "") or None, + "source_chunk_path": str(chunk.get("path") or "") or None, + "file_path": file_path, + "sort_order": _sort_order(source=source, chunk=chunk), + "metadata": metadata, + "asset_url": _asset_url(source=source, file_path=file_path), + "created_at": None, + } + + +def _sort_order( + *, + source: DemoSourceDefinition, + chunk: dict[str, Any], +) -> int: + try: + return _load_source_chunks(source).index(chunk) + except ValueError: + return 0 + + +def _metadata(chunk: dict[str, Any]) -> dict[str, Any]: + metadata = chunk.get("metadata") + return dict(metadata) if isinstance(metadata, dict) else {} + + +def _is_media_chunk(chunk: dict[str, Any]) -> bool: + return _normalize_chunk_type(chunk.get("type")) in {"image", "table"} + + +def _canonical_chunk_id( + source: DemoSourceDefinition, + chunk: dict[str, Any], +) -> str: + return f"{source.demo_source_id}:{chunk['chunk_id']}" + + +def _asset_url( + *, + source: DemoSourceDefinition, + file_path: str | None, +) -> str | None: + if not file_path: + return None + encoded_path = quote(file_path, safe="/") + return f"/api/v1/demo/sources/{source.demo_source_id}/assets/{encoded_path}" + + +def _normalize_chunk_type(value: object) -> str: + raw = str(value or "").strip().split("\n", 1)[0].lower() + return raw if raw in {"text", "image", "table"} else "text" + + +def _first_string(*values: object) -> str | None: + for value in values: + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _normalize_text(value: str) -> str: + return " ".join(value.lower().split()) + + +@lru_cache(maxsize=8) +def _load_source_chunks(source: DemoSourceDefinition) -> tuple[dict[str, Any], ...]: + chunks_path = (_DATA_ROOT / source.asset_directory) / "chunks.json" + with chunks_path.open("r", encoding="utf-8") as file: + payload = json.load(file) + + chunks = payload.get("chunks") if isinstance(payload, dict) else None + if not isinstance(chunks, list): + return () + + return tuple( + dict(chunk) + for chunk in chunks + if isinstance(chunk, dict) and isinstance(chunk.get("chunk_id"), str) + ) diff --git a/apps/api/app/services/document_ingestion_service.py b/apps/api/app/services/document_ingestion_service.py new file mode 100644 index 000000000..12f69a5b6 --- /dev/null +++ b/apps/api/app/services/document_ingestion_service.py @@ -0,0 +1,605 @@ +from __future__ import annotations + +import os +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import cast +from urllib.parse import urlparse + +from app.repositories.job_repository import JobRepository +from app.services.job_document_scope_service import ( + find_active_job_for_document, + is_active_document_job_unique_violation, + raise_document_ingestion_conflict, + resolve_effective_document_scope, +) +from app.services.job_read_service import check_job_permission +from app.services.job_response_projection import to_job_status_value +from app.services.knowledge.kb_orchestrator import KBOrchestrator +from app.services.rate_limit.data_structures import CurrentUser +from app.services.rate_limit.dependencies import enforce_job_creation_capacity +from app.services.state_machine import JobStateMachine +from fastapi import Request +from loguru import logger +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + ConflictException, + JobOperationException, + NotFoundException, + PermissionDeniedException, + RateLimitException, + UnavailableException, + ValidationException, +) +from shared.core.exceptions.webhook_exceptions import WebhookConfigException +from shared.core.state_machine.states import JobStatus +from shared.models.database.job import Job +from shared.models.schemas.job import ConfirmUploadRequest, JobCreate, JobResponse +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.redis import JobInfoRedisService, RedisServiceFactory +from shared.services.redis.job_metadata_service import JobMetadataService +from shared.services.storage.file_upload_service import FileUploadService +from shared.utils.url_file_type import resolve_file_extension_async +from shared.utils.url_security import validate_http_url_and_resolve_ip_async + +JOB_TYPE_KB_MANAGEMENT = "kb_management" +JobMetadata = dict[str, object] +UploadHeaders = dict[str, str] + + +@dataclass(frozen=True) +class ResolvedDocumentIngestionScope: + job_metadata: JobMetadata + document_id: str + namespace: str + + +def _get_supported_formats() -> str: + return ", ".join(sorted(settings.get_supported_extensions())) + + +def _is_supported_file_name(file_name: str) -> bool: + if not file_name: + return False + file_extension = os.path.splitext(file_name)[1].lower() + return file_extension in settings.get_supported_extensions() + + +def _build_job_response( + *, + job_id: str, + job: Job, + source_type: str, + data_id: str | None, + namespace: str | None = None, + document_id: str | None = None, + upload_url: str | None = None, + upload_headers: UploadHeaders | None = None, + expires_in: int | None = None, +) -> JobResponse: + return JobResponse( + job_id=job_id, + status=to_job_status_value(job.status), + source_type=source_type, + data_id=data_id, + namespace=namespace, + document_id=document_id, + created_at=job.created_at, + upload_url=upload_url, + upload_headers=upload_headers, + expires_in=expires_in, + ) + + +def _resolve_url_source_file_name(*, source_url: str, file_extension: str) -> str: + parsed_url = urlparse(source_url) + url_basename = str(os.path.basename(parsed_url.path)) + if url_basename and os.path.splitext(url_basename)[1].lower() == file_extension: + return url_basename + if url_basename: + return f"{url_basename}{file_extension}" + return f"url_file_{uuid.uuid4().hex[:8]}{file_extension}" + + +def _schedule_url_upload(*, job_id: str, source_url: str, user_id: str) -> None: + from shared.core.celery_app import get_celery_app + + celery_app = get_celery_app() + upload_url_file_task = celery_app.signature( + "app.core.tasks.kb_tasks.upload_url_file_task" + ) + upload_url_file_task.apply_async( + args=[job_id, source_url, user_id], + kwargs={"job_type": JOB_TYPE_KB_MANAGEMENT}, + ) + + +async def _transition_job_to_uploaded( + db: AsyncSession, + *, + job_id: str, + trigger: str = "manual_upload_completed", +) -> None: + state_machine = JobStateMachine() + await state_machine.transition( + db, + job_id, + JobStatus.PENDING.value, + trigger, + None, + "system", + ) + + +async def _start_job_workflow( + db: AsyncSession, + *, + job_id: str, + job_type: str, + source_type: str, + user_id: str, + file_path: str | None = None, + file_url: str | None = None, +) -> None: + if job_type == JOB_TYPE_KB_MANAGEMENT: + orchestrator = KBOrchestrator() + await orchestrator.start_workflow( + db=db, + job_id=job_id, + source_type=source_type, + file_path=file_path, + file_url=file_url, + user_id=user_id, + ) + return + + raise ValidationException( + user_message="Unsupported job type", + violations=[ + { + "field": "job_type", + "description": f"Job type '{job_type}' is not supported", + } + ], + ) + + +class DocumentIngestionService: + def __init__( + self, + *, + job_repository: JobRepository | None = None, + file_upload_service: FileUploadService | None = None, + ) -> None: + self._job_repository = job_repository or JobRepository() + self._file_upload_service = file_upload_service or FileUploadService() + + async def create_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + current_user: CurrentUser, + request: Request, + ) -> JobResponse: + try: + job_id = f"job_{uuid.uuid4().hex[:12]}" + await self._validate_create_payload(payload) + scope = await self._resolve_scope( + db, + payload=payload, + current_user=current_user, + ) + + await enforce_job_creation_capacity( + request=request, + db=db, + current_user=current_user, + ) + + if payload.source_type == "file": + return await self._create_file_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + scope=scope, + ) + return await self._create_url_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + scope=scope, + ) + except NotFoundException: + raise + except ValidationException: + raise + except ConflictException: + raise + except WebhookConfigException: + raise + except (RateLimitException, UnavailableException): + raise + except JobOperationException: + raise + except Exception as exc: + logger.error(f"Failed to create job: {exc}") + raise JobOperationException( + internal_message=f"Job creation failed: {str(exc)}" + ) + + async def confirm_upload( + self, + db: AsyncSession, + *, + job_id: str, + request_payload: ConfirmUploadRequest | None, + user_id: str, + ) -> dict[str, str]: + del request_payload + + try: + job = await self._job_repository.get_job_by_id(db, job_id) + check_job_permission(job, user_id, job_id) + assert job is not None + + logger.info(f"Confirm upload - Job {job_id} current status: {job.status}") + if job.status not in [JobStatus.PENDING.value, JobStatus.WAITING_FILE.value]: + logger.info(f"Job {job_id} already processed, status: {job.status}") + return {"message": "Job status already updated"} + + if not job.s3_key: + raise ValidationException( + user_message="Job is missing S3 key information", + violations=[ + { + "field": "s3_key", + "description": "S3 key not set for this job", + } + ], + ) + + file_info = await self._file_upload_service.verify_s3_file_exists(job.s3_key) + if not bool(file_info.get("exists")): + raise ValidationException( + user_message="S3 file does not exist, please upload the file first", + violations=[ + {"field": "file", "description": "File not found in S3"} + ], + ) + + await _transition_job_to_uploaded(db, job_id=job_id) + await _start_job_workflow( + db=db, + job_id=job_id, + job_type=job.job_type, + source_type="file", + user_id=user_id, + ) + return {"message": "File upload confirmed; processing started"} + except NotFoundException: + raise + except PermissionDeniedException: + raise + except ValidationException: + raise + except Exception as exc: + logger.error(f"Failed to confirm upload: {exc}") + raise JobOperationException( + internal_message=f"Failed to confirm upload: {str(exc)}" + ) + + async def _validate_create_payload(self, payload: JobCreate) -> None: + if payload.source_type == "file" and not payload.file_name: + raise ValidationException( + user_message="file_name is required when source_type is 'file'", + violations=[ + { + "field": "file_name", + "description": "Required for file source type", + } + ], + ) + if payload.source_type == "url" and not payload.source_url: + raise ValidationException( + user_message="source_url is required when source_type is 'url'", + violations=[ + { + "field": "source_url", + "description": "Required for url source type", + } + ], + ) + + if payload.webhook and payload.webhook.url: + validation_result = await validate_http_url_and_resolve_ip_async( + payload.webhook.url, + ) + if not validation_result.is_valid: + raise WebhookConfigException( + user_message="Invalid webhook URL", + internal_message=( + "Webhook validation failed: " + f"{validation_result.error_message}" + ), + ) + + if ( + payload.source_type == "file" + and payload.file_name + and not _is_supported_file_name(payload.file_name) + ): + supported_formats = _get_supported_formats() + raise ValidationException( + user_message=( + "Unsupported file type. Supported formats: " + f"{supported_formats}" + ), + violations=[ + {"field": "file_name", "description": "File type not supported"} + ], + ) + + if payload.source_type == "url": + assert payload.source_url is not None + file_extension = await resolve_file_extension_async(payload.source_url) + if not file_extension: + supported_formats = _get_supported_formats() + raise ValidationException( + user_message=( + "Unsupported URL file type. Supported formats: " + f"{supported_formats}" + ), + violations=[ + { + "field": "source_url", + "description": "URL file type not supported", + } + ], + ) + + async def _resolve_scope( + self, + db: AsyncSession, + *, + payload: JobCreate, + current_user: CurrentUser, + ) -> ResolvedDocumentIngestionScope: + job_metadata = cast(JobMetadata, JobMetadataHelper.create_from_request(payload)) + requested_document_id = cast(str | None, job_metadata.get("document_id")) + if requested_document_id: + active_job = await find_active_job_for_document( + db, + user_id=current_user.user_id, + document_id=requested_document_id, + ) + if active_job is not None: + raise_document_ingestion_conflict( + document_id=requested_document_id, + active_job_id=active_job.job_id, + ) + + ( + effective_document_id, + effective_namespace, + ) = await resolve_effective_document_scope( + db, + user_id=current_user.user_id, + document_id=requested_document_id, + requested_namespace=cast(str | None, payload.namespace), + ) + + if not requested_document_id: + active_job = await find_active_job_for_document( + db, + user_id=current_user.user_id, + document_id=effective_document_id, + ) + if active_job is not None: + raise_document_ingestion_conflict( + document_id=effective_document_id, + active_job_id=active_job.job_id, + ) + + job_metadata["document_id"] = effective_document_id + job_metadata["namespace"] = effective_namespace + return ResolvedDocumentIngestionScope( + job_metadata=job_metadata, + document_id=effective_document_id, + namespace=effective_namespace, + ) + + async def _create_waiting_job( + self, + db: AsyncSession, + *, + job_id: str, + user_id: str, + source_type: str, + webhook_url: str | None, + job_metadata: JobMetadata, + s3_key: str, + document_id: str, + ) -> Job: + try: + job = await self._job_repository.create_job( + db=db, + job_id=job_id, + user_id=user_id, + job_type=JOB_TYPE_KB_MANAGEMENT, + source_type=source_type, + file_path=None, + webhook_url=webhook_url, + metadata=job_metadata, + initial_state=JobStatus.WAITING_FILE.value, + s3_key=s3_key, + ) + except IntegrityError as exc: + if is_active_document_job_unique_violation(exc): + raise_document_ingestion_conflict(document_id=document_id) + raise + + if job is None: + raise JobOperationException( + internal_message="Failed to create job in database" + ) + return job + + async def _cache_job_creation_state( + self, + *, + job_id: str, + s3_key: str, + user_id: str, + webhook_enabled: bool, + source_type: str, + job_metadata: JobMetadata, + ) -> None: + redis_service = RedisServiceFactory.get_service() + metadata_service = JobMetadataService(redis_service) + await metadata_service.save_metadata(job_id, job_metadata) + + job_info_service = JobInfoRedisService(redis_service) + job_info: dict[str, object] = { + "job_id": job_id, + "s3_key": s3_key, + "user_id": user_id, + "webhook_enabled": webhook_enabled, + "job_type": JOB_TYPE_KB_MANAGEMENT, + "source_type": source_type, + "created_at": datetime.now(timezone.utc).isoformat(), + } + await job_info_service.save_job_info(job_id, job_info) + + async def _create_file_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + scope: ResolvedDocumentIngestionScope, + ) -> JobResponse: + assert payload.file_name is not None + file_extension = os.path.splitext(payload.file_name)[1] + s3_key = f"uploads/{job_id}{file_extension}" + scope.job_metadata["source_file_name"] = payload.file_name + scope.job_metadata["source_type"] = "file" + + job = await self._create_waiting_job( + db, + job_id=job_id, + user_id=current_user.user_id, + source_type="file", + webhook_url=payload.webhook.url if payload.webhook else None, + job_metadata=scope.job_metadata, + s3_key=s3_key, + document_id=scope.document_id, + ) + + upload_info = await self._file_upload_service.generate_upload_url( + job_id, + file_extension, + ) + upload_url = cast(str, upload_info["upload_url"]) + upload_headers = cast(UploadHeaders, upload_info["upload_headers"]) + expires_in = cast(int, upload_info["expires_in"]) + + await self._cache_job_creation_state( + job_id=job_id, + s3_key=s3_key, + user_id=current_user.user_id, + webhook_enabled=bool(payload.webhook and payload.webhook.url), + source_type="file", + job_metadata=scope.job_metadata, + ) + + logger.info(f"Job {job_id} upload_url returned to client: {upload_url}") + return _build_job_response( + job_id=job_id, + job=job, + source_type="file", + data_id=payload.data_id, + namespace=scope.namespace, + upload_url=upload_url, + upload_headers=upload_headers, + expires_in=expires_in, + ) + + async def _create_url_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + scope: ResolvedDocumentIngestionScope, + ) -> JobResponse: + assert payload.source_url is not None + file_extension = await resolve_file_extension_async(payload.source_url) + if not file_extension: + supported_formats = _get_supported_formats() + raise ValidationException( + user_message=( + "Unsupported URL file type. Supported formats: " + f"{supported_formats}" + ), + violations=[ + { + "field": "source_url", + "description": "URL file type not supported", + } + ], + ) + + source_file_name = _resolve_url_source_file_name( + source_url=payload.source_url, + file_extension=file_extension, + ) + s3_key = f"uploads/{job_id}{file_extension}" + scope.job_metadata.update( + { + "source_file_name": source_file_name, + "source_url": payload.source_url, + "source_type": "url", + } + ) + + job = await self._create_waiting_job( + db, + job_id=job_id, + user_id=current_user.user_id, + source_type="url", + webhook_url=payload.webhook.url if payload.webhook else None, + job_metadata=scope.job_metadata, + s3_key=s3_key, + document_id=scope.document_id, + ) + + await self._cache_job_creation_state( + job_id=job_id, + s3_key=s3_key, + user_id=current_user.user_id, + webhook_enabled=bool(payload.webhook and payload.webhook.url), + source_type="url", + job_metadata=scope.job_metadata, + ) + _schedule_url_upload( + job_id=job_id, + source_url=payload.source_url, + user_id=current_user.user_id, + ) + + return _build_job_response( + job_id=job_id, + job=job, + source_type="url", + data_id=payload.data_id, + namespace=scope.namespace, + ) diff --git a/apps/api/app/services/job_creation_service.py b/apps/api/app/services/job_creation_service.py deleted file mode 100644 index c5ed2b284..000000000 --- a/apps/api/app/services/job_creation_service.py +++ /dev/null @@ -1,463 +0,0 @@ -from __future__ import annotations - -import os -import uuid -from datetime import datetime, timezone -from typing import Optional, cast -from urllib.parse import urlparse - -from app.repositories.job_repository import JobRepository -from app.services.job_document_scope_service import ( - find_active_job_for_document, - is_active_document_job_unique_violation, - raise_document_ingestion_conflict, - resolve_effective_document_scope, -) -from app.services.rate_limit.data_structures import CurrentUser -from fastapi import Request -from loguru import logger -from sqlalchemy.exc import IntegrityError -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - ConflictException, - JobOperationException, - NotFoundException, - RateLimitException, - UnavailableException, - ValidationException, -) -from shared.core.exceptions.webhook_exceptions import WebhookConfigException -from shared.core.state_machine.states import JobStatus -from shared.models.schemas.job import JobCreate, JobResponse -from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.redis import JobInfoRedisService, RedisServiceFactory -from shared.services.redis.job_metadata_service import JobMetadataService -from shared.services.storage.file_upload_service import FileUploadService -from shared.utils.url_file_type import resolve_file_extension_async -from shared.utils.url_security import validate_http_url_and_resolve_ip_async - -JOB_TYPE_KB_MANAGEMENT = "kb_management" - - -def get_supported_formats() -> str: - return ", ".join(sorted(settings.get_supported_extensions())) - - -def validate_file_type(file_name: str) -> bool: - if not file_name: - return False - file_extension = os.path.splitext(file_name)[1].lower() - return file_extension in settings.get_supported_extensions() - - -def create_job_response( - job_id: str, - job, - source_type: str, - data_id: Optional[str], - namespace: Optional[str] = None, - document_id: Optional[str] = None, - upload_url: Optional[str] = None, - upload_headers: Optional[dict] = None, - expires_in: Optional[int] = None, -) -> JobResponse: - return JobResponse( - job_id=job_id, - status=job.status, - source_type=source_type, - data_id=data_id, - namespace=namespace, - document_id=document_id, - created_at=job.created_at, - upload_url=upload_url, - upload_headers=upload_headers, - expires_in=expires_in, - ) - - -async def _validate_create_job_payload(payload: JobCreate) -> None: - if payload.source_type == "file" and not payload.file_name: - raise ValidationException( - user_message="file_name is required when source_type is 'file'", - violations=[ - { - "field": "file_name", - "description": "Required for file source type", - } - ], - ) - if payload.source_type == "url" and not payload.source_url: - raise ValidationException( - user_message="source_url is required when source_type is 'url'", - violations=[ - { - "field": "source_url", - "description": "Required for url source type", - } - ], - ) - - if payload.webhook and payload.webhook.url: - validation_result = await validate_http_url_and_resolve_ip_async( - payload.webhook.url, - ) - if not validation_result.is_valid: - raise WebhookConfigException( - user_message="Invalid webhook URL", - internal_message=f"Webhook validation failed: {validation_result.error_message}", - ) - - if ( - payload.source_type == "file" - and payload.file_name - and not validate_file_type(payload.file_name) - ): - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported file type. Supported formats: {supported_formats}", - violations=[ - {"field": "file_name", "description": "File type not supported"} - ], - ) - - if payload.source_type == "url": - assert payload.source_url is not None - file_extension = await resolve_file_extension_async(payload.source_url) - if not file_extension: - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported URL file type. Supported formats: {supported_formats}", - violations=[ - { - "field": "source_url", - "description": "URL file type not supported", - } - ], - ) - - -async def _resolve_job_metadata( - db: AsyncSession, - *, - payload: JobCreate, - current_user: CurrentUser, -) -> tuple[dict, str, str]: - job_metadata = JobMetadataHelper.create_from_request(payload) - requested_document_id = cast(Optional[str], job_metadata.get("document_id")) - if requested_document_id: - active_job = await find_active_job_for_document( - db, - user_id=current_user.user_id, - document_id=requested_document_id, - ) - if active_job is not None: - raise_document_ingestion_conflict( - document_id=requested_document_id, - active_job_id=active_job.job_id, - ) - ( - effective_document_id, - effective_namespace, - ) = await resolve_effective_document_scope( - db, - user_id=current_user.user_id, - document_id=requested_document_id, - requested_namespace=cast(Optional[str], payload.namespace), - ) - if not requested_document_id: - active_job = await find_active_job_for_document( - db, - user_id=current_user.user_id, - document_id=effective_document_id, - ) - if active_job is not None: - raise_document_ingestion_conflict( - document_id=effective_document_id, - active_job_id=active_job.job_id, - ) - job_metadata["document_id"] = effective_document_id - job_metadata["namespace"] = effective_namespace - return job_metadata, effective_document_id, effective_namespace - - -async def _cache_job_creation_state( - *, - job_id: str, - s3_key: str, - user_id: str, - webhook_enabled: bool, - source_type: str, - job_metadata: dict, -) -> None: - redis_service = RedisServiceFactory.get_service() - metadata_service = JobMetadataService(redis_service) - await metadata_service.save_metadata(job_id, job_metadata) - - job_info_service = JobInfoRedisService(redis_service) - job_info = { - "job_id": job_id, - "s3_key": s3_key, - "user_id": user_id, - "webhook_enabled": webhook_enabled, - "job_type": JOB_TYPE_KB_MANAGEMENT, - "source_type": source_type, - "created_at": datetime.now(timezone.utc).isoformat(), - } - await job_info_service.save_job_info(job_id, job_info) - - -async def _create_waiting_job( - db: AsyncSession, - *, - job_id: str, - user_id: str, - source_type: str, - webhook_url: str | None, - job_metadata: dict, - s3_key: str, - effective_document_id: str, -): - job_repo = JobRepository() - try: - return await job_repo.create_job( - db=db, - job_id=job_id, - user_id=user_id, - job_type=JOB_TYPE_KB_MANAGEMENT, - source_type=source_type, - file_path=None, - webhook_url=webhook_url, - metadata=job_metadata, - initial_state=JobStatus.WAITING_FILE.value, - s3_key=s3_key, - ) - except IntegrityError as exc: - if is_active_document_job_unique_violation(exc): - raise_document_ingestion_conflict(document_id=effective_document_id) - raise - - -def _resolve_url_source_file_name(*, source_url: str, file_extension: str) -> str: - parsed_url = urlparse(source_url) - url_basename = str(os.path.basename(parsed_url.path)) - if url_basename and os.path.splitext(url_basename)[1].lower() == file_extension: - return url_basename - if url_basename: - return f"{url_basename}{file_extension}" - return f"url_file_{uuid.uuid4().hex[:8]}{file_extension}" - - -def _schedule_url_upload(*, job_id: str, source_url: str, user_id: str) -> None: - from shared.core.celery_app import get_celery_app - - celery_app = get_celery_app() - upload_url_file_task = celery_app.signature( - "app.core.tasks.kb_tasks.upload_url_file_task" - ) - upload_url_file_task.apply_async( - args=[job_id, source_url, user_id], - kwargs={ - "job_type": JOB_TYPE_KB_MANAGEMENT, - }, - ) - - -async def _create_file_job( - db: AsyncSession, - *, - payload: JobCreate, - job_id: str, - current_user: CurrentUser, - job_metadata: dict, - effective_document_id: str, - effective_namespace: str, -) -> JobResponse: - assert payload.file_name is not None - file_extension = os.path.splitext(payload.file_name)[1] - s3_key = f"uploads/{job_id}{file_extension}" - job_metadata["source_file_name"] = payload.file_name - job_metadata["source_type"] = "file" - - job = await _create_waiting_job( - db, - job_id=job_id, - user_id=current_user.user_id, - source_type="file", - webhook_url=payload.webhook.url if payload.webhook else None, - job_metadata=job_metadata, - s3_key=s3_key, - effective_document_id=effective_document_id, - ) - if not job: - raise JobOperationException( - internal_message="Failed to create job in database" - ) - - upload_service = FileUploadService() - upload_info = await upload_service.generate_upload_url(job_id, file_extension) - - await _cache_job_creation_state( - job_id=job_id, - s3_key=s3_key, - user_id=current_user.user_id, - webhook_enabled=bool(payload.webhook and payload.webhook.url), - source_type="file", - job_metadata=job_metadata, - ) - - logger.info( - f"Job {job_id} upload_url returned to client: {upload_info['upload_url']}" - ) - return create_job_response( - job_id=job_id, - job=job, - source_type="file", - data_id=payload.data_id, - namespace=effective_namespace, - upload_url=upload_info["upload_url"], - upload_headers=upload_info["upload_headers"], - expires_in=upload_info["expires_in"], - ) - - -async def _create_url_job( - db: AsyncSession, - *, - payload: JobCreate, - job_id: str, - current_user: CurrentUser, - job_metadata: dict, - effective_document_id: str, - effective_namespace: str, -) -> JobResponse: - assert payload.source_url is not None - file_extension = await resolve_file_extension_async(payload.source_url) - if not file_extension: - supported_formats = get_supported_formats() - raise ValidationException( - user_message=f"Unsupported URL file type. Supported formats: {supported_formats}", - violations=[ - { - "field": "source_url", - "description": "URL file type not supported", - } - ], - ) - - source_file_name = _resolve_url_source_file_name( - source_url=payload.source_url, - file_extension=file_extension, - ) - s3_key = f"uploads/{job_id}{file_extension}" - job_metadata.update( - { - "source_file_name": source_file_name, - "source_url": payload.source_url, - "source_type": "url", - } - ) - - job = await _create_waiting_job( - db, - job_id=job_id, - user_id=current_user.user_id, - source_type="url", - webhook_url=payload.webhook.url if payload.webhook else None, - job_metadata=job_metadata, - s3_key=s3_key, - effective_document_id=effective_document_id, - ) - if not job: - raise JobOperationException( - internal_message="Failed to create URL job in database" - ) - - await _cache_job_creation_state( - job_id=job_id, - s3_key=s3_key, - user_id=current_user.user_id, - webhook_enabled=bool(payload.webhook and payload.webhook.url), - source_type="url", - job_metadata=job_metadata, - ) - _schedule_url_upload( - job_id=job_id, - source_url=payload.source_url, - user_id=current_user.user_id, - ) - - return create_job_response( - job_id=job_id, - job=job, - source_type="url", - data_id=payload.data_id, - namespace=effective_namespace, - ) - - -async def create_job_from_request( - db: AsyncSession, - *, - payload: JobCreate, - current_user: CurrentUser, - enforce_capacity, - request: Request, -) -> JobResponse: - try: - job_id = f"job_{uuid.uuid4().hex[:12]}" - await _validate_create_job_payload(payload) - ( - job_metadata, - effective_document_id, - effective_namespace, - ) = await _resolve_job_metadata( - db, - payload=payload, - current_user=current_user, - ) - - await enforce_capacity( - request=request, - db=db, - current_user=current_user, - ) - - if payload.source_type == "file": - return await _create_file_job( - db, - payload=payload, - job_id=job_id, - current_user=current_user, - job_metadata=job_metadata, - effective_document_id=effective_document_id, - effective_namespace=effective_namespace, - ) - return await _create_url_job( - db, - payload=payload, - job_id=job_id, - current_user=current_user, - job_metadata=job_metadata, - effective_document_id=effective_document_id, - effective_namespace=effective_namespace, - ) - - except NotFoundException: - raise - except ValidationException: - raise - except ConflictException: - raise - except WebhookConfigException: - raise - except (RateLimitException, UnavailableException): - raise - except JobOperationException: - raise - except Exception as exc: - logger.error(f"Failed to create job: {exc}") - raise JobOperationException( - internal_message=f"Job creation failed: {str(exc)}" - ) diff --git a/apps/api/app/services/job_upload_confirmation_service.py b/apps/api/app/services/job_upload_confirmation_service.py deleted file mode 100644 index 421915d58..000000000 --- a/apps/api/app/services/job_upload_confirmation_service.py +++ /dev/null @@ -1,124 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -from app.repositories.job_repository import JobRepository -from app.services.job_read_service import check_job_permission -from app.services.knowledge.kb_orchestrator import KBOrchestrator -from app.services.state_machine import JobStateMachine -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.exceptions.domain_exceptions import ( - JobOperationException, - NotFoundException, - PermissionDeniedException, - ValidationException, -) -from shared.core.state_machine.states import JobStatus -from shared.models.schemas.job import ConfirmUploadRequest -from shared.services.storage.file_upload_service import FileUploadService - - -async def transition_to_uploaded( - db: AsyncSession, - job_id: str, - trigger: str = "manual_upload_completed", -) -> None: - state_machine = JobStateMachine() - await state_machine.transition( - db, job_id, JobStatus.PENDING.value, trigger, None, "system" - ) - - -async def start_workflow_for_job( - db: AsyncSession, - job_id: str, - job_type: str, - source_type: str, - user_id: str, - file_path: Optional[str] = None, - file_url: Optional[str] = None, -) -> None: - if job_type == "kb_management": - orchestrator = KBOrchestrator() - await orchestrator.start_workflow( - db=db, - job_id=job_id, - source_type=source_type, - file_path=file_path, - file_url=file_url, - user_id=user_id, - ) - return - - raise ValidationException( - user_message="Unsupported job type", - violations=[ - { - "field": "job_type", - "description": f"Job type '{job_type}' is not supported", - } - ], - ) - - -async def confirm_job_upload( - db: AsyncSession, - *, - job_id: str, - request: ConfirmUploadRequest | None, - user_id: str, -) -> dict[str, str]: - del request - - try: - job_repo = JobRepository() - job = await job_repo.get_job_by_id(db, job_id) - check_job_permission(job, user_id, job_id) - assert job is not None - - logger.info(f"Confirm upload - Job {job_id} current status: {job.status}") - if job.status not in [JobStatus.PENDING.value, JobStatus.WAITING_FILE.value]: - logger.info(f"Job {job_id} already processed, status: {job.status}") - return {"message": "Job status already updated"} - - if not job.s3_key: - raise ValidationException( - user_message="Job is missing S3 key information", - violations=[ - {"field": "s3_key", "description": "S3 key not set for this job"} - ], - ) - - upload_service = FileUploadService() - file_info = await upload_service.verify_s3_file_exists(job.s3_key) - - if not file_info.get("exists"): - raise ValidationException( - user_message="S3 file does not exist, please upload the file first", - violations=[{"field": "file", "description": "File not found in S3"}], - ) - - await transition_to_uploaded(db, job_id) - await start_workflow_for_job( - db=db, - job_id=job_id, - job_type=job.job_type, - source_type="file", - user_id=user_id, - ) - - return {"message": "File upload confirmed; processing started"} - - except NotFoundException: - raise - except PermissionDeniedException: - raise - except ValidationException: - raise - except Exception as exc: - logger.error(f"Failed to confirm upload: {exc}") - raise JobOperationException( - internal_message=f"Failed to confirm upload: {str(exc)}" - ) diff --git a/apps/api/app/services/rate_limit/dependencies.py b/apps/api/app/services/rate_limit/dependencies.py index c9a697b74..888f9a570 100644 --- a/apps/api/app/services/rate_limit/dependencies.py +++ b/apps/api/app/services/rate_limit/dependencies.py @@ -1,192 +1,24 @@ -""" -FastAPI dependencies for the rate-limit layer. +"""FastAPI adapters for the Job Admission module.""" -Dependency chain (outermost -> innermost): - require_billing_limits -> with_current_user -> get_current_user_id - -> get_db - -``with_current_user`` resolves the user's billing tier through TierService and -enforces the matched system limit (Layer 0). - -``require_billing_limits`` enforces billing RPM (Layer 1) when billing is -enabled and yields control to the route handler. Concurrency (Layer 2) and -daily quota (Layer 3) are enforced just before insert in the create-job route -only when billing is enabled. -""" - -import math -from fnmatch import fnmatch from typing import AsyncGenerator from app.core.dependencies import get_current_user_id -from app.services.rate_limit.config import ( - CONCURRENCY_RETRY_AFTER_SECONDS, - RateLimitConfig, -) -from app.services.rate_limit.data_structures import CurrentUser, TierLimits -from app.services.rate_limit.limiter import RateLimiter -from app.services.rate_limit.system_limit import find_system_rule -from app.services.rate_limit.tier_service import TierService +from app.services.rate_limit.data_structures import CurrentUser +from app.services.rate_limit.job_admission_service import JobAdmissionService from fastapi import Depends, Request -from loguru import logger -from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import settings from shared.core.database import get_db -from shared.core.exceptions.domain_exceptions import ( - PermissionDeniedException, - RateLimitException, - UnavailableException, -) -from shared.core.logging import log_context -from shared.core.state_machine.states import JobStatus -from shared.models.database.job import Job -from shared.models.database.user_balance import UserBalance - -_ACTIVE_JOB_STATES: tuple[str, ...] = ( - JobStatus.WAITING_FILE.value, - JobStatus.PENDING.value, - JobStatus.RUNNING.value, - JobStatus.CONVERTING.value, -) -_GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS: tuple[str, ...] = ( - "/v1/jobs", - "/v1/jobs/*", - "/v1/billing/credits", - "/v1/retrieval/query", - "/v1/documents", - "/v1/documents/*", - "/mcp", -) -_GUEST_API_KEY_REQUIRED_PERMISSION: str = ( - "jobs_documents_retrieval_mcp_or_billing_credits" -) -_GUEST_API_KEY_SCOPE_MESSAGE: str = ( - "Guest API keys can only access job, document, retrieval, MCP query, " - "and billing credits APIs" -) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _get_route_path(request: Request) -> str: - """Return the request path without the application's root_path prefix.""" - scope_path: str = request.scope.get("path", request.url.path) - root_path: str = request.scope.get("root_path", "") - if root_path and scope_path.startswith(root_path): - return scope_path[len(root_path) :] - return scope_path - - -def _get_route_limit_identifier(request: Request) -> str: - """Return a stable identifier for route-scoped system limits.""" - route = request.scope.get("route") - route_path = getattr(route, "path", None) - if isinstance(route_path, str) and route_path: - return route_path - - route_path_format = getattr(route, "path_format", None) - if isinstance(route_path_format, str) and route_path_format: - return route_path_format - - return _get_route_path(request) - - -def _normalize_route_path(route_path: str) -> str: - """Normalize guest route checks across slash-redirect variants.""" - normalized_path = route_path.rstrip("/") - return normalized_path or "/" - - -def _is_guest_api_key_route_allowed(route_path: str) -> bool: - """Return whether a guest API key may access the given route.""" - normalized_path = _normalize_route_path(route_path) - return any( - fnmatch(normalized_path, pattern) - for pattern in _GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS - ) - - -def _enforce_guest_api_key_scope(request: Request, user_tier: str) -> None: - """Reject guest API keys outside the guest-allowed API surface.""" - if user_tier != "guest": - return - - route_path = _get_route_path(request) - if _is_guest_api_key_route_allowed(route_path): - return - - raise PermissionDeniedException( - user_message=_GUEST_API_KEY_SCOPE_MESSAGE, - required_permission=_GUEST_API_KEY_REQUIRED_PERMISSION, - ) - - -# --------------------------------------------------------------------------- -# with_current_user -- Layer 0 (matched system limit) -# --------------------------------------------------------------------------- - +_job_admission_service = JobAdmissionService() async def with_current_user( request: Request, user_id: str = Depends(get_current_user_id), ) -> AsyncGenerator[CurrentUser, None]: - """Resolve the current user tier and enforce the matched system limit. - - Steps: - 1. ``get_current_user_id`` already authenticated the user (401 - on failure). - 2. Resolve ``user_tier`` through ``TierService.get_tier(user_id)``. - 3. If ``RATE_LIMIT_ENABLED=false`` is set, return immediately. - 4. Check the matched system limit via the rate limiter (fail-open on - Redis error). - """ - user_tier = await TierService.get_tier(user_id) - _enforce_guest_api_key_scope(request, user_tier) - current_user = CurrentUser(user_id=user_id, user_tier=user_tier) - - with log_context(user_id=user_id): - # -- Global rate-limit switch -- - config = RateLimitConfig.get_instance() - if not config.is_enabled: - yield current_user - return - - # -- Layer 0: matched system limit -- - try: - route_path = _get_route_path(request) - rule = find_system_rule(request.method, route_path, config.system_rules) - limiter = RateLimiter(config) - await limiter.check_system_limit( - identifier=user_id, - limit=rule.limit, - matched_pattern=rule.api_pattern, - period=rule.period, - ) - except RateLimitException: - raise - except Exception as exc: - # Fail-open: log and let the request through. - logger.warning( - "rate_limit: Redis error during system limit check, " - "failing open for user_id={}, error={}", - user_id, - exc, - ) - - yield current_user - - -# --------------------------------------------------------------------------- -# require_billing_limits -- Layer 1 -# (billing RPM) -# --------------------------------------------------------------------------- - -_RETRY_AFTER_SECONDS: int = 15 + current_user = await _job_admission_service.resolve_current_user( + request=request, + user_id=user_id, + ) + yield current_user async def require_billing_limits( @@ -194,91 +26,14 @@ async def require_billing_limits( current_user: CurrentUser = Depends(with_current_user), _db: AsyncSession = Depends(get_db), ) -> AsyncGenerator[CurrentUser, None]: - """Enforce billing RPM (Layer 1) around the route handler. - - This is an async-generator (yield) dependency so that teardown logic - can run after the route handler completes. - - Layer enforced before yield: - 1. Billing RPM -- per-user requests-per-minute - - Layers enforced inside route just before insert: - 2. Non-terminal jobs concurrency -- max pending/running jobs - 3. Daily quota (free tier only) -- hard daily cap - - When ``BILLING_ENABLED=false``, this yields after identity and system - route limiting. Otherwise, Redis failures raise 503 because billing - enforcement must not be silently skipped. - """ - if not settings.BILLING_ENABLED: - yield current_user - return - - config = RateLimitConfig.get_instance() - if not config.is_enabled: - yield current_user - return - - tier_limits: TierLimits | None = config.tier_map.get(current_user.user_tier) - if tier_limits is None: - logger.error( - "rate_limit: no tier config for tier='{}', user_id={}", - current_user.user_tier, - current_user.user_id, - ) - raise UnavailableException( - internal_message=(f"Missing tier config for tier={current_user.user_tier}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - # -- Layer 1: billing RPM -- - limiter = RateLimiter(config) - try: - await limiter.check_billing_rpm(current_user.user_id, tier_limits.rpm_limit) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in billing RPM check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - + del request + del _db + await _job_admission_service.enforce_billing_limits(current_user=current_user) yield current_user async def require_route_system_limit(request: Request) -> None: - """Apply the matched system limit to the current route using a route key. - - Prefer the framework route template so paths with different parameters - share the same budget bucket. If no explicit rule matches, the default - system rule still protects the route with the wider fallback budget. - Fail closed when Redis or limiter state is unavailable. - """ - config = RateLimitConfig.get_instance() - if not config.is_enabled: - return - - route_path = _get_route_path(request) - route_identifier = _get_route_limit_identifier(request) - rule = find_system_rule(request.method, route_path, config.system_rules) - limiter = RateLimiter(config) - try: - await limiter.check_system_limit( - identifier=route_identifier, - limit=rule.limit, - matched_pattern=rule.api_pattern, - period=rule.period, - use_global_key=True, - ) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in route system limit: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - limit=rule.limit, - period=rule.period, - ) + await _job_admission_service.enforce_route_system_limit(request=request) async def enforce_job_creation_capacity( @@ -286,138 +41,8 @@ async def enforce_job_creation_capacity( db: AsyncSession, current_user: CurrentUser, ) -> None: - """Enforce Layers 2-3 immediately before job insert.""" - if not settings.BILLING_ENABLED: - return - - config = RateLimitConfig.get_instance() - if not config.is_enabled: - return - - tier_limits = config.tier_map.get(current_user.user_tier) - if tier_limits is None: - logger.error( - "rate_limit: no tier config for tier='{}', user_id={}", - current_user.user_tier, - current_user.user_id, - ) - raise UnavailableException( - internal_message=(f"Missing tier config for tier={current_user.user_tier}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - limiter = RateLimiter(config) - - # -- Layer 2: non-terminal jobs concurrency (DB-locked) -- - if tier_limits.max_concurrent_jobs != -1: - try: - await _acquire_user_concurrency_lock(db, current_user.user_id) - active_jobs = await _count_non_terminal_jobs(db, current_user.user_id) - if active_jobs >= tier_limits.max_concurrent_jobs: - retry_after_seconds = _compute_concurrency_retry_after_seconds( - base_retry_after_seconds=CONCURRENCY_RETRY_AFTER_SECONDS, - rpm_limit=tier_limits.rpm_limit, - ) - exc = RateLimitException( - retry_after=retry_after_seconds, - limit=tier_limits.max_concurrent_jobs, - period="concurrent", - user_message=( - f"Too many concurrent requests " - f"({active_jobs}/{tier_limits.max_concurrent_jobs} active). " - f"Please retry after {retry_after_seconds} seconds." - ), - internal_message=( - "Concurrency limit exceeded: " - f"user_id={current_user.user_id}, " - f"active_jobs={active_jobs}, " - f"limit={tier_limits.max_concurrent_jobs}, " - f"retry_after={retry_after_seconds}s" - ), - ) - exc.details.update( - { - "active_jobs": active_jobs, - "available_slots": max( - 0, tier_limits.max_concurrent_jobs - active_jobs - ), - } - ) - raise exc - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"DB error in concurrency check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - # -- Layer 3: daily quota -- - if tier_limits.daily_quota != -1: - try: - await limiter.check_daily_quota( - current_user.user_id, tier_limits.daily_quota - ) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in daily quota check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - -# --------------------------------------------------------------------------- -# Internal helpers -# --------------------------------------------------------------------------- - - -async def _acquire_user_concurrency_lock( - db: AsyncSession, - user_id: str, -) -> None: - """Acquire a per-user row lock to serialize concurrent job creation. - - Locks the UserBalance row instead of User to avoid contention with - unrelated operations (profile updates, etc.) that may also lock User. - """ - result = await db.execute( - select(UserBalance.user_id) - .where(UserBalance.user_id == user_id) - .with_for_update() - ) - if result.scalar_one_or_none() is None: - raise RateLimitException( - internal_message=f"UserBalance row not found for user_id={user_id}" - ) - - -def _compute_concurrency_retry_after_seconds( - base_retry_after_seconds: int, - rpm_limit: int, -) -> int: - """ - Compute Retry-After hint for concurrency rejections. - - Concurrency has no deterministic reset timestamp, so we provide a - conservative client hint: - - floor: configured base retry (currently 30s) - - if billing RPM is finite, also respect one request spacing - (ceil(60 / rpm_limit)) to reduce immediate repeated 429s - """ - if rpm_limit <= 0: - return base_retry_after_seconds - return max(base_retry_after_seconds, int(math.ceil(60 / rpm_limit))) - - -async def _count_non_terminal_jobs( - db: AsyncSession, - user_id: str, -) -> int: - """Count non-terminal jobs for a user in the current transaction.""" - result = await db.execute( - select(func.count(Job.job_id)) - .where(Job.user_id == user_id) - .where(Job.status.in_(_ACTIVE_JOB_STATES)) + del request + await _job_admission_service.enforce_job_creation_capacity( + db=db, + current_user=current_user, ) - return int(result.scalar_one() or 0) diff --git a/apps/api/app/services/rate_limit/job_admission_service.py b/apps/api/app/services/rate_limit/job_admission_service.py new file mode 100644 index 000000000..ff410b46e --- /dev/null +++ b/apps/api/app/services/rate_limit/job_admission_service.py @@ -0,0 +1,356 @@ +from __future__ import annotations + +import math +from fnmatch import fnmatch + +from app.services.rate_limit.config import ( + CONCURRENCY_RETRY_AFTER_SECONDS, + RateLimitConfig, +) +from app.services.rate_limit.data_structures import CurrentUser, TierLimits +from app.services.rate_limit.limiter import RateLimiter +from app.services.rate_limit.system_limit import find_system_rule +from app.services.rate_limit.tier_service import TierService +from fastapi import Request +from loguru import logger +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + PermissionDeniedException, + RateLimitException, + UnavailableException, +) +from shared.core.logging import log_context +from shared.core.state_machine.states import JobStatus +from shared.models.database.job import Job +from shared.models.database.user_balance import UserBalance + +_ACTIVE_JOB_STATES: tuple[str, ...] = ( + JobStatus.WAITING_FILE.value, + JobStatus.PENDING.value, + JobStatus.RUNNING.value, + JobStatus.CONVERTING.value, +) +_GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS: tuple[str, ...] = ( + "/v1/jobs", + "/v1/jobs/*", + "/v1/billing/credits", + "/v1/retrieval/query", + "/v1/documents", + "/v1/documents/*", + "/mcp", +) +_GUEST_API_KEY_REQUIRED_PERMISSION: str = ( + "jobs_documents_retrieval_mcp_or_billing_credits" +) +_GUEST_API_KEY_SCOPE_MESSAGE: str = ( + "Guest API keys can only access job, document, retrieval, MCP query, " + "and billing credits APIs" +) +_RETRY_AFTER_SECONDS: int = 15 + + +class JobAdmissionService: + async def resolve_current_user( + self, + *, + request: Request, + user_id: str, + ) -> CurrentUser: + user_tier = await TierService.get_tier(user_id) + self._enforce_guest_api_key_scope(request=request, user_tier=user_tier) + current_user = CurrentUser(user_id=user_id, user_tier=user_tier) + + with log_context(user_id=user_id): + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return current_user + + try: + await self._check_user_system_limit( + request=request, + config=config, + user_id=user_id, + ) + except RateLimitException: + raise + except Exception as exc: + logger.warning( + "rate_limit: Redis error during system limit check, " + "failing open for user_id={}, error={}", + user_id, + exc, + ) + + return current_user + + async def enforce_billing_limits( + self, + *, + current_user: CurrentUser, + ) -> None: + if not settings.BILLING_ENABLED: + return + + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + tier_limits = self._require_tier_limits( + config=config, + current_user=current_user, + ) + + limiter = RateLimiter(config) + try: + await limiter.check_billing_rpm( + current_user.user_id, + tier_limits.rpm_limit, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in billing RPM check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + async def enforce_route_system_limit(self, *, request: Request) -> None: + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + route_path = self._get_route_path(request) + route_identifier = self._get_route_limit_identifier(request) + rule = find_system_rule(request.method, route_path, config.system_rules) + limiter = RateLimiter(config) + + try: + await limiter.check_system_limit( + identifier=route_identifier, + limit=rule.limit, + matched_pattern=rule.api_pattern, + period=rule.period, + use_global_key=True, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in route system limit: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + limit=rule.limit, + period=rule.period, + ) + + async def enforce_job_creation_capacity( + self, + *, + db: AsyncSession, + current_user: CurrentUser, + ) -> None: + if not settings.BILLING_ENABLED: + return + + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + tier_limits = self._require_tier_limits( + config=config, + current_user=current_user, + ) + limiter = RateLimiter(config) + + if tier_limits.max_concurrent_jobs != -1: + try: + await self._acquire_user_concurrency_lock( + db=db, + user_id=current_user.user_id, + ) + active_jobs = await self._count_non_terminal_jobs( + db=db, + user_id=current_user.user_id, + ) + if active_jobs >= tier_limits.max_concurrent_jobs: + retry_after_seconds = self._compute_concurrency_retry_after_seconds( + base_retry_after_seconds=CONCURRENCY_RETRY_AFTER_SECONDS, + rpm_limit=tier_limits.rpm_limit, + ) + exc = RateLimitException( + retry_after=retry_after_seconds, + limit=tier_limits.max_concurrent_jobs, + period="concurrent", + user_message=( + f"Too many concurrent requests " + f"({active_jobs}/{tier_limits.max_concurrent_jobs} active). " + f"Please retry after {retry_after_seconds} seconds." + ), + internal_message=( + "Concurrency limit exceeded: " + f"user_id={current_user.user_id}, " + f"active_jobs={active_jobs}, " + f"limit={tier_limits.max_concurrent_jobs}, " + f"retry_after={retry_after_seconds}s" + ), + ) + exc.details.update( + { + "active_jobs": active_jobs, + "available_slots": max( + 0, + tier_limits.max_concurrent_jobs - active_jobs, + ), + } + ) + raise exc + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"DB error in concurrency check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + if tier_limits.daily_quota != -1: + try: + await limiter.check_daily_quota( + current_user.user_id, + tier_limits.daily_quota, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in daily quota check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + async def _check_user_system_limit( + self, + *, + request: Request, + config: RateLimitConfig, + user_id: str, + ) -> None: + route_path = self._get_route_path(request) + rule = find_system_rule(request.method, route_path, config.system_rules) + limiter = RateLimiter(config) + await limiter.check_system_limit( + identifier=user_id, + limit=rule.limit, + matched_pattern=rule.api_pattern, + period=rule.period, + ) + + def _require_tier_limits( + self, + *, + config: RateLimitConfig, + current_user: CurrentUser, + ) -> TierLimits: + tier_limits = config.tier_map.get(current_user.user_tier) + if tier_limits is None: + logger.error( + "rate_limit: no tier config for tier='{}', user_id={}", + current_user.user_tier, + current_user.user_id, + ) + raise UnavailableException( + internal_message=( + f"Missing tier config for tier={current_user.user_tier}" + ), + retry_after=_RETRY_AFTER_SECONDS, + ) + return tier_limits + + def _get_route_path(self, request: Request) -> str: + scope_path = request.scope.get("path", request.url.path) + root_path = request.scope.get("root_path", "") + if isinstance(scope_path, str) and isinstance(root_path, str): + if root_path and scope_path.startswith(root_path): + return scope_path[len(root_path) :] + return scope_path + return request.url.path + + def _get_route_limit_identifier(self, request: Request) -> str: + route = request.scope.get("route") + route_path = getattr(route, "path", None) + if isinstance(route_path, str) and route_path: + return route_path + + route_path_format = getattr(route, "path_format", None) + if isinstance(route_path_format, str) and route_path_format: + return route_path_format + + return self._get_route_path(request) + + def _normalize_route_path(self, route_path: str) -> str: + normalized_path = route_path.rstrip("/") + return normalized_path or "/" + + def _is_guest_api_key_route_allowed(self, route_path: str) -> bool: + normalized_path = self._normalize_route_path(route_path) + return any( + fnmatch(normalized_path, pattern) + for pattern in _GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS + ) + + def _enforce_guest_api_key_scope( + self, + *, + request: Request, + user_tier: str, + ) -> None: + if user_tier != "guest": + return + + route_path = self._get_route_path(request) + if self._is_guest_api_key_route_allowed(route_path): + return + + raise PermissionDeniedException( + user_message=_GUEST_API_KEY_SCOPE_MESSAGE, + required_permission=_GUEST_API_KEY_REQUIRED_PERMISSION, + ) + + async def _acquire_user_concurrency_lock( + self, + *, + db: AsyncSession, + user_id: str, + ) -> None: + result = await db.execute( + select(UserBalance.user_id) + .where(UserBalance.user_id == user_id) + .with_for_update() + ) + if result.scalar_one_or_none() is None: + raise RateLimitException( + internal_message=f"UserBalance row not found for user_id={user_id}" + ) + + async def _count_non_terminal_jobs( + self, + *, + db: AsyncSession, + user_id: str, + ) -> int: + result = await db.execute( + select(func.count(Job.job_id)) + .where(Job.user_id == user_id) + .where(Job.status.in_(_ACTIVE_JOB_STATES)) + ) + return int(result.scalar_one() or 0) + + def _compute_concurrency_retry_after_seconds( + self, + *, + base_retry_after_seconds: int, + rpm_limit: int, + ) -> int: + if rpm_limit <= 0: + return base_retry_after_seconds + return max(base_retry_after_seconds, int(math.ceil(60 / rpm_limit))) diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index a2bc658cb..c1111e342 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -453,7 +453,7 @@ async def create_checkout_session_for_credits_package( async with developer_api_client_factory() as api_client: billing_service_module = importlib.import_module( - "app.services.billing.billing_app_service" + "app.services.billing.billing_workflow_service" ) monkeypatch.setattr( billing_service_module, @@ -498,7 +498,7 @@ async def create_payment_intent( async with developer_api_client_factory() as api_client: billing_service_module = importlib.import_module( - "app.services.billing.billing_app_service" + "app.services.billing.billing_workflow_service" ) monkeypatch.setattr( billing_service_module, diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index 8cfa4993c..b1eda40b5 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -1002,7 +1002,7 @@ async def _fake_start_workflow_for_job( ) async with developer_api_client_factory() as api_client: - import app.services.job_upload_confirmation_service as upload_confirmation_service + import app.services.document_ingestion_service as document_ingestion_service import shared.services.storage.file_upload_service as file_upload_service_module monkeypatch.setattr( @@ -1011,8 +1011,8 @@ async def _fake_start_workflow_for_job( _fake_verify_s3_file_exists, ) monkeypatch.setattr( - upload_confirmation_service, - "start_workflow_for_job", + document_ingestion_service, + "_start_job_workflow", _fake_start_workflow_for_job, ) From a2edca9f52a56d261cb64980d7783f3598947fe8 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 02:36:34 +0800 Subject: [PATCH 03/40] refactor: split stripe billing workflows --- CONTEXT.md | 10 + .../billing/billing_workflow_service.py | 26 +- .../billing/stripe_purchase_service.py | 200 +++++ .../app/services/billing/stripe_service.py | 766 ------------------ .../billing/stripe_webhook_service.py | 600 ++++++++++++++ .../tests/contract/test_billing_contract.py | 61 +- 6 files changed, 881 insertions(+), 782 deletions(-) create mode 100644 apps/api/app/services/billing/stripe_purchase_service.py delete mode 100644 apps/api/app/services/billing/stripe_service.py create mode 100644 apps/api/app/services/billing/stripe_webhook_service.py diff --git a/CONTEXT.md b/CONTEXT.md index 0c956ac22..db2f07453 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -77,6 +77,16 @@ Job Result, Document, and Document Chunk records. The credits purchase, checkout, webhook handling, refund reconciliation, and tier refresh flows. +### Stripe Purchase + +The Billing Workflow adapter that creates Stripe payment intents and checkout +sessions for credits purchases. + +### Stripe Webhook Reconciliation + +The Billing Workflow adapter that verifies Stripe events and reconciles credits, +payment records, and refunds. + ### Guest API Key A guest-tier API key with a restricted route surface. diff --git a/apps/api/app/services/billing/billing_workflow_service.py b/apps/api/app/services/billing/billing_workflow_service.py index 51dbd44bb..bf55ae6f7 100644 --- a/apps/api/app/services/billing/billing_workflow_service.py +++ b/apps/api/app/services/billing/billing_workflow_service.py @@ -3,7 +3,8 @@ from typing import Optional from app.services.billing.price_config_service import PriceConfigService -from app.services.billing.stripe_service import StripeService +from app.services.billing.stripe_purchase_service import StripePurchaseService +from app.services.billing.stripe_webhook_service import StripeWebhookService from pydantic import BaseModel from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession @@ -52,11 +53,11 @@ async def buy_credits( request: BuyCreditsRequest, user_id: str, ) -> PaymentIntentResponse: - stripe_service = self._create_stripe_service() + stripe_purchase_service = self._create_stripe_purchase_service() try: amount_cny = request.credits_amount * 0.02 amount_cents = int(amount_cny * 100) - payment_intent = await stripe_service.create_payment_intent( + payment_intent = await stripe_purchase_service.create_payment_intent( user_id=user_id, amount=amount_cents, credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, @@ -229,7 +230,7 @@ async def buy_credits_package( request: BuyCreditsPackageRequest, user_id: str, ) -> CheckoutSessionResponse: - stripe_service = self._create_stripe_service() + stripe_purchase_service = self._create_stripe_purchase_service() try: result = await db.execute(select(User.email).where(User.id == user_id)) user_email = result.scalar_one_or_none() @@ -238,7 +239,7 @@ async def buy_credits_package( success_url = f"{frontend_url}/billing?success=true&type=credits_package" cancel_url = f"{frontend_url}/billing?canceled=true" - checkout_url = await stripe_service.create_checkout_session_for_credits_package( + checkout_url = await stripe_purchase_service.create_credits_package_checkout_session( db=db, user_id=user_id, price_id=request.price_id, @@ -264,20 +265,27 @@ async def handle_stripe_webhook( payload: bytes, stripe_signature: str | None, ) -> dict[str, object]: - stripe_service = self._create_stripe_service() + stripe_webhook_service = self._create_stripe_webhook_service() try: if not stripe_signature: raise StripeServiceException( internal_message="Missing stripe-signature header" ) - return await stripe_service.handle_webhook(db, payload, stripe_signature) + return await stripe_webhook_service.handle_webhook( + db, + payload=payload, + sig_header=stripe_signature, + ) except Exception as exc: raise StripeServiceException( internal_message=f"Failed to handle webhook: {str(exc)}" ) - def _create_stripe_service(self) -> StripeService: - return StripeService() + def _create_stripe_purchase_service(self) -> StripePurchaseService: + return StripePurchaseService() + + def _create_stripe_webhook_service(self) -> StripeWebhookService: + return StripeWebhookService() async def _load_total_parse_micro_credits_used( self, diff --git a/apps/api/app/services/billing/stripe_purchase_service.py b/apps/api/app/services/billing/stripe_purchase_service.py new file mode 100644 index 000000000..c5ea9bedd --- /dev/null +++ b/apps/api/app/services/billing/stripe_purchase_service.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +from typing import Any + +import stripe +from app.services.billing.price_config_service import PriceConfigService +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + StripeServiceException, + SystemSettingMissingException, + ValidationException, +) +from shared.core.logging import logger +from shared.repositories.credits_repository import CreditsRepository +from shared.services.billing import CreditsService + + +class StripePurchaseService: + def __init__( + self, + *, + price_config_service: PriceConfigService | None = None, + credits_repository: CreditsRepository | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._configure_stripe_api() + self._price_config_service = price_config_service or PriceConfigService() + self._credits_repository = credits_repository or CreditsRepository() + self._credits_service = credits_service or CreditsService() + + async def create_credits_package_checkout_session( + self, + db: AsyncSession, + *, + user_id: str, + price_id: str, + success_url: str, + cancel_url: str, + quantity: int, + email: str | None = None, + ) -> str: + try: + config = await self._price_config_service.get_price_config(db, price_id) + if not config.is_credits_package(): + raise ValidationException( + user_message="Invalid price configuration", + violations=[ + { + "field": "price_id", + "description": f"Price ID {price_id} is not a credits package", + } + ], + ) + + customer_id = await self._resolve_customer_id( + db, + user_id=user_id, + email=email, + ) + metadata = { + "user_id": str(user_id), + "price_id": str(price_id), + "type": "credits_package", + "credits_amount": ( + str(config.credits_amount) if config.credits_amount else None + ), + "quantity": str(quantity), + } + session_params: dict[str, Any] = { + "customer": customer_id, + "customer_update": {"address": "auto"}, + "client_reference_id": str(user_id), + "line_items": [ + { + "price": price_id, + "quantity": quantity, + } + ], + "mode": "payment", + "success_url": success_url, + "cancel_url": cancel_url, + "metadata": metadata, + "payment_intent_data": {"metadata": metadata}, + "allow_promotion_codes": True, + "adaptive_pricing": {"enabled": False}, + "billing_address_collection": "required", + } + session = stripe.checkout.Session.create(**session_params) + await db.commit() + return str(session.url or "") + except stripe.StripeError as exc: + logger.error(f"Stripe credits checkout session failed: {exc}") + raise StripeServiceException( + internal_message=f"Stripe credits checkout session failed: {exc}" + ) + + async def create_payment_intent( + self, + *, + user_id: str, + amount: int, + credits_amount: int, + currency: str = "usd", + ) -> dict[str, str]: + try: + intent = stripe.PaymentIntent.create( + amount=amount, + currency=currency, + automatic_payment_methods={"enabled": True}, + metadata={ + "user_id": user_id, + "type": "credits", + "credits_amount": str(credits_amount), + }, + ) + return { + "client_secret": str(intent.client_secret or ""), + "payment_intent_id": str(intent.id), + } + except stripe.StripeError as exc: + logger.error(f"Failed to create payment intent: {exc}") + raise StripeServiceException( + internal_message=f"Stripe payment intent creation failed: {exc}" + ) + + async def _resolve_customer_id( + self, + db: AsyncSession, + *, + user_id: str, + email: str | None, + ) -> str: + await self._credits_service.ensure_user_initialized(db, user_id) + user_balance = await self._credits_repository.get_user_balance(db, user_id) + if not user_balance: + raise ValidationException( + user_message="Failed to initialize user balance", + violations=[ + { + "field": "user_id", + "description": f"Failed to initialize user balance for {user_id}", + } + ], + ) + + customer_id = user_balance.stripe_customer_id + if customer_id: + return customer_id + + customer_id = self._find_existing_customer_id(email=email) + if customer_id is None: + customer_id = self._create_customer(user_id=user_id, email=email) + + user_balance.stripe_customer_id = customer_id + return customer_id + + def _find_existing_customer_id( + self, + *, + email: str | None, + ) -> str | None: + if not email: + return None + + existing_customers = stripe.Customer.list(email=email, limit=1) + if existing_customers.data: + return str(existing_customers.data[0].id) + return None + + def _create_customer( + self, + *, + user_id: str, + email: str | None, + ) -> str: + if not email: + raise ValidationException( + user_message="Email required for first-time payment", + violations=[ + { + "field": "email", + "description": "Email is required to create a billing profile", + } + ], + ) + + customer = stripe.Customer.create( + email=email, + metadata={"user_id": str(user_id)}, + ) + return str(customer.id) + + def _configure_stripe_api(self) -> None: + if not settings.STRIPE_SECRET_KEY: + raise SystemSettingMissingException( + internal_message="Stripe API key not configured (STRIPE_SECRET_KEY)" + ) + stripe.api_key = settings.STRIPE_SECRET_KEY diff --git a/apps/api/app/services/billing/stripe_service.py b/apps/api/app/services/billing/stripe_service.py deleted file mode 100644 index be1792a2e..000000000 --- a/apps/api/app/services/billing/stripe_service.py +++ /dev/null @@ -1,766 +0,0 @@ -"""Stripe payment service.""" - -from typing import Any, Dict, Optional -from uuid import UUID - -import stripe -from app.repositories.payment_record_repository import PaymentRecordRepository -from app.services.billing.price_config_service import PriceConfigService -from app.services.rate_limit.tier_service import TierService -from sqlalchemy import func, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - AuthException, - KnowhereException, - StripeServiceException, - SystemSettingMissingException, - ValidationException, -) -from shared.core.logging import logger -from shared.models.database.payment_record import PaymentRecord -from shared.repositories.credits_repository import CreditsRepository -from shared.services.billing import CreditsService -from shared.utils.utc_now import utc_now_naive - - -class StripeService: - """Stripe payment service.""" - - def __init__(self): - if not settings.STRIPE_SECRET_KEY: - raise SystemSettingMissingException( - internal_message="Stripe API key not configured (STRIPE_SECRET_KEY)" - ) - stripe.api_key = settings.STRIPE_SECRET_KEY - self.credits_repo = CreditsRepository() - self.payment_record_repo = PaymentRecordRepository() - self.price_config_service = PriceConfigService() - self.credits_service = CreditsService() - - async def create_checkout_session( - self, - db: AsyncSession, - user_id: str, - plan_id: str, - success_url: str, - cancel_url: str, - ) -> str: - """Create a Stripe Checkout session for a subscription.""" - try: - # Load the Stripe price ID for the requested plan from the database. - price_id = await self.price_config_service.get_plan_price_id(db, plan_id) - - session = stripe.checkout.Session.create( - line_items=[ - { - "price": price_id, - "quantity": 1, - } - ], - mode="subscription", - success_url=success_url, - cancel_url=cancel_url, - metadata={ - "user_id": user_id, - "plan_id": plan_id, - "type": "subscription", - }, - allow_promotion_codes=True, - # Disable Adaptive Pricing to prevent currency switcher from hiding Alipay. - # Alipay handles USD→CNY conversion internally for customers. - adaptive_pricing={"enabled": False}, - ) - return str(session.url or "") - except stripe.StripeError as e: - logger.error(f"Failed to create subscription checkout session: {e}") - raise StripeServiceException( - internal_message=f"Stripe checkout session creation failed: {e}" - ) - - async def create_checkout_session_for_credits_package( - self, - db: AsyncSession, - user_id: str, - price_id: str, - success_url: str, - cancel_url: str, - quantity: int, - email: Optional[str] = None, - ) -> str: - """Create a Stripe Checkout session for a credits package.""" - try: - # Validate that the selected price configuration exists. - config = await self.price_config_service.get_price_config(db, price_id) - if not config.is_credits_package(): - raise ValidationException( - user_message="Invalid price configuration", - violations=[ - { - "field": "price_id", - "description": f"Price ID {price_id} is not a credits package", - } - ], - ) - - # Ensure user is initialized (UserBalance exists) - await self.credits_service.ensure_user_initialized(db, user_id) - - user_balance = await self.credits_repo.get_user_balance(db, user_id) - if not user_balance: - # Should not happen after ensure_user_initialized - raise ValidationException( - user_message="Failed to initialize user balance", - violations=[ - { - "field": "user_id", - "description": f"Failed to initialize user balance for {user_id}", - } - ], - ) - - customer_id = user_balance.stripe_customer_id - - if not customer_id: - # Reuse an existing Stripe customer when the email already exists. - if email: - existing_customers = stripe.Customer.list(email=email, limit=1) - if existing_customers.data: - customer_id = existing_customers.data[0].id - - if not customer_id: - # Create a new Stripe customer when no existing record matches. - if not email: - # For new customers, we prefer having an email. - # If no email provided, we can't create a good customer record. - # But technically Stripe allows it. - # Better: Require email for new billing profiles. - raise ValidationException( - user_message="Email required for first-time payment", - violations=[ - { - "field": "email", - "description": "Email is required to create a billing profile", - } - ], - ) - - customer_params = { - "email": email, - "metadata": {"user_id": str(user_id)}, - } - # Username is not available without User model, omit it. - - customer = stripe.Customer.create(**customer_params) - customer_id = customer.id - - user_balance.stripe_customer_id = customer_id - - # Keep metadata values as strings so refunds can recover the user ID later. - metadata = { - "user_id": str(user_id), - "price_id": str(price_id), - "type": "credits_package", - "credits_amount": ( - str(config.credits_amount) if config.credits_amount else None - ), - "quantity": str(quantity), - } - - session_params: Dict[str, Any] = { - "customer": customer_id, - "customer_update": {"address": "auto"}, - "client_reference_id": str(user_id), - "line_items": [ - { - "price": price_id, - "quantity": quantity, - } - ], - "mode": "payment", # One-time payment. - "success_url": success_url, - "cancel_url": cancel_url, - "metadata": metadata, - # Copy metadata onto the PaymentIntent and Charge for refund handling. - "payment_intent_data": { - "metadata": metadata, - }, - # Collect more customer information for later reconciliation. - "allow_promotion_codes": True, - # Disable Adaptive Pricing to prevent currency switcher from hiding Alipay. - # Alipay handles USD→CNY conversion internally for customers. - "adaptive_pricing": {"enabled": False}, - # Require a billing address so Checkout syncs it to the customer record. - "billing_address_collection": "required", - } - - session = stripe.checkout.Session.create(**session_params) - - await db.commit() - - return str(session.url or "") - except stripe.StripeError as e: - logger.error(f"Stripe credits checkout session failed: {e}") - raise StripeServiceException( - internal_message=f"Stripe credits checkout session failed: {e}" - ) - - async def create_payment_intent( - self, user_id: str, amount: int, credits_amount: int, currency: str = "usd" - ) -> Dict[str, Any]: - """Create a PaymentIntent for a credits purchase.""" - try: - intent = stripe.PaymentIntent.create( - amount=amount, # amount in cents - currency=currency, - automatic_payment_methods={"enabled": True}, - metadata={ - "user_id": user_id, - "type": "credits", - "credits_amount": str(credits_amount), - }, - ) - return { - "client_secret": intent.client_secret, - "payment_intent_id": intent.id, - } - except stripe.StripeError as e: - logger.error(f"Failed to create payment intent: {e}") - raise StripeServiceException( - internal_message=f"Stripe payment intent creation failed: {e}" - ) - - async def handle_webhook( - self, db: AsyncSession, payload: bytes, sig_header: str - ) -> Dict[str, Any]: - """Handle a Stripe webhook payload.""" - try: - event = stripe.Webhook.construct_event( - payload, sig_header, settings.STRIPE_WEBHOOK_SECRET - ) - return await self._process_webhook_event(db, event) - except ValueError as e: - logger.error(f"Invalid payload: {e}") - raise ValidationException( - user_message="Invalid webhook payload", - violations=[ - {"field": "payload", "description": "Webhook payload is malformed"} - ], - ) - except stripe.SignatureVerificationError as e: - logger.error(f"Invalid signature: {e}") - raise AuthException( - user_message="Invalid webhook signature", - internal_message=f"Webhook signature verification failed: {e}", - ) - - async def _process_webhook_event( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Dispatch an incoming Stripe webhook event.""" - event_type = event["type"] - - if event_type == "checkout.session.completed": - return await self._handle_checkout_completed(db, event) - elif event_type == "payment_intent.succeeded": - return await self._handle_payment_intent_succeeded(db, event) - elif event_type == "invoice.payment_succeeded": - return await self._handle_payment_succeeded(db, event) - elif event_type == "customer.subscription.deleted": - return await self._handle_subscription_deleted(db, event) - elif event_type == "charge.refunded": - return await self._handle_charge_refunded(db, event) - else: - return {"status": "ignored", "event_type": event_type} - - async def _handle_checkout_completed( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle a completed Checkout session.""" - session = event["data"]["object"] - session_id = session["id"] - mode = session.get("mode") - metadata = session.get("metadata", {}) - user_id = metadata.get("user_id") - payment_type = metadata.get("type") - quantity = int(metadata.get("quantity", 1)) - - if not user_id: - logger.warning( - f"Checkout session {session_id} is missing user_id metadata; likely a test event, skipping" - ) - return { - "status": "ignored", - "message": "Missing user_id metadata (likely test event)", - "checkout_session_id": session_id, - "event_type": "checkout.session.completed", - } - - # Skip work that was already processed for this Checkout session. - if await self.payment_record_repo.is_processed( - db, checkout_session_id=session_id - ): - logger.info(f"Checkout session {session_id} already processed, skipping") - return { - "status": "ignored", - "message": "Already processed", - "checkout_session_id": session_id, - } - - # Seed audit metadata for the payment record. - payment_metadata = { - "session_id": session_id, - "stripe_session": session, # Full session payload for debugging and audits. - } - - # Create the pending payment record before side effects run. - payment_record = PaymentRecord( - checkout_session_id=session_id, - user_id=user_id, - payment_type=payment_type or "unknown", - amount_cents=session.get("amount_total", 0), - currency=session.get("currency", "cny").upper(), - status="pending", - extra_metadata=payment_metadata, - ) - db.add(payment_record) - await db.flush() # Get the database ID without committing yet. - - try: - if mode == "payment" and payment_type == "credits_package": - # Credits package purchase flow. - price_id = metadata.get("price_id") - - if not price_id: - logger.error(f"Incomplete Credits pack info: price_id={price_id}") - return {"status": "error", "message": "Missing price_id"} - - # Load the credits amount and product metadata from the price config. - price_config = await self.price_config_service.get_price_config( - db, price_id - ) - credits_amount = price_config.credits_amount * quantity - if credits_amount is None: - logger.error( - f"Credits amount is not configured for price ID {price_id}" - ) - return { - "status": "error", - "message": "Credits amount not configured", - } - - # Attach purchased product details to the payment record. - product_description = f"Credits pack - {credits_amount} Credits" - if price_config.extra_metadata and price_config.extra_metadata.get( - "description" - ): - product_description = price_config.extra_metadata.get("description") - - payment_record.extra_metadata = { - **payment_metadata, - "product_description": product_description, - "price_id": price_id, - "credits_amount": credits_amount, - "product_metadata": price_config.extra_metadata - or {}, # Product metadata from the price config. - } - - # Grant the purchased credits to the user balance. - await self.credits_service.add_credits( - session=db, - user_id=user_id, - amount=credits_amount, - reason=f"Purchase credits pack: {product_description}", - stripe_payment_id=session.get("payment_intent"), - ) - - # Mark the payment record as completed. - payment_record.status = "succeeded" - payment_record.credits_amount = credits_amount - payment_record.processed_at = utc_now_naive() - - await TierService.refresh_tier(user_id, db) - await db.commit() - await db.refresh(payment_record) - - logger.info( - f"Credits pack purchase succeeded: user_id={user_id}, credits={credits_amount}, price_id={price_id}" - ) - return { - "status": "success", - "event_type": "checkout.session.completed", - "user_id": user_id, - "credits_amount": credits_amount, - "payment_type": "credits_package", - } - else: - logger.warning( - f"Unknown payment type: mode={mode}, type={payment_type}" - ) - return {"status": "ignored", "message": "Unknown payment type"} - - except KnowhereException: - raise - except Exception as e: - logger.error( - f"Failed to process checkout.session.completed: {e}", exc_info=True - ) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": str(e), - } - await db.commit() - raise StripeServiceException( - internal_message=f"Failed to process checkout.session.completed: {str(e)}", - original_exception=e, - ) - - async def _handle_payment_intent_succeeded( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle a successful PaymentIntent for a credits purchase.""" - payment_intent = event["data"]["object"] - payment_intent_id = payment_intent["id"] - metadata = payment_intent.get("metadata", {}) - user_id = metadata.get("user_id") - payment_type = metadata.get("type") - - if payment_type != "credits": - logger.info( - f"PaymentIntent {payment_intent_id} is not a Credits payment, skipping" - ) - return {"status": "ignored", "payment_intent_id": payment_intent_id} - - if not user_id: - logger.warning( - f"PaymentIntent {payment_intent_id} is missing user_id metadata; likely a test event, skipping" - ) - return { - "status": "ignored", - "message": "Missing user_id metadata (likely test event)", - "payment_intent_id": payment_intent_id, - } - - # Skip work that was already processed for this PaymentIntent. - if await self.payment_record_repo.is_processed( - db, payment_intent_id=payment_intent_id - ): - logger.info( - f"PaymentIntent {payment_intent_id} already processed, skipping" - ) - return { - "status": "ignored", - "message": "Already processed", - "payment_intent_id": payment_intent_id, - } - - # Seed audit metadata for the payment record. - payment_metadata = { - "payment_intent_id": payment_intent_id, - "stripe_payment_intent": payment_intent, # Full PaymentIntent payload for debugging and audits. - } - - # Create the pending payment record before side effects run. - payment_record = PaymentRecord( - payment_intent_id=payment_intent_id, - user_id=user_id, - payment_type="credits_package", - amount_cents=payment_intent.get("amount", 0), - currency=payment_intent.get("currency", "cny").upper(), - status="pending", - extra_metadata=payment_metadata, - ) - db.add(payment_record) - await db.flush() # Get the database ID without committing yet. - - try: - # Read the purchased credits amount from metadata. - credits_amount_str = metadata.get("credits_amount") - if not credits_amount_str: - logger.error( - f"PaymentIntent {payment_intent_id} is missing credits_amount" - ) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": "Missing credits_amount", - } - await db.commit() - return {"status": "error", "message": "Missing credits_amount"} - - credits_amount = int(credits_amount_str) - - # Attach purchased product details to the payment record. - payment_record.extra_metadata = { - **payment_metadata, - "product_description": f"Credits package - {credits_amount} Credits", - "credits_amount": credits_amount, - "payment_method": "payment_intent", # Marks this purchase as PaymentIntent-based. - } - - # Amount validation can be layered in here if needed later. - payment_intent.get("amount", 0) - - # Grant the purchased credits to the user balance. - await self.credits_service.add_credits( - session=db, - user_id=user_id, - amount=credits_amount, - reason=f"buy credits - {credits_amount} Credits", - stripe_payment_id=payment_intent_id, - ) - - # Mark the payment record as completed. - payment_record.status = "succeeded" - payment_record.credits_amount = credits_amount - payment_record.processed_at = utc_now_naive() - - await TierService.refresh_tier(user_id, db) - await db.commit() - await db.refresh(payment_record) - - logger.info( - f"buy credits success: user_id={user_id}, credits={credits_amount}, payment_intent_id={payment_intent_id}" - ) - return { - "status": "success", - "event_type": "payment_intent.succeeded", - "user_id": user_id, - "credits_amount": credits_amount, - "payment_type": "credits_package", - } - - except Exception as e: - logger.error(f"Failed to process Credits purchase: {e}", exc_info=True) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": str(e), - } - await db.commit() - raise StripeServiceException( - internal_message=f"Failed to process Credits purchase: {str(e)}", - original_exception=e, - ) - - async def _handle_payment_succeeded( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle a successful subscription renewal payment.""" - invoice = event["data"]["object"] - subscription_id = invoice.get("subscription") - - if not subscription_id: - logger.warning("Invoice is missing subscription ID") - return {"status": "ignored", "message": "Missing subscription_id"} - - return {"status": "ignored", "message": "Subscription renewal not implemented"} - - async def _handle_subscription_deleted( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle subscription deletion events.""" - subscription = event["data"]["object"] - stripe_subscription_id = subscription["id"] - - try: - # Subscription management not yet implemented - logger.warning( - f"Local subscription record not found: stripe_subscription_id={stripe_subscription_id}" - ) - - return {"status": "success", "subscription_id": stripe_subscription_id} - except KnowhereException: - raise - except Exception as e: - logger.error( - f"Failed to process customer.subscription.deleted: {e}", exc_info=True - ) - raise StripeServiceException( - internal_message=f"Failed to process customer.subscription.deleted: {str(e)}", - original_exception=e, - ) - - async def _handle_charge_refunded( - self, db: AsyncSession, event: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle refund events, including manual refunds from the Stripe dashboard.""" - charge = event["data"]["object"] - charge_id = charge.get("id") - refund_items = (charge.get("refunds", {}) or {}).get("data", []) or [] - latest_refund = refund_items[-1] if refund_items else None - - payment_intent_id = charge.get("payment_intent") - refund_id = latest_refund.get("id") if latest_refund else None - - currency = (charge.get("currency") or "cny").upper() - - # Use a stable idempotency key derived from the refund or charge identifier. - idempotency_key = refund_id or f"{charge_id}-refund" - - # Recover the original payment record to reuse billing context such as user_id. - original_record = None - if payment_intent_id: - original_record = await self.payment_record_repo.get_by_payment_intent_id( - db, payment_intent_id - ) - - metadata = charge.get("metadata") or {} - user_id = metadata.get("user_id") or (getattr(original_record, "user_id", None)) - payment_type = ( - metadata.get("type") - or (getattr(original_record, "payment_type", None)) - or "refund" - ) - - if not user_id: - logger.error( - f"Refund event is missing user_id; cannot record refund: charge_id={charge_id}" - ) - return { - "status": "error", - "message": "Missing user_id for refund", - "event_type": "charge.refunded", - } - - # Normalize user_id to UUID before using it in SQL filters. - if user_id and isinstance(user_id, str): - try: - user_id = UUID(user_id) - except ValueError: - logger.error(f"Invalid user_id format: {user_id}") - return { - "status": "error", - "message": "Invalid user_id format", - "event_type": "charge.refunded", - } - - user_id_str = str(user_id) - - # Compute the incremental refund amount from the cumulative Stripe total. - # total_refund_amount_cents already includes the current refund event. - total_refund_amount_cents = charge.get("amount_refunded") or 0 - - # Load previously recorded refund totals for the same payment flow. - origin_total_refund_amount_cents = 0 - - # Sum historical refund records that use the same synthetic refund key. - query = ( - select(func.sum(PaymentRecord.amount_cents)) - .where(PaymentRecord.payment_intent_id == idempotency_key) - .where(PaymentRecord.user_id == user_id) - .where( - PaymentRecord.amount_cents - < 0 # Refund rows are stored as negative amounts. - ) - ) - result = await db.execute(query) - # Sum negative refund amounts and convert back to a positive total. - origin_total_refund_amount_cents = abs(result.scalar() or 0) - - refund_amount_cents = ( - total_refund_amount_cents - origin_total_refund_amount_cents - ) - if refund_amount_cents <= 0: - # The refund has already been processed; keep this path idempotent. - logger.info( - f"Refund already processed, skipping: charge_id={charge_id}, refund_id={refund_id}" - ) - return { - "status": "success", - "event_type": "charge.refunded", - "message": "Already processed", - "user_id": user_id, - "refund_id": refund_id, - } - - # Translate the refunded cash amount back into credits using price metadata. - credits_refunded = None - price_id = metadata.get("price_id") or ( - getattr(original_record, "extra_metadata", {}) or {} - ).get("price_id") - if price_id: - try: - price_cfg = await self.price_config_service.get_price_config( - db, price_id - ) - if price_cfg and price_cfg.amount_cents: - credits_refunded = -int( - price_cfg.credits_amount - * abs(refund_amount_cents) - / abs(price_cfg.amount_cents) # credits_amount * quantity - ) - except Exception as e: - logger.warning( - f"Failed to calculate refunded Credits, price_id={price_id}: {e}" - ) - credits_refunded = None - - # Fall back to the original payment record ratio when price metadata is unavailable. - if ( - credits_refunded is None - and original_record - and original_record.credits_amount - and original_record.amount_cents - ): - credits_refunded = -int( - abs(original_record.credits_amount) - * abs(refund_amount_cents) - / abs(original_record.amount_cents) - ) - - # Apply the credit adjustment to the user balance when needed. - if credits_refunded is not None and credits_refunded < 0: - await self.credits_service.add_credits( - session=db, - user_id=user_id_str, - amount=credits_refunded, - reason="Refund adjustment", - transaction_type="refund", - transaction_metadata={"refund_id": refund_id, "charge_id": charge_id}, - ) - - refund_metadata = { - "refund_id": refund_id, - "charge_id": charge_id, - "original_payment_intent_id": payment_intent_id, - "original_payment_record_id": getattr(original_record, "id", None), - "reason": (latest_refund or {}).get("reason"), - "balance_transaction": (latest_refund or {}).get("balance_transaction"), - } - - refund_record = PaymentRecord( - payment_intent_id=idempotency_key, - user_id=user_id, - payment_type=payment_type, - amount_cents=-abs(refund_amount_cents), - currency=currency, - status="succeeded", - credits_amount=credits_refunded, - plan_id=getattr(original_record, "plan_id", None), - stripe_subscription_id=getattr( - original_record, "stripe_subscription_id", None - ), - processed_at=utc_now_naive(), - extra_metadata=refund_metadata, - ) - - db.add(refund_record) - await db.commit() - await db.refresh(refund_record) - - logger.info( - f"Refund record created: user_id={user_id}, amount_cents={refund_record.amount_cents}, " - f"refund_id={refund_id}, charge_id={charge_id}" - ) - - return { - "status": "success", - "event_type": "charge.refunded", - "user_id": user_id, - "refund_amount_cents": abs(refund_amount_cents), - "payment_intent_id": payment_intent_id, - "refund_id": refund_id, - } diff --git a/apps/api/app/services/billing/stripe_webhook_service.py b/apps/api/app/services/billing/stripe_webhook_service.py new file mode 100644 index 000000000..7da961fdb --- /dev/null +++ b/apps/api/app/services/billing/stripe_webhook_service.py @@ -0,0 +1,600 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any, TypeAlias +from uuid import UUID + +import stripe +from app.repositories.payment_record_repository import PaymentRecordRepository +from app.services.billing.price_config_service import PriceConfigService +from app.services.rate_limit.tier_service import TierService +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + AuthException, + KnowhereException, + StripeServiceException, + SystemSettingMissingException, + ValidationException, +) +from shared.core.logging import logger +from shared.models.database.payment_record import PaymentRecord +from shared.services.billing import CreditsService +from shared.utils.utc_now import utc_now_naive + +StripeEvent: TypeAlias = dict[str, Any] +StripeWebhookHandler: TypeAlias = Callable[ + [AsyncSession, StripeEvent], Awaitable[dict[str, object]] +] + + +class StripeWebhookService: + def __init__( + self, + *, + payment_record_repository: PaymentRecordRepository | None = None, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._configure_stripe_api() + self._payment_record_repository = ( + payment_record_repository or PaymentRecordRepository() + ) + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + + async def handle_webhook( + self, + db: AsyncSession, + *, + payload: bytes, + sig_header: str, + ) -> dict[str, object]: + try: + event = stripe.Webhook.construct_event( + payload, + sig_header, + settings.STRIPE_WEBHOOK_SECRET, + ) + return await self._dispatch_event(db, event) + except ValueError as exc: + logger.error(f"Invalid payload: {exc}") + raise ValidationException( + user_message="Invalid webhook payload", + violations=[ + {"field": "payload", "description": "Webhook payload is malformed"} + ], + ) + except stripe.SignatureVerificationError as exc: + logger.error(f"Invalid signature: {exc}") + raise AuthException( + user_message="Invalid webhook signature", + internal_message=f"Webhook signature verification failed: {exc}", + ) + + async def _dispatch_event( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + event_type = str(event["type"]) + handler = self._event_handlers().get(event_type) + if handler is None: + return {"status": "ignored", "event_type": event_type} + return await handler(db, event) + + def _event_handlers(self) -> dict[str, StripeWebhookHandler]: + return { + "checkout.session.completed": self._handle_checkout_completed, + "payment_intent.succeeded": self._handle_payment_intent_succeeded, + "invoice.payment_succeeded": self._handle_payment_succeeded, + "customer.subscription.deleted": self._handle_subscription_deleted, + "charge.refunded": self._handle_charge_refunded, + } + + async def _handle_checkout_completed( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + session = event["data"]["object"] + session_id = str(session["id"]) + mode = session.get("mode") + metadata = session.get("metadata", {}) + user_id = metadata.get("user_id") + payment_type = metadata.get("type") + quantity = int(metadata.get("quantity", 1)) + + if not user_id: + logger.warning( + f"Checkout session {session_id} is missing user_id metadata; likely a test event, skipping" + ) + return { + "status": "ignored", + "message": "Missing user_id metadata (likely test event)", + "checkout_session_id": session_id, + "event_type": "checkout.session.completed", + } + + if await self._payment_record_repository.is_processed( + db, + checkout_session_id=session_id, + ): + logger.info(f"Checkout session {session_id} already processed, skipping") + return { + "status": "ignored", + "message": "Already processed", + "checkout_session_id": session_id, + } + + payment_metadata = { + "session_id": session_id, + "stripe_session": session, + } + payment_record = PaymentRecord( + checkout_session_id=session_id, + user_id=user_id, + payment_type=payment_type or "unknown", + amount_cents=session.get("amount_total", 0), + currency=session.get("currency", "cny").upper(), + status="pending", + extra_metadata=payment_metadata, + ) + db.add(payment_record) + await db.flush() + + try: + if mode != "payment" or payment_type != "credits_package": + logger.warning(f"Unknown payment type: mode={mode}, type={payment_type}") + return {"status": "ignored", "message": "Unknown payment type"} + + price_id = metadata.get("price_id") + if not price_id: + logger.error(f"Incomplete Credits pack info: price_id={price_id}") + return {"status": "error", "message": "Missing price_id"} + + price_config = await self._price_config_service.get_price_config(db, price_id) + configured_credits_amount = price_config.credits_amount + if configured_credits_amount is None: + logger.error( + f"Credits amount is not configured for price ID {price_id}" + ) + return { + "status": "error", + "message": "Credits amount not configured", + } + credits_amount = configured_credits_amount * quantity + + product_description = f"Credits pack - {credits_amount} Credits" + if price_config.extra_metadata and price_config.extra_metadata.get( + "description" + ): + product_description = str( + price_config.extra_metadata.get("description") + ) + + payment_record.extra_metadata = { + **payment_metadata, + "product_description": product_description, + "price_id": price_id, + "credits_amount": credits_amount, + "product_metadata": price_config.extra_metadata or {}, + } + await self._credits_service.add_credits( + session=db, + user_id=user_id, + amount=credits_amount, + reason=f"Purchase credits pack: {product_description}", + stripe_payment_id=session.get("payment_intent"), + ) + payment_record.status = "succeeded" + payment_record.credits_amount = credits_amount + payment_record.processed_at = utc_now_naive() + + await TierService.refresh_tier(user_id, db) + await db.commit() + await db.refresh(payment_record) + + logger.info( + f"Credits pack purchase succeeded: user_id={user_id}, credits={credits_amount}, price_id={price_id}" + ) + return { + "status": "success", + "event_type": "checkout.session.completed", + "user_id": user_id, + "credits_amount": credits_amount, + "payment_type": "credits_package", + } + except KnowhereException: + raise + except Exception as exc: + logger.error( + f"Failed to process checkout.session.completed: {exc}", + exc_info=True, + ) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": str(exc), + } + await db.commit() + raise StripeServiceException( + internal_message=( + "Failed to process checkout.session.completed: " + f"{str(exc)}" + ), + original_exception=exc, + ) + + async def _handle_payment_intent_succeeded( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + payment_intent = event["data"]["object"] + payment_intent_id = str(payment_intent["id"]) + metadata = payment_intent.get("metadata", {}) + user_id = metadata.get("user_id") + payment_type = metadata.get("type") + + if payment_type != "credits": + logger.info( + f"PaymentIntent {payment_intent_id} is not a Credits payment, skipping" + ) + return {"status": "ignored", "payment_intent_id": payment_intent_id} + + if not user_id: + logger.warning( + f"PaymentIntent {payment_intent_id} is missing user_id metadata; likely a test event, skipping" + ) + return { + "status": "ignored", + "message": "Missing user_id metadata (likely test event)", + "payment_intent_id": payment_intent_id, + } + + if await self._payment_record_repository.is_processed( + db, + payment_intent_id=payment_intent_id, + ): + logger.info( + f"PaymentIntent {payment_intent_id} already processed, skipping" + ) + return { + "status": "ignored", + "message": "Already processed", + "payment_intent_id": payment_intent_id, + } + + payment_metadata = { + "payment_intent_id": payment_intent_id, + "stripe_payment_intent": payment_intent, + } + payment_record = PaymentRecord( + payment_intent_id=payment_intent_id, + user_id=user_id, + payment_type="credits_package", + amount_cents=payment_intent.get("amount", 0), + currency=payment_intent.get("currency", "cny").upper(), + status="pending", + extra_metadata=payment_metadata, + ) + db.add(payment_record) + await db.flush() + + try: + credits_amount_str = metadata.get("credits_amount") + if not credits_amount_str: + logger.error( + f"PaymentIntent {payment_intent_id} is missing credits_amount" + ) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": "Missing credits_amount", + } + await db.commit() + return {"status": "error", "message": "Missing credits_amount"} + + credits_amount = int(credits_amount_str) + payment_record.extra_metadata = { + **payment_metadata, + "product_description": f"Credits package - {credits_amount} Credits", + "credits_amount": credits_amount, + "payment_method": "payment_intent", + } + await self._credits_service.add_credits( + session=db, + user_id=user_id, + amount=credits_amount, + reason=f"buy credits - {credits_amount} Credits", + stripe_payment_id=payment_intent_id, + ) + payment_record.status = "succeeded" + payment_record.credits_amount = credits_amount + payment_record.processed_at = utc_now_naive() + + await TierService.refresh_tier(user_id, db) + await db.commit() + await db.refresh(payment_record) + + logger.info( + f"buy credits success: user_id={user_id}, credits={credits_amount}, payment_intent_id={payment_intent_id}" + ) + return { + "status": "success", + "event_type": "payment_intent.succeeded", + "user_id": user_id, + "credits_amount": credits_amount, + "payment_type": "credits_package", + } + except Exception as exc: + logger.error(f"Failed to process Credits purchase: {exc}", exc_info=True) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": str(exc), + } + await db.commit() + raise StripeServiceException( + internal_message=f"Failed to process Credits purchase: {str(exc)}", + original_exception=exc, + ) + + async def _handle_payment_succeeded( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + del db + invoice = event["data"]["object"] + subscription_id = invoice.get("subscription") + + if not subscription_id: + logger.warning("Invoice is missing subscription ID") + return {"status": "ignored", "message": "Missing subscription_id"} + + return {"status": "ignored", "message": "Subscription renewal not implemented"} + + async def _handle_subscription_deleted( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + del db + subscription = event["data"]["object"] + stripe_subscription_id = str(subscription["id"]) + + try: + logger.warning( + "Local subscription record not found: " + f"stripe_subscription_id={stripe_subscription_id}" + ) + return {"status": "success", "subscription_id": stripe_subscription_id} + except KnowhereException: + raise + except Exception as exc: + logger.error( + f"Failed to process customer.subscription.deleted: {exc}", + exc_info=True, + ) + raise StripeServiceException( + internal_message=( + "Failed to process customer.subscription.deleted: " + f"{str(exc)}" + ), + original_exception=exc, + ) + + async def _handle_charge_refunded( + self, + db: AsyncSession, + event: StripeEvent, + ) -> dict[str, object]: + charge = event["data"]["object"] + charge_id = charge.get("id") + refund_items = (charge.get("refunds", {}) or {}).get("data", []) or [] + latest_refund = refund_items[-1] if refund_items else None + + payment_intent_id = charge.get("payment_intent") + refund_id = latest_refund.get("id") if latest_refund else None + currency = (charge.get("currency") or "cny").upper() + idempotency_key = refund_id or f"{charge_id}-refund" + + original_record = None + if payment_intent_id: + original_record = await self._payment_record_repository.get_by_payment_intent_id( + db, + payment_intent_id, + ) + + metadata = charge.get("metadata") or {} + user_id = metadata.get("user_id") or ( + getattr(original_record, "user_id", None) + ) + payment_type = ( + metadata.get("type") + or getattr(original_record, "payment_type", None) + or "refund" + ) + + if not user_id: + logger.error( + f"Refund event is missing user_id; cannot record refund: charge_id={charge_id}" + ) + return { + "status": "error", + "message": "Missing user_id for refund", + "event_type": "charge.refunded", + } + + normalized_user_id = self._normalize_user_id(user_id) + if normalized_user_id is None: + logger.error(f"Invalid user_id format: {user_id}") + return { + "status": "error", + "message": "Invalid user_id format", + "event_type": "charge.refunded", + } + + user_id_str = str(normalized_user_id) + total_refund_amount_cents = charge.get("amount_refunded") or 0 + origin_total_refund_amount_cents = await self._load_recorded_refund_amount( + db, + payment_intent_id=idempotency_key, + user_id=normalized_user_id, + ) + + refund_amount_cents = ( + total_refund_amount_cents - origin_total_refund_amount_cents + ) + if refund_amount_cents <= 0: + logger.info( + f"Refund already processed, skipping: charge_id={charge_id}, refund_id={refund_id}" + ) + return { + "status": "success", + "event_type": "charge.refunded", + "message": "Already processed", + "user_id": normalized_user_id, + "refund_id": refund_id, + } + + credits_refunded = await self._calculate_refunded_credits( + db, + metadata=metadata, + original_record=original_record, + refund_amount_cents=refund_amount_cents, + ) + + if credits_refunded is not None and credits_refunded < 0: + await self._credits_service.add_credits( + session=db, + user_id=user_id_str, + amount=credits_refunded, + reason="Refund adjustment", + transaction_type="refund", + transaction_metadata={"refund_id": refund_id, "charge_id": charge_id}, + ) + + refund_metadata = { + "refund_id": refund_id, + "charge_id": charge_id, + "original_payment_intent_id": payment_intent_id, + "original_payment_record_id": getattr(original_record, "id", None), + "reason": (latest_refund or {}).get("reason"), + "balance_transaction": (latest_refund or {}).get("balance_transaction"), + } + refund_record = PaymentRecord( + payment_intent_id=idempotency_key, + user_id=normalized_user_id, + payment_type=payment_type, + amount_cents=-abs(refund_amount_cents), + currency=currency, + status="succeeded", + credits_amount=credits_refunded, + plan_id=getattr(original_record, "plan_id", None), + stripe_subscription_id=getattr( + original_record, + "stripe_subscription_id", + None, + ), + processed_at=utc_now_naive(), + extra_metadata=refund_metadata, + ) + db.add(refund_record) + await db.commit() + await db.refresh(refund_record) + + logger.info( + f"Refund record created: user_id={normalized_user_id}, amount_cents={refund_record.amount_cents}, " + f"refund_id={refund_id}, charge_id={charge_id}" + ) + return { + "status": "success", + "event_type": "charge.refunded", + "user_id": normalized_user_id, + "refund_amount_cents": abs(refund_amount_cents), + "payment_intent_id": payment_intent_id, + "refund_id": refund_id, + } + + async def _load_recorded_refund_amount( + self, + db: AsyncSession, + *, + payment_intent_id: str, + user_id: UUID, + ) -> int: + result = await db.execute( + select(func.sum(PaymentRecord.amount_cents)) + .where(PaymentRecord.payment_intent_id == payment_intent_id) + .where(PaymentRecord.user_id == user_id) + .where(PaymentRecord.amount_cents < 0) + ) + return int(abs(result.scalar() or 0)) + + async def _calculate_refunded_credits( + self, + db: AsyncSession, + *, + metadata: dict[str, Any], + original_record: PaymentRecord | None, + refund_amount_cents: int, + ) -> int | None: + credits_refunded: int | None = None + price_id = metadata.get("price_id") or ( + getattr(original_record, "extra_metadata", {}) or {} + ).get("price_id") + if price_id: + try: + price_cfg = await self._price_config_service.get_price_config( + db, + price_id, + ) + if price_cfg and price_cfg.amount_cents and price_cfg.credits_amount: + credits_refunded = -int( + price_cfg.credits_amount + * abs(refund_amount_cents) + / abs(price_cfg.amount_cents) + ) + except Exception as exc: + logger.warning( + f"Failed to calculate refunded Credits, price_id={price_id}: {exc}" + ) + credits_refunded = None + + if ( + credits_refunded is None + and original_record + and original_record.credits_amount + and original_record.amount_cents + ): + credits_refunded = -int( + abs(original_record.credits_amount) + * abs(refund_amount_cents) + / abs(original_record.amount_cents) + ) + + return credits_refunded + + def _normalize_user_id( + self, + user_id: str | UUID, + ) -> UUID | None: + if isinstance(user_id, UUID): + return user_id + + try: + return UUID(user_id) + except ValueError: + return None + + def _configure_stripe_api(self) -> None: + if not settings.STRIPE_SECRET_KEY: + raise SystemSettingMissingException( + internal_message="Stripe API key not configured (STRIPE_SECRET_KEY)" + ) + stripe.api_key = settings.STRIPE_SECRET_KEY diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index c1111e342..de7e3cc04 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -432,8 +432,8 @@ async def test_should_return_a_checkout_url_when_buying_a_credit_package( ], monkeypatch: MonkeyPatch, ) -> None: - class FakeStripeService: - async def create_checkout_session_for_credits_package( + class FakeStripePurchaseService: + async def create_credits_package_checkout_session( self, db, user_id: str, @@ -457,8 +457,8 @@ async def create_checkout_session_for_credits_package( ) monkeypatch.setattr( billing_service_module, - "StripeService", - FakeStripeService, + "StripePurchaseService", + FakeStripePurchaseService, ) response = await api_client.post( "/api/v1/billing/buy-credits-package", @@ -479,7 +479,7 @@ async def test_should_return_a_payment_intent_payload_when_buying_credits( ], monkeypatch: MonkeyPatch, ) -> None: - class FakeStripeService: + class FakeStripePurchaseService: async def create_payment_intent( self, user_id: str, @@ -502,8 +502,8 @@ async def create_payment_intent( ) monkeypatch.setattr( billing_service_module, - "StripeService", - FakeStripeService, + "StripePurchaseService", + FakeStripePurchaseService, ) response = await api_client.post( "/api/v1/billing/buy-credits", @@ -515,3 +515,50 @@ async def create_payment_intent( "client_secret": "pi_contract_secret", "payment_intent_id": "pi_contract_id", } + + +@pytest.mark.asyncio +async def test_should_delegate_the_webhook_endpoint_to_the_stripe_webhook_service( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + class FakeStripeWebhookService: + async def handle_webhook( + self, + db, + *, + payload: bytes, + sig_header: str, + ) -> dict[str, object]: + del db + assert payload == b'{"type":"payment_intent.succeeded"}' + assert sig_header == "sig_contract" + return { + "status": "success", + "event_type": "payment_intent.succeeded", + "user_id": "local-dev-user", + } + + async with developer_api_client_factory() as api_client: + billing_service_module = importlib.import_module( + "app.services.billing.billing_workflow_service" + ) + monkeypatch.setattr( + billing_service_module, + "StripeWebhookService", + FakeStripeWebhookService, + ) + response = await api_client.post( + "/api/v1/billing/webhook", + content=b'{"type":"payment_intent.succeeded"}', + headers={"stripe-signature": "sig_contract"}, + ) + + assert response.status_code == 200 + assert response.json() == { + "status": "success", + "event_type": "payment_intent.succeeded", + "user_id": "local-dev-user", + } From 3fc99ce3781d697c551c6262571e7e119f6cb708 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 02:46:45 +0800 Subject: [PATCH 04/40] refactor: split api key workflows --- CONTEXT.md | 16 + apps/api/app/api/v1/routes/api_key.py | 38 +- apps/api/app/core/dependencies.py | 8 +- .../auth/api_key_authentication_service.py | 178 +++++++++ .../auth/api_key_management_service.py | 217 +++++++++++ apps/api/app/services/auth/api_key_service.py | 362 ------------------ .../guest/guest_registration_service.py | 4 +- 7 files changed, 436 insertions(+), 387 deletions(-) create mode 100644 apps/api/app/services/auth/api_key_authentication_service.py create mode 100644 apps/api/app/services/auth/api_key_management_service.py delete mode 100644 apps/api/app/services/auth/api_key_service.py diff --git a/CONTEXT.md b/CONTEXT.md index db2f07453..e83408e5e 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -77,6 +77,16 @@ Job Result, Document, and Document Chunk records. The credits purchase, checkout, webhook handling, refund reconciliation, and tier refresh flows. +### API Key Authentication + +The auth-time workflow that validates API keys, reads and writes the API-key +cache, and schedules best-effort last-used updates. + +### API Key Management + +The user-facing workflow that creates, lists, reads, revokes, and toggles API +keys. + ### Stripe Purchase The Billing Workflow adapter that creates Stripe payment intents and checkout @@ -169,6 +179,12 @@ exceptions. - `app/repositories/payment_record_repository.py` - shared billing modules in `packages/shared-python/shared/services/billing/*` +### API Key Management + +- `app/api/v1/routes/api_key.py` +- `app/services/auth/*` +- `app/repositories/api_key_repository.py` + ### Webhook Management - `app/api/v1/routes/webhook.py` diff --git a/apps/api/app/api/v1/routes/api_key.py b/apps/api/app/api/v1/routes/api_key.py index 370521f7b..2dbcefce4 100644 --- a/apps/api/app/api/v1/routes/api_key.py +++ b/apps/api/app/api/v1/routes/api_key.py @@ -2,7 +2,7 @@ API key management endpoints. """ -from app.services.auth.api_key_service import APIKeyService +from app.services.auth.api_key_management_service import APIKeyManagementService from app.services.rate_limit.dependencies import ( CurrentUser, with_current_user, @@ -25,6 +25,7 @@ ) router = APIRouter(tags=["API Key Management"]) +_api_key_management_service = APIKeyManagementService() @router.post("/create", summary="Create an API key") @@ -34,10 +35,8 @@ async def create_api_key( db: AsyncSession = Depends(get_db), ): """Create an API key.""" - api_key_service = APIKeyService.get_instance() - try: - api_key = await api_key_service.create_api_key( + api_key = await _api_key_management_service.create_api_key( session=db, user_id=current_user.user_id, name=request.name, @@ -68,11 +67,10 @@ async def list_api_keys( db: AsyncSession = Depends(get_db), ): """List API keys for the current user.""" - api_key_service = APIKeyService.get_instance() - try: - api_keys_data = await api_key_service.list_user_api_keys( - db, current_user.user_id + api_keys_data = await _api_key_management_service.list_user_api_keys( + db, + user_id=current_user.user_id, ) api_keys = [ @@ -104,11 +102,11 @@ async def revoke_api_key( db: AsyncSession = Depends(get_db), ): """Revoke an API key.""" - api_key_service = APIKeyService.get_instance() - try: - await api_key_service.revoke_api_key( - session=db, api_key_id=request.api_key_id, user_id=current_user.user_id + await _api_key_management_service.revoke_api_key( + session=db, + api_key_id=request.api_key_id, + user_id=current_user.user_id, ) return {"message": "API key revoked"} @@ -129,11 +127,11 @@ async def get_api_key( db: AsyncSession = Depends(get_db), ): """Get details for a single API key.""" - api_key_service = APIKeyService.get_instance() - try: - api_key = await api_key_service.get_api_key( - db, current_user.user_id, api_key_id + api_key = await _api_key_management_service.get_api_key( + db, + user_id=current_user.user_id, + api_key_id=api_key_id, ) if not api_key: raise NotFoundException( @@ -167,11 +165,11 @@ async def toggle_api_key( db: AsyncSession = Depends(get_db), ): """Enable or disable an API key.""" - api_key_service = APIKeyService.get_instance() - try: - success = await api_key_service.toggle_api_key( - db, current_user.user_id, api_key_id + success = await _api_key_management_service.toggle_api_key( + db, + user_id=current_user.user_id, + api_key_id=api_key_id, ) if success: return {"message": "API key status updated"} diff --git a/apps/api/app/core/dependencies.py b/apps/api/app/core/dependencies.py index bf56d39c8..153177cd4 100644 --- a/apps/api/app/core/dependencies.py +++ b/apps/api/app/core/dependencies.py @@ -3,7 +3,9 @@ from typing import Any import jwt -from app.services.auth.api_key_service import APIKeyService +from app.services.auth.api_key_authentication_service import ( + APIKeyAuthenticationService, +) from fastapi import Depends, Header, Request from jwt import PyJWKClient from loguru import logger @@ -27,6 +29,7 @@ # Cached PyJWKClient instance _jwks_client: PyJWKClient | None = None _jwks_client_lock = threading.Lock() +_api_key_authentication_service = APIKeyAuthenticationService() def _get_jwks_client() -> PyJWKClient: @@ -145,8 +148,7 @@ async def get_current_user_id( # Mode 1: API Key verification (for external clients) if is_api_key_token(token): - api_key_service = APIKeyService.get_instance() - user_id = await api_key_service.validate_api_key(db, token) + user_id = await _api_key_authentication_service.validate_api_key(db, token) if user_id: return user_id diff --git a/apps/api/app/services/auth/api_key_authentication_service.py b/apps/api/app/services/auth/api_key_authentication_service.py new file mode 100644 index 000000000..67e677247 --- /dev/null +++ b/apps/api/app/services/auth/api_key_authentication_service.py @@ -0,0 +1,178 @@ +"""API key authentication workflow.""" + +from __future__ import annotations + +import asyncio +import json +from datetime import datetime, timezone + +from app.repositories.api_key_repository import APIKeyRepository +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import redis_pool_manager +from shared.core.database import get_db_context +from shared.services.redis.redis_service import RedisService +from shared.utils.api_keys import hash_api_key + +_API_KEY_USER_CACHE_TTL_SECONDS: int = 3600 + + +class APIKeyAuthenticationService: + """Validate API keys and maintain the API-key auth cache.""" + + def __init__( + self, + *, + repository: APIKeyRepository | None = None, + ) -> None: + self._repository = repository or APIKeyRepository() + + async def validate_api_key( + self, + session: AsyncSession, + api_key: str, + ) -> str | None: + """Validate an API key and return the owning user ID.""" + key_hash = hash_api_key(api_key) + redis_service = redis_pool_manager.get_redis_service() + cached_user_id = await self._get_cached_user_id(redis_service, key_hash) + if cached_user_id is not None: + return cached_user_id + + api_key_record = await self._repository.get_by_key_hash(session, key_hash) + if not api_key_record or not api_key_record.is_valid(): + return None + + self._schedule_last_used_update(str(api_key_record.id)) + user_id = str(api_key_record.user_id) + await self._set_cached_user_id( + redis_service, + key_hash, + user_id, + self._resolve_api_key_cache_ttl_seconds(api_key_record.expires_at), + ) + return user_id + + async def invalidate_api_key_user_cache( + self, + *, + user_id: str, + api_key_hash: str, + ) -> None: + """Remove one API-key auth cache entry.""" + await self._invalidate_cached_api_key_user_id( + redis_pool_manager.get_redis_service(), + user_id, + api_key_hash, + ) + + @staticmethod + def _get_user_id_key(api_key_hash: str) -> str: + return f"api-key:user-id:{api_key_hash}" + + @staticmethod + def _get_user_api_keys_key(user_id: str) -> str: + return f"api-key:user-hashes:{user_id}" + + async def _get_cached_user_id( + self, + redis_service: RedisService, + api_key_hash: str, + ) -> str | None: + try: + raw_user_id = await redis_service.get(self._get_user_id_key(api_key_hash)) + return self._coerce_user_id(raw_user_id) + except Exception: + logger.warning("api_key_authentication: failed to read API-key user cache") + return None + + async def _set_cached_user_id( + self, + redis_service: RedisService, + api_key_hash: str, + user_id: str, + ttl_seconds: int, + ) -> None: + effective_ttl_seconds = min(_API_KEY_USER_CACHE_TTL_SECONDS, ttl_seconds) + user_id_key = self._get_user_id_key(api_key_hash) + user_api_keys_key = self._get_user_api_keys_key(user_id) + + try: + await redis_service.set(user_id_key, user_id, ttl=effective_ttl_seconds) + await redis_service.sadd(user_api_keys_key, api_key_hash) + reverse_ttl_seconds = await redis_service.ttl(user_api_keys_key) + if ( + reverse_ttl_seconds in (-2, -1) + or reverse_ttl_seconds < effective_ttl_seconds + ): + await redis_service.expire(user_api_keys_key, effective_ttl_seconds) + except Exception: + logger.warning( + "api_key_authentication: failed to write API-key user cache for user_id={}", + user_id, + ) + + async def _invalidate_cached_api_key_user_id( + self, + redis_service: RedisService, + user_id: str, + api_key_hash: str, + ) -> None: + try: + await redis_service.delete(self._get_user_id_key(api_key_hash)) + await redis_service.srem(self._get_user_api_keys_key(user_id), api_key_hash) + except Exception: + logger.warning( + "api_key_authentication: failed to invalidate API-key cache for user_id={}", + user_id, + ) + + def _coerce_user_id(self, raw_user_id: object) -> str | None: + if isinstance(raw_user_id, str): + try: + parsed_user_id: object = json.loads(raw_user_id) + except json.JSONDecodeError: + return raw_user_id + else: + parsed_user_id = raw_user_id + + if isinstance(parsed_user_id, str): + return parsed_user_id + + if isinstance(parsed_user_id, dict): + legacy_user_id = parsed_user_id.get("user_id") + if isinstance(legacy_user_id, str): + return legacy_user_id + + return None + + def _resolve_api_key_cache_ttl_seconds(self, expires_at: datetime | None) -> int: + if expires_at is None: + return _API_KEY_USER_CACHE_TTL_SECONDS + + expires_at_utc = expires_at + if expires_at_utc.tzinfo is None: + expires_at_utc = expires_at_utc.replace(tzinfo=timezone.utc) + + now = datetime.now(timezone.utc) + remaining_seconds = int((expires_at_utc - now).total_seconds()) + return max(1, min(_API_KEY_USER_CACHE_TTL_SECONDS, remaining_seconds)) + + def _schedule_last_used_update(self, api_key_id: str) -> None: + try: + asyncio.create_task( + self._update_last_used_best_effort(api_key_id), + name=f"api_key_last_used:{api_key_id}", + ) + except Exception as exc: + logger.warning( + f"Failed to schedule API key last-used update (ignored): {exc}" + ) + + async def _update_last_used_best_effort(self, api_key_id: str) -> None: + try: + async with get_db_context() as db: + await self._repository.update_last_used(db, api_key_id) + except Exception as exc: + logger.warning(f"Failed to update API key last-used time (ignored): {exc}") diff --git a/apps/api/app/services/auth/api_key_management_service.py b/apps/api/app/services/auth/api_key_management_service.py new file mode 100644 index 000000000..0adbfee94 --- /dev/null +++ b/apps/api/app/services/auth/api_key_management_service.py @@ -0,0 +1,217 @@ +"""API key management workflow.""" + +from __future__ import annotations + +from datetime import datetime +from typing import TypedDict + +from app.repositories.api_key_repository import APIKeyRepository +from app.services.auth.api_key_authentication_service import ( + APIKeyAuthenticationService, +) +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + APIKeyOperationException, + KnowhereException, + NotFoundException, + ValidationException, +) +from shared.models.database.api_key import APIKey +from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key + + +class APIKeyListItem(TypedDict): + id: str + name: str + api_key: str + enabled_modules: list[str] | None + is_active: bool + created_at: datetime + last_used_at: datetime | None + expires_at: datetime | None + + +class APIKeyManagementService: + """Create, list, read, revoke, and toggle API keys.""" + + def __init__( + self, + *, + repository: APIKeyRepository | None = None, + authentication_service: APIKeyAuthenticationService | None = None, + ) -> None: + self._repository = repository or APIKeyRepository() + self._authentication_service = ( + authentication_service or APIKeyAuthenticationService() + ) + + async def create_api_key( + self, + session: AsyncSession, + *, + user_id: str, + name: str, + enabled_modules: list[str] | None = None, + expires_at: datetime | None = None, + ) -> str: + key_count = await self._repository.count_by_user(session, user_id) + if key_count >= 10: + raise ValidationException( + user_message="Maximum API Key limit reached (10)", + violations=[ + { + "field": "api_keys", + "description": "User has reached the maximum API Key limit", + } + ], + ) + + existing_key = await self._repository.get_by_user_and_name( + session, + user_id, + name, + ) + if existing_key: + raise ValidationException( + user_message="API Key name already exists", + violations=[ + { + "field": "name", + "description": f"An API Key with name '{name}' already exists", + } + ], + ) + + api_key = generate_api_key() + api_key_record = APIKey( + user_id=user_id, + key_hash=hash_api_key(api_key), + key_mask=mask_api_key(api_key), + name=name, + enabled_modules=enabled_modules or ["all"], + expires_at=expires_at, + ) + await self._repository.create(session, api_key_record) + return api_key + + async def revoke_api_key( + self, + session: AsyncSession, + *, + api_key_id: str, + user_id: str, + ) -> bool: + logger.info(f"Revoking API key: api_key_id={api_key_id}, user_id={user_id}") + api_key = await self._repository.get_by_id(session, api_key_id) + + if not api_key: + logger.warning("API key does not exist") + raise NotFoundException( + resource="APIKey", + resource_id=api_key_id, + internal_message="API Key not found", + ) + + if str(api_key.user_id) != user_id: + logger.warning( + f"User ID mismatch: api_key.user_id={api_key.user_id}, user_id={user_id}" + ) + raise NotFoundException( + resource="APIKey", + resource_id=api_key_id, + internal_message="API Key not found or does not belong to user", + ) + + success = await self._repository.delete_by_id(session, api_key_id) + logger.info(f"Delete result: {success}") + + if success: + await session.commit() + logger.info("Transaction committed") + await self._authentication_service.invalidate_api_key_user_cache( + user_id=user_id, + api_key_hash=api_key.key_hash, + ) + + return success + + async def list_user_api_keys( + self, + session: AsyncSession, + *, + user_id: str, + ) -> list[APIKeyListItem]: + api_keys = await self._repository.get_unexpired_by_user_id(session, user_id) + return [ + { + "id": str(api_key.id), + "name": api_key.name, + "api_key": api_key.key_mask + or f"sk_{api_key.id[:8]}••••••••••••••••••••••••••••••••••••••••", + "enabled_modules": api_key.enabled_modules, + "is_active": api_key.is_active, + "created_at": api_key.created_at, + "last_used_at": api_key.last_used_at, + "expires_at": api_key.expires_at, + } + for api_key in api_keys + ] + + async def get_api_key( + self, + session: AsyncSession, + *, + user_id: str, + api_key_id: str, + ) -> APIKey | None: + try: + api_key = await self._repository.get(session, api_key_id) + if api_key and str(api_key.user_id) == user_id: + return api_key + return None + except KnowhereException: + raise + except Exception as exc: + logger.error(f"Failed to get API key: {exc}") + raise APIKeyOperationException( + internal_message=f"Failed to get API key: {str(exc)}", + original_exception=exc, + ) + + async def toggle_api_key( + self, + session: AsyncSession, + *, + user_id: str, + api_key_id: str, + ) -> bool: + try: + api_key = await self._repository.get(session, api_key_id) + if not api_key or str(api_key.user_id) != user_id: + return False + + api_key.is_active = not api_key.is_active + await session.commit() + await session.refresh(api_key) + + if not api_key.is_active: + await self._authentication_service.invalidate_api_key_user_cache( + user_id=user_id, + api_key_hash=api_key.key_hash, + ) + + logger.info( + f"API key status toggled successfully: {api_key_id}, new_status={api_key.is_active}" + ) + return True + except KnowhereException: + raise + except Exception as exc: + logger.error(f"Failed to toggle API key status: {exc}") + await session.rollback() + raise APIKeyOperationException( + internal_message=f"Failed to toggle API key status: {str(exc)}", + original_exception=exc, + ) diff --git a/apps/api/app/services/auth/api_key_service.py b/apps/api/app/services/auth/api_key_service.py deleted file mode 100644 index f2b22e4da..000000000 --- a/apps/api/app/services/auth/api_key_service.py +++ /dev/null @@ -1,362 +0,0 @@ -"""API key management service.""" - -from __future__ import annotations - -import asyncio -import json -from datetime import datetime, timezone -from typing import List, Optional - -from app.repositories.api_key_repository import APIKeyRepository -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.config import redis_pool_manager -from shared.core.database import get_db_context -from shared.core.exceptions.domain_exceptions import ( - APIKeyOperationException, - KnowhereException, - NotFoundException, - ValidationException, -) -from shared.models.database.api_key import APIKey -from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key - -from shared.services.redis.redis_service import RedisService - -_API_KEY_USER_CACHE_TTL_SECONDS: int = 3600 - - -class APIKeyService: - """API key management service.""" - - _instance: "APIKeyService | None" = None - - def __new__(cls) -> "APIKeyService": - """Return the singleton API-key service object.""" - if cls._instance is None: - cls._instance = super().__new__(cls) - return cls._instance - - def __init__(self) -> None: - if hasattr(self, "repository"): - return - self.repository = APIKeyRepository() - - @classmethod - def get_instance(cls) -> "APIKeyService": - """Return the singleton API-key service instance.""" - return cls() - - def _mask_api_key(self, api_key: str) -> str: - """Mask an API key, exposing only the first 8 and last 4 characters.""" - return mask_api_key(api_key) - - async def create_api_key( - self, - session: AsyncSession, - user_id: str, - name: str, - enabled_modules: Optional[List[str]] = None, - expires_at: Optional[datetime] = None, - ) -> str: - """Create an API key.""" - key_count = await self.repository.count_by_user(session, user_id) - if key_count >= 10: - raise ValidationException( - user_message="Maximum API Key limit reached (10)", - violations=[ - { - "field": "api_keys", - "description": "User has reached the maximum API Key limit", - } - ], - ) - - existing_key = await self.repository.get_by_user_and_name( - session, user_id, name - ) - if existing_key: - raise ValidationException( - user_message="API Key name already exists", - violations=[ - { - "field": "name", - "description": f"An API Key with name '{name}' already exists", - } - ], - ) - - api_key = generate_api_key() - key_hash = hash_api_key(api_key) - key_mask = mask_api_key(api_key) - - api_key_record = APIKey( - user_id=user_id, - key_hash=key_hash, - key_mask=key_mask, - name=name, - enabled_modules=enabled_modules or ["all"], - expires_at=expires_at, - ) - - await self.repository.create(session, api_key_record) - - return api_key - - async def validate_api_key( - self, session: AsyncSession, api_key: str - ) -> Optional[str]: - """Validate API key against DB, return user_id or None.""" - key_hash: str = hash_api_key(api_key) - cached_user_id = await self._get_cached_user_id( - redis_pool_manager.get_redis_service(), - key_hash, - ) - if cached_user_id is not None: - return cached_user_id - - api_key_record = await self.repository.get_by_key_hash(session, key_hash) - if not api_key_record or not api_key_record.is_valid(): - return None - - self._schedule_last_used_update(str(api_key_record.id)) - user_id = str(api_key_record.user_id) - await self._set_cached_user_id( - redis_pool_manager.get_redis_service(), - key_hash, - user_id, - self._resolve_api_key_cache_ttl_seconds(api_key_record.expires_at), - ) - return user_id - - @staticmethod - def _get_user_id_key(api_key_hash: str) -> str: - """Return the Redis key for an API-key hash to user ID lookup.""" - return f"api-key:user-id:{api_key_hash}" - - @staticmethod - def _get_user_api_keys_key(user_id: str) -> str: - """Return the Redis reverse-index key for a user's API-key hashes.""" - return f"api-key:user-hashes:{user_id}" - - async def _get_cached_user_id( - self, - redis_service: RedisService, - api_key_hash: str, - ) -> str | None: - """Return cached API-key user ID or None on miss/cache failure.""" - try: - raw_user_id = await redis_service.get(self._get_user_id_key(api_key_hash)) - return self._coerce_user_id(raw_user_id) - except Exception: - logger.warning("api_key_service: failed to read API-key user cache") - return None - - async def _set_cached_user_id( - self, - redis_service: RedisService, - api_key_hash: str, - user_id: str, - ttl_seconds: int, - ) -> None: - """Cache a validated API-key to user ID lookup.""" - effective_ttl_seconds = min(_API_KEY_USER_CACHE_TTL_SECONDS, ttl_seconds) - user_id_key = self._get_user_id_key(api_key_hash) - user_api_keys_key = self._get_user_api_keys_key(user_id) - - try: - await redis_service.set(user_id_key, user_id, ttl=effective_ttl_seconds) - await redis_service.sadd(user_api_keys_key, api_key_hash) - reverse_ttl_seconds = await redis_service.ttl(user_api_keys_key) - if ( - reverse_ttl_seconds in (-2, -1) - or reverse_ttl_seconds < effective_ttl_seconds - ): - await redis_service.expire(user_api_keys_key, effective_ttl_seconds) - except Exception: - logger.warning( - "api_key_service: failed to write API-key user cache for user_id={}", - user_id, - ) - - async def _invalidate_cached_api_key_user_id( - self, - redis_service: RedisService, - user_id: str, - api_key_hash: str, - ) -> None: - """Delete one API-key to user ID cache entry.""" - try: - await redis_service.delete(self._get_user_id_key(api_key_hash)) - await redis_service.srem(self._get_user_api_keys_key(user_id), api_key_hash) - except Exception: - logger.warning( - "api_key_service: failed to invalidate API-key cache for user_id={}", - user_id, - ) - - def _coerce_user_id(self, raw_user_id: object) -> str | None: - """Return a typed user ID from current or legacy Redis values.""" - if isinstance(raw_user_id, str): - try: - parsed_user_id: object = json.loads(raw_user_id) - except json.JSONDecodeError: - return raw_user_id - else: - parsed_user_id = raw_user_id - - if isinstance(parsed_user_id, str): - return parsed_user_id - - if isinstance(parsed_user_id, dict): - legacy_user_id = parsed_user_id.get("user_id") - if isinstance(legacy_user_id, str): - return legacy_user_id - - return None - - def _resolve_api_key_cache_ttl_seconds(self, expires_at: datetime | None) -> int: - """Resolve cache TTL for an API-key lookup without exceeding key expiry.""" - if expires_at is None: - return _API_KEY_USER_CACHE_TTL_SECONDS - - expires_at_utc = expires_at - if expires_at_utc.tzinfo is None: - expires_at_utc = expires_at_utc.replace(tzinfo=timezone.utc) - - now = datetime.now(timezone.utc) - remaining_seconds = int((expires_at_utc - now).total_seconds()) - return max(1, min(_API_KEY_USER_CACHE_TTL_SECONDS, remaining_seconds)) - - async def revoke_api_key( - self, session: AsyncSession, api_key_id: str, user_id: str - ) -> bool: - """Revoke an API key by deleting it directly.""" - logger.info(f"Revoking API key: api_key_id={api_key_id}, user_id={user_id}") - - api_key = await self.repository.get_by_id(session, api_key_id) - - if not api_key: - logger.warning("API key does not exist") - raise NotFoundException( - resource="APIKey", - resource_id=api_key_id, - internal_message="API Key not found", - ) - - if str(api_key.user_id) != user_id: - logger.warning( - f"User ID mismatch: api_key.user_id={api_key.user_id}, user_id={user_id}" - ) - raise NotFoundException( - resource="APIKey", - resource_id=api_key_id, - internal_message="API Key not found or does not belong to user", - ) - - success = await self.repository.delete_by_id(session, api_key_id) - logger.info(f"Delete result: {success}") - - if success: - await session.commit() - logger.info("Transaction committed") - await self._invalidate_cached_api_key_user_id( - redis_pool_manager.get_redis_service(), - user_id, - api_key.key_hash, - ) - - return success - - async def list_user_api_keys( - self, session: AsyncSession, user_id: str - ) -> List[dict]: - """List a user's API keys, including disabled ones that are still valid.""" - api_keys = await self.repository.get_unexpired_by_user_id(session, user_id) - return [ - { - "id": str(api_key.id), - "name": api_key.name, - "api_key": api_key.key_mask - or f"sk_{api_key.id[:8]}••••••••••••••••••••••••••••••••••••••••", - "enabled_modules": api_key.enabled_modules, - "is_active": api_key.is_active, - "created_at": api_key.created_at, - "last_used_at": api_key.last_used_at, - "expires_at": api_key.expires_at, - } - for api_key in api_keys - ] - - def _schedule_last_used_update(self, api_key_id: str) -> None: - """Schedule a best-effort background update for api_keys.last_used_at.""" - try: - asyncio.create_task( - self._update_last_used_best_effort(api_key_id), - name=f"api_key_last_used:{api_key_id}", - ) - except Exception as e: - logger.warning( - f"Failed to schedule API key last-used update (ignored): {e}" - ) - - async def _update_last_used_best_effort(self, api_key_id: str) -> None: - """Best-effort async update; failures are logged but never propagated.""" - try: - async with get_db_context() as db: - await self.repository.update_last_used(db, api_key_id) - except Exception as e: - logger.warning(f"Failed to update API key last-used time (ignored): {e}") - - async def get_api_key( - self, session: AsyncSession, user_id: str, api_key_id: str - ) -> Optional[APIKey]: - """Get a single API key for a user.""" - try: - api_key = await self.repository.get(session, api_key_id) - if api_key and api_key.user_id == user_id: - return api_key - return None - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to get API key: {e}") - raise APIKeyOperationException( - internal_message=f"Failed to get API key: {str(e)}", - original_exception=e, - ) - - async def toggle_api_key( - self, session: AsyncSession, user_id: str, api_key_id: str - ) -> bool: - """Enable or disable an API key.""" - try: - api_key = await self.repository.get(session, api_key_id) - if not api_key or str(api_key.user_id) != user_id: - return False - - api_key.is_active = not api_key.is_active - await session.commit() - await session.refresh(api_key) - - if not api_key.is_active: - await self._invalidate_cached_api_key_user_id( - redis_pool_manager.get_redis_service(), - user_id, - api_key.key_hash, - ) - - logger.info( - f"API key status toggled successfully: {api_key_id}, new_status={api_key.is_active}" - ) - return True - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to toggle API key status: {e}") - await session.rollback() - raise APIKeyOperationException( - internal_message=f"Failed to toggle API key status: {str(e)}", - original_exception=e, - ) diff --git a/apps/api/app/services/guest/guest_registration_service.py b/apps/api/app/services/guest/guest_registration_service.py index 58069fac2..4f189807e 100644 --- a/apps/api/app/services/guest/guest_registration_service.py +++ b/apps/api/app/services/guest/guest_registration_service.py @@ -149,8 +149,8 @@ async def _create_api_key_without_commit( ) -> str: """Generate an API key record and flush (but do not commit). - This avoids the internal commit inside APIKeyService.create_api_key() - which would make the key durable before the device row is inserted. + This keeps the guest API key row in the same transaction as the device + row so guest registration stays atomic. """ from shared.models.database.api_key import APIKey From 5439a9284e7fc8f8a47a7400100b3fa00b3a597f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 02:50:26 +0800 Subject: [PATCH 05/40] refactor: extract stripe refund reconciliation --- .../stripe_refund_reconciliation_service.py | 234 ++++++++++++++++++ .../billing/stripe_webhook_service.py | 212 ++-------------- 2 files changed, 248 insertions(+), 198 deletions(-) create mode 100644 apps/api/app/services/billing/stripe_refund_reconciliation_service.py diff --git a/apps/api/app/services/billing/stripe_refund_reconciliation_service.py b/apps/api/app/services/billing/stripe_refund_reconciliation_service.py new file mode 100644 index 000000000..d770195e1 --- /dev/null +++ b/apps/api/app/services/billing/stripe_refund_reconciliation_service.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from typing import Any +from uuid import UUID + +from app.repositories.payment_record_repository import PaymentRecordRepository +from app.services.billing.price_config_service import PriceConfigService +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.logging import logger +from shared.models.database.payment_record import PaymentRecord +from shared.services.billing import CreditsService +from shared.utils.utc_now import utc_now_naive + + +class StripeRefundReconciliationService: + def __init__( + self, + *, + payment_record_repository: PaymentRecordRepository | None = None, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._payment_record_repository = ( + payment_record_repository or PaymentRecordRepository() + ) + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + + async def reconcile_charge_refund( + self, + db: AsyncSession, + *, + event: dict[str, Any], + ) -> dict[str, object]: + charge = event["data"]["object"] + charge_id = charge.get("id") + refund_items = (charge.get("refunds", {}) or {}).get("data", []) or [] + latest_refund = refund_items[-1] if refund_items else None + + payment_intent_id = charge.get("payment_intent") + refund_id = latest_refund.get("id") if latest_refund else None + currency = (charge.get("currency") or "cny").upper() + idempotency_key = refund_id or f"{charge_id}-refund" + + original_record = None + if payment_intent_id: + original_record = await self._payment_record_repository.get_by_payment_intent_id( + db, + payment_intent_id, + ) + + metadata = charge.get("metadata") or {} + user_id = metadata.get("user_id") or ( + getattr(original_record, "user_id", None) + ) + payment_type = ( + metadata.get("type") + or getattr(original_record, "payment_type", None) + or "refund" + ) + + if not user_id: + logger.error( + f"Refund event is missing user_id; cannot record refund: charge_id={charge_id}" + ) + return { + "status": "error", + "message": "Missing user_id for refund", + "event_type": "charge.refunded", + } + + normalized_user_id = self._normalize_user_id(user_id) + if normalized_user_id is None: + logger.error(f"Invalid user_id format: {user_id}") + return { + "status": "error", + "message": "Invalid user_id format", + "event_type": "charge.refunded", + } + + user_id_str = str(normalized_user_id) + total_refund_amount_cents = charge.get("amount_refunded") or 0 + origin_total_refund_amount_cents = await self._load_recorded_refund_amount( + db, + payment_intent_id=idempotency_key, + user_id=normalized_user_id, + ) + + refund_amount_cents = ( + total_refund_amount_cents - origin_total_refund_amount_cents + ) + if refund_amount_cents <= 0: + logger.info( + f"Refund already processed, skipping: charge_id={charge_id}, refund_id={refund_id}" + ) + return { + "status": "success", + "event_type": "charge.refunded", + "message": "Already processed", + "user_id": normalized_user_id, + "refund_id": refund_id, + } + + credits_refunded = await self._calculate_refunded_credits( + db, + metadata=metadata, + original_record=original_record, + refund_amount_cents=refund_amount_cents, + ) + + if credits_refunded is not None and credits_refunded < 0: + await self._credits_service.add_credits( + session=db, + user_id=user_id_str, + amount=credits_refunded, + reason="Refund adjustment", + transaction_type="refund", + transaction_metadata={"refund_id": refund_id, "charge_id": charge_id}, + ) + + refund_metadata = { + "refund_id": refund_id, + "charge_id": charge_id, + "original_payment_intent_id": payment_intent_id, + "original_payment_record_id": getattr(original_record, "id", None), + "reason": (latest_refund or {}).get("reason"), + "balance_transaction": (latest_refund or {}).get("balance_transaction"), + } + refund_record = PaymentRecord( + payment_intent_id=idempotency_key, + user_id=normalized_user_id, + payment_type=payment_type, + amount_cents=-abs(refund_amount_cents), + currency=currency, + status="succeeded", + credits_amount=credits_refunded, + plan_id=getattr(original_record, "plan_id", None), + stripe_subscription_id=getattr( + original_record, + "stripe_subscription_id", + None, + ), + processed_at=utc_now_naive(), + extra_metadata=refund_metadata, + ) + db.add(refund_record) + await db.commit() + await db.refresh(refund_record) + + logger.info( + f"Refund record created: user_id={normalized_user_id}, amount_cents={refund_record.amount_cents}, " + f"refund_id={refund_id}, charge_id={charge_id}" + ) + return { + "status": "success", + "event_type": "charge.refunded", + "user_id": normalized_user_id, + "refund_amount_cents": abs(refund_amount_cents), + "payment_intent_id": payment_intent_id, + "refund_id": refund_id, + } + + async def _load_recorded_refund_amount( + self, + db: AsyncSession, + *, + payment_intent_id: str, + user_id: UUID, + ) -> int: + result = await db.execute( + select(func.sum(PaymentRecord.amount_cents)) + .where(PaymentRecord.payment_intent_id == payment_intent_id) + .where(PaymentRecord.user_id == user_id) + .where(PaymentRecord.amount_cents < 0) + ) + return int(abs(result.scalar() or 0)) + + async def _calculate_refunded_credits( + self, + db: AsyncSession, + *, + metadata: dict[str, Any], + original_record: PaymentRecord | None, + refund_amount_cents: int, + ) -> int | None: + credits_refunded: int | None = None + price_id = metadata.get("price_id") or ( + getattr(original_record, "extra_metadata", {}) or {} + ).get("price_id") + if price_id: + try: + price_cfg = await self._price_config_service.get_price_config( + db, + price_id, + ) + if price_cfg and price_cfg.amount_cents and price_cfg.credits_amount: + credits_refunded = -int( + price_cfg.credits_amount + * abs(refund_amount_cents) + / abs(price_cfg.amount_cents) + ) + except Exception as exc: + logger.warning( + f"Failed to calculate refunded Credits, price_id={price_id}: {exc}" + ) + credits_refunded = None + + if ( + credits_refunded is None + and original_record + and original_record.credits_amount + and original_record.amount_cents + ): + credits_refunded = -int( + abs(original_record.credits_amount) + * abs(refund_amount_cents) + / abs(original_record.amount_cents) + ) + + return credits_refunded + + def _normalize_user_id( + self, + user_id: str | UUID, + ) -> UUID | None: + if isinstance(user_id, UUID): + return user_id + + try: + return UUID(user_id) + except ValueError: + return None diff --git a/apps/api/app/services/billing/stripe_webhook_service.py b/apps/api/app/services/billing/stripe_webhook_service.py index 7da961fdb..6a7aa2848 100644 --- a/apps/api/app/services/billing/stripe_webhook_service.py +++ b/apps/api/app/services/billing/stripe_webhook_service.py @@ -2,13 +2,14 @@ from collections.abc import Awaitable, Callable from typing import Any, TypeAlias -from uuid import UUID import stripe from app.repositories.payment_record_repository import PaymentRecordRepository from app.services.billing.price_config_service import PriceConfigService +from app.services.billing.stripe_refund_reconciliation_service import ( + StripeRefundReconciliationService, +) from app.services.rate_limit.tier_service import TierService -from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from shared.core.config import settings @@ -37,6 +38,7 @@ def __init__( payment_record_repository: PaymentRecordRepository | None = None, price_config_service: PriceConfigService | None = None, credits_service: CreditsService | None = None, + refund_reconciliation_service: StripeRefundReconciliationService | None = None, ) -> None: self._configure_stripe_api() self._payment_record_repository = ( @@ -44,6 +46,14 @@ def __init__( ) self._price_config_service = price_config_service or PriceConfigService() self._credits_service = credits_service or CreditsService() + self._refund_reconciliation_service = ( + refund_reconciliation_service + or StripeRefundReconciliationService( + payment_record_repository=self._payment_record_repository, + price_config_service=self._price_config_service, + credits_service=self._credits_service, + ) + ) async def handle_webhook( self, @@ -393,204 +403,10 @@ async def _handle_charge_refunded( db: AsyncSession, event: StripeEvent, ) -> dict[str, object]: - charge = event["data"]["object"] - charge_id = charge.get("id") - refund_items = (charge.get("refunds", {}) or {}).get("data", []) or [] - latest_refund = refund_items[-1] if refund_items else None - - payment_intent_id = charge.get("payment_intent") - refund_id = latest_refund.get("id") if latest_refund else None - currency = (charge.get("currency") or "cny").upper() - idempotency_key = refund_id or f"{charge_id}-refund" - - original_record = None - if payment_intent_id: - original_record = await self._payment_record_repository.get_by_payment_intent_id( - db, - payment_intent_id, - ) - - metadata = charge.get("metadata") or {} - user_id = metadata.get("user_id") or ( - getattr(original_record, "user_id", None) - ) - payment_type = ( - metadata.get("type") - or getattr(original_record, "payment_type", None) - or "refund" - ) - - if not user_id: - logger.error( - f"Refund event is missing user_id; cannot record refund: charge_id={charge_id}" - ) - return { - "status": "error", - "message": "Missing user_id for refund", - "event_type": "charge.refunded", - } - - normalized_user_id = self._normalize_user_id(user_id) - if normalized_user_id is None: - logger.error(f"Invalid user_id format: {user_id}") - return { - "status": "error", - "message": "Invalid user_id format", - "event_type": "charge.refunded", - } - - user_id_str = str(normalized_user_id) - total_refund_amount_cents = charge.get("amount_refunded") or 0 - origin_total_refund_amount_cents = await self._load_recorded_refund_amount( + return await self._refund_reconciliation_service.reconcile_charge_refund( db, - payment_intent_id=idempotency_key, - user_id=normalized_user_id, - ) - - refund_amount_cents = ( - total_refund_amount_cents - origin_total_refund_amount_cents + event=event, ) - if refund_amount_cents <= 0: - logger.info( - f"Refund already processed, skipping: charge_id={charge_id}, refund_id={refund_id}" - ) - return { - "status": "success", - "event_type": "charge.refunded", - "message": "Already processed", - "user_id": normalized_user_id, - "refund_id": refund_id, - } - - credits_refunded = await self._calculate_refunded_credits( - db, - metadata=metadata, - original_record=original_record, - refund_amount_cents=refund_amount_cents, - ) - - if credits_refunded is not None and credits_refunded < 0: - await self._credits_service.add_credits( - session=db, - user_id=user_id_str, - amount=credits_refunded, - reason="Refund adjustment", - transaction_type="refund", - transaction_metadata={"refund_id": refund_id, "charge_id": charge_id}, - ) - - refund_metadata = { - "refund_id": refund_id, - "charge_id": charge_id, - "original_payment_intent_id": payment_intent_id, - "original_payment_record_id": getattr(original_record, "id", None), - "reason": (latest_refund or {}).get("reason"), - "balance_transaction": (latest_refund or {}).get("balance_transaction"), - } - refund_record = PaymentRecord( - payment_intent_id=idempotency_key, - user_id=normalized_user_id, - payment_type=payment_type, - amount_cents=-abs(refund_amount_cents), - currency=currency, - status="succeeded", - credits_amount=credits_refunded, - plan_id=getattr(original_record, "plan_id", None), - stripe_subscription_id=getattr( - original_record, - "stripe_subscription_id", - None, - ), - processed_at=utc_now_naive(), - extra_metadata=refund_metadata, - ) - db.add(refund_record) - await db.commit() - await db.refresh(refund_record) - - logger.info( - f"Refund record created: user_id={normalized_user_id}, amount_cents={refund_record.amount_cents}, " - f"refund_id={refund_id}, charge_id={charge_id}" - ) - return { - "status": "success", - "event_type": "charge.refunded", - "user_id": normalized_user_id, - "refund_amount_cents": abs(refund_amount_cents), - "payment_intent_id": payment_intent_id, - "refund_id": refund_id, - } - - async def _load_recorded_refund_amount( - self, - db: AsyncSession, - *, - payment_intent_id: str, - user_id: UUID, - ) -> int: - result = await db.execute( - select(func.sum(PaymentRecord.amount_cents)) - .where(PaymentRecord.payment_intent_id == payment_intent_id) - .where(PaymentRecord.user_id == user_id) - .where(PaymentRecord.amount_cents < 0) - ) - return int(abs(result.scalar() or 0)) - - async def _calculate_refunded_credits( - self, - db: AsyncSession, - *, - metadata: dict[str, Any], - original_record: PaymentRecord | None, - refund_amount_cents: int, - ) -> int | None: - credits_refunded: int | None = None - price_id = metadata.get("price_id") or ( - getattr(original_record, "extra_metadata", {}) or {} - ).get("price_id") - if price_id: - try: - price_cfg = await self._price_config_service.get_price_config( - db, - price_id, - ) - if price_cfg and price_cfg.amount_cents and price_cfg.credits_amount: - credits_refunded = -int( - price_cfg.credits_amount - * abs(refund_amount_cents) - / abs(price_cfg.amount_cents) - ) - except Exception as exc: - logger.warning( - f"Failed to calculate refunded Credits, price_id={price_id}: {exc}" - ) - credits_refunded = None - - if ( - credits_refunded is None - and original_record - and original_record.credits_amount - and original_record.amount_cents - ): - credits_refunded = -int( - abs(original_record.credits_amount) - * abs(refund_amount_cents) - / abs(original_record.amount_cents) - ) - - return credits_refunded - - def _normalize_user_id( - self, - user_id: str | UUID, - ) -> UUID | None: - if isinstance(user_id, UUID): - return user_id - - try: - return UUID(user_id) - except ValueError: - return None def _configure_stripe_api(self) -> None: if not settings.STRIPE_SECRET_KEY: From 84072fe530e0c23c602174db73a22f030c6aa2a3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 13:58:44 +0800 Subject: [PATCH 06/40] refactor: split document ingestion workflows --- ...document_ingestion_confirmation_service.py | 138 ++++++ .../document_ingestion_creation_service.py | 321 ++++++++++++++ .../services/document_ingestion_service.py | 394 ++---------------- .../contract/test_job_creation_contract.py | 4 +- 4 files changed, 485 insertions(+), 372 deletions(-) create mode 100644 apps/api/app/services/document_ingestion_confirmation_service.py create mode 100644 apps/api/app/services/document_ingestion_creation_service.py diff --git a/apps/api/app/services/document_ingestion_confirmation_service.py b/apps/api/app/services/document_ingestion_confirmation_service.py new file mode 100644 index 000000000..c9947e205 --- /dev/null +++ b/apps/api/app/services/document_ingestion_confirmation_service.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from app.repositories.job_repository import JobRepository +from app.services.job_read_service import check_job_permission +from app.services.knowledge.kb_orchestrator import KBOrchestrator +from app.services.state_machine import JobStateMachine +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + JobOperationException, + NotFoundException, + PermissionDeniedException, + ValidationException, +) +from shared.core.state_machine.states import JobStatus +from shared.services.storage.file_upload_service import FileUploadService + +_JOB_TYPE_KB_MANAGEMENT = "kb_management" + + +class DocumentIngestionConfirmationService: + def __init__( + self, + *, + job_repository: JobRepository | None = None, + file_upload_service: FileUploadService | None = None, + ) -> None: + self._job_repository = job_repository or JobRepository() + self._file_upload_service = file_upload_service or FileUploadService() + + async def confirm_upload( + self, + db: AsyncSession, + *, + job_id: str, + user_id: str, + ) -> dict[str, str]: + try: + job = await self._job_repository.get_job_by_id(db, job_id) + check_job_permission(job, user_id, job_id) + assert job is not None + + logger.info(f"Confirm upload - Job {job_id} current status: {job.status}") + if job.status not in [JobStatus.PENDING.value, JobStatus.WAITING_FILE.value]: + logger.info(f"Job {job_id} already processed, status: {job.status}") + return {"message": "Job status already updated"} + + if not job.s3_key: + raise ValidationException( + user_message="Job is missing S3 key information", + violations=[ + { + "field": "s3_key", + "description": "S3 key not set for this job", + } + ], + ) + + file_info = await self._file_upload_service.verify_s3_file_exists(job.s3_key) + if not bool(file_info.get("exists")): + raise ValidationException( + user_message="S3 file does not exist, please upload the file first", + violations=[ + {"field": "file", "description": "File not found in S3"} + ], + ) + + await _transition_job_to_uploaded(db, job_id=job_id) + await _start_job_workflow( + db=db, + job_id=job_id, + job_type=job.job_type, + source_type="file", + user_id=user_id, + ) + return {"message": "File upload confirmed; processing started"} + except NotFoundException: + raise + except PermissionDeniedException: + raise + except ValidationException: + raise + except Exception as exc: + logger.error(f"Failed to confirm upload: {exc}") + raise JobOperationException( + internal_message=f"Failed to confirm upload: {str(exc)}" + ) + + +async def _transition_job_to_uploaded( + db: AsyncSession, + *, + job_id: str, + trigger: str = "manual_upload_completed", +) -> None: + state_machine = JobStateMachine() + await state_machine.transition( + db, + job_id, + JobStatus.PENDING.value, + trigger, + None, + "system", + ) + + +async def _start_job_workflow( + db: AsyncSession, + *, + job_id: str, + job_type: str, + source_type: str, + user_id: str, + file_path: str | None = None, + file_url: str | None = None, +) -> None: + if job_type == _JOB_TYPE_KB_MANAGEMENT: + orchestrator = KBOrchestrator() + await orchestrator.start_workflow( + db=db, + job_id=job_id, + source_type=source_type, + file_path=file_path, + file_url=file_url, + user_id=user_id, + ) + return + + raise ValidationException( + user_message="Unsupported job type", + violations=[ + { + "field": "job_type", + "description": f"Job type '{job_type}' is not supported", + } + ], + ) diff --git a/apps/api/app/services/document_ingestion_creation_service.py b/apps/api/app/services/document_ingestion_creation_service.py new file mode 100644 index 000000000..c5c8c06d7 --- /dev/null +++ b/apps/api/app/services/document_ingestion_creation_service.py @@ -0,0 +1,321 @@ +from __future__ import annotations + +import os +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import cast +from urllib.parse import urlparse + +from app.repositories.job_repository import JobRepository +from app.services.job_document_scope_service import ( + is_active_document_job_unique_violation, + raise_document_ingestion_conflict, +) +from app.services.job_response_projection import to_job_status_value +from app.services.rate_limit.data_structures import CurrentUser +from loguru import logger +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + JobOperationException, + ValidationException, +) +from shared.core.state_machine.states import JobStatus +from shared.models.database.job import Job +from shared.models.schemas.job import JobCreate, JobResponse +from shared.services.redis import JobInfoRedisService, RedisServiceFactory +from shared.services.redis.job_metadata_service import JobMetadataService +from shared.services.storage.file_upload_service import FileUploadService +from shared.utils.url_file_type import resolve_file_extension_async + +_JOB_TYPE_KB_MANAGEMENT = "kb_management" +JobMetadata = dict[str, object] +UploadHeaders = dict[str, str] + + +@dataclass(frozen=True) +class ResolvedDocumentIngestionScope: + job_metadata: JobMetadata + document_id: str + namespace: str + + +class DocumentIngestionCreationService: + def __init__( + self, + *, + job_repository: JobRepository | None = None, + file_upload_service: FileUploadService | None = None, + ) -> None: + self._job_repository = job_repository or JobRepository() + self._file_upload_service = file_upload_service or FileUploadService() + + async def create_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + scope: ResolvedDocumentIngestionScope, + ) -> JobResponse: + if payload.source_type == "file": + return await self._create_file_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + scope=scope, + ) + return await self._create_url_job( + db, + payload=payload, + job_id=job_id, + current_user=current_user, + scope=scope, + ) + + async def _create_waiting_job( + self, + db: AsyncSession, + *, + job_id: str, + user_id: str, + source_type: str, + webhook_url: str | None, + job_metadata: JobMetadata, + s3_key: str, + document_id: str, + ) -> Job: + try: + job = await self._job_repository.create_job( + db=db, + job_id=job_id, + user_id=user_id, + job_type=_JOB_TYPE_KB_MANAGEMENT, + source_type=source_type, + file_path=None, + webhook_url=webhook_url, + metadata=job_metadata, + initial_state=JobStatus.WAITING_FILE.value, + s3_key=s3_key, + ) + except IntegrityError as exc: + if is_active_document_job_unique_violation(exc): + raise_document_ingestion_conflict(document_id=document_id) + raise + + if job is None: + raise JobOperationException( + internal_message="Failed to create job in database" + ) + return job + + async def _cache_job_creation_state( + self, + *, + job_id: str, + s3_key: str, + user_id: str, + webhook_enabled: bool, + source_type: str, + job_metadata: JobMetadata, + ) -> None: + redis_service = RedisServiceFactory.get_service() + metadata_service = JobMetadataService(redis_service) + await metadata_service.save_metadata(job_id, job_metadata) + + job_info_service = JobInfoRedisService(redis_service) + job_info: dict[str, object] = { + "job_id": job_id, + "s3_key": s3_key, + "user_id": user_id, + "webhook_enabled": webhook_enabled, + "job_type": _JOB_TYPE_KB_MANAGEMENT, + "source_type": source_type, + "created_at": datetime.now(timezone.utc).isoformat(), + } + await job_info_service.save_job_info(job_id, job_info) + + async def _create_file_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + scope: ResolvedDocumentIngestionScope, + ) -> JobResponse: + assert payload.file_name is not None + file_extension = os.path.splitext(payload.file_name)[1] + s3_key = f"uploads/{job_id}{file_extension}" + scope.job_metadata["source_file_name"] = payload.file_name + scope.job_metadata["source_type"] = "file" + + job = await self._create_waiting_job( + db, + job_id=job_id, + user_id=current_user.user_id, + source_type="file", + webhook_url=payload.webhook.url if payload.webhook else None, + job_metadata=scope.job_metadata, + s3_key=s3_key, + document_id=scope.document_id, + ) + + upload_info = await self._file_upload_service.generate_upload_url( + job_id, + file_extension, + ) + upload_url = cast(str, upload_info["upload_url"]) + upload_headers = cast(UploadHeaders, upload_info["upload_headers"]) + expires_in = cast(int, upload_info["expires_in"]) + + await self._cache_job_creation_state( + job_id=job_id, + s3_key=s3_key, + user_id=current_user.user_id, + webhook_enabled=bool(payload.webhook and payload.webhook.url), + source_type="file", + job_metadata=scope.job_metadata, + ) + + logger.info(f"Job {job_id} upload_url returned to client: {upload_url}") + return _build_job_response( + job_id=job_id, + job=job, + source_type="file", + data_id=payload.data_id, + namespace=scope.namespace, + upload_url=upload_url, + upload_headers=upload_headers, + expires_in=expires_in, + ) + + async def _create_url_job( + self, + db: AsyncSession, + *, + payload: JobCreate, + job_id: str, + current_user: CurrentUser, + scope: ResolvedDocumentIngestionScope, + ) -> JobResponse: + assert payload.source_url is not None + file_extension = await resolve_file_extension_async(payload.source_url) + if not file_extension: + raise ValidationException( + user_message=( + "Unsupported URL file type. Supported formats: " + f"{_get_supported_formats()}" + ), + violations=[ + { + "field": "source_url", + "description": "URL file type not supported", + } + ], + ) + + source_file_name = _resolve_url_source_file_name( + source_url=payload.source_url, + file_extension=file_extension, + ) + s3_key = f"uploads/{job_id}{file_extension}" + scope.job_metadata.update( + { + "source_file_name": source_file_name, + "source_url": payload.source_url, + "source_type": "url", + } + ) + + job = await self._create_waiting_job( + db, + job_id=job_id, + user_id=current_user.user_id, + source_type="url", + webhook_url=payload.webhook.url if payload.webhook else None, + job_metadata=scope.job_metadata, + s3_key=s3_key, + document_id=scope.document_id, + ) + + await self._cache_job_creation_state( + job_id=job_id, + s3_key=s3_key, + user_id=current_user.user_id, + webhook_enabled=bool(payload.webhook and payload.webhook.url), + source_type="url", + job_metadata=scope.job_metadata, + ) + _schedule_url_upload( + job_id=job_id, + source_url=payload.source_url, + user_id=current_user.user_id, + ) + + return _build_job_response( + job_id=job_id, + job=job, + source_type="url", + data_id=payload.data_id, + namespace=scope.namespace, + ) + + +def _get_supported_formats() -> str: + from shared.core.config import settings + + return ", ".join(sorted(settings.get_supported_extensions())) + + +def _build_job_response( + *, + job_id: str, + job: Job, + source_type: str, + data_id: str | None, + namespace: str | None = None, + document_id: str | None = None, + upload_url: str | None = None, + upload_headers: UploadHeaders | None = None, + expires_in: int | None = None, +) -> JobResponse: + return JobResponse( + job_id=job_id, + status=to_job_status_value(job.status), + source_type=source_type, + data_id=data_id, + namespace=namespace, + document_id=document_id, + created_at=job.created_at, + upload_url=upload_url, + upload_headers=upload_headers, + expires_in=expires_in, + ) + + +def _resolve_url_source_file_name(*, source_url: str, file_extension: str) -> str: + parsed_url = urlparse(source_url) + url_basename = str(os.path.basename(parsed_url.path)) + if url_basename and os.path.splitext(url_basename)[1].lower() == file_extension: + return url_basename + if url_basename: + return f"{url_basename}{file_extension}" + return f"url_file_{uuid.uuid4().hex[:8]}{file_extension}" + + +def _schedule_url_upload(*, job_id: str, source_url: str, user_id: str) -> None: + from shared.core.celery_app import get_celery_app + + celery_app = get_celery_app() + upload_url_file_task = celery_app.signature( + "app.core.tasks.kb_tasks.upload_url_file_task" + ) + upload_url_file_task.apply_async( + args=[job_id, source_url, user_id], + kwargs={"job_type": _JOB_TYPE_KB_MANAGEMENT}, + ) diff --git a/apps/api/app/services/document_ingestion_service.py b/apps/api/app/services/document_ingestion_service.py index 12f69a5b6..6bb29e35d 100644 --- a/apps/api/app/services/document_ingestion_service.py +++ b/apps/api/app/services/document_ingestion_service.py @@ -2,27 +2,24 @@ import os import uuid -from dataclasses import dataclass -from datetime import datetime, timezone from typing import cast -from urllib.parse import urlparse -from app.repositories.job_repository import JobRepository +from app.services.document_ingestion_confirmation_service import ( + DocumentIngestionConfirmationService, +) +from app.services.document_ingestion_creation_service import ( + DocumentIngestionCreationService, + ResolvedDocumentIngestionScope, +) from app.services.job_document_scope_service import ( find_active_job_for_document, - is_active_document_job_unique_violation, raise_document_ingestion_conflict, resolve_effective_document_scope, ) -from app.services.job_read_service import check_job_permission -from app.services.job_response_projection import to_job_status_value -from app.services.knowledge.kb_orchestrator import KBOrchestrator from app.services.rate_limit.data_structures import CurrentUser from app.services.rate_limit.dependencies import enforce_job_creation_capacity -from app.services.state_machine import JobStateMachine from fastapi import Request from loguru import logger -from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from shared.core.config import settings @@ -36,147 +33,25 @@ ValidationException, ) from shared.core.exceptions.webhook_exceptions import WebhookConfigException -from shared.core.state_machine.states import JobStatus -from shared.models.database.job import Job from shared.models.schemas.job import ConfirmUploadRequest, JobCreate, JobResponse from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.redis import JobInfoRedisService, RedisServiceFactory -from shared.services.redis.job_metadata_service import JobMetadataService -from shared.services.storage.file_upload_service import FileUploadService from shared.utils.url_file_type import resolve_file_extension_async from shared.utils.url_security import validate_http_url_and_resolve_ip_async -JOB_TYPE_KB_MANAGEMENT = "kb_management" JobMetadata = dict[str, object] -UploadHeaders = dict[str, str] - - -@dataclass(frozen=True) -class ResolvedDocumentIngestionScope: - job_metadata: JobMetadata - document_id: str - namespace: str - - -def _get_supported_formats() -> str: - return ", ".join(sorted(settings.get_supported_extensions())) - - -def _is_supported_file_name(file_name: str) -> bool: - if not file_name: - return False - file_extension = os.path.splitext(file_name)[1].lower() - return file_extension in settings.get_supported_extensions() - - -def _build_job_response( - *, - job_id: str, - job: Job, - source_type: str, - data_id: str | None, - namespace: str | None = None, - document_id: str | None = None, - upload_url: str | None = None, - upload_headers: UploadHeaders | None = None, - expires_in: int | None = None, -) -> JobResponse: - return JobResponse( - job_id=job_id, - status=to_job_status_value(job.status), - source_type=source_type, - data_id=data_id, - namespace=namespace, - document_id=document_id, - created_at=job.created_at, - upload_url=upload_url, - upload_headers=upload_headers, - expires_in=expires_in, - ) - - -def _resolve_url_source_file_name(*, source_url: str, file_extension: str) -> str: - parsed_url = urlparse(source_url) - url_basename = str(os.path.basename(parsed_url.path)) - if url_basename and os.path.splitext(url_basename)[1].lower() == file_extension: - return url_basename - if url_basename: - return f"{url_basename}{file_extension}" - return f"url_file_{uuid.uuid4().hex[:8]}{file_extension}" - - -def _schedule_url_upload(*, job_id: str, source_url: str, user_id: str) -> None: - from shared.core.celery_app import get_celery_app - - celery_app = get_celery_app() - upload_url_file_task = celery_app.signature( - "app.core.tasks.kb_tasks.upload_url_file_task" - ) - upload_url_file_task.apply_async( - args=[job_id, source_url, user_id], - kwargs={"job_type": JOB_TYPE_KB_MANAGEMENT}, - ) - - -async def _transition_job_to_uploaded( - db: AsyncSession, - *, - job_id: str, - trigger: str = "manual_upload_completed", -) -> None: - state_machine = JobStateMachine() - await state_machine.transition( - db, - job_id, - JobStatus.PENDING.value, - trigger, - None, - "system", - ) - - -async def _start_job_workflow( - db: AsyncSession, - *, - job_id: str, - job_type: str, - source_type: str, - user_id: str, - file_path: str | None = None, - file_url: str | None = None, -) -> None: - if job_type == JOB_TYPE_KB_MANAGEMENT: - orchestrator = KBOrchestrator() - await orchestrator.start_workflow( - db=db, - job_id=job_id, - source_type=source_type, - file_path=file_path, - file_url=file_url, - user_id=user_id, - ) - return - - raise ValidationException( - user_message="Unsupported job type", - violations=[ - { - "field": "job_type", - "description": f"Job type '{job_type}' is not supported", - } - ], - ) class DocumentIngestionService: def __init__( self, *, - job_repository: JobRepository | None = None, - file_upload_service: FileUploadService | None = None, + creation_service: DocumentIngestionCreationService | None = None, + confirmation_service: DocumentIngestionConfirmationService | None = None, ) -> None: - self._job_repository = job_repository or JobRepository() - self._file_upload_service = file_upload_service or FileUploadService() + self._creation_service = creation_service or DocumentIngestionCreationService() + self._confirmation_service = ( + confirmation_service or DocumentIngestionConfirmationService() + ) async def create_job( self, @@ -201,15 +76,7 @@ async def create_job( current_user=current_user, ) - if payload.source_type == "file": - return await self._create_file_job( - db, - payload=payload, - job_id=job_id, - current_user=current_user, - scope=scope, - ) - return await self._create_url_job( + return await self._creation_service.create_job( db, payload=payload, job_id=job_id, @@ -245,44 +112,11 @@ async def confirm_upload( del request_payload try: - job = await self._job_repository.get_job_by_id(db, job_id) - check_job_permission(job, user_id, job_id) - assert job is not None - - logger.info(f"Confirm upload - Job {job_id} current status: {job.status}") - if job.status not in [JobStatus.PENDING.value, JobStatus.WAITING_FILE.value]: - logger.info(f"Job {job_id} already processed, status: {job.status}") - return {"message": "Job status already updated"} - - if not job.s3_key: - raise ValidationException( - user_message="Job is missing S3 key information", - violations=[ - { - "field": "s3_key", - "description": "S3 key not set for this job", - } - ], - ) - - file_info = await self._file_upload_service.verify_s3_file_exists(job.s3_key) - if not bool(file_info.get("exists")): - raise ValidationException( - user_message="S3 file does not exist, please upload the file first", - violations=[ - {"field": "file", "description": "File not found in S3"} - ], - ) - - await _transition_job_to_uploaded(db, job_id=job_id) - await _start_job_workflow( + return await self._confirmation_service.confirm_upload( db=db, job_id=job_id, - job_type=job.job_type, - source_type="file", user_id=user_id, ) - return {"message": "File upload confirmed; processing started"} except NotFoundException: raise except PermissionDeniedException: @@ -335,11 +169,10 @@ async def _validate_create_payload(self, payload: JobCreate) -> None: and payload.file_name and not _is_supported_file_name(payload.file_name) ): - supported_formats = _get_supported_formats() raise ValidationException( user_message=( "Unsupported file type. Supported formats: " - f"{supported_formats}" + f"{_get_supported_formats()}" ), violations=[ {"field": "file_name", "description": "File type not supported"} @@ -350,11 +183,10 @@ async def _validate_create_payload(self, payload: JobCreate) -> None: assert payload.source_url is not None file_extension = await resolve_file_extension_async(payload.source_url) if not file_extension: - supported_formats = _get_supported_formats() raise ValidationException( user_message=( "Unsupported URL file type. Supported formats: " - f"{supported_formats}" + f"{_get_supported_formats()}" ), violations=[ { @@ -415,191 +247,13 @@ async def _resolve_scope( namespace=effective_namespace, ) - async def _create_waiting_job( - self, - db: AsyncSession, - *, - job_id: str, - user_id: str, - source_type: str, - webhook_url: str | None, - job_metadata: JobMetadata, - s3_key: str, - document_id: str, - ) -> Job: - try: - job = await self._job_repository.create_job( - db=db, - job_id=job_id, - user_id=user_id, - job_type=JOB_TYPE_KB_MANAGEMENT, - source_type=source_type, - file_path=None, - webhook_url=webhook_url, - metadata=job_metadata, - initial_state=JobStatus.WAITING_FILE.value, - s3_key=s3_key, - ) - except IntegrityError as exc: - if is_active_document_job_unique_violation(exc): - raise_document_ingestion_conflict(document_id=document_id) - raise - - if job is None: - raise JobOperationException( - internal_message="Failed to create job in database" - ) - return job - async def _cache_job_creation_state( - self, - *, - job_id: str, - s3_key: str, - user_id: str, - webhook_enabled: bool, - source_type: str, - job_metadata: JobMetadata, - ) -> None: - redis_service = RedisServiceFactory.get_service() - metadata_service = JobMetadataService(redis_service) - await metadata_service.save_metadata(job_id, job_metadata) - - job_info_service = JobInfoRedisService(redis_service) - job_info: dict[str, object] = { - "job_id": job_id, - "s3_key": s3_key, - "user_id": user_id, - "webhook_enabled": webhook_enabled, - "job_type": JOB_TYPE_KB_MANAGEMENT, - "source_type": source_type, - "created_at": datetime.now(timezone.utc).isoformat(), - } - await job_info_service.save_job_info(job_id, job_info) - - async def _create_file_job( - self, - db: AsyncSession, - *, - payload: JobCreate, - job_id: str, - current_user: CurrentUser, - scope: ResolvedDocumentIngestionScope, - ) -> JobResponse: - assert payload.file_name is not None - file_extension = os.path.splitext(payload.file_name)[1] - s3_key = f"uploads/{job_id}{file_extension}" - scope.job_metadata["source_file_name"] = payload.file_name - scope.job_metadata["source_type"] = "file" - - job = await self._create_waiting_job( - db, - job_id=job_id, - user_id=current_user.user_id, - source_type="file", - webhook_url=payload.webhook.url if payload.webhook else None, - job_metadata=scope.job_metadata, - s3_key=s3_key, - document_id=scope.document_id, - ) - - upload_info = await self._file_upload_service.generate_upload_url( - job_id, - file_extension, - ) - upload_url = cast(str, upload_info["upload_url"]) - upload_headers = cast(UploadHeaders, upload_info["upload_headers"]) - expires_in = cast(int, upload_info["expires_in"]) - - await self._cache_job_creation_state( - job_id=job_id, - s3_key=s3_key, - user_id=current_user.user_id, - webhook_enabled=bool(payload.webhook and payload.webhook.url), - source_type="file", - job_metadata=scope.job_metadata, - ) - - logger.info(f"Job {job_id} upload_url returned to client: {upload_url}") - return _build_job_response( - job_id=job_id, - job=job, - source_type="file", - data_id=payload.data_id, - namespace=scope.namespace, - upload_url=upload_url, - upload_headers=upload_headers, - expires_in=expires_in, - ) - - async def _create_url_job( - self, - db: AsyncSession, - *, - payload: JobCreate, - job_id: str, - current_user: CurrentUser, - scope: ResolvedDocumentIngestionScope, - ) -> JobResponse: - assert payload.source_url is not None - file_extension = await resolve_file_extension_async(payload.source_url) - if not file_extension: - supported_formats = _get_supported_formats() - raise ValidationException( - user_message=( - "Unsupported URL file type. Supported formats: " - f"{supported_formats}" - ), - violations=[ - { - "field": "source_url", - "description": "URL file type not supported", - } - ], - ) - - source_file_name = _resolve_url_source_file_name( - source_url=payload.source_url, - file_extension=file_extension, - ) - s3_key = f"uploads/{job_id}{file_extension}" - scope.job_metadata.update( - { - "source_file_name": source_file_name, - "source_url": payload.source_url, - "source_type": "url", - } - ) - - job = await self._create_waiting_job( - db, - job_id=job_id, - user_id=current_user.user_id, - source_type="url", - webhook_url=payload.webhook.url if payload.webhook else None, - job_metadata=scope.job_metadata, - s3_key=s3_key, - document_id=scope.document_id, - ) +def _get_supported_formats() -> str: + return ", ".join(sorted(settings.get_supported_extensions())) - await self._cache_job_creation_state( - job_id=job_id, - s3_key=s3_key, - user_id=current_user.user_id, - webhook_enabled=bool(payload.webhook and payload.webhook.url), - source_type="url", - job_metadata=scope.job_metadata, - ) - _schedule_url_upload( - job_id=job_id, - source_url=payload.source_url, - user_id=current_user.user_id, - ) - return _build_job_response( - job_id=job_id, - job=job, - source_type="url", - data_id=payload.data_id, - namespace=scope.namespace, - ) +def _is_supported_file_name(file_name: str) -> bool: + if not file_name: + return False + file_extension = os.path.splitext(file_name)[1].lower() + return file_extension in settings.get_supported_extensions() diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index b1eda40b5..33d324acc 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -1002,7 +1002,7 @@ async def _fake_start_workflow_for_job( ) async with developer_api_client_factory() as api_client: - import app.services.document_ingestion_service as document_ingestion_service + import app.services.document_ingestion_confirmation_service as document_ingestion_confirmation_service import shared.services.storage.file_upload_service as file_upload_service_module monkeypatch.setattr( @@ -1011,7 +1011,7 @@ async def _fake_start_workflow_for_job( _fake_verify_s3_file_exists, ) monkeypatch.setattr( - document_ingestion_service, + document_ingestion_confirmation_service, "_start_job_workflow", _fake_start_workflow_for_job, ) From dd0b17818f1f041b8f65853dbddab0d300d4f928 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 14:13:53 +0800 Subject: [PATCH 07/40] refactor: split job admission policies --- CONTEXT.md | 10 + .../job_admission_capacity_service.py | 206 +++++++++++ .../job_admission_route_policy_service.py | 129 +++++++ .../rate_limit/job_admission_service.py | 325 ++---------------- 4 files changed, 373 insertions(+), 297 deletions(-) create mode 100644 apps/api/app/services/rate_limit/job_admission_capacity_service.py create mode 100644 apps/api/app/services/rate_limit/job_admission_route_policy_service.py diff --git a/CONTEXT.md b/CONTEXT.md index e83408e5e..aec24b1e7 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -53,6 +53,16 @@ The policy checks that must pass before a new Job is created: authentication, guest scope, system limits, billing RPM, concurrent job limits, and daily quota. +### Job Admission Route Policy + +The route-aware part of Job Admission that enforces guest API key scope and +system limits. + +### Job Admission Capacity + +The quota-aware part of Job Admission that enforces billing RPM, concurrent +jobs, and daily quota. + ### Publication The shared workflow that turns parsed chunks into Documents, Document Sections, diff --git a/apps/api/app/services/rate_limit/job_admission_capacity_service.py b/apps/api/app/services/rate_limit/job_admission_capacity_service.py new file mode 100644 index 000000000..fb29bb7ad --- /dev/null +++ b/apps/api/app/services/rate_limit/job_admission_capacity_service.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +import math + +from app.services.rate_limit.config import ( + CONCURRENCY_RETRY_AFTER_SECONDS, + RateLimitConfig, +) +from app.services.rate_limit.data_structures import CurrentUser, TierLimits +from app.services.rate_limit.limiter import RateLimiter +from loguru import logger +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + RateLimitException, + UnavailableException, +) +from shared.core.state_machine.states import JobStatus +from shared.models.database.job import Job +from shared.models.database.user_balance import UserBalance + +_ACTIVE_JOB_STATES: tuple[str, ...] = ( + JobStatus.WAITING_FILE.value, + JobStatus.PENDING.value, + JobStatus.RUNNING.value, + JobStatus.CONVERTING.value, +) +_RETRY_AFTER_SECONDS: int = 15 + + +class JobAdmissionCapacityService: + async def enforce_billing_limits( + self, + *, + current_user: CurrentUser, + ) -> None: + if not settings.BILLING_ENABLED: + return + + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + tier_limits = self._require_tier_limits( + config=config, + current_user=current_user, + ) + limiter = RateLimiter(config) + + try: + await limiter.check_billing_rpm( + current_user.user_id, + tier_limits.rpm_limit, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in billing RPM check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + async def enforce_job_creation_capacity( + self, + *, + db: AsyncSession, + current_user: CurrentUser, + ) -> None: + if not settings.BILLING_ENABLED: + return + + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + tier_limits = self._require_tier_limits( + config=config, + current_user=current_user, + ) + limiter = RateLimiter(config) + + if tier_limits.max_concurrent_jobs != -1: + try: + await self._acquire_user_concurrency_lock( + db=db, + user_id=current_user.user_id, + ) + active_jobs = await self._count_non_terminal_jobs( + db=db, + user_id=current_user.user_id, + ) + if active_jobs >= tier_limits.max_concurrent_jobs: + retry_after_seconds = self._compute_concurrency_retry_after_seconds( + base_retry_after_seconds=CONCURRENCY_RETRY_AFTER_SECONDS, + rpm_limit=tier_limits.rpm_limit, + ) + exc = RateLimitException( + retry_after=retry_after_seconds, + limit=tier_limits.max_concurrent_jobs, + period="concurrent", + user_message=( + f"Too many concurrent requests " + f"({active_jobs}/{tier_limits.max_concurrent_jobs} active). " + f"Please retry after {retry_after_seconds} seconds." + ), + internal_message=( + "Concurrency limit exceeded: " + f"user_id={current_user.user_id}, " + f"active_jobs={active_jobs}, " + f"limit={tier_limits.max_concurrent_jobs}, " + f"retry_after={retry_after_seconds}s" + ), + ) + exc.details.update( + { + "active_jobs": active_jobs, + "available_slots": max( + 0, + tier_limits.max_concurrent_jobs - active_jobs, + ), + } + ) + raise exc + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"DB error in concurrency check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + if tier_limits.daily_quota != -1: + try: + await limiter.check_daily_quota( + current_user.user_id, + tier_limits.daily_quota, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in daily quota check: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + ) + + def _require_tier_limits( + self, + *, + config: RateLimitConfig, + current_user: CurrentUser, + ) -> TierLimits: + tier_limits = config.tier_map.get(current_user.user_tier) + if tier_limits is None: + logger.error( + "rate_limit: no tier config for tier='{}', user_id={}", + current_user.user_tier, + current_user.user_id, + ) + raise UnavailableException( + internal_message=( + f"Missing tier config for tier={current_user.user_tier}" + ), + retry_after=_RETRY_AFTER_SECONDS, + ) + return tier_limits + + async def _acquire_user_concurrency_lock( + self, + *, + db: AsyncSession, + user_id: str, + ) -> None: + result = await db.execute( + select(UserBalance.user_id) + .where(UserBalance.user_id == user_id) + .with_for_update() + ) + if result.scalar_one_or_none() is None: + raise RateLimitException( + internal_message=f"UserBalance row not found for user_id={user_id}" + ) + + async def _count_non_terminal_jobs( + self, + *, + db: AsyncSession, + user_id: str, + ) -> int: + result = await db.execute( + select(func.count(Job.job_id)) + .where(Job.user_id == user_id) + .where(Job.status.in_(_ACTIVE_JOB_STATES)) + ) + return int(result.scalar_one() or 0) + + def _compute_concurrency_retry_after_seconds( + self, + *, + base_retry_after_seconds: int, + rpm_limit: int, + ) -> int: + if rpm_limit <= 0: + return base_retry_after_seconds + return max(base_retry_after_seconds, int(math.ceil(60 / rpm_limit))) diff --git a/apps/api/app/services/rate_limit/job_admission_route_policy_service.py b/apps/api/app/services/rate_limit/job_admission_route_policy_service.py new file mode 100644 index 000000000..d455f9927 --- /dev/null +++ b/apps/api/app/services/rate_limit/job_admission_route_policy_service.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from fnmatch import fnmatch + +from app.services.rate_limit.config import RateLimitConfig +from app.services.rate_limit.limiter import RateLimiter +from app.services.rate_limit.system_limit import find_system_rule +from fastapi import Request + +from shared.core.exceptions.domain_exceptions import ( + PermissionDeniedException, + RateLimitException, + UnavailableException, +) + +_GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS: tuple[str, ...] = ( + "/v1/jobs", + "/v1/jobs/*", + "/v1/billing/credits", + "/v1/retrieval/query", + "/v1/documents", + "/v1/documents/*", + "/mcp", +) +_GUEST_API_KEY_REQUIRED_PERMISSION: str = ( + "jobs_documents_retrieval_mcp_or_billing_credits" +) +_GUEST_API_KEY_SCOPE_MESSAGE: str = ( + "Guest API keys can only access job, document, retrieval, MCP query, " + "and billing credits APIs" +) +_RETRY_AFTER_SECONDS: int = 15 + + +class JobAdmissionRoutePolicyService: + async def enforce_user_system_limit( + self, + *, + request: Request, + config: RateLimitConfig, + user_id: str, + ) -> None: + route_path = self._get_route_path(request) + rule = find_system_rule(request.method, route_path, config.system_rules) + limiter = RateLimiter(config) + await limiter.check_system_limit( + identifier=user_id, + limit=rule.limit, + matched_pattern=rule.api_pattern, + period=rule.period, + ) + + async def enforce_route_system_limit(self, *, request: Request) -> None: + config = RateLimitConfig.get_instance() + if not config.is_enabled: + return + + route_path = self._get_route_path(request) + route_identifier = self._get_route_limit_identifier(request) + rule = find_system_rule(request.method, route_path, config.system_rules) + limiter = RateLimiter(config) + + try: + await limiter.check_system_limit( + identifier=route_identifier, + limit=rule.limit, + matched_pattern=rule.api_pattern, + period=rule.period, + use_global_key=True, + ) + except RateLimitException: + raise + except Exception as exc: + raise UnavailableException( + internal_message=(f"Redis error in route system limit: {exc}"), + retry_after=_RETRY_AFTER_SECONDS, + limit=rule.limit, + period=rule.period, + ) + + def enforce_guest_api_key_scope( + self, + *, + request: Request, + user_tier: str, + ) -> None: + if user_tier != "guest": + return + + route_path = self._get_route_path(request) + if self._is_guest_api_key_route_allowed(route_path): + return + + raise PermissionDeniedException( + user_message=_GUEST_API_KEY_SCOPE_MESSAGE, + required_permission=_GUEST_API_KEY_REQUIRED_PERMISSION, + ) + + def _get_route_path(self, request: Request) -> str: + scope_path = request.scope.get("path", request.url.path) + root_path = request.scope.get("root_path", "") + if isinstance(scope_path, str) and isinstance(root_path, str): + if root_path and scope_path.startswith(root_path): + return scope_path[len(root_path) :] + return scope_path + return request.url.path + + def _get_route_limit_identifier(self, request: Request) -> str: + route = request.scope.get("route") + route_path = getattr(route, "path", None) + if isinstance(route_path, str) and route_path: + return route_path + + route_path_format = getattr(route, "path_format", None) + if isinstance(route_path_format, str) and route_path_format: + return route_path_format + + return self._get_route_path(request) + + def _normalize_route_path(self, route_path: str) -> str: + normalized_path = route_path.rstrip("/") + return normalized_path or "/" + + def _is_guest_api_key_route_allowed(self, route_path: str) -> bool: + normalized_path = self._normalize_route_path(route_path) + return any( + fnmatch(normalized_path, pattern) + for pattern in _GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS + ) diff --git a/apps/api/app/services/rate_limit/job_admission_service.py b/apps/api/app/services/rate_limit/job_admission_service.py index ff410b46e..554b8a36d 100644 --- a/apps/api/app/services/rate_limit/job_admission_service.py +++ b/apps/api/app/services/rate_limit/job_admission_service.py @@ -1,58 +1,34 @@ from __future__ import annotations -import math -from fnmatch import fnmatch - -from app.services.rate_limit.config import ( - CONCURRENCY_RETRY_AFTER_SECONDS, - RateLimitConfig, +from app.services.rate_limit.config import RateLimitConfig +from app.services.rate_limit.data_structures import CurrentUser +from app.services.rate_limit.job_admission_capacity_service import ( + JobAdmissionCapacityService, +) +from app.services.rate_limit.job_admission_route_policy_service import ( + JobAdmissionRoutePolicyService, ) -from app.services.rate_limit.data_structures import CurrentUser, TierLimits -from app.services.rate_limit.limiter import RateLimiter -from app.services.rate_limit.system_limit import find_system_rule from app.services.rate_limit.tier_service import TierService from fastapi import Request from loguru import logger -from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - PermissionDeniedException, - RateLimitException, - UnavailableException, -) +from shared.core.exceptions.domain_exceptions import RateLimitException from shared.core.logging import log_context -from shared.core.state_machine.states import JobStatus -from shared.models.database.job import Job -from shared.models.database.user_balance import UserBalance - -_ACTIVE_JOB_STATES: tuple[str, ...] = ( - JobStatus.WAITING_FILE.value, - JobStatus.PENDING.value, - JobStatus.RUNNING.value, - JobStatus.CONVERTING.value, -) -_GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS: tuple[str, ...] = ( - "/v1/jobs", - "/v1/jobs/*", - "/v1/billing/credits", - "/v1/retrieval/query", - "/v1/documents", - "/v1/documents/*", - "/mcp", -) -_GUEST_API_KEY_REQUIRED_PERMISSION: str = ( - "jobs_documents_retrieval_mcp_or_billing_credits" -) -_GUEST_API_KEY_SCOPE_MESSAGE: str = ( - "Guest API keys can only access job, document, retrieval, MCP query, " - "and billing credits APIs" -) -_RETRY_AFTER_SECONDS: int = 15 class JobAdmissionService: + def __init__( + self, + *, + route_policy_service: JobAdmissionRoutePolicyService | None = None, + capacity_service: JobAdmissionCapacityService | None = None, + ) -> None: + self._route_policy_service = ( + route_policy_service or JobAdmissionRoutePolicyService() + ) + self._capacity_service = capacity_service or JobAdmissionCapacityService() + async def resolve_current_user( self, *, @@ -60,7 +36,10 @@ async def resolve_current_user( user_id: str, ) -> CurrentUser: user_tier = await TierService.get_tier(user_id) - self._enforce_guest_api_key_scope(request=request, user_tier=user_tier) + self._route_policy_service.enforce_guest_api_key_scope( + request=request, + user_tier=user_tier, + ) current_user = CurrentUser(user_id=user_id, user_tier=user_tier) with log_context(user_id=user_id): @@ -69,7 +48,7 @@ async def resolve_current_user( return current_user try: - await self._check_user_system_limit( + await self._route_policy_service.enforce_user_system_limit( request=request, config=config, user_id=user_id, @@ -91,59 +70,10 @@ async def enforce_billing_limits( *, current_user: CurrentUser, ) -> None: - if not settings.BILLING_ENABLED: - return - - config = RateLimitConfig.get_instance() - if not config.is_enabled: - return - - tier_limits = self._require_tier_limits( - config=config, - current_user=current_user, - ) - - limiter = RateLimiter(config) - try: - await limiter.check_billing_rpm( - current_user.user_id, - tier_limits.rpm_limit, - ) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in billing RPM check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) + await self._capacity_service.enforce_billing_limits(current_user=current_user) async def enforce_route_system_limit(self, *, request: Request) -> None: - config = RateLimitConfig.get_instance() - if not config.is_enabled: - return - - route_path = self._get_route_path(request) - route_identifier = self._get_route_limit_identifier(request) - rule = find_system_rule(request.method, route_path, config.system_rules) - limiter = RateLimiter(config) - - try: - await limiter.check_system_limit( - identifier=route_identifier, - limit=rule.limit, - matched_pattern=rule.api_pattern, - period=rule.period, - use_global_key=True, - ) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in route system limit: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - limit=rule.limit, - period=rule.period, - ) + await self._route_policy_service.enforce_route_system_limit(request=request) async def enforce_job_creation_capacity( self, @@ -151,206 +81,7 @@ async def enforce_job_creation_capacity( db: AsyncSession, current_user: CurrentUser, ) -> None: - if not settings.BILLING_ENABLED: - return - - config = RateLimitConfig.get_instance() - if not config.is_enabled: - return - - tier_limits = self._require_tier_limits( - config=config, + await self._capacity_service.enforce_job_creation_capacity( + db=db, current_user=current_user, ) - limiter = RateLimiter(config) - - if tier_limits.max_concurrent_jobs != -1: - try: - await self._acquire_user_concurrency_lock( - db=db, - user_id=current_user.user_id, - ) - active_jobs = await self._count_non_terminal_jobs( - db=db, - user_id=current_user.user_id, - ) - if active_jobs >= tier_limits.max_concurrent_jobs: - retry_after_seconds = self._compute_concurrency_retry_after_seconds( - base_retry_after_seconds=CONCURRENCY_RETRY_AFTER_SECONDS, - rpm_limit=tier_limits.rpm_limit, - ) - exc = RateLimitException( - retry_after=retry_after_seconds, - limit=tier_limits.max_concurrent_jobs, - period="concurrent", - user_message=( - f"Too many concurrent requests " - f"({active_jobs}/{tier_limits.max_concurrent_jobs} active). " - f"Please retry after {retry_after_seconds} seconds." - ), - internal_message=( - "Concurrency limit exceeded: " - f"user_id={current_user.user_id}, " - f"active_jobs={active_jobs}, " - f"limit={tier_limits.max_concurrent_jobs}, " - f"retry_after={retry_after_seconds}s" - ), - ) - exc.details.update( - { - "active_jobs": active_jobs, - "available_slots": max( - 0, - tier_limits.max_concurrent_jobs - active_jobs, - ), - } - ) - raise exc - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"DB error in concurrency check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - if tier_limits.daily_quota != -1: - try: - await limiter.check_daily_quota( - current_user.user_id, - tier_limits.daily_quota, - ) - except RateLimitException: - raise - except Exception as exc: - raise UnavailableException( - internal_message=(f"Redis error in daily quota check: {exc}"), - retry_after=_RETRY_AFTER_SECONDS, - ) - - async def _check_user_system_limit( - self, - *, - request: Request, - config: RateLimitConfig, - user_id: str, - ) -> None: - route_path = self._get_route_path(request) - rule = find_system_rule(request.method, route_path, config.system_rules) - limiter = RateLimiter(config) - await limiter.check_system_limit( - identifier=user_id, - limit=rule.limit, - matched_pattern=rule.api_pattern, - period=rule.period, - ) - - def _require_tier_limits( - self, - *, - config: RateLimitConfig, - current_user: CurrentUser, - ) -> TierLimits: - tier_limits = config.tier_map.get(current_user.user_tier) - if tier_limits is None: - logger.error( - "rate_limit: no tier config for tier='{}', user_id={}", - current_user.user_tier, - current_user.user_id, - ) - raise UnavailableException( - internal_message=( - f"Missing tier config for tier={current_user.user_tier}" - ), - retry_after=_RETRY_AFTER_SECONDS, - ) - return tier_limits - - def _get_route_path(self, request: Request) -> str: - scope_path = request.scope.get("path", request.url.path) - root_path = request.scope.get("root_path", "") - if isinstance(scope_path, str) and isinstance(root_path, str): - if root_path and scope_path.startswith(root_path): - return scope_path[len(root_path) :] - return scope_path - return request.url.path - - def _get_route_limit_identifier(self, request: Request) -> str: - route = request.scope.get("route") - route_path = getattr(route, "path", None) - if isinstance(route_path, str) and route_path: - return route_path - - route_path_format = getattr(route, "path_format", None) - if isinstance(route_path_format, str) and route_path_format: - return route_path_format - - return self._get_route_path(request) - - def _normalize_route_path(self, route_path: str) -> str: - normalized_path = route_path.rstrip("/") - return normalized_path or "/" - - def _is_guest_api_key_route_allowed(self, route_path: str) -> bool: - normalized_path = self._normalize_route_path(route_path) - return any( - fnmatch(normalized_path, pattern) - for pattern in _GUEST_API_KEY_ALLOWED_ROUTE_PATTERNS - ) - - def _enforce_guest_api_key_scope( - self, - *, - request: Request, - user_tier: str, - ) -> None: - if user_tier != "guest": - return - - route_path = self._get_route_path(request) - if self._is_guest_api_key_route_allowed(route_path): - return - - raise PermissionDeniedException( - user_message=_GUEST_API_KEY_SCOPE_MESSAGE, - required_permission=_GUEST_API_KEY_REQUIRED_PERMISSION, - ) - - async def _acquire_user_concurrency_lock( - self, - *, - db: AsyncSession, - user_id: str, - ) -> None: - result = await db.execute( - select(UserBalance.user_id) - .where(UserBalance.user_id == user_id) - .with_for_update() - ) - if result.scalar_one_or_none() is None: - raise RateLimitException( - internal_message=f"UserBalance row not found for user_id={user_id}" - ) - - async def _count_non_terminal_jobs( - self, - *, - db: AsyncSession, - user_id: str, - ) -> int: - result = await db.execute( - select(func.count(Job.job_id)) - .where(Job.user_id == user_id) - .where(Job.status.in_(_ACTIVE_JOB_STATES)) - ) - return int(result.scalar_one() or 0) - - def _compute_concurrency_retry_after_seconds( - self, - *, - base_retry_after_seconds: int, - rpm_limit: int, - ) -> int: - if rpm_limit <= 0: - return base_retry_after_seconds - return max(base_retry_after_seconds, int(math.ceil(60 / rpm_limit))) From 67418bda01c3c95286c436fa2725bf65fc143491 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 14:28:37 +0800 Subject: [PATCH 08/40] refactor: extract stripe credits settlement --- CONTEXT.md | 5 + .../stripe_credits_settlement_service.py | 283 ++++++++++++++++++ .../billing/stripe_webhook_service.py | 252 +--------------- 3 files changed, 304 insertions(+), 236 deletions(-) create mode 100644 apps/api/app/services/billing/stripe_credits_settlement_service.py diff --git a/CONTEXT.md b/CONTEXT.md index aec24b1e7..539fe5f9f 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -102,6 +102,11 @@ keys. The Billing Workflow adapter that creates Stripe payment intents and checkout sessions for credits purchases. +### Stripe Credits Settlement + +The Billing Workflow adapter that settles successful Stripe checkout and +payment-intent events into credits, payment records, and tier refreshes. + ### Stripe Webhook Reconciliation The Billing Workflow adapter that verifies Stripe events and reconciles credits, diff --git a/apps/api/app/services/billing/stripe_credits_settlement_service.py b/apps/api/app/services/billing/stripe_credits_settlement_service.py new file mode 100644 index 000000000..9091dd6e6 --- /dev/null +++ b/apps/api/app/services/billing/stripe_credits_settlement_service.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +from typing import Any + +from app.repositories.payment_record_repository import PaymentRecordRepository +from app.services.billing.price_config_service import PriceConfigService +from app.services.rate_limit.tier_service import TierService +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + KnowhereException, + StripeServiceException, +) +from shared.models.database.payment_record import PaymentRecord +from shared.services.billing import CreditsService +from shared.utils.utc_now import utc_now_naive + + +class StripeCreditsSettlementService: + def __init__( + self, + *, + payment_record_repository: PaymentRecordRepository | None = None, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._payment_record_repository = ( + payment_record_repository or PaymentRecordRepository() + ) + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + + async def handle_checkout_completed( + self, + db: AsyncSession, + *, + event: dict[str, Any], + ) -> dict[str, object]: + session = event["data"]["object"] + session_id = str(session["id"]) + mode = session.get("mode") + metadata = session.get("metadata", {}) + user_id = metadata.get("user_id") + payment_type = metadata.get("type") + quantity = int(metadata.get("quantity", 1)) + + if not user_id: + logger.warning( + f"Checkout session {session_id} is missing user_id metadata; likely a test event, skipping" + ) + return { + "status": "ignored", + "message": "Missing user_id metadata (likely test event)", + "checkout_session_id": session_id, + "event_type": "checkout.session.completed", + } + + if await self._payment_record_repository.is_processed( + db, + checkout_session_id=session_id, + ): + logger.info(f"Checkout session {session_id} already processed, skipping") + return { + "status": "ignored", + "message": "Already processed", + "checkout_session_id": session_id, + } + + payment_metadata = { + "session_id": session_id, + "stripe_session": session, + } + payment_record = PaymentRecord( + checkout_session_id=session_id, + user_id=user_id, + payment_type=payment_type or "unknown", + amount_cents=session.get("amount_total", 0), + currency=session.get("currency", "cny").upper(), + status="pending", + extra_metadata=payment_metadata, + ) + db.add(payment_record) + await db.flush() + + try: + if mode != "payment" or payment_type != "credits_package": + logger.warning(f"Unknown payment type: mode={mode}, type={payment_type}") + return {"status": "ignored", "message": "Unknown payment type"} + + price_id = metadata.get("price_id") + if not price_id: + logger.error(f"Incomplete Credits pack info: price_id={price_id}") + return {"status": "error", "message": "Missing price_id"} + + price_config = await self._price_config_service.get_price_config(db, price_id) + configured_credits_amount = price_config.credits_amount + if configured_credits_amount is None: + logger.error( + f"Credits amount is not configured for price ID {price_id}" + ) + return { + "status": "error", + "message": "Credits amount not configured", + } + credits_amount = configured_credits_amount * quantity + + product_description = f"Credits pack - {credits_amount} Credits" + if price_config.extra_metadata and price_config.extra_metadata.get( + "description" + ): + product_description = str( + price_config.extra_metadata.get("description") + ) + + payment_record.extra_metadata = { + **payment_metadata, + "product_description": product_description, + "price_id": price_id, + "credits_amount": credits_amount, + "product_metadata": price_config.extra_metadata or {}, + } + await self._credits_service.add_credits( + session=db, + user_id=user_id, + amount=credits_amount, + reason=f"Purchase credits pack: {product_description}", + stripe_payment_id=session.get("payment_intent"), + ) + payment_record.status = "succeeded" + payment_record.credits_amount = credits_amount + payment_record.processed_at = utc_now_naive() + + await TierService.refresh_tier(user_id, db) + await db.commit() + await db.refresh(payment_record) + + logger.info( + f"Credits pack purchase succeeded: user_id={user_id}, credits={credits_amount}, price_id={price_id}" + ) + return { + "status": "success", + "event_type": "checkout.session.completed", + "user_id": user_id, + "credits_amount": credits_amount, + "payment_type": "credits_package", + } + except KnowhereException: + raise + except Exception as exc: + logger.error( + f"Failed to process checkout.session.completed: {exc}", + exc_info=True, + ) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": str(exc), + } + await db.commit() + raise StripeServiceException( + internal_message=( + "Failed to process checkout.session.completed: " + f"{str(exc)}" + ), + original_exception=exc, + ) + + async def handle_payment_intent_succeeded( + self, + db: AsyncSession, + *, + event: dict[str, Any], + ) -> dict[str, object]: + payment_intent = event["data"]["object"] + payment_intent_id = str(payment_intent["id"]) + metadata = payment_intent.get("metadata", {}) + user_id = metadata.get("user_id") + payment_type = metadata.get("type") + + if payment_type != "credits": + logger.info( + f"PaymentIntent {payment_intent_id} is not a Credits payment, skipping" + ) + return {"status": "ignored", "payment_intent_id": payment_intent_id} + + if not user_id: + logger.warning( + f"PaymentIntent {payment_intent_id} is missing user_id metadata; likely a test event, skipping" + ) + return { + "status": "ignored", + "message": "Missing user_id metadata (likely test event)", + "payment_intent_id": payment_intent_id, + } + + if await self._payment_record_repository.is_processed( + db, + payment_intent_id=payment_intent_id, + ): + logger.info( + f"PaymentIntent {payment_intent_id} already processed, skipping" + ) + return { + "status": "ignored", + "message": "Already processed", + "payment_intent_id": payment_intent_id, + } + + payment_metadata = { + "payment_intent_id": payment_intent_id, + "stripe_payment_intent": payment_intent, + } + payment_record = PaymentRecord( + payment_intent_id=payment_intent_id, + user_id=user_id, + payment_type="credits_package", + amount_cents=payment_intent.get("amount", 0), + currency=payment_intent.get("currency", "cny").upper(), + status="pending", + extra_metadata=payment_metadata, + ) + db.add(payment_record) + await db.flush() + + try: + credits_amount_str = metadata.get("credits_amount") + if not credits_amount_str: + logger.error( + f"PaymentIntent {payment_intent_id} is missing credits_amount" + ) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": "Missing credits_amount", + } + await db.commit() + return {"status": "error", "message": "Missing credits_amount"} + + credits_amount = int(credits_amount_str) + payment_record.extra_metadata = { + **payment_metadata, + "product_description": f"Credits package - {credits_amount} Credits", + "credits_amount": credits_amount, + "payment_method": "payment_intent", + } + await self._credits_service.add_credits( + session=db, + user_id=user_id, + amount=credits_amount, + reason=f"buy credits - {credits_amount} Credits", + stripe_payment_id=payment_intent_id, + ) + payment_record.status = "succeeded" + payment_record.credits_amount = credits_amount + payment_record.processed_at = utc_now_naive() + + await TierService.refresh_tier(user_id, db) + await db.commit() + await db.refresh(payment_record) + + logger.info( + f"buy credits success: user_id={user_id}, credits={credits_amount}, payment_intent_id={payment_intent_id}" + ) + return { + "status": "success", + "event_type": "payment_intent.succeeded", + "user_id": user_id, + "credits_amount": credits_amount, + "payment_type": "credits_package", + } + except Exception as exc: + logger.error(f"Failed to process Credits purchase: {exc}", exc_info=True) + payment_record.status = "failed" + payment_record.extra_metadata = { + **(payment_record.extra_metadata or {}), + "error": str(exc), + } + await db.commit() + raise StripeServiceException( + internal_message=f"Failed to process Credits purchase: {str(exc)}", + original_exception=exc, + ) diff --git a/apps/api/app/services/billing/stripe_webhook_service.py b/apps/api/app/services/billing/stripe_webhook_service.py index 6a7aa2848..41581d287 100644 --- a/apps/api/app/services/billing/stripe_webhook_service.py +++ b/apps/api/app/services/billing/stripe_webhook_service.py @@ -5,11 +5,13 @@ import stripe from app.repositories.payment_record_repository import PaymentRecordRepository +from app.services.billing.stripe_credits_settlement_service import ( + StripeCreditsSettlementService, +) from app.services.billing.price_config_service import PriceConfigService from app.services.billing.stripe_refund_reconciliation_service import ( StripeRefundReconciliationService, ) -from app.services.rate_limit.tier_service import TierService from sqlalchemy.ext.asyncio import AsyncSession from shared.core.config import settings @@ -21,9 +23,7 @@ ValidationException, ) from shared.core.logging import logger -from shared.models.database.payment_record import PaymentRecord from shared.services.billing import CreditsService -from shared.utils.utc_now import utc_now_naive StripeEvent: TypeAlias = dict[str, Any] StripeWebhookHandler: TypeAlias = Callable[ @@ -38,6 +38,7 @@ def __init__( payment_record_repository: PaymentRecordRepository | None = None, price_config_service: PriceConfigService | None = None, credits_service: CreditsService | None = None, + credits_settlement_service: StripeCreditsSettlementService | None = None, refund_reconciliation_service: StripeRefundReconciliationService | None = None, ) -> None: self._configure_stripe_api() @@ -46,6 +47,14 @@ def __init__( ) self._price_config_service = price_config_service or PriceConfigService() self._credits_service = credits_service or CreditsService() + self._credits_settlement_service = ( + credits_settlement_service + or StripeCreditsSettlementService( + payment_record_repository=self._payment_record_repository, + price_config_service=self._price_config_service, + credits_service=self._credits_service, + ) + ) self._refund_reconciliation_service = ( refund_reconciliation_service or StripeRefundReconciliationService( @@ -109,249 +118,20 @@ async def _handle_checkout_completed( db: AsyncSession, event: StripeEvent, ) -> dict[str, object]: - session = event["data"]["object"] - session_id = str(session["id"]) - mode = session.get("mode") - metadata = session.get("metadata", {}) - user_id = metadata.get("user_id") - payment_type = metadata.get("type") - quantity = int(metadata.get("quantity", 1)) - - if not user_id: - logger.warning( - f"Checkout session {session_id} is missing user_id metadata; likely a test event, skipping" - ) - return { - "status": "ignored", - "message": "Missing user_id metadata (likely test event)", - "checkout_session_id": session_id, - "event_type": "checkout.session.completed", - } - - if await self._payment_record_repository.is_processed( + return await self._credits_settlement_service.handle_checkout_completed( db, - checkout_session_id=session_id, - ): - logger.info(f"Checkout session {session_id} already processed, skipping") - return { - "status": "ignored", - "message": "Already processed", - "checkout_session_id": session_id, - } - - payment_metadata = { - "session_id": session_id, - "stripe_session": session, - } - payment_record = PaymentRecord( - checkout_session_id=session_id, - user_id=user_id, - payment_type=payment_type or "unknown", - amount_cents=session.get("amount_total", 0), - currency=session.get("currency", "cny").upper(), - status="pending", - extra_metadata=payment_metadata, + event=event, ) - db.add(payment_record) - await db.flush() - - try: - if mode != "payment" or payment_type != "credits_package": - logger.warning(f"Unknown payment type: mode={mode}, type={payment_type}") - return {"status": "ignored", "message": "Unknown payment type"} - - price_id = metadata.get("price_id") - if not price_id: - logger.error(f"Incomplete Credits pack info: price_id={price_id}") - return {"status": "error", "message": "Missing price_id"} - - price_config = await self._price_config_service.get_price_config(db, price_id) - configured_credits_amount = price_config.credits_amount - if configured_credits_amount is None: - logger.error( - f"Credits amount is not configured for price ID {price_id}" - ) - return { - "status": "error", - "message": "Credits amount not configured", - } - credits_amount = configured_credits_amount * quantity - - product_description = f"Credits pack - {credits_amount} Credits" - if price_config.extra_metadata and price_config.extra_metadata.get( - "description" - ): - product_description = str( - price_config.extra_metadata.get("description") - ) - - payment_record.extra_metadata = { - **payment_metadata, - "product_description": product_description, - "price_id": price_id, - "credits_amount": credits_amount, - "product_metadata": price_config.extra_metadata or {}, - } - await self._credits_service.add_credits( - session=db, - user_id=user_id, - amount=credits_amount, - reason=f"Purchase credits pack: {product_description}", - stripe_payment_id=session.get("payment_intent"), - ) - payment_record.status = "succeeded" - payment_record.credits_amount = credits_amount - payment_record.processed_at = utc_now_naive() - - await TierService.refresh_tier(user_id, db) - await db.commit() - await db.refresh(payment_record) - - logger.info( - f"Credits pack purchase succeeded: user_id={user_id}, credits={credits_amount}, price_id={price_id}" - ) - return { - "status": "success", - "event_type": "checkout.session.completed", - "user_id": user_id, - "credits_amount": credits_amount, - "payment_type": "credits_package", - } - except KnowhereException: - raise - except Exception as exc: - logger.error( - f"Failed to process checkout.session.completed: {exc}", - exc_info=True, - ) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": str(exc), - } - await db.commit() - raise StripeServiceException( - internal_message=( - "Failed to process checkout.session.completed: " - f"{str(exc)}" - ), - original_exception=exc, - ) async def _handle_payment_intent_succeeded( self, db: AsyncSession, event: StripeEvent, ) -> dict[str, object]: - payment_intent = event["data"]["object"] - payment_intent_id = str(payment_intent["id"]) - metadata = payment_intent.get("metadata", {}) - user_id = metadata.get("user_id") - payment_type = metadata.get("type") - - if payment_type != "credits": - logger.info( - f"PaymentIntent {payment_intent_id} is not a Credits payment, skipping" - ) - return {"status": "ignored", "payment_intent_id": payment_intent_id} - - if not user_id: - logger.warning( - f"PaymentIntent {payment_intent_id} is missing user_id metadata; likely a test event, skipping" - ) - return { - "status": "ignored", - "message": "Missing user_id metadata (likely test event)", - "payment_intent_id": payment_intent_id, - } - - if await self._payment_record_repository.is_processed( + return await self._credits_settlement_service.handle_payment_intent_succeeded( db, - payment_intent_id=payment_intent_id, - ): - logger.info( - f"PaymentIntent {payment_intent_id} already processed, skipping" - ) - return { - "status": "ignored", - "message": "Already processed", - "payment_intent_id": payment_intent_id, - } - - payment_metadata = { - "payment_intent_id": payment_intent_id, - "stripe_payment_intent": payment_intent, - } - payment_record = PaymentRecord( - payment_intent_id=payment_intent_id, - user_id=user_id, - payment_type="credits_package", - amount_cents=payment_intent.get("amount", 0), - currency=payment_intent.get("currency", "cny").upper(), - status="pending", - extra_metadata=payment_metadata, + event=event, ) - db.add(payment_record) - await db.flush() - - try: - credits_amount_str = metadata.get("credits_amount") - if not credits_amount_str: - logger.error( - f"PaymentIntent {payment_intent_id} is missing credits_amount" - ) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": "Missing credits_amount", - } - await db.commit() - return {"status": "error", "message": "Missing credits_amount"} - - credits_amount = int(credits_amount_str) - payment_record.extra_metadata = { - **payment_metadata, - "product_description": f"Credits package - {credits_amount} Credits", - "credits_amount": credits_amount, - "payment_method": "payment_intent", - } - await self._credits_service.add_credits( - session=db, - user_id=user_id, - amount=credits_amount, - reason=f"buy credits - {credits_amount} Credits", - stripe_payment_id=payment_intent_id, - ) - payment_record.status = "succeeded" - payment_record.credits_amount = credits_amount - payment_record.processed_at = utc_now_naive() - - await TierService.refresh_tier(user_id, db) - await db.commit() - await db.refresh(payment_record) - - logger.info( - f"buy credits success: user_id={user_id}, credits={credits_amount}, payment_intent_id={payment_intent_id}" - ) - return { - "status": "success", - "event_type": "payment_intent.succeeded", - "user_id": user_id, - "credits_amount": credits_amount, - "payment_type": "credits_package", - } - except Exception as exc: - logger.error(f"Failed to process Credits purchase: {exc}", exc_info=True) - payment_record.status = "failed" - payment_record.extra_metadata = { - **(payment_record.extra_metadata or {}), - "error": str(exc), - } - await db.commit() - raise StripeServiceException( - internal_message=f"Failed to process Credits purchase: {str(exc)}", - original_exception=exc, - ) async def _handle_payment_succeeded( self, From 2bc351cb3f4d7049475aca17ec94c49fd2c95678 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 15:20:17 +0800 Subject: [PATCH 09/40] refactor: package document ingestion workflow --- CONTEXT.md | 7 ++++--- apps/api/app/api/v1/routes/jobs.py | 2 +- apps/api/app/services/document_ingestion/__init__.py | 3 +++ .../confirmation_service.py} | 0 .../creation_service.py} | 2 +- .../scope_service.py} | 2 +- .../service.py} | 6 +++--- apps/api/tests/contract/test_job_creation_contract.py | 2 +- 8 files changed, 14 insertions(+), 10 deletions(-) create mode 100644 apps/api/app/services/document_ingestion/__init__.py rename apps/api/app/services/{document_ingestion_confirmation_service.py => document_ingestion/confirmation_service.py} (100%) rename apps/api/app/services/{document_ingestion_creation_service.py => document_ingestion/creation_service.py} (99%) rename apps/api/app/services/{job_document_scope_service.py => document_ingestion/scope_service.py} (98%) rename apps/api/app/services/{document_ingestion_service.py => document_ingestion/service.py} (97%) diff --git a/CONTEXT.md b/CONTEXT.md index 539fe5f9f..a4bdca7a4 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -161,9 +161,10 @@ exceptions. ### Document Ingestion - `app/api/v1/routes/jobs.py` -- `app/services/job_creation_service.py` -- `app/services/job_upload_confirmation_service.py` -- `app/services/job_document_scope_service.py` +- `app/services/document_ingestion/service.py` +- `app/services/document_ingestion/creation_service.py` +- `app/services/document_ingestion/confirmation_service.py` +- `app/services/document_ingestion/scope_service.py` - `app/repositories/job_repository.py` ### Job Admission diff --git a/apps/api/app/api/v1/routes/jobs.py b/apps/api/app/api/v1/routes/jobs.py index befada76c..38b3b61cc 100644 --- a/apps/api/app/api/v1/routes/jobs.py +++ b/apps/api/app/api/v1/routes/jobs.py @@ -7,7 +7,7 @@ from datetime import datetime from typing import Optional -from app.services.document_ingestion_service import DocumentIngestionService +from app.services.document_ingestion import DocumentIngestionService from app.services.job_read_service import ( get_job_result_for_user, list_jobs_for_user, diff --git a/apps/api/app/services/document_ingestion/__init__.py b/apps/api/app/services/document_ingestion/__init__.py new file mode 100644 index 000000000..0e67b8017 --- /dev/null +++ b/apps/api/app/services/document_ingestion/__init__.py @@ -0,0 +1,3 @@ +from app.services.document_ingestion.service import DocumentIngestionService + +__all__ = ["DocumentIngestionService"] diff --git a/apps/api/app/services/document_ingestion_confirmation_service.py b/apps/api/app/services/document_ingestion/confirmation_service.py similarity index 100% rename from apps/api/app/services/document_ingestion_confirmation_service.py rename to apps/api/app/services/document_ingestion/confirmation_service.py diff --git a/apps/api/app/services/document_ingestion_creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py similarity index 99% rename from apps/api/app/services/document_ingestion_creation_service.py rename to apps/api/app/services/document_ingestion/creation_service.py index c5c8c06d7..19baabfa8 100644 --- a/apps/api/app/services/document_ingestion_creation_service.py +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -8,7 +8,7 @@ from urllib.parse import urlparse from app.repositories.job_repository import JobRepository -from app.services.job_document_scope_service import ( +from app.services.document_ingestion.scope_service import ( is_active_document_job_unique_violation, raise_document_ingestion_conflict, ) diff --git a/apps/api/app/services/job_document_scope_service.py b/apps/api/app/services/document_ingestion/scope_service.py similarity index 98% rename from apps/api/app/services/job_document_scope_service.py rename to apps/api/app/services/document_ingestion/scope_service.py index 80625fe87..34390799d 100644 --- a/apps/api/app/services/job_document_scope_service.py +++ b/apps/api/app/services/document_ingestion/scope_service.py @@ -1,5 +1,5 @@ """ -Document-scope rules used by job creation/update flows. +Document-scope rules used by document-ingestion workflows. """ from __future__ import annotations diff --git a/apps/api/app/services/document_ingestion_service.py b/apps/api/app/services/document_ingestion/service.py similarity index 97% rename from apps/api/app/services/document_ingestion_service.py rename to apps/api/app/services/document_ingestion/service.py index 6bb29e35d..022ccc340 100644 --- a/apps/api/app/services/document_ingestion_service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -4,14 +4,14 @@ import uuid from typing import cast -from app.services.document_ingestion_confirmation_service import ( +from app.services.document_ingestion.confirmation_service import ( DocumentIngestionConfirmationService, ) -from app.services.document_ingestion_creation_service import ( +from app.services.document_ingestion.creation_service import ( DocumentIngestionCreationService, ResolvedDocumentIngestionScope, ) -from app.services.job_document_scope_service import ( +from app.services.document_ingestion.scope_service import ( find_active_job_for_document, raise_document_ingestion_conflict, resolve_effective_document_scope, diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index 33d324acc..32c52ff51 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -1002,7 +1002,7 @@ async def _fake_start_workflow_for_job( ) async with developer_api_client_factory() as api_client: - import app.services.document_ingestion_confirmation_service as document_ingestion_confirmation_service + import app.services.document_ingestion.confirmation_service as document_ingestion_confirmation_service import shared.services.storage.file_upload_service as file_upload_service_module monkeypatch.setattr( From a6f2631c815b95fad7a54d62e80a3c1fb75245d8 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 15:50:19 +0800 Subject: [PATCH 10/40] refactor: package job read workflows --- CONTEXT.md | 12 ++++++++++++ apps/api/app/api/v1/routes/jobs.py | 2 +- .../document_ingestion/confirmation_service.py | 2 +- .../services/document_ingestion/creation_service.py | 2 +- apps/api/app/services/jobs/__init__.py | 11 +++++++++++ .../{job_read_service.py => jobs/read_service.py} | 2 +- .../result_projection.py} | 4 +++- apps/api/tests/contract/test_job_read_contract.py | 8 ++++++++ 8 files changed, 38 insertions(+), 5 deletions(-) create mode 100644 apps/api/app/services/jobs/__init__.py rename apps/api/app/services/{job_read_service.py => jobs/read_service.py} (99%) rename apps/api/app/services/{job_response_projection.py => jobs/result_projection.py} (98%) diff --git a/CONTEXT.md b/CONTEXT.md index a4bdca7a4..bafd49c76 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -29,6 +29,11 @@ URL ingestion, or demo source materialization. The terminal artifact record attached to a Job. It stores delivery metadata, result bundle references, and the revision that publication uses. +### Job Read + +The workflow that lists a User's Jobs and projects one Job into the public Job +Result response shape. + ### Document The retrieval-visible knowledge object produced from a Job Result after @@ -167,6 +172,13 @@ exceptions. - `app/services/document_ingestion/scope_service.py` - `app/repositories/job_repository.py` +### Job Read + +- `app/api/v1/routes/jobs.py` +- `app/services/jobs/read_service.py` +- `app/services/jobs/result_projection.py` +- `app/repositories/job_repository.py` + ### Job Admission - `app/services/rate_limit/*` diff --git a/apps/api/app/api/v1/routes/jobs.py b/apps/api/app/api/v1/routes/jobs.py index 38b3b61cc..e65c7c6c4 100644 --- a/apps/api/app/api/v1/routes/jobs.py +++ b/apps/api/app/api/v1/routes/jobs.py @@ -8,7 +8,7 @@ from typing import Optional from app.services.document_ingestion import DocumentIngestionService -from app.services.job_read_service import ( +from app.services.jobs import ( get_job_result_for_user, list_jobs_for_user, ) diff --git a/apps/api/app/services/document_ingestion/confirmation_service.py b/apps/api/app/services/document_ingestion/confirmation_service.py index c9947e205..970520d61 100644 --- a/apps/api/app/services/document_ingestion/confirmation_service.py +++ b/apps/api/app/services/document_ingestion/confirmation_service.py @@ -1,7 +1,7 @@ from __future__ import annotations from app.repositories.job_repository import JobRepository -from app.services.job_read_service import check_job_permission +from app.services.jobs import check_job_permission from app.services.knowledge.kb_orchestrator import KBOrchestrator from app.services.state_machine import JobStateMachine from loguru import logger diff --git a/apps/api/app/services/document_ingestion/creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py index 19baabfa8..1e418360c 100644 --- a/apps/api/app/services/document_ingestion/creation_service.py +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -12,7 +12,7 @@ is_active_document_job_unique_violation, raise_document_ingestion_conflict, ) -from app.services.job_response_projection import to_job_status_value +from app.services.jobs.result_projection import to_job_status_value from app.services.rate_limit.data_structures import CurrentUser from loguru import logger from sqlalchemy.exc import IntegrityError diff --git a/apps/api/app/services/jobs/__init__.py b/apps/api/app/services/jobs/__init__.py new file mode 100644 index 000000000..fcdc2ccdc --- /dev/null +++ b/apps/api/app/services/jobs/__init__.py @@ -0,0 +1,11 @@ +from app.services.jobs.read_service import ( + check_job_permission, + get_job_result_for_user, + list_jobs_for_user, +) + +__all__ = [ + "check_job_permission", + "get_job_result_for_user", + "list_jobs_for_user", +] diff --git a/apps/api/app/services/job_read_service.py b/apps/api/app/services/jobs/read_service.py similarity index 99% rename from apps/api/app/services/job_read_service.py rename to apps/api/app/services/jobs/read_service.py index 4b68d3397..70deacf5d 100644 --- a/apps/api/app/services/job_read_service.py +++ b/apps/api/app/services/jobs/read_service.py @@ -5,7 +5,7 @@ from typing import Optional from app.repositories.job_repository import JobRepository -from app.services.job_response_projection import ( +from app.services.jobs.result_projection import ( build_job_result_response, to_job_status_value, ) diff --git a/apps/api/app/services/job_response_projection.py b/apps/api/app/services/jobs/result_projection.py similarity index 98% rename from apps/api/app/services/job_response_projection.py rename to apps/api/app/services/jobs/result_projection.py index 9563206bb..308f58116 100644 --- a/apps/api/app/services/job_response_projection.py +++ b/apps/api/app/services/jobs/result_projection.py @@ -109,7 +109,9 @@ def _resolve_duration_seconds(job: Any) -> float | None: return None -async def _resolve_result_delivery(job: Any) -> tuple[dict[str, Any] | None, str | None, datetime]: +async def _resolve_result_delivery( + job: Any, +) -> tuple[dict[str, Any] | None, str | None, datetime]: job_result = job.job_result result_url = None result = None diff --git a/apps/api/tests/contract/test_job_read_contract.py b/apps/api/tests/contract/test_job_read_contract.py index 136d20f89..b913d2a6c 100644 --- a/apps/api/tests/contract/test_job_read_contract.py +++ b/apps/api/tests/contract/test_job_read_contract.py @@ -1,3 +1,4 @@ +import importlib from collections.abc import Callable from contextlib import AbstractAsyncContextManager from datetime import datetime, timedelta, timezone @@ -187,3 +188,10 @@ async def test_should_forbid_access_to_a_job_owned_by_another_user( assert error["code"] == "PERMISSION_DENIED" assert error["message"] == "You don't have permission to access this job" assert "details" not in error + + +def test_job_read_workflow_module_should_be_importable() -> None: + job_read_service = importlib.import_module("app.services.jobs.read_service") + + assert callable(job_read_service.list_jobs_for_user) + assert callable(job_read_service.get_job_result_for_user) From b783850de13cda30ba4b25b6bc9b508c27ae3f4b Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 16:32:59 +0800 Subject: [PATCH 11/40] refactor: remove GitHub flow test file --- apps/docs/test_github_flow.md | 4 ---- 1 file changed, 4 deletions(-) delete mode 100644 apps/docs/test_github_flow.md diff --git a/apps/docs/test_github_flow.md b/apps/docs/test_github_flow.md deleted file mode 100644 index 84fd82cd5..000000000 --- a/apps/docs/test_github_flow.md +++ /dev/null @@ -1,4 +0,0 @@ -# GitHub Flow Test - -This is a test file to verify the GitHub workflow simulation (Issue -> Branch -> PR -> Merge). -It is safe to ignore or delete this file later. From fa2589672fb5ec963865b2fc0861b9dc14fbc4c3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 21:16:23 +0800 Subject: [PATCH 12/40] refactor: split worker ingestion and parser orchestration --- apps/worker/app/core/tasks/kb_tasks.py | 2 +- apps/worker/app/services/common/__init__.py | 33 +- apps/worker/app/services/common/kb_utils.py | 406 ------------ .../services/document_ingestion/__init__.py | 1 + .../job_state_gate.py} | 15 +- .../page_estimator.py | 148 ++--- .../services/document_ingestion/service.py | 586 ++++++++++++++++++ .../document_ingestion/workspace.py} | 21 +- .../services/document_parser/atlas_parser.py | 7 +- .../document_parser/dataframe_helpers.py | 109 ++++ .../services/document_parser/doc_parser.py | 11 +- .../services/document_parser/html_parser.py | 2 +- .../services/document_parser/identifiers.py | 14 + .../services/document_parser/image_parser.py | 7 +- .../services/document_parser/layout_parser.py | 5 +- .../app/services/document_parser/md_parser.py | 9 +- .../document_parser/orchestration/__init__.py | 1 + .../orchestration/parse_session.py | 189 ++++++ .../orchestration/postprocess.py | 70 +++ .../orchestration/route_parse.py | 199 ++++++ .../services/document_parser/parse_service.py | 456 +------------- .../services/document_parser/path_helpers.py | 122 ++++ .../services/document_parser/pptx_parser.py | 2 +- .../services/document_parser/table_parser.py | 10 +- .../services/document_parser/text_helpers.py | 67 ++ .../services/document_parser/toc_parser.py | 5 +- apps/worker/app/services/workload/__init__.py | 8 +- .../services/workload/parse_job_service.py | 467 -------------- .../contract/test_parse_task_contract.py | 2 +- 29 files changed, 1474 insertions(+), 1500 deletions(-) delete mode 100755 apps/worker/app/services/common/kb_utils.py create mode 100644 apps/worker/app/services/document_ingestion/__init__.py rename apps/worker/app/services/{common/job_start_service.py => document_ingestion/job_state_gate.py} (84%) rename apps/worker/app/services/{workload => document_ingestion}/page_estimator.py (53%) create mode 100644 apps/worker/app/services/document_ingestion/service.py rename apps/worker/app/{core/tasks/task_utils.py => services/document_ingestion/workspace.py} (86%) create mode 100644 apps/worker/app/services/document_parser/dataframe_helpers.py create mode 100644 apps/worker/app/services/document_parser/identifiers.py create mode 100644 apps/worker/app/services/document_parser/orchestration/__init__.py create mode 100644 apps/worker/app/services/document_parser/orchestration/parse_session.py create mode 100644 apps/worker/app/services/document_parser/orchestration/postprocess.py create mode 100644 apps/worker/app/services/document_parser/orchestration/route_parse.py create mode 100644 apps/worker/app/services/document_parser/path_helpers.py create mode 100644 apps/worker/app/services/document_parser/text_helpers.py delete mode 100644 apps/worker/app/services/workload/parse_job_service.py diff --git a/apps/worker/app/core/tasks/kb_tasks.py b/apps/worker/app/core/tasks/kb_tasks.py index 5a4416a90..13e9d67f7 100644 --- a/apps/worker/app/core/tasks/kb_tasks.py +++ b/apps/worker/app/core/tasks/kb_tasks.py @@ -7,7 +7,7 @@ # Base task class from app.core.tasks.base_task import KBBaseTask -from app.services.workload.parse_job_service import parse_uploaded_file_job +from app.services.document_ingestion.service import parse_uploaded_file_job from app.services.workload.url_upload_service import upload_url_file from loguru import logger diff --git a/apps/worker/app/services/common/__init__.py b/apps/worker/app/services/common/__init__.py index ca66e606f..ffb53d823 100644 --- a/apps/worker/app/services/common/__init__.py +++ b/apps/worker/app/services/common/__init__.py @@ -1,4 +1,4 @@ -"""Common worker services, including reusable knowledge-base helpers.""" +"""Common worker services.""" from shared.utils.device_utils import check_internet from shared.utils.file_utils import clean_file, path_handle @@ -10,39 +10,8 @@ tokenize2stw_remove, ) -from .kb_utils import ( - find_images, - find_matches_parsing, - flatten_dic2paths, - flatten_list, - gen_str_codes, - get_str_time, - html2txt, - merge_df, - process_dup_paths_df, - process_path_texts, - remove_spaces, - restore_graph_by_paths, - traverse_dict, -) - __all__ = [ - # From kb_utils "count_cn_en", - "find_images", - "find_matches_parsing", - "flatten_dic2paths", - "flatten_list", - "gen_str_codes", - "get_str_time", - "html2txt", - "merge_df", - "process_dup_paths_df", - "process_path_texts", - "remove_spaces", - "restore_graph_by_paths", - "traverse_dict", - # From shared-python "check_internet", "clean_file", "min_max_normalize", diff --git a/apps/worker/app/services/common/kb_utils.py b/apps/worker/app/services/common/kb_utils.py deleted file mode 100755 index 43389a56f..000000000 --- a/apps/worker/app/services/common/kb_utils.py +++ /dev/null @@ -1,406 +0,0 @@ -import os -import re -import uuid -from datetime import datetime - -import pandas as pd -from bs4 import BeautifulSoup - -from shared.core.config import settings -from shared.utils.chunk_refs import extract_chunk_refs -from shared.utils.file_utils import path_handle -from shared.utils.text_utils import _CN_EN_NUM_RE - -SUMMARY_PATH_MARKERS: tuple[str, ...] = ("summary", "\u6458\u8981\u603b\u7ed3") - - -def gen_str_codes(input_string): - """Generate a UUID5 code from a string.""" - namespace = uuid.NAMESPACE_DNS - return str(uuid.uuid5(namespace, input_string)) - - -def get_str_time(): - """Get the current time as a string.""" - now = datetime.now() - return now.strftime("%Y-%m-%d %H:%M:%S") - - -def find_images(folder_path): - """Find image files inside a folder tree.""" - image_extensions = {".png", ".jpg", ".jpeg"} - image_files = [] - - for _, _, files in os.walk(folder_path): - files.sort() - for file in files: - if os.path.splitext(file)[1].lower() in image_extensions: - image_files.append(file) - return image_files - - -def find_matches_parsing(content, path): - """Parse table and image markers from content.""" - matches = extract_chunk_refs(content) - if len(matches) == 0: - match_type = "PTXT" - else: - match_type = "\n".join((["PTXT"] + matches)) - - split_char = settings.SPLIT_CHAR or ";" - if any( - f"{split_char}{summary_marker}" in path - for summary_marker in SUMMARY_PATH_MARKERS - ): - parent_path = path.split(split_char)[-2] - match_type = "SUMMARY_" + parent_path + "_SUMMARY" - return match_type - - -def flatten_list(nested_list): - """Flatten a nested list.""" - flat_list = [] - for item in nested_list: - if isinstance(item, list): - flat_list.extend(flatten_list(item)) - else: - flat_list.append(item) - return flat_list - - -def flatten_dic2paths(d, current_path=None, result=None): - """Flatten a nested dict into path strings.""" - if result is None: - result = [] - if current_path is None: - current_path = [] - - for key, value in d.items(): - if not isinstance(key, str): - continue - new_path = current_path + [key] - if isinstance(value, dict) and value: - flatten_dic2paths(value, new_path, result) - else: - split_char = settings.SPLIT_CHAR or ";" - result.append(split_char.join(new_path)) - return result - - -def merge_df(input_df): - """Merge DataFrame rows that share the same path.""" - dfs_by_path = list(input_df.groupby("path", sort=False)) - processed_dfs = [] - - for key, df in dfs_by_path: - content_to_merge = [] - types_to_merge = [] - total_length = 0 - - for i, row in df.iterrows(): - content_to_merge.append(row["content"]) - types_to_merge.extend(row["type"].split("\n")) - total_length += len(row["content"]) - - content_to_merge = "\n".join(content_to_merge) - temp_merge_df = pd.DataFrame( - [ - { - "content": content_to_merge, - "type": "\n".join(list(set(types_to_merge))), - "path": key, - "length": total_length, - "know_id": gen_str_codes(content_to_merge), - } - ] - ) - processed_dfs.append(temp_merge_df) - - final_df = pd.concat(processed_dfs, axis=0, ignore_index=True) - return final_df - - -def process_path_texts(path_, last=50): - """Normalize path text for downstream use.""" - temp_path = path_handle(path_, mode="sanitize") - if temp_path is None: - return "" - return "_".join(temp_path.split(os.sep))[:last] - - -def process_dup_paths_df(df): - """ - de-duplicate kbs dataframe for final output - - Args: - df: initial dataframe after all heading stacking - - Returns: - Dataframe without any duplicate paths - """ - if "path" not in df.columns: - return df - - split_char = settings.SPLIT_CHAR or "/" - - # Step 1: detect if there are any duplicated paths - dup_mask = df["path"].duplicated(keep=False) - if not dup_mask.any(): - return df - - # Step 2: record ids of duplicated paths as a mapping - path_occurrences = {} # path -> list of row indices - for idx, path in enumerate(df["path"]): - if path not in path_occurrences: - path_occurrences[path] = [] - path_occurrences[path].append(idx) - - # path_renames: row_index -> new_path (recording rows renamed) - # parent_rename_map: original_path -> {row_index: new_path} - path_renames = {} - parent_rename_map = {} - - for path, indices in path_occurrences.items(): - if len(indices) > 1: # only process duplicated paths - parent_rename_map[path] = {} - for occurrence, idx in enumerate(indices): - if occurrence == 0: - # keep the first appearance as it is - path_renames[idx] = path - else: - # add suffix to subsequent appearances - new_path = f"{path}_{occurrence + 1}" - path_renames[idx] = new_path - parent_rename_map[path][idx] = new_path - - # Step 3: process all rows, update paths - new_paths = [] - - for idx, row in df.iterrows(): - path = row["path"] - - # Check whether this row itself needs renaming. - new_path = path_renames.get(idx, path) - path_parts = new_path.split(split_char) - - # Check whether this row is under a renamed parent path. - for parent_path, rename_info in parent_rename_map.items(): - parent_parts = parent_path.split(split_char) - - # Check whether the current path starts with that parent path. - if ( - len(path_parts) > len(parent_parts) - and path_parts[: len(parent_parts)] == parent_parts - ): - # Find the nearest renamed parent path that appears earlier. - matching_parent_idx = None - for parent_idx in sorted(rename_info.keys(), reverse=True): - if parent_idx < idx: - matching_parent_idx = parent_idx - break - - if matching_parent_idx is not None: - renamed_parent = rename_info[matching_parent_idx] - renamed_parent_parts = renamed_parent.split(split_char) - new_path_parts = ( - renamed_parent_parts + path_parts[len(parent_parts) :] - ) - new_path = split_char.join(new_path_parts) - break - new_paths.append(new_path) - - df = df.copy() - df["path"] = new_paths - return df - - -def remove_spaces(text, handle_punctuation=False): - """Remove spaces between Chinese chars while keeping English word spacing.""" - if handle_punctuation: - punctuation = ( - r"""!"#$%&'()*+,-./:;<=>?@[\]^_`{|}~，。、【】《》？；：''""（）…—-！""" - ) - res_text = re.sub(f"[{re.escape(punctuation)}]", "", text) - else: - pattern = re.compile(r"([\u4e00-\u9fff])\s+|(?<=\s)([\u4e00-\u9fff])") - - def replacer(match): - return match.group(1) or match.group(2) - - res_text = pattern.sub(replacer, text) - - res_text = re.sub(r"\s+", " ", res_text) - return res_text.strip() - - -def traverse_dict(d, parent=None): - """Traverse a dictionary and generate description text.""" - dic_texts = [] - for key, value in d.items(): - if value: - child_keys = ", ".join(value.keys()) - text = f"'{key}' includes {child_keys}" - dic_texts.append(text) - dic_texts.extend(traverse_dict(value, key)) - return dic_texts - - -def restore_graph_by_paths(paths): - """Rebuild a graph structure from path strings.""" - root_dict = {} - split_char = settings.SPLIT_CHAR or ";" - for path in paths: - nodes = path.split(split_char) - current_dict = root_dict - for node in nodes: - if node not in current_dict: - current_dict[node] = {} - current_dict = current_dict[node] - dic_texts = traverse_dict(root_dict) - return root_dict, dic_texts - - -def html2txt(html_text): - """Convert HTML into plain text.""" - soup = BeautifulSoup(html_text, "html.parser") - text = soup.get_text() - return text - - -def normalize_md(s: str) -> str: - """Normalize markdown string for comparison - - Removes heading markers (###) and whitespace, converts to lowercase. - Used for TOC keyword matching. - """ - s = re.sub(r"^\s*#+\s*", "", s) - s = re.sub(r"\s+", "", s) - return s.lower() - - -# --------------------------------------------------------------------------- -# truncate_text (character-based) — KEPT for table-cell display callers in -# doc_parser.py and html_parser.py where a per-character limit is intentional. -# Do NOT use for heading / semantic text truncation; use truncate_text_by_tokens. -# --------------------------------------------------------------------------- -def truncate_text(text: str, start_limit: int, end_limit: int) -> str: - """Truncate text by raw character count, keeping start and end parts. - - Intended for short display values (table headers, file names, etc.) where - a fixed character budget is appropriate. For heading / semantic text where - English words must not be split mid-word, use ``truncate_text_by_tokens``. - - Args: - text: Text to truncate. - start_limit: Number of characters to keep from start. - end_limit: Number of characters to keep from end (0 = no tail). - - Returns: - Truncated text with '...' in the middle if it exceeds the limits. - """ - text = str(text) - total_limit = start_limit + end_limit - if len(text) <= total_limit: - return text - start_part = text[:start_limit] - end_part = text[-end_limit:] if end_limit > 0 else "" - return f"{start_part}...{end_part}" - - -# --------------------------------------------------------------------------- -# Language detection & language-aware token truncation -# --------------------------------------------------------------------------- - -_CN_CHAR_RE = re.compile(r"[\u4e00-\u9fff]") - -EN_START_LIMIT = 15 # token budget for English-dominant headings -CN_RATIO_THRESHOLD = 0.3 # if ≥30 % of tokens are Chinese chars → "Chinese" - - -def detect_primary_lang(text: str) -> str: - """Detect whether *text* is primarily Chinese or English/other. - - Uses the semantic tokens already defined by ``_CN_EN_NUM_RE`` - (Chinese chars, English word runs, number groups). If Chinese - characters account for at least ``CN_RATIO_THRESHOLD`` of all - tokens the text is classified as ``'zh'``; otherwise ``'en'``. - - Args: - text: Input text (heading or any short string). - - Returns: - ``'zh'`` for Chinese-dominant text, ``'en'`` otherwise. - """ - if not text: - return "en" - tokens = _CN_EN_NUM_RE.findall(text) - if not tokens: - return "en" - cn_count = sum(1 for t in tokens if _CN_CHAR_RE.fullmatch(t)) - return "zh" if (cn_count / len(tokens)) >= CN_RATIO_THRESHOLD else "en" - - -def count_cn_en(text: str) -> int: - """Count semantic Chinese/English/number tokens in a string.""" - return len(_CN_EN_NUM_RE.findall(str(text))) - - -def truncate_text_by_tokens( - text: str, - start_limit: int, - end_limit: int, - lang_aware: bool = True, -) -> str: - """Truncate text by semantic token count, preserving whole words. - - Uses the same token definition as ``count_cn_en``: - - - each Chinese character = 1 token - - each run of English letters = 1 token - - each number group = 1 token - - punctuation and whitespace are excluded from the count but - preserved in the output up to the split point. - - When *lang_aware* is ``True`` (default), the function auto-detects - whether the text is English-dominant and caps ``start_limit`` at - ``EN_START_LIMIT`` (15) in that case. Chinese-dominant text keeps - the caller-supplied ``start_limit`` (typically 30). This prevents - over-long English heading chunks while still allowing a generous - budget for dense Chinese text. - - Cut points are placed *after* the last character of the start - token and *before* the first character of the first tail token, - so no word is ever split in the middle. - - Args: - text: Text to truncate. - start_limit: Max tokens to keep from the start. When - *lang_aware* is True and the text is English-dominant, - this is silently capped at ``EN_START_LIMIT``. - end_limit: Max tokens to keep from the end (0 = no tail). - lang_aware: When True, auto-detect language and apply a tighter - budget for English text. Set to False to use raw limits. - - Returns: - Truncated text with ``'...'`` in the middle when the token - count exceeds ``start_limit + end_limit``. Returns the - original text unchanged when the count is within the budget. - """ - text = str(text) - matches = list(_CN_EN_NUM_RE.finditer(text)) - total = len(matches) - - if lang_aware and total > 0: - lang = detect_primary_lang(text) - if lang == "en": - start_limit = min(start_limit, EN_START_LIMIT) - - if total <= start_limit + end_limit: - return text - # Cut position: end of the start_limit-th token - cut_start = matches[start_limit - 1].end() if start_limit > 0 else 0 - # Tail position: start of the (total - end_limit)-th token - cut_end = matches[total - end_limit].start() if end_limit > 0 else len(text) - if cut_start >= cut_end: - return text - return text[:cut_start] + "..." + text[cut_end:] diff --git a/apps/worker/app/services/document_ingestion/__init__.py b/apps/worker/app/services/document_ingestion/__init__.py new file mode 100644 index 000000000..191bf4665 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/__init__.py @@ -0,0 +1 @@ +"""Worker-side Document Ingestion modules.""" diff --git a/apps/worker/app/services/common/job_start_service.py b/apps/worker/app/services/document_ingestion/job_state_gate.py similarity index 84% rename from apps/worker/app/services/common/job_start_service.py rename to apps/worker/app/services/document_ingestion/job_state_gate.py index 150b852c4..5cd5d067c 100644 --- a/apps/worker/app/services/common/job_start_service.py +++ b/apps/worker/app/services/document_ingestion/job_state_gate.py @@ -1,9 +1,9 @@ """ -Worker-side job start gating. +Worker-side Document Ingestion state gate. Keeps the worker-specific policy for when a parse task is allowed to move a -job into ``running`` while delegating the actual state transition to the -shared sync state machine service. +Job into ``running`` while delegating the actual transition to the shared sync +state machine service. """ from typing import Any @@ -24,11 +24,7 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: - """Transition a job from pending to running before parse execution. - - Returns ``True`` when parsing should proceed. Returns ``False`` when the - job is already terminal and the task should skip quietly. - """ + """Transition a Job from pending to running before parse execution.""" state_machine = SyncStateMachineService(redis_service) with get_sync_db_context() as db: @@ -45,9 +41,6 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: current_state = job.status if current_state == JobStatus.RUNNING.value: - # Likely a broker redelivery while the original worker is still - # processing. Let the caller proceed to RedisJobLock, which gates - # actual execution. logger.info( f"Job already running (likely redelivery), deferring to lock: {job_id}" ) diff --git a/apps/worker/app/services/workload/page_estimator.py b/apps/worker/app/services/document_ingestion/page_estimator.py similarity index 53% rename from apps/worker/app/services/workload/page_estimator.py rename to apps/worker/app/services/document_ingestion/page_estimator.py index a2dfb3799..ee681d4c4 100644 --- a/apps/worker/app/services/workload/page_estimator.py +++ b/apps/worker/app/services/document_ingestion/page_estimator.py @@ -1,16 +1,11 @@ """ -Page Estimator Service +Page estimator for worker-side Document Ingestion billing. Calculates page counts for billing based on: - PDF: Physical page count from metadata - PPTX: Slide count - Text-based (DOC, DOCX, TXT, MD, JSON): Word-based estimation using count_cn_en - Spreadsheet-based (XLS, XLSX): Row-based estimation - -Supported file types: -- .pdf, .doc, .docx, .pptx, .xls, .xlsx -- .txt, .md, .json, .fragment -- .png, .jpg, .jpeg """ import math @@ -20,68 +15,48 @@ from shared.core.logging import logger from shared.utils.text_utils import count_cn_en -# Constants for page estimation -WORDS_PER_PAGE = 500 # Chinese chars + English words + numbers per page -ROWS_PER_PAGE = 50 # For Excel files +WORDS_PER_PAGE = 500 +ROWS_PER_PAGE = 50 class PageEstimator: - """ - Estimates page count for billing purposes. - - Billing Logic: - - PDF: Physical page count from metadata - - PPTX: Slide count - - DOC/DOCX/TXT/MD/JSON: Word-based estimation (count_cn_en / 500) - - XLS/XLSX: Row-based estimation (rows / 50) - - Images: 1 page per image - """ + """Estimate page count for billing purposes.""" @classmethod def estimate(cls, file_path: str) -> int: - """ - Estimate page count for a file. - - Args: - file_path: Path to the file - - Returns: - Estimated page count (minimum 1) - """ + """Estimate the billable page count for a file.""" path = Path(file_path) suffix = path.suffix.lower() try: if suffix == ".pdf": return cls._estimate_pdf(file_path) - elif suffix == ".pptx": + if suffix == ".pptx": return cls._estimate_pptx(file_path) - elif suffix == ".doc": + if suffix == ".doc": return cls._estimate_doc(file_path) - elif suffix == ".docx": + if suffix == ".docx": return cls._estimate_docx(file_path) - elif suffix == ".xls": + if suffix == ".xls": return cls._estimate_xls(file_path) - elif suffix == ".xlsx": + if suffix == ".xlsx": return cls._estimate_xlsx(file_path) - elif suffix in [".txt", ".md", ".json", ".fragment"]: + if suffix in [".txt", ".md", ".json", ".fragment"]: return cls._estimate_text(file_path) - elif suffix in [".png", ".jpg", ".jpeg"]: - return 1 # Image = 1 page - else: - logger.warning( - f"Unknown file type for billing: {suffix}, defaulting to 1 page" - ) + if suffix in [".png", ".jpg", ".jpeg"]: return 1 - except Exception as e: - logger.error(f"Error estimating pages for {file_path}: {e}") - return 1 # Fallback to minimum charge + + logger.warning( + f"Unknown file type for billing: {suffix}, defaulting to 1 page" + ) + return 1 + except Exception as exc: + logger.error(f"Error estimating pages for {file_path}: {exc}") + return 1 @classmethod def _estimate_pdf(cls, file_path: str) -> int: - """ - Estimate pages for PDF using physical page count from metadata. - """ + """Estimate pages for PDF using physical page count.""" try: from pypdf import PdfReader @@ -90,124 +65,107 @@ def _estimate_pdf(cls, file_path: str) -> int: except ImportError: logger.warning("pypdf not installed, defaulting to 1 page") return 1 - except Exception as e: - logger.error(f"PDF estimation error: {e}") + except Exception as exc: + logger.error(f"PDF estimation error: {exc}") return 1 @classmethod def _estimate_pptx(cls, file_path: str) -> int: - """ - Estimate pages for PPTX using slide count. - """ + """Estimate pages for PPTX using slide count.""" try: from pptx import Presentation - prs = Presentation(file_path) - return max(1, len(prs.slides)) + presentation = Presentation(file_path) + return max(1, len(presentation.slides)) except ImportError: logger.warning("python-pptx not installed, defaulting to 1 page") return 1 - except Exception as e: - logger.error(f"PPTX estimation error: {e}") + except Exception as exc: + logger.error(f"PPTX estimation error: {exc}") return 1 @classmethod def _estimate_docx(cls, file_path: str) -> int: - """ - Estimate pages for DOCX using word-based counting. - """ + """Estimate pages for DOCX using word-based counting.""" try: from docx import Document - doc = Document(file_path) + document = Document(file_path) total_text = "" - # Collect text from paragraphs - for para in doc.paragraphs: - total_text += para.text + " " + for paragraph in document.paragraphs: + total_text += paragraph.text + " " - # Collect text from tables - for table in doc.tables: + for table in document.tables: for row in table.rows: for cell in row.cells: total_text += cell.text + " " word_count = count_cn_en(total_text) return max(1, math.ceil(word_count / WORDS_PER_PAGE)) - except ImportError: logger.warning("python-docx not installed") return 1 - except Exception as e: - logger.error(f"DOCX estimation error: {e}") + except Exception as exc: + logger.error(f"DOCX estimation error: {exc}") return 1 @classmethod def _estimate_doc(cls, file_path: str) -> int: - """ - Estimate pages for DOC by converting it to DOCX first. - """ + """Estimate pages for DOC by converting it to DOCX first.""" try: from app.services.document_parser.legacy_converter import doc_to_docx with tempfile.TemporaryDirectory(prefix="page-estimator-doc-") as temp_dir: converted_path, _ = doc_to_docx(file_path, temp_dir) return cls._estimate_docx(converted_path) - except Exception as e: - logger.error(f"DOC estimation error: {e}") + except Exception as exc: + logger.error(f"DOC estimation error: {exc}") return 1 @classmethod def _estimate_xlsx(cls, file_path: str) -> int: - """ - Estimate pages for XLSX using row count. - """ + """Estimate pages for XLSX using row count.""" try: import pandas as pd - xlsx = pd.ExcelFile(file_path) + workbook = pd.ExcelFile(file_path) total_rows = 0 - for sheet_name in xlsx.sheet_names: - df = pd.read_excel(xlsx, sheet_name=sheet_name) - total_rows += len(df) + for sheet_name in workbook.sheet_names: + dataframe = pd.read_excel(workbook, sheet_name=sheet_name) + total_rows += len(dataframe) return max(1, math.ceil(total_rows / ROWS_PER_PAGE)) - except ImportError: logger.warning("pandas not installed") return 1 - except Exception as e: - logger.error(f"XLSX estimation error: {e}") + except Exception as exc: + logger.error(f"XLSX estimation error: {exc}") return 1 @classmethod def _estimate_xls(cls, file_path: str) -> int: - """ - Estimate pages for XLS by converting it to XLSX first. - """ + """Estimate pages for XLS by converting it to XLSX first.""" try: from app.services.document_parser.legacy_converter import xls_to_xlsx with tempfile.TemporaryDirectory(prefix="page-estimator-xls-") as temp_dir: converted_path, _ = xls_to_xlsx(file_path, temp_dir) return cls._estimate_xlsx(converted_path) - except Exception as e: - logger.error(f"XLS estimation error: {e}") + except Exception as exc: + logger.error(f"XLS estimation error: {exc}") return 1 @classmethod def _estimate_text(cls, file_path: str) -> int: - """ - Estimate pages for text files using word-based counting. - """ + """Estimate pages for text-like files using word-based counting.""" try: - with open(file_path, "r", encoding="utf-8", errors="ignore") as f: - content = f.read() + with open(file_path, "r", encoding="utf-8", errors="ignore") as file: + content = file.read() word_count = count_cn_en(content) return max(1, math.ceil(word_count / WORDS_PER_PAGE)) - - except Exception as e: - logger.error(f"Text estimation error: {e}") + except Exception as exc: + logger.error(f"Text estimation error: {exc}") return 1 diff --git a/apps/worker/app/services/document_ingestion/service.py b/apps/worker/app/services/document_ingestion/service.py new file mode 100644 index 000000000..625464a41 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/service.py @@ -0,0 +1,586 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any + +import pandas as pd +from app.services.connect_builder.summary_builder import ( + build_section_summary_lookup, + enrich_doc_nav_summaries, + ensure_doc_nav_json, + load_nav_top_summary, +) +from app.services.document_ingestion.job_state_gate import mark_job_running +from app.services.document_ingestion.page_estimator import PageEstimator +from app.services.document_ingestion.workspace import ( + cleanup_task_workspace, + create_task_workspace, + download_s3_file_to_temp, +) +from app.services.document_parser.stage_profiler import stage_timer +from app.services.storage.sync_storage_service import ( + generate_download_url, + verify_s3_file_exists, +) +from loguru import logger +from sqlalchemy import select + +from shared.core.config import settings +from shared.core.database_sync import get_sync_db_context +from shared.core.exceptions.domain_exceptions import ( + InsufficientCreditsException, + NotFoundException, + ValidationException, + WorkerHandlingException, +) +from shared.models.database.job import Job +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.billing.work_billing_service import WorkBillingService +from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks +from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service +from shared.services.redis.distributed_lock import RedisJobLock +from shared.services.redis.redis_sync_service import ( + SyncJobInfoRedisService, + SyncJobMetadataService, + SyncRedisServiceFactory, +) +from shared.services.storage.result_storage import get_result_storage +from shared.services.storage.zip_result_service import ZipResultService + + +@dataclass(frozen=True) +class _ParseJobContext: + job_metadata: dict[str, object] + job_user_id: str | None + metadata_service: SyncJobMetadataService + redis_service: Any + s3_key: str + + +@dataclass(frozen=True) +class _ParseJobBillingSnapshot: + billing_amount_micro_dollars: int + billing_credits: float + billing_status: str + + +def parse_uploaded_file_job(job_id: str, user_id: str | None) -> dict[str, object]: + """Run worker-side Document Ingestion for an uploaded file Job.""" + logger.info(f"Parse started: job_id={job_id}, user_id={user_id}") + lifecycle_service = get_sync_job_lifecycle_service() + + redis_service = SyncRedisServiceFactory.get_service() + job_context = _load_parse_job_context(job_id, user_id, redis_service) + _assert_source_file_within_size_limit(job_context.s3_key) + + should_process = mark_job_running(job_id, job_context.redis_service) + if not should_process: + logger.warning(f"Skipping parse_task for inactive job: job_id={job_id}") + return { + "status": "skipped", + "job_id": job_id, + "reason": "job_already_terminal", + } + + with RedisJobLock(job_context.redis_service, job_id): + task_workspace_dir, input_dir, output_dir = _prepare_task_workspace(job_id) + try: + return _run_parse_job( + job_id=job_id, + job_context=job_context, + lifecycle_service=lifecycle_service, + input_dir=input_dir, + output_dir=output_dir, + task_workspace_dir=task_workspace_dir, + ) + finally: + cleanup_task_workspace(task_workspace_dir) + + raise WorkerHandlingException( + user_message="We could not complete document processing", + internal_message=f"Parse workflow exited without a result for job_id={job_id}", + ) + + +def _load_parse_job_context( + job_id: str, + requested_user_id: str | None, + redis_service: Any, +) -> _ParseJobContext: + job_info_service = SyncJobInfoRedisService(redis_service) + job_info = job_info_service.get_job_info(job_id) + + if not job_info: + logger.warning( + f"JobInfo not found in Redis for job_id={job_id}; falling back to database" + ) + with get_sync_db_context() as fallback_db: + job_row = fallback_db.execute( + select(Job).where(Job.job_id == job_id) + ).scalar_one_or_none() + + if not job_row or not job_row.s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id=job_id, + internal_message="job info not found in Redis or database", + ) + + s3_key: str = job_row.s3_key + job_user_id: str | None = ( + str(job_row.user_id) if job_row.user_id else requested_user_id + ) + logger.info(f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}") + else: + raw_s3_key = job_info.get("s3_key") + if not isinstance(raw_s3_key, str) or not raw_s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id="s3_key", + internal_message="Missing s3_key in job_info", + ) + + s3_key = raw_s3_key + raw_job_user_id = job_info.get("user_id") + job_user_id = ( + raw_job_user_id if isinstance(raw_job_user_id, str) else requested_user_id + ) + + metadata_service = SyncJobMetadataService(redis_service) + raw_job_metadata = metadata_service.get_metadata(job_id) + if not isinstance(raw_job_metadata, dict) or not raw_job_metadata: + raise NotFoundException( + resource="JobMetadata", + resource_id=job_id, + internal_message=f"Job metadata not found for job_id={job_id}", + ) + + return _ParseJobContext( + job_metadata=dict(raw_job_metadata), + job_user_id=job_user_id, + metadata_service=metadata_service, + redis_service=redis_service, + s3_key=s3_key, + ) + + +def _assert_source_file_within_size_limit(s3_key: str) -> None: + file_info = verify_s3_file_exists(s3_key) + if not file_info.get("exists"): + raise NotFoundException( + resource="S3File", + resource_id=s3_key, + internal_message=f"S3 file not found: {s3_key}", + ) + + logger.info(f"S3 file verified: {s3_key}") + + file_size = file_info.get("size", 0) + file_extension = os.path.splitext(s3_key)[1].lower() + if file_size > settings.MAX_FILE_SIZE: + limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) + raise ValidationException( + user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", + violations=[ + { + "field": "file_size", + "description": ( + f"Size {file_size} bytes exceeds limit of " + f"{settings.MAX_FILE_SIZE} bytes" + ), + } + ], + ) + + +def _prepare_task_workspace(job_id: str) -> tuple[str, str, str]: + task_workspace_dir = create_task_workspace(job_id) + input_dir = os.path.join(task_workspace_dir, "input") + output_dir = os.path.join(task_workspace_dir, "output") + os.makedirs(input_dir, exist_ok=True) + os.makedirs(output_dir, exist_ok=True) + logger.info( + f"Task workspace ready: job_id={job_id}, workspace={task_workspace_dir}" + ) + return task_workspace_dir, input_dir, output_dir + + +def _run_parse_job( + *, + job_id: str, + job_context: _ParseJobContext, + lifecycle_service: Any, + input_dir: str, + output_dir: str, + task_workspace_dir: str, +) -> dict[str, object]: + lifecycle_service.update_progress(job_id, progress=10, message="Parsing document...") + + filename = JobMetadataHelper.get_field(job_context.job_metadata, "source_file_name") + file_ext = os.path.splitext(job_context.s3_key)[1].lower() if job_context.s3_key else "" + file_url = generate_download_url( + job_context.s3_key, + settings.S3_BUCKET_NAME, + )["download_url"] + local_temp_path = download_s3_file_to_temp(file_url, file_ext, input_dir) + logger.info(f"File downloaded: job_id={job_id}, local_path={local_temp_path}") + + from app.services.document_parser.internal_parse_name import ( + prepare_internal_parse_input, + ) + from app.services.document_parser import parse_service + + prepared_parse_input = prepare_internal_parse_input( + local_temp_path, + filename, + fallback_ext=file_ext, + prefer_fallback_ext=True, + ) + internal_parse_name = prepared_parse_input.internal_filename + local_temp_path = prepared_parse_input.file_path + logger.info( + f"File prepared for parsing: job_id={job_id}, " + f"internal_filename={internal_parse_name}, local_path={local_temp_path}" + ) + + page_count = PageEstimator.estimate(local_temp_path) + logger.info(f"Workload estimation: job_id={job_id}, page_count={page_count}") + + processing_started_at = datetime.now(timezone.utc) + billing_snapshot = _charge_parse_job_pages( + job_id=job_id, + filename=filename, + job_user_id=job_context.job_user_id, + page_count=page_count, + ) + _record_processing_start( + job_id=job_id, + job_context=job_context, + billing_snapshot=billing_snapshot, + page_count=page_count, + processing_started_at=processing_started_at, + ) + + doc_type = JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "doc_type", + "auto", + ) + logger.info( + f"Start parse: job_id={job_id}, filename={filename}, " + f"internal_filename={internal_parse_name}, type={doc_type}" + ) + + with stage_timer( + "worker.parse.document", + job_id=job_id, + filename=filename, + doc_type=doc_type, + ): + add_dir, parsed_contents_df = parse_service.checkerboard_inject_parse( + file_full_path=local_temp_path, + filename=filename, + output_dir=output_dir, + job_id=job_id, + internal_output_filename=internal_parse_name, + kb_dir=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "kb_dir", + "Default_Root", + ), + doc_type=doc_type, + smart_title_parse=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "smart_title_parse", + True, + ), + summary_image=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_image", + True, + ), + summary_table=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_table", + True, + ), + summary_txt=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_txt", + True, + ), + add_frag_desc=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "add_frag_desc", + "", + ), + s3_key=job_context.s3_key, + ) + + logger.info( + "File parsing completed: " + f"job_id={job_id}, add_dir={add_dir}, " + f"chunks={len(parsed_contents_df) if parsed_contents_df is not None else 0}" + ) + + if parsed_contents_df is None: + raise WorkerHandlingException( + user_message="We could not extract content from your file", + internal_message="File parsing failed, no content returned from parser", + ) + + if parsed_contents_df.empty: + logger.warning( + f"No content returned from file parsing: job_id={job_id}, filename={filename}" + ) + + lifecycle_service.update_progress( + job_id, + progress=30, + message="Parse completed, preparing chunks...", + ) + chunks = dataframe_to_chunks(parsed_contents_df) + + lifecycle_service.update_progress( + job_id, + progress=70, + message="Chunks ready, generating zip...", + ) + logger.info(f"Chunks prepared: job_id={job_id}, count={len(chunks)}") + + return _finalize_parse_job_success( + add_dir=add_dir, + chunks=chunks, + job_context=job_context, + job_id=job_id, + lifecycle_service=lifecycle_service, + parsed_contents_df=parsed_contents_df, + processing_started_at=processing_started_at, + task_workspace_dir=task_workspace_dir, + ) + + +def _charge_parse_job_pages( + *, + job_id: str, + filename: str | None, + job_user_id: str | None, + page_count: int, +) -> _ParseJobBillingSnapshot: + if not job_user_id: + raise NotFoundException( + resource="JobInfo", + resource_id="user_id", + internal_message=f"Missing user_id in job info for job_id={job_id}", + ) + + billing_service = WorkBillingService() + billing_filename = filename or "" + billing_status = "skipped" + billing_amount_micro_dollars = 0 + billing_credits = 0.0 + + with get_sync_db_context() as db: + job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()) + job = job_result.scalar_one_or_none() + + if job and getattr(job, "billing_status", "") == "charged": + logger.info(f"Job already charged: {job_id}") + billing_status = "charged" + billing_amount_micro_dollars = int(job.credits_charged or 0) + billing_credits = billing_amount_micro_dollars / 1_000_000 + else: + try: + billing_result = billing_service.charge_for_pages( + session=db, + user_id=job_user_id, + page_count=page_count, + filename=billing_filename, + ) + except InsufficientCreditsException: + logger.warning(f"Billing failed: job_id={job_id}, user_id={job_user_id}") + billing_amount = billing_service.estimate_page_charge( + page_count=page_count + ) + if job: + job.page_count = page_count + job.credits_charged = billing_amount.amount_micro_dollars + job.billing_status = "billing_failed" + db.commit() + + raise InsufficientCreditsException( + user_message=( + "Insufficient credits to process this document " + f"({page_count} pages required, cost: " + f"{billing_amount.credits})." + ), + required_credits=billing_amount.credits, + internal_message=( + f"job_id={job_id}, user_id={job_user_id}, " + f"page_count={page_count}" + ), + ) + + billing_status = billing_result.billing_status + billing_amount_micro_dollars = billing_result.amount_micro_dollars + billing_credits = billing_result.credits + if job: + job.page_count = page_count + job.credits_charged = billing_amount_micro_dollars + job.billing_status = billing_status + + return _ParseJobBillingSnapshot( + billing_amount_micro_dollars=billing_amount_micro_dollars, + billing_credits=billing_credits, + billing_status=billing_status, + ) + + +def _record_processing_start( + *, + job_id: str, + job_context: _ParseJobContext, + billing_snapshot: _ParseJobBillingSnapshot, + page_count: int, + processing_started_at: datetime, +) -> None: + metadata_updates = { + "page_count": page_count, + "billing_status": billing_snapshot.billing_status, + "billing_amount_micro_dollars": billing_snapshot.billing_amount_micro_dollars, + "billing_credits": billing_snapshot.billing_credits, + "processing_started_at": processing_started_at.isoformat(), + } + job_context.metadata_service.update_metadata(job_id, metadata_updates) + job_context.job_metadata.update(metadata_updates) + + +def _finalize_parse_job_success( + *, + add_dir: str, + chunks: list[dict[str, Any]], + job_context: _ParseJobContext, + job_id: str, + lifecycle_service: Any, + parsed_contents_df: pd.DataFrame, + processing_started_at: datetime, + task_workspace_dir: str, +) -> dict[str, object]: + source_file_name = JobMetadataHelper.get_field( + job_context.job_metadata, + "source_file_name", + ) or JobMetadataHelper.get_field(job_context.job_metadata, "source_url") + if isinstance(source_file_name, str) and "/" in source_file_name: + source_file_name = os.path.basename(source_file_name) + + document_top_summary = "" + section_summaries: dict[str, str] = {} + if add_dir and source_file_name: + if "path" in parsed_contents_df.columns: + ensure_doc_nav_json( + str(add_dir), + chunks, + source_file_name=str(source_file_name), + ) + try: + kb_dir_for_enrich = os.path.dirname(str(add_dir)) + summary_use_llm = JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_use_llm", + False, + ) + enrich_doc_nav_summaries( + kb_dir_for_enrich, + source_file=str(source_file_name), + use_llm=summary_use_llm, + ) + section_summaries = build_section_summary_lookup(str(add_dir)) + except Exception as exc: + logger.warning(f"doc_nav enrichment failed (non-fatal): {exc}") + document_top_summary = load_nav_top_summary(str(add_dir), str(source_file_name)) + + if document_top_summary: + for chunk in chunks: + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + chunk["metadata"] = metadata + metadata["document_top_summary"] = document_top_summary + + lifecycle_service.update_progress( + job_id, + progress=80, + message="Generating ZIP package...", + ) + processing_completed_at = datetime.now(timezone.utc) + processing_timing_updates = { + "processing_completed_at": processing_completed_at.isoformat(), + "processing_duration_ms": max( + 0, + int((processing_completed_at - processing_started_at).total_seconds() * 1000), + ), + } + job_context.metadata_service.update_metadata(job_id, processing_timing_updates) + job_context.job_metadata.update(processing_timing_updates) + + data_id = JobMetadataHelper.get_field(job_context.job_metadata, "data_id") + zip_service = ZipResultService() + zip_file_path, checksum, statistics, zip_size = zip_service.generate_zip_package( + job_id=job_id, + chunks=chunks, + add_dir=str(add_dir) if add_dir else "", + source_file_name=source_file_name, + data_id=data_id, + job_metadata=job_context.job_metadata, + parsed_df=parsed_contents_df, + temp_dir=task_workspace_dir, + ) + del statistics + + checksum_value = ( + checksum.get("value", "") + if isinstance(checksum, dict) + else (checksum or "") + ) + + lifecycle_service.update_progress( + job_id, + progress=90, + message="Uploading results to S3...", + ) + result_bundle = get_result_storage().upload( + job_id=job_id, + result_dir=str(add_dir) if add_dir else "", + zip_file_path=zip_file_path, + ) + result_s3_key = result_bundle.zip_key + stored_count = 0 + + lifecycle_service.update_progress(job_id, progress=100, message="Task complete!") + lifecycle_service.finalize_job_success( + job_id=job_id, + chunks=chunks, + result_s3_key=result_s3_key, + checksum=checksum_value, + zip_size=zip_size, + stored_count=stored_count, + delivery_mode="url", + section_summaries=section_summaries, + ) + + logger.info( + f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" + ) + + return { + "status": "success", + "job_id": job_id, + "add_dir": None, + "vectors_count": 0, + "contents_count": len(parsed_contents_df), + "stored_count": stored_count, + "delivery_mode": "url", + "result_s3_key": result_s3_key, + } diff --git a/apps/worker/app/core/tasks/task_utils.py b/apps/worker/app/services/document_ingestion/workspace.py similarity index 86% rename from apps/worker/app/core/tasks/task_utils.py rename to apps/worker/app/services/document_ingestion/workspace.py index dafe3dd0e..80217df1f 100644 --- a/apps/worker/app/core/tasks/task_utils.py +++ b/apps/worker/app/services/document_ingestion/workspace.py @@ -1,3 +1,5 @@ +"""Task-scoped workspace helpers for worker-side Document Ingestion.""" + import os import shutil import tempfile @@ -28,10 +30,7 @@ def cleanup_temp_file(file_path: str | None) -> None: def cleanup_task_workspace(workspace_dir: str | None) -> bool: """Best-effort cleanup for a task-scoped temporary workspace.""" - if not workspace_dir: - return False - - if not os.path.isdir(workspace_dir): + if not workspace_dir or not os.path.isdir(workspace_dir): return False try: @@ -71,15 +70,17 @@ def create_task_workspace(job_id: str) -> str: def download_s3_file_to_temp(file_url: str, file_ext: str, temp_dir: str) -> str: - """Download the source file from object storage into a task workspace file.""" - local_temp_path = None + """Download the source file from object storage into the task workspace.""" + local_temp_path: str | None = None try: os.makedirs(temp_dir, exist_ok=True) with tempfile.NamedTemporaryFile( - delete=False, suffix=file_ext, dir=temp_dir - ) as tmp_file: - local_temp_path = tmp_file.name + delete=False, + suffix=file_ext, + dir=temp_dir, + ) as temp_file: + local_temp_path = temp_file.name with requests.get( file_url, timeout=120, @@ -89,7 +90,7 @@ def download_s3_file_to_temp(file_url: str, file_ext: str, temp_dir: str) -> str response.raise_for_status() for chunk in response.iter_content(chunk_size=65536): if chunk: - tmp_file.write(chunk) + temp_file.write(chunk) except requests.RequestException as exc: cleanup_temp_file(local_temp_path) raise StorageServiceException( diff --git a/apps/worker/app/services/document_parser/atlas_parser.py b/apps/worker/app/services/document_parser/atlas_parser.py index a379d9c51..2d95b1b5a 100644 --- a/apps/worker/app/services/document_parser/atlas_parser.py +++ b/apps/worker/app/services/document_parser/atlas_parser.py @@ -19,11 +19,8 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import pandas as pd -from app.services.common.kb_utils import ( - gen_str_codes, - get_str_time, - process_dup_paths_df, -) +from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.identifiers import gen_str_codes, get_str_time from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker from app.services.document_parser.toc_parser import detect_tocs_in_texts from loguru import logger diff --git a/apps/worker/app/services/document_parser/dataframe_helpers.py b/apps/worker/app/services/document_parser/dataframe_helpers.py new file mode 100644 index 000000000..2173286b3 --- /dev/null +++ b/apps/worker/app/services/document_parser/dataframe_helpers.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import os +from typing import Any + +import pandas as pd + +from app.services.document_parser.identifiers import gen_str_codes + + +def flatten_list(nested_list: list[Any]) -> list[Any]: + """Flatten a nested list.""" + flat_list: list[Any] = [] + for item in nested_list: + if isinstance(item, list): + flat_list.extend(flatten_list(item)) + else: + flat_list.append(item) + return flat_list + + +def merge_df(input_df: pd.DataFrame) -> pd.DataFrame: + """Merge DataFrame rows that share the same path.""" + dfs_by_path = list(input_df.groupby("path", sort=False)) + processed_dfs: list[pd.DataFrame] = [] + + for key, df in dfs_by_path: + content_to_merge: list[str] = [] + types_to_merge: list[str] = [] + total_length = 0 + + for _, row in df.iterrows(): + content_to_merge.append(str(row["content"])) + types_to_merge.extend(str(row["type"]).split("\n")) + total_length += len(str(row["content"])) + + processed_dfs.append( + pd.DataFrame( + [ + { + "content": "\n".join(content_to_merge), + "type": "\n".join(sorted(set(types_to_merge))), + "path": key, + "length": total_length, + "know_id": gen_str_codes("\n".join(content_to_merge)), + } + ] + ) + ) + + return pd.concat(processed_dfs, axis=0, ignore_index=True) + + +def process_dup_paths_df(df: pd.DataFrame) -> pd.DataFrame: + """De-duplicate KB DataFrame paths for final output.""" + if "path" not in df.columns: + return df + + split_char = os.getenv("SPLIT_CHAR", "/") + dup_mask = df["path"].duplicated(keep=False) + if not dup_mask.any(): + return df + + path_occurrences: dict[str, list[int]] = {} + for idx, path in enumerate(df["path"]): + path_occurrences.setdefault(str(path), []).append(idx) + + path_renames: dict[int, str] = {} + parent_rename_map: dict[str, dict[int, str]] = {} + + for path, indices in path_occurrences.items(): + if len(indices) > 1: + parent_rename_map[path] = {} + for occurrence, idx in enumerate(indices): + if occurrence == 0: + path_renames[idx] = path + else: + new_path = f"{path}_{occurrence + 1}" + path_renames[idx] = new_path + parent_rename_map[path][idx] = new_path + + new_paths: list[str] = [] + for idx, row in df.iterrows(): + row_index = int(str(idx)) + path = str(row["path"]) + new_path = path_renames.get(row_index, path) + path_parts = new_path.split(split_char) + + for parent_path, rename_info in parent_rename_map.items(): + parent_parts = parent_path.split(split_char) + if len(path_parts) > len(parent_parts) and path_parts[: len(parent_parts)] == parent_parts: + matching_parent_idx = None + for parent_idx in sorted(rename_info.keys(), reverse=True): + if parent_idx < row_index: + matching_parent_idx = parent_idx + break + if matching_parent_idx is not None: + renamed_parent = rename_info[matching_parent_idx] + renamed_parent_parts = renamed_parent.split(split_char) + new_path = split_char.join( + renamed_parent_parts + path_parts[len(parent_parts) :] + ) + break + + new_paths.append(new_path) + + df = df.copy() + df["path"] = new_paths + return df diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index a253100c2..8e50ab15e 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -6,11 +6,10 @@ import zipfile import pandas as pd -from app.services.common.kb_utils import ( +from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.path_helpers import ( find_matches_parsing, - gen_str_codes, - get_str_time, - process_dup_paths_df, process_path_texts, remove_spaces, ) @@ -74,7 +73,7 @@ def _find_img_context(headings_stack, max_chars=100): Returns: The nearest valid text context, or empty string if none found """ - from app.services.common.kb_utils import truncate_text + from app.services.document_parser.text_helpers import truncate_text try: content_list = headings_stack[-1].get("content", []) @@ -236,7 +235,7 @@ def _first_cols_rows(table_block, max_items=10, max_chars=20): Returns: Tuple of (first_row_text, first_col_text) with ' | ' as separator """ - from app.services.common.kb_utils import truncate_text + from app.services.document_parser.text_helpers import truncate_text first_row_text = "" first_col_text = "" diff --git a/apps/worker/app/services/document_parser/html_parser.py b/apps/worker/app/services/document_parser/html_parser.py index c94fbfce8..8c65a4ec7 100644 --- a/apps/worker/app/services/document_parser/html_parser.py +++ b/apps/worker/app/services/document_parser/html_parser.py @@ -410,7 +410,7 @@ def first_cols_rows_html(html_str, max_items=10, max_chars=20): Returns: Tuple of (first_row_text, first_col_text) with ' | ' as separator """ - from app.services.common.kb_utils import truncate_text + from app.services.document_parser.text_helpers import truncate_text soup = BeautifulSoup(html_str, "html.parser") table = soup.find("table") diff --git a/apps/worker/app/services/document_parser/identifiers.py b/apps/worker/app/services/document_parser/identifiers.py new file mode 100644 index 000000000..c8d177c3a --- /dev/null +++ b/apps/worker/app/services/document_parser/identifiers.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +import uuid +from datetime import datetime + + +def gen_str_codes(input_string: str) -> str: + """Generate a UUID5 code from a string.""" + return str(uuid.uuid5(uuid.NAMESPACE_DNS, input_string)) + + +def get_str_time() -> str: + """Get the current time as a string.""" + return datetime.now().strftime("%Y-%m-%d %H:%M:%S") diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/image_parser.py index dc8ffb541..70c3c232a 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/image_parser.py @@ -8,11 +8,8 @@ from pathlib import Path import pandas as pd -from app.services.common.kb_utils import ( - gen_str_codes, - get_str_time, - process_dup_paths_df, -) +from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.identifiers import gen_str_codes, get_str_time from loguru import logger from PIL import Image diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py index f308e7f5a..6aa83ba63 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/layout_parser.py @@ -6,10 +6,7 @@ import gevent import pandas as pd -from app.services.common.kb_utils import ( - count_cn_en, - truncate_text_by_tokens, -) +from app.services.document_parser.text_helpers import count_cn_en, truncate_text_by_tokens from app.services.document_parser.stage_profiler import stage_timer from app.services.document_parser.table_parser import df2md from docx.oxml.ns import qn diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py index ab85751bc..33e7357af 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/md_parser.py @@ -7,12 +7,9 @@ import gevent import pandas as pd -from app.services.common.kb_utils import ( - find_matches_parsing, - gen_str_codes, - get_str_time, - process_dup_paths_df, -) +from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.path_helpers import find_matches_parsing from app.services.document_parser.html_parser import ( first_cols_rows_html, merge_html_tables, diff --git a/apps/worker/app/services/document_parser/orchestration/__init__.py b/apps/worker/app/services/document_parser/orchestration/__init__.py new file mode 100644 index 000000000..6f12d8497 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/__init__.py @@ -0,0 +1 @@ +"""Document parser orchestration modules.""" diff --git a/apps/worker/app/services/document_parser/orchestration/parse_session.py b/apps/worker/app/services/document_parser/orchestration/parse_session.py new file mode 100644 index 000000000..44a58e343 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any + +from app.services.document_parser.atlas_classifier import classify_atlas_with_vlm +from app.services.document_parser.doc_profiler import profile_document +from app.services.document_parser.stage_profiler import stage_timer +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ValidationException +from shared.utils.file_utils import path_handle + +PDF_PAGE_LIMIT = 600 + + +@dataclass(frozen=True) +class ParseSession: + base_llm_paras: dict[str, object] + base_url: str + file_full_path: str + filename: str + fragment_content: str + full_output_dir: str + internal_output_filename: str + job_id: str | None + kb_dir: str + output_dir: str + profile: Any + relative_root: str + s3_key: str | None + + +def build_parse_session( + *, + add_frag_desc: str, + base_url: str, + doc_type: str, + file_full_path: str, + filename: str, + fragment_content: str, + internal_output_filename: str, + job_id: str | None, + kb_dir: str, + llm_histories: int, + output_dir: str, + s3_key: str | None, + smart_title_parse: bool, + stopwords: list[str] | None, + summary_image: bool, + summary_table: bool, + summary_txt: bool, +) -> ParseSession: + """Build the parser routing session from explicit parse inputs.""" + base_llm_paras = { + "llm_histories": llm_histories, + "smart_title_parse": smart_title_parse, + "summary_image": summary_image, + "summary_table": summary_table, + "summary_txt": summary_txt, + "stopwords": stopwords, + "doc_type": doc_type, + "frag_desc": add_frag_desc, + "model_name": settings.NORMOL_MODEL, + "hierarchy_model_name": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, + } + + logger.debug(f"baseurl: {base_url}") + logger.debug(f"file_full_path: {file_full_path}") + + relative_root, full_output_dir = _resolve_output_paths( + filename=filename, + internal_output_filename=internal_output_filename, + kb_dir=kb_dir, + output_dir=output_dir, + ) + logger.debug(f"relative_root: {relative_root}") + logger.debug(f"full_output_dir: {full_output_dir}") + + with stage_timer("document.profile", filename=filename): + profile = profile_document(file_full_path, internal_output_filename) + logger.info(f"📋 DocProfile: {profile.summary()}") + logger.debug(f"📋 Reasoning: {profile.reasoning}") + + if profile.atlas_candidate and profile.doc_category not in ("atlas", "ppt_converted"): + logger.info(f"🔍 Atlas candidate detected, running VLM visual check for {filename}") + with stage_timer("document.atlas_vlm_check", filename=filename): + vlm_is_atlas = classify_atlas_with_vlm(file_full_path) + if vlm_is_atlas: + profile.doc_category = "atlas" + profile.reasoning += " | vlm_confirmed_atlas=True" + logger.info(f"✅ VLM confirmed atlas for {filename}") + else: + profile.reasoning += " | vlm_confirmed_atlas=False" + logger.info(f"ℹ️ VLM rejected atlas for {filename}, routing as generic") + + if profile.file_type == "pdf" and profile.page_count > PDF_PAGE_LIMIT: + raise ValidationException( + user_message=( + f"Document too large: {profile.page_count} pages exceeds the {PDF_PAGE_LIMIT}-page limit. " + "Please split the document and upload in smaller batches." + ), + violations=[ + { + "field": "page_count", + "description": f"PDF has {profile.page_count} pages, limit is {PDF_PAGE_LIMIT}", + } + ], + ) + + if profile.doc_category == "atlas": + filename, internal_output_filename, relative_root, full_output_dir = _rename_atlas_output( + filename=filename, + internal_output_filename=internal_output_filename, + kb_dir=kb_dir, + output_dir=output_dir, + ) + logger.info(f"📐 Atlas output renamed: {filename}") + + return ParseSession( + base_llm_paras=base_llm_paras, + base_url=base_url, + file_full_path=file_full_path, + filename=filename, + fragment_content=fragment_content, + full_output_dir=full_output_dir, + internal_output_filename=internal_output_filename, + job_id=job_id, + kb_dir=kb_dir, + output_dir=output_dir, + profile=profile, + relative_root=relative_root, + s3_key=s3_key, + ) + + +def _rename_atlas_output( + *, + filename: str, + internal_output_filename: str, + kb_dir: str, + output_dir: str, +) -> tuple[str, str, str, str]: + name_base, _ = os.path.splitext(filename) + internal_name_base, _ = os.path.splitext(internal_output_filename) + atlas_filename = name_base + ".atlas" + atlas_internal_filename = internal_name_base + ".atlas" + relative_root, full_output_dir = _resolve_output_paths( + filename=atlas_filename, + internal_output_filename=atlas_internal_filename, + kb_dir=kb_dir, + output_dir=output_dir, + ) + return atlas_filename, atlas_internal_filename, relative_root, full_output_dir + + +def _resolve_output_paths( + *, + filename: str, + internal_output_filename: str, + kb_dir: str, + output_dir: str, +) -> tuple[str, str]: + split_char = settings.SPLIT_CHAR or "/" + kb_dir_parts = kb_dir.split(split_char) + + if filename and "images" not in kb_dir_parts: + relative_root = "/".join(kb_dir_parts + [filename]) + else: + relative_root = "/".join(kb_dir_parts) + + if internal_output_filename and "images" not in kb_dir_parts: + internal_relative_root = "/".join(kb_dir_parts + [internal_output_filename]) + else: + internal_relative_root = "/".join(kb_dir_parts) + + full_output_dir = os.path.join( + output_dir, + internal_relative_root.replace("/", os.sep), + ) + sanitized_output_dir = path_handle(full_output_dir, mode="sanitize") + if not isinstance(sanitized_output_dir, str) or not sanitized_output_dir: + raise ValueError(f"Failed to sanitize parser output directory: {full_output_dir}") + os.makedirs(sanitized_output_dir, exist_ok=True) + + logger.debug(f"internal_relative_root: {internal_relative_root}") + return relative_root, sanitized_output_dir diff --git a/apps/worker/app/services/document_parser/orchestration/postprocess.py b/apps/worker/app/services/document_parser/orchestration/postprocess.py new file mode 100644 index 000000000..9cebc7c12 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/postprocess.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import os +import re + +import pandas as pd +from app.services.document_parser.image_compressor import ( + apply_rename_map_to_dataframe, + compress_output_images, +) +from app.services.document_parser.stage_profiler import stage_timer +from loguru import logger + + +def apply_parse_postprocess( + output_dir: str, + parsed_df: pd.DataFrame | None, +) -> pd.DataFrame | None: + """Apply output cleanup and image compression after parsing.""" + logger.debug(f"full_output_dir: {output_dir}") + + with stage_timer("document.cleanup_unreferenced_images", output_dir=output_dir): + cleanup_unreferenced_images(output_dir) + + with stage_timer("document.compress_images", output_dir=output_dir): + compress_stats = compress_output_images(output_dir) + if compress_stats.processed > 0: + logger.info( + f"📦 Image compression: {compress_stats.processed} processed " + f"({compress_stats.converted_png_to_jpg} PNG→JPG, " + f"{compress_stats.resized} resized), " + f"{compress_stats.bytes_before / 1024 / 1024:.1f}MB → " + f"{compress_stats.bytes_after / 1024 / 1024:.1f}MB" + ) + if compress_stats.rename_map and parsed_df is not None: + return apply_rename_map_to_dataframe(parsed_df, compress_stats.rename_map) + + return parsed_df + + +def cleanup_unreferenced_images(output_dir: str) -> int: + """Remove UUID-named images that are not referenced by final parsed output.""" + image_dir = os.path.join(output_dir, "images") + if not os.path.isdir(image_dir): + return 0 + + uuid_pattern = re.compile( + r"^[a-f0-9]{64}\.(?:jpg|jpeg|png|gif|webp)$", + re.IGNORECASE, + ) + removed_count = 0 + + for filename in os.listdir(image_dir): + if not uuid_pattern.match(filename): + continue + + file_path = os.path.join(image_dir, filename) + try: + os.remove(file_path) + removed_count += 1 + logger.debug(f"Removed unreferenced image: {filename}") + except OSError as exc: + logger.warning(f"Failed to remove {filename}: {exc}") + + if removed_count > 0: + logger.info( + f"Cleaned up {removed_count} unreferenced UUID-named images from {image_dir}" + ) + + return removed_count diff --git a/apps/worker/app/services/document_parser/orchestration/route_parse.py b/apps/worker/app/services/document_parser/orchestration/route_parse.py new file mode 100644 index 000000000..faeda3dee --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/route_parse.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import os + +import pandas as pd + +from app.services.document_parser.orchestration.parse_session import ParseSession +from shared.core.exceptions.domain_exceptions import ValidationException + +SUPPORTED_FILE_TYPES: tuple[str, ...] = ( + ".txt", + ".fragment", + ".png", + ".jpg", + ".jpeg", + ".pdf", + ".doc", + ".docx", + ".xls", + ".xlsx", + ".pptx", + ".md", + ".json", +) + + +def route_document_parse(session: ParseSession) -> tuple[str, pd.DataFrame | None]: + """Route a parser session to the correct adapter and return its output.""" + file_path_lower = session.file_full_path.lower() + + if ".fragment" in file_path_lower: + from app.services.document_parser.fragment_parser import parse_fragment + + full_output_dir, _relative_root, parsed_df = parse_fragment( + session.fragment_content, + filename=session.filename, + output_dir=session.output_dir, + kb_dir=session.kb_dir, + base_llm_paras=session.base_llm_paras, + ) + return full_output_dir, parsed_df + + if ".txt" in file_path_lower: + from app.services.document_parser.md_parser import parse_md + from app.services.document_parser.txt_parser import parse_texts + + text_lines = parse_texts(file_path=session.file_full_path, baseurl=session.base_url) + parsed_df = parse_md( + session.full_output_dir, + source_type="md", + md_lines=text_lines, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + if any(extension in file_path_lower for extension in (".png", ".jpg", ".jpeg")): + from app.services.document_parser.image_parser import parse_image + + parsed_df = parse_image( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + baseurl=session.base_url, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + if ".pdf" in file_path_lower: + from app.services.document_parser.pdf_parser import parse_pdfs + + parsed_df = parse_pdfs( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + base_llm_paras=session.base_llm_paras, + profile=session.profile, + relative_root=session.relative_root, + s3_key=session.s3_key, + ) + return session.full_output_dir, parsed_df + + if ".doc" in file_path_lower and ".docx" not in file_path_lower: + from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx + from app.services.document_parser.legacy_converter import doc_to_docx + + converted_docx_path, _ = doc_to_docx( + session.file_full_path, + outdir=session.full_output_dir, + ) + parsed_structure, dataframe_list = parse_docx( + converted_docx_path, + session.base_llm_paras, + session.full_output_dir, + session.filename, + session.base_url, + relative_root=session.relative_root, + ) + parsed_df = convert_doc2dics( + parsed_structure, + dataframe_list, + session.full_output_dir, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + if ".docx" in file_path_lower: + from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx + + parsed_structure, dataframe_list = parse_docx( + session.file_full_path, + session.base_llm_paras, + session.full_output_dir, + session.filename, + session.base_url, + relative_root=session.relative_root, + ) + parsed_df = convert_doc2dics( + parsed_structure, + dataframe_list, + session.full_output_dir, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + if ".xls" in file_path_lower and ".xlsx" not in file_path_lower: + from app.services.document_parser.legacy_converter import xls_to_xlsx + from app.services.document_parser.table_parser import parse_xlsx + + converted_xlsx_path, _ = xls_to_xlsx( + session.file_full_path, + outdir=session.full_output_dir, + ) + parsed_df = parse_xlsx( + converted_xlsx_path, + session.filename, + session.full_output_dir, + session.base_url, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + if ".xlsx" in file_path_lower: + from app.services.document_parser.table_parser import parse_xlsx + + parsed_df = parse_xlsx( + session.file_full_path, + session.filename, + session.full_output_dir, + session.base_url, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + if ".pptx" in file_path_lower: + from app.services.document_parser.pptx_parser import parse_pptx + + parsed_df = parse_pptx( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + base_llm_paras=session.base_llm_paras, + strategy="to_pdf_api", + job_id=session.job_id, + relative_root=session.relative_root, + baseurl=session.base_url, + ) + return session.full_output_dir, parsed_df + + if ".md" in file_path_lower: + from app.services.document_parser.md_parser import parse_md + + parsed_df = parse_md( + session.full_output_dir, + source_type="md", + file_path=session.file_full_path, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + if ".json" in file_path_lower: + return session.full_output_dir, None + + file_ext = os.path.splitext(session.file_full_path)[1].lower() + raise ValidationException( + user_message=f"Unsupported file type: {file_ext}", + violations=[ + { + "field": "file_type", + "description": f"Must be one of: {', '.join(SUPPORTED_FILE_TYPES)}", + } + ], + ) diff --git a/apps/worker/app/services/document_parser/parse_service.py b/apps/worker/app/services/document_parser/parse_service.py index ad5058e47..482147099 100644 --- a/apps/worker/app/services/document_parser/parse_service.py +++ b/apps/worker/app/services/document_parser/parse_service.py @@ -1,68 +1,10 @@ -# pyright: reportArgumentType=false, reportReturnType=false -""" -main parsing service -""" - -import os -import re +"""Stable parser seam backed by dedicated orchestration modules.""" import pandas as pd -from app.services.document_parser.atlas_classifier import classify_atlas_with_vlm - -# document_parser imports -from app.services.document_parser.doc_profiler import profile_document -from app.services.document_parser.stage_profiler import stage_timer -from loguru import logger - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - ValidationException, -) -from shared.utils.file_utils import path_handle - - -def cleanup_unreferenced_images(output_dir: str) -> int: - """ - Clean up unreferenced UUID-named images from the images directory. - - After document parsing (PDF, DOCX, PPTX, etc.), the images/ directory may contain: - 1. Processed images: renamed with semantic names like 'image-0-xxx.jpg' - 2. Unreferenced images: UUID-named (64-char hex) that were parsed as tables/formulas - - This function removes the unreferenced UUID-named images to reduce final package size. - Args: - output_dir: The full output directory path - - Returns: - Number of files removed - """ - img_dir = os.path.join(output_dir, "images") - if not os.path.isdir(img_dir): - return 0 - - # UUID pattern: 64 hex characters followed by image extension - uuid_pattern = re.compile( - r"^[a-f0-9]{64}\.(?:jpg|jpeg|png|gif|webp)$", re.IGNORECASE - ) - - removed_count = 0 - for filename in os.listdir(img_dir): - if uuid_pattern.match(filename): - file_path = os.path.join(img_dir, filename) - try: - os.remove(file_path) - removed_count += 1 - logger.debug(f"Removed unreferenced image: {filename}") - except OSError as e: - logger.warning(f"Failed to remove {filename}: {e}") - - if removed_count > 0: - logger.info( - f"Cleaned up {removed_count} unreferenced UUID-named images from {img_dir}" - ) - - return removed_count +from app.services.document_parser.orchestration.parse_session import build_parse_session +from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess +from app.services.document_parser.orchestration.route_parse import route_document_parse def checkerboard_inject_parse( @@ -84,374 +26,26 @@ def checkerboard_inject_parse( fragment_content: str = "", s3_key: str | None = None, ) -> tuple[str, pd.DataFrame | None]: - """ - main parsing function - - Args: - file_full_path: source file path (local or URL) - filename: file name - output_dir: output directory (absolute path, caller provides) - kb_dir: sub-directory name - llm_histories: retained for downstream LLM settings - smart_title_parse: enable smart heading parsing - summary_image: enable image summaries - summary_table: enable table summaries - summary_txt: enable text summaries - stopwords: optional stopword list - doc_type: parser document type hint - add_frag_desc: extra fragment description - base_url: optional source base URL - fragment_content: raw fragment content - job_id: optional job identifier used for parser artifacts - internal_output_filename: normalized internal folder name for on-disk output - s3_key: optional S3 key for downstream parsers - - Returns: - tuple: (output_dir, parsed_df) - - output_dir: directory path after parsing - - parsed_df: parsed content DataFrame - """ - # Build base_llm_paras from explicit parameters - base_llm_paras = { - "llm_histories": llm_histories, - "smart_title_parse": smart_title_parse, - "summary_image": summary_image, - "summary_table": summary_table, - "summary_txt": summary_txt, - "stopwords": stopwords, - "doc_type": doc_type, - "frag_desc": add_frag_desc, - "model_name": settings.NORMOL_MODEL, - "hierarchy_model_name": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, - } - - baseurl = base_url - - logger.debug(f"baseurl: {baseurl}") - logger.debug(f"file_full_path: {file_full_path}") - - # ========== Path handling ========== - split_char = settings.SPLIT_CHAR or "/" - - # Develop relative root path for chunk path field - kb_dir_parts = kb_dir.split(split_char) - if filename and "images" not in kb_dir_parts: - relative_root = "/".join(kb_dir_parts + [filename]) - else: - relative_root = "/".join(kb_dir_parts) - - if internal_output_filename and "images" not in kb_dir_parts: - internal_relative_root = "/".join(kb_dir_parts + [internal_output_filename]) - else: - internal_relative_root = "/".join(kb_dir_parts) - - # Develop full output directory (output_dir + relative_root) - full_output_dir = os.path.join( - output_dir, internal_relative_root.replace("/", os.sep) - ) - full_output_dir = path_handle(full_output_dir, mode="sanitize") - os.makedirs(full_output_dir, exist_ok=True) - - logger.debug(f"relative_root: {relative_root}") - logger.debug(f"internal_relative_root: {internal_relative_root}") - logger.debug(f"full_output_dir: {full_output_dir}") - - file_path_lower = file_full_path.lower() - parsed_df = None - - # ── Agentic Profiler: classify document before routing ── - with stage_timer("document.profile", filename=filename): - profile = profile_document(file_full_path, internal_output_filename) - logger.info(f"📋 DocProfile: {profile.summary()}") - logger.debug(f"📋 Reasoning: {profile.reasoning}") - - # ── VLM second-pass: confirm atlas_candidate with visual check ── - # Heuristics can miss atlases that have a rich OCR text layer on top of - # scanned drawing pages (avg_text_density too high). VLM sees the actual - # page layout and makes the final call. - if profile.atlas_candidate and profile.doc_category not in ( - "atlas", - "ppt_converted", - ): - logger.info( - f"🔍 Atlas candidate detected, running VLM visual check for {filename}" - ) - with stage_timer("document.atlas_vlm_check", filename=filename): - vlm_is_atlas = classify_atlas_with_vlm(file_full_path) - if vlm_is_atlas: - profile.doc_category = "atlas" - profile.reasoning += " | vlm_confirmed_atlas=True" - logger.info(f"✅ VLM confirmed atlas for {filename}") - else: - profile.reasoning += " | vlm_confirmed_atlas=False" - logger.info(f"ℹ️ VLM rejected atlas for {filename}, routing as generic") - - # ── Page count guard: reject oversized PDFs before routing ── - PDF_PAGE_LIMIT = 600 - if profile and profile.file_type == "pdf" and profile.page_count > PDF_PAGE_LIMIT: - raise ValidationException( - user_message=( - f"Document too large: {profile.page_count} pages exceeds the {PDF_PAGE_LIMIT}-page limit. " - f"Please split the document and upload in smaller batches." - ), - violations=[ - { - "field": "page_count", - "description": f"PDF has {profile.page_count} pages, limit is {PDF_PAGE_LIMIT}", - } - ], - ) - - # Atlas routing: rename output folder from .pdf → .atlas for easy filtering - if profile and profile.doc_category == "atlas": - name_base, _ = os.path.splitext(filename) - internal_name_base, _ = os.path.splitext(internal_output_filename) - filename = name_base + ".atlas" - internal_output_filename = internal_name_base + ".atlas" - relative_root = "/".join(kb_dir_parts + [filename]) - internal_relative_root = "/".join(kb_dir_parts + [internal_output_filename]) - full_output_dir = os.path.join( - output_dir, internal_relative_root.replace("/", os.sep) - ) - full_output_dir = path_handle(full_output_dir, mode="sanitize") - os.makedirs(full_output_dir, exist_ok=True) - logger.info(f"📐 Atlas output renamed: {filename}") - - if ".fragment" in file_path_lower: - logger.debug("file type is fragment") - from app.services.document_parser.fragment_parser import parse_fragment - - full_output_dir, relative_root, parsed_df = parse_fragment( - fragment_content, - filename=filename, - output_dir=output_dir, - kb_dir=kb_dir, - base_llm_paras=base_llm_paras, - ) - - elif ".txt" in file_path_lower: - logger.debug("file type is txt") - from app.services.document_parser.md_parser import parse_md - from app.services.document_parser.txt_parser import parse_texts - - txt_lines = parse_texts(file_path=file_full_path, baseurl=baseurl) - parsed_df = parse_md( - full_output_dir, - source_type="md", - md_lines=txt_lines, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ( - ".png" in file_path_lower - or ".jpg" in file_path_lower - or ".jpeg" in file_path_lower - ): - logger.debug("file type is image") - from app.services.document_parser.image_parser import parse_image - - parsed_df = parse_image( - file_full_path, - filename=filename, - output_dir=full_output_dir, - baseurl=baseurl, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".pdf" in file_path_lower: - logger.debug("file type is pdf") - from app.services.document_parser.pdf_parser import parse_pdfs - - if filename and file_full_path: - parsed_df = parse_pdfs( - file_full_path, - filename=filename, - output_dir=full_output_dir, - base_llm_paras=base_llm_paras, - profile=profile, - relative_root=relative_root, - s3_key=s3_key, - ) - - elif ".doc" in file_path_lower and ".docx" not in file_path_lower: - logger.debug("file type is doc") - from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx - from app.services.document_parser.legacy_converter import doc_to_docx - - if filename and file_full_path: - converted_docx_path, _ = doc_to_docx(file_full_path, outdir=full_output_dir) - parsed_structure, df_list = parse_docx( - converted_docx_path, - base_llm_paras, - full_output_dir, - filename, - baseurl, - relative_root=relative_root, - ) - parsed_df = convert_doc2dics( - parsed_structure, - df_list, - full_output_dir, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".docx" in file_path_lower: - logger.debug("file type is docx") - from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx - - if filename and file_full_path: - parsed_structure, df_list = parse_docx( - file_full_path, - base_llm_paras, - full_output_dir, - filename, - baseurl, - relative_root=relative_root, - ) - parsed_df = convert_doc2dics( - parsed_structure, - df_list, - full_output_dir, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".xls" in file_path_lower and ".xlsx" not in file_path_lower: - logger.debug("file type is xls") - from app.services.document_parser.legacy_converter import xls_to_xlsx - from app.services.document_parser.table_parser import parse_xlsx - - if filename and file_full_path: - converted_xlsx_path, _ = xls_to_xlsx(file_full_path, outdir=full_output_dir) - parsed_df = parse_xlsx( - converted_xlsx_path, - filename, - full_output_dir, - baseurl, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".xlsx" in file_path_lower: - logger.debug("file type is xlsx") - from app.services.document_parser.table_parser import parse_xlsx - - if filename and file_full_path: - parsed_df = parse_xlsx( - file_full_path, - filename, - full_output_dir, - baseurl, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".pptx" in file_path_lower: - logger.debug("file type is pptx") - from app.services.document_parser.pptx_parser import parse_pptx - - if filename and file_full_path: - # ====== iLoveAPI PPTX → PDF → MinerU (default production route) ====== - parsed_df = parse_pptx( - file_full_path, - filename=filename, - output_dir=full_output_dir, - base_llm_paras=base_llm_paras, - strategy="to_pdf_api", - job_id=job_id, - relative_root=relative_root, - baseurl=baseurl, - ) - - # ====== [EXPERIMENTAL] Directly send PPTX to MinerU via parse_pdfs ====== - # Uncomment the block below (and comment out parse_pptx above) to bypass iLoveAPI - # from app.services.document_parser.pdf_parser import parse_pdfs - # parsed_df = parse_pdfs( - # file_full_path, - # filename=filename, - # output_dir=full_output_dir, - # base_llm_paras=base_llm_paras, - # profile=profile, - # relative_root=relative_root, - # s3_key=s3_key - # ) - - elif ".md" in file_path_lower: - logger.debug("file type is md") - from app.services.document_parser.md_parser import parse_md - - if filename and file_full_path: - parsed_df = parse_md( - full_output_dir, - source_type="md", - file_path=file_full_path, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - ) - - elif ".json" in file_path_lower: - logger.debug("file type is json") - # JSON parsing not yet implemented - - else: - # Unsupported file type - file_ext = os.path.splitext(file_full_path)[1].lower() - supported_types = [ - ".txt", - ".fragment", - ".png", - ".jpg", - ".jpeg", - ".pdf", - ".doc", - ".docx", - ".xls", - ".xlsx", - ".pptx", - ".md", - ".json", - ] - raise ValidationException( - user_message=f"Unsupported file type: {file_ext}", - violations=[ - { - "field": "file_type", - "description": f"Must be one of: {', '.join(supported_types)}", - } - ], - ) - - logger.debug(f"full_output_dir: {full_output_dir}") - - # Post-processing: clean up unreferenced UUID-named images - with stage_timer( - "document.cleanup_unreferenced_images", output_dir=full_output_dir - ): - cleanup_unreferenced_images(full_output_dir) - - # Post-processing: compress output images (PNG→JPEG, resize oversized) - from app.services.document_parser.image_compressor import ( - apply_rename_map_to_dataframe, - compress_output_images, + """Run the stable parser seam using dedicated orchestration modules.""" + session = build_parse_session( + add_frag_desc=add_frag_desc, + base_url=base_url, + doc_type=doc_type, + file_full_path=file_full_path, + filename=filename, + fragment_content=fragment_content, + internal_output_filename=internal_output_filename, + job_id=job_id, + kb_dir=kb_dir, + llm_histories=llm_histories, + output_dir=output_dir, + s3_key=s3_key, + smart_title_parse=smart_title_parse, + stopwords=stopwords, + summary_image=summary_image, + summary_table=summary_table, + summary_txt=summary_txt, ) - - with stage_timer("document.compress_images", output_dir=full_output_dir): - compress_stats = compress_output_images(full_output_dir) - if compress_stats.processed > 0: - logger.info( - f"📦 Image compression: {compress_stats.processed} processed " - f"({compress_stats.converted_png_to_jpg} PNG→JPG, " - f"{compress_stats.resized} resized), " - f"{compress_stats.bytes_before / 1024 / 1024:.1f}MB → " - f"{compress_stats.bytes_after / 1024 / 1024:.1f}MB" - ) - # Update DataFrame references when PNG→JPG conversions occurred - if compress_stats.rename_map and parsed_df is not None: - parsed_df = apply_rename_map_to_dataframe( - parsed_df, compress_stats.rename_map - ) - + full_output_dir, parsed_df = route_document_parse(session) + parsed_df = apply_parse_postprocess(full_output_dir, parsed_df) return full_output_dir, parsed_df diff --git a/apps/worker/app/services/document_parser/path_helpers.py b/apps/worker/app/services/document_parser/path_helpers.py new file mode 100644 index 000000000..41965e125 --- /dev/null +++ b/apps/worker/app/services/document_parser/path_helpers.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import os +import re +from typing import Any + +from bs4 import BeautifulSoup +from shared.utils.chunk_refs import extract_chunk_refs +from shared.utils.file_utils import path_handle + +SUMMARY_PATH_MARKERS: tuple[str, ...] = ("summary", "\u6458\u8981\u603b\u7ed3") + + +def find_images(folder_path: str) -> list[str]: + """Find image files inside a folder tree.""" + image_extensions = {".png", ".jpg", ".jpeg"} + image_files: list[str] = [] + + for _, _, files in os.walk(folder_path): + files.sort() + for file in files: + if os.path.splitext(file)[1].lower() in image_extensions: + image_files.append(file) + return image_files + + +def find_matches_parsing(content: str, path: str) -> str: + """Parse table and image markers from content.""" + matches = extract_chunk_refs(content) + match_type = "PTXT" if len(matches) == 0 else "\n".join((["PTXT"] + matches)) + + split_char = os.getenv("SPLIT_CHAR", "/") + if any( + f"{split_char}{summary_marker}" in path + for summary_marker in SUMMARY_PATH_MARKERS + ): + parent_path = path.split(split_char)[-2] + match_type = "SUMMARY_" + parent_path + "_SUMMARY" + return match_type + + +def flatten_dic2paths( + d: dict[str, Any], + current_path: list[str] | None = None, + result: list[str] | None = None, +) -> list[str]: + """Flatten a nested dict into path strings.""" + if result is None: + result = [] + if current_path is None: + current_path = [] + + for key, value in d.items(): + if not isinstance(key, str): + continue + new_path = current_path + [key] + if isinstance(value, dict) and value: + flatten_dic2paths(value, new_path, result) + else: + split_char = os.getenv("SPLIT_CHAR", "/") + result.append(split_char.join(new_path)) + return result + + +def process_path_texts(path_: str, last: int = 50) -> str: + """Normalize path text for downstream use.""" + temp_path = path_handle(path_, mode="sanitize") + if not isinstance(temp_path, str) or temp_path == "": + return "" + return "_".join(temp_path.split(os.sep))[:last] + + +def remove_spaces(text: str, handle_punctuation: bool = False) -> str: + """Remove spaces between Chinese chars while keeping English word spacing.""" + if handle_punctuation: + punctuation = ( + r"""!"#$%&'()*+,-./:;<=>?@[\]^_`{|}~，。、【】《》？；：''""（）…—-！""" + ) + res_text = re.sub(f"[{re.escape(punctuation)}]", "", text) + else: + pattern = re.compile(r"([\u4e00-\u9fff])\s+|(?<=\s)([\u4e00-\u9fff])") + + def replacer(match: re.Match[str]) -> str: + return match.group(1) or match.group(2) + + res_text = pattern.sub(replacer, text) + + res_text = re.sub(r"\s+", " ", res_text) + return res_text.strip() + + +def traverse_dict(d: dict[str, Any], parent: str | None = None) -> list[str]: + """Traverse a dictionary and generate description text.""" + dic_texts: list[str] = [] + for key, value in d.items(): + if value: + child_keys = ", ".join(value.keys()) + text = f"'{key}' includes {child_keys}" + dic_texts.append(text) + dic_texts.extend(traverse_dict(value, key)) + return dic_texts + + +def restore_graph_by_paths(paths: list[str]) -> tuple[dict[str, Any], list[str]]: + """Rebuild a graph structure from path strings.""" + root_dict: dict[str, Any] = {} + split_char = os.getenv("SPLIT_CHAR", "/") + for path in paths: + nodes = path.split(split_char) + current_dict = root_dict + for node in nodes: + if node not in current_dict: + current_dict[node] = {} + current_dict = current_dict[node] + dic_texts = traverse_dict(root_dict) + return root_dict, dic_texts + + +def html2txt(html_text: str) -> str: + """Convert HTML into plain text.""" + soup = BeautifulSoup(html_text, "html.parser") + return soup.get_text() diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/pptx_parser.py index dde943110..408918655 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/pptx_parser.py @@ -6,7 +6,7 @@ import jwt import requests -from app.services.common.kb_utils import find_images +from app.services.document_parser.path_helpers import find_images from app.services.document_parser.legacy_converter import ( _convert_with_libreoffice, ) diff --git a/apps/worker/app/services/document_parser/table_parser.py b/apps/worker/app/services/document_parser/table_parser.py index 74860b867..0916eab57 100755 --- a/apps/worker/app/services/document_parser/table_parser.py +++ b/apps/worker/app/services/document_parser/table_parser.py @@ -11,13 +11,9 @@ import numpy as np import openpyxl import pandas as pd -from app.services.common.kb_utils import ( - flatten_dic2paths, - gen_str_codes, - get_str_time, - process_dup_paths_df, - remove_spaces, -) +from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.path_helpers import flatten_dic2paths, remove_spaces from app.services.document_parser.html_parser import df2html from bs4 import BeautifulSoup from loguru import logger diff --git a/apps/worker/app/services/document_parser/text_helpers.py b/apps/worker/app/services/document_parser/text_helpers.py new file mode 100644 index 000000000..b11fdc772 --- /dev/null +++ b/apps/worker/app/services/document_parser/text_helpers.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import re + +from shared.utils.text_utils import _CN_EN_NUM_RE + +_CN_CHAR_RE = re.compile(r"[\u4e00-\u9fff]") +EN_START_LIMIT = 15 +CN_RATIO_THRESHOLD = 0.3 + + +def normalize_md(text: str) -> str: + """Normalize markdown string for comparison.""" + text = re.sub(r"^\s*#+\s*", "", text) + text = re.sub(r"\s+", "", text) + return text.lower() + + +def truncate_text(text: str, start_limit: int, end_limit: int) -> str: + """Truncate text by raw character count, keeping start and end parts.""" + text = str(text) + total_limit = start_limit + end_limit + if len(text) <= total_limit: + return text + start_part = text[:start_limit] + end_part = text[-end_limit:] if end_limit > 0 else "" + return f"{start_part}...{end_part}" + + +def detect_primary_lang(text: str) -> str: + """Detect whether text is primarily Chinese or English/other.""" + if not text: + return "en" + tokens = _CN_EN_NUM_RE.findall(text) + if not tokens: + return "en" + cn_count = sum(1 for token in tokens if _CN_CHAR_RE.fullmatch(token)) + return "zh" if (cn_count / len(tokens)) >= CN_RATIO_THRESHOLD else "en" + + +def count_cn_en(text: str) -> int: + """Count semantic Chinese/English/number tokens in a string.""" + return len(_CN_EN_NUM_RE.findall(str(text))) + + +def truncate_text_by_tokens( + text: str, + start_limit: int, + end_limit: int, + lang_aware: bool = True, +) -> str: + """Truncate text by semantic token count, preserving whole words.""" + text = str(text) + matches = list(_CN_EN_NUM_RE.finditer(text)) + total = len(matches) + + if lang_aware and total > 0 and detect_primary_lang(text) == "en": + start_limit = min(start_limit, EN_START_LIMIT) + + if total <= start_limit + end_limit: + return text + + cut_start = matches[start_limit - 1].end() if start_limit > 0 else 0 + cut_end = matches[total - end_limit].start() if end_limit > 0 else len(text) + if cut_start >= cut_end: + return text + return text[:cut_start] + "..." + text[cut_end:] diff --git a/apps/worker/app/services/document_parser/toc_parser.py b/apps/worker/app/services/document_parser/toc_parser.py index e192bea40..c885f8b4c 100644 --- a/apps/worker/app/services/document_parser/toc_parser.py +++ b/apps/worker/app/services/document_parser/toc_parser.py @@ -15,10 +15,7 @@ import gevent import pandas as pd -from app.services.common.kb_utils import ( - normalize_md, - truncate_text_by_tokens, -) +from app.services.document_parser.text_helpers import normalize_md, truncate_text_by_tokens from app.services.document_parser.layout_parser import ( hiearchy_llm, judge_by_conditions, diff --git a/apps/worker/app/services/workload/__init__.py b/apps/worker/app/services/workload/__init__.py index b7df5ee6b..3ad3226c0 100644 --- a/apps/worker/app/services/workload/__init__.py +++ b/apps/worker/app/services/workload/__init__.py @@ -1,7 +1 @@ -""" -Billing services for the worker. -""" - -from .page_estimator import PageEstimator - -__all__ = ["PageEstimator"] +"""Worker workload adapters.""" diff --git a/apps/worker/app/services/workload/parse_job_service.py b/apps/worker/app/services/workload/parse_job_service.py deleted file mode 100644 index f8b6d453f..000000000 --- a/apps/worker/app/services/workload/parse_job_service.py +++ /dev/null @@ -1,467 +0,0 @@ -from __future__ import annotations - -import os -from datetime import datetime, timezone - -import pandas as pd -from app.core.tasks.task_utils import ( - cleanup_task_workspace, - create_task_workspace, - download_s3_file_to_temp, -) -from app.services.common.job_start_service import mark_job_running -from app.services.connect_builder.summary_builder import ( - build_section_summary_lookup, - enrich_doc_nav_summaries, - ensure_doc_nav_json, - load_nav_top_summary, -) -from app.services.document_parser.stage_profiler import stage_timer -from app.services.storage.sync_storage_service import ( - generate_download_url, - verify_s3_file_exists, -) -from app.services.workload.page_estimator import PageEstimator -from loguru import logger -from sqlalchemy import select - -from shared.core.config import settings -from shared.core.database_sync import get_sync_db_context -from shared.core.exceptions.domain_exceptions import ( - InsufficientCreditsException, - NotFoundException, - ValidationException, - WorkerHandlingException, -) -from shared.models.database.job import Job -from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.billing.work_billing_service import WorkBillingService -from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks -from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service -from shared.services.redis.distributed_lock import RedisJobLock -from shared.services.redis.redis_sync_service import ( - SyncJobInfoRedisService, - SyncJobMetadataService, - SyncRedisServiceFactory, -) -from shared.services.storage.result_storage import get_result_storage -from shared.services.storage.zip_result_service import ZipResultService - - -def parse_uploaded_file_job(job_id: str, user_id: str | None) -> dict[str, object]: - logger.info(f"Parse started: job_id={job_id}, user_id={user_id}") - lifecycle_service = get_sync_job_lifecycle_service() - - redis_service = SyncRedisServiceFactory.get_service() - job_info_service = SyncJobInfoRedisService(redis_service) - job_info = job_info_service.get_job_info(job_id) - - if not job_info: - logger.warning( - f"JobInfo not found in Redis for job_id={job_id}; falling back to database" - ) - with get_sync_db_context() as fallback_db: - job_row = fallback_db.execute( - select(Job).where(Job.job_id == job_id) - ).scalar_one_or_none() - - if not job_row or not job_row.s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id=job_id, - internal_message="job info not found in Redis or database", - ) - - s3_key: str = job_row.s3_key - job_user_id: str | None = str(job_row.user_id) if job_row.user_id else user_id - logger.info( - f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}" - ) - else: - raw_s3_key = job_info.get("s3_key") - if not isinstance(raw_s3_key, str) or not raw_s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id="s3_key", - internal_message="Missing s3_key in job_info", - ) - - s3_key = raw_s3_key - raw_job_user_id = job_info.get("user_id") - job_user_id = raw_job_user_id if isinstance(raw_job_user_id, str) else user_id - - file_info = verify_s3_file_exists(s3_key) - if not file_info.get("exists"): - raise NotFoundException( - resource="S3File", - resource_id=s3_key, - internal_message=f"S3 file not found: {s3_key}", - ) - - logger.info(f"S3 file verified: {s3_key}") - - file_size = file_info.get("size", 0) - file_extension = os.path.splitext(s3_key)[1].lower() - - if file_size > settings.MAX_FILE_SIZE: - limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) - raise ValidationException( - user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", - violations=[ - { - "field": "file_size", - "description": ( - f"Size {file_size} bytes exceeds limit of " - f"{settings.MAX_FILE_SIZE} bytes" - ), - } - ], - ) - - metadata_service = SyncJobMetadataService(redis_service) - job_metadata = metadata_service.get_metadata(job_id) - if not job_metadata: - raise NotFoundException( - resource="JobMetadata", - resource_id=job_id, - internal_message=f"Job metadata not found for job_id={job_id}", - ) - - should_process = mark_job_running(job_id, redis_service) - if not should_process: - logger.warning(f"Skipping parse_task for inactive job: job_id={job_id}") - return { - "status": "skipped", - "job_id": job_id, - "reason": "job_already_terminal", - } - - with RedisJobLock(redis_service, job_id): - task_workspace_dir = create_task_workspace(job_id) - input_dir = os.path.join(task_workspace_dir, "input") - output_dir = os.path.join(task_workspace_dir, "output") - os.makedirs(input_dir, exist_ok=True) - os.makedirs(output_dir, exist_ok=True) - logger.info( - f"Task workspace ready: job_id={job_id}, workspace={task_workspace_dir}" - ) - - try: - lifecycle_service.update_progress( - job_id, progress=10, message="Parsing document..." - ) - - file_url_response = generate_download_url(s3_key, settings.S3_BUCKET_NAME) - file_url = file_url_response["download_url"] - - filename = JobMetadataHelper.get_field(job_metadata, "source_file_name") - - file_ext = os.path.splitext(s3_key)[1].lower() if s3_key else "" - local_temp_path = download_s3_file_to_temp(file_url, file_ext, input_dir) - - logger.info( - f"File downloaded: job_id={job_id}, local_path={local_temp_path}" - ) - - from app.services.document_parser.internal_parse_name import ( - prepare_internal_parse_input, - ) - from app.services.document_parser.parse_service import ( - checkerboard_inject_parse, - ) - - prepared_parse_input = prepare_internal_parse_input( - local_temp_path, - filename, - fallback_ext=file_ext, - prefer_fallback_ext=True, - ) - internal_parse_name = prepared_parse_input.internal_filename - local_temp_path = prepared_parse_input.file_path - logger.info( - f"File prepared for parsing: job_id={job_id}, " - f"internal_filename={internal_parse_name}, local_path={local_temp_path}" - ) - - page_count = PageEstimator.estimate(local_temp_path) - logger.info( - f"Workload estimation: job_id={job_id}, page_count={page_count}" - ) - - processing_started_at = datetime.now(timezone.utc) - - if not job_user_id: - raise NotFoundException( - resource="JobInfo", - resource_id="user_id", - internal_message=f"Missing user_id in job info for job_id={job_id}", - ) - - billing_service = WorkBillingService() - billing_status = "skipped" - billing_amount_micro_dollars = 0 - billing_credits = 0.0 - with get_sync_db_context() as db: - job_result = db.execute( - select(Job).where(Job.job_id == job_id).with_for_update() - ) - job = job_result.scalar_one_or_none() - - if job and getattr(job, "billing_status", "") == "charged": - logger.info(f"Job already charged: {job_id}") - billing_status = "charged" - billing_amount_micro_dollars = int(job.credits_charged or 0) - billing_credits = billing_amount_micro_dollars / 1_000_000 - else: - try: - billing_result = billing_service.charge_for_pages( - session=db, - user_id=job_user_id, - page_count=page_count, - filename=filename, - ) - except InsufficientCreditsException: - logger.warning( - f"Billing failed: job_id={job_id}, user_id={job_user_id}" - ) - billing_amount = billing_service.estimate_page_charge( - page_count=page_count - ) - if job: - job.page_count = page_count - job.credits_charged = billing_amount.amount_micro_dollars - job.billing_status = "billing_failed" - db.commit() - - raise InsufficientCreditsException( - user_message=( - "Insufficient credits to process this document " - f"({page_count} pages required, cost: " - f"{billing_amount.credits})." - ), - required_credits=billing_amount.credits, - internal_message=( - f"job_id={job_id}, user_id={job_user_id}, " - f"page_count={page_count}" - ), - ) - - billing_status = billing_result.billing_status - billing_amount_micro_dollars = billing_result.amount_micro_dollars - billing_credits = billing_result.credits - if job: - job.page_count = page_count - job.credits_charged = billing_amount_micro_dollars - job.billing_status = billing_status - - metadata_updates = { - "page_count": page_count, - "billing_status": billing_status, - "billing_amount_micro_dollars": billing_amount_micro_dollars, - "billing_credits": billing_credits, - "processing_started_at": processing_started_at.isoformat(), - } - metadata_service.update_metadata(job_id, metadata_updates) - job_metadata.update(metadata_updates) - - doc_type = JobMetadataHelper.get_parsing_param( - job_metadata, "doc_type", "auto" - ) - logger.info( - f"Start parse: job_id={job_id}, filename={filename}, " - f"internal_filename={internal_parse_name}, type={doc_type}" - ) - - with stage_timer( - "worker.parse.document", - job_id=job_id, - filename=filename, - doc_type=doc_type, - ): - add_dir, add_contents_df = checkerboard_inject_parse( - file_full_path=local_temp_path, - filename=filename, - output_dir=output_dir, - job_id=job_id, - internal_output_filename=internal_parse_name, - kb_dir=JobMetadataHelper.get_parsing_param( - job_metadata, "kb_dir", "Default_Root" - ), - doc_type=doc_type, - smart_title_parse=JobMetadataHelper.get_parsing_param( - job_metadata, "smart_title_parse", True - ), - summary_image=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_image", True - ), - summary_table=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_table", True - ), - summary_txt=JobMetadataHelper.get_parsing_param( - job_metadata, "summary_txt", True - ), - add_frag_desc=JobMetadataHelper.get_parsing_param( - job_metadata, "add_frag_desc", "" - ), - s3_key=s3_key, - ) - parsed_contents_df: pd.DataFrame | None = add_contents_df - - logger.info( - "File parsing completed: " - f"job_id={job_id}, add_dir={add_dir}, " - f"chunks={len(parsed_contents_df) if parsed_contents_df is not None else 0}" - ) - - if parsed_contents_df is None: - raise WorkerHandlingException( - user_message="We could not extract content from your file", - internal_message="File parsing failed, no content returned from parser", - ) - - if parsed_contents_df.empty: - logger.warning( - f"No content returned from file parsing: job_id={job_id}, filename={filename}" - ) - - lifecycle_service.update_progress( - job_id, progress=30, message="Parse completed, preparing chunks..." - ) - - chunks = dataframe_to_chunks(parsed_contents_df) - - lifecycle_service.update_progress( - job_id, progress=70, message="Chunks ready, generating zip..." - ) - logger.info(f"Chunks prepared: job_id={job_id}, count={len(chunks)}") - - source_file_name = JobMetadataHelper.get_field( - job_metadata, "source_file_name" - ) or JobMetadataHelper.get_field(job_metadata, "source_url") - if isinstance(source_file_name, str) and "/" in source_file_name: - source_file_name = os.path.basename(source_file_name) - - document_top_summary = "" - section_summaries: dict[str, str] = {} - if add_dir and source_file_name: - if add_contents_df is not None and "path" in add_contents_df.columns: - ensure_doc_nav_json( - str(add_dir), - chunks, - source_file_name=str(source_file_name), - ) - try: - kb_dir_for_enrich = os.path.dirname(str(add_dir)) - summary_use_llm = JobMetadataHelper.get_parsing_param( - job_metadata, "summary_use_llm", False - ) - enrich_doc_nav_summaries( - kb_dir_for_enrich, - source_file=str(source_file_name), - use_llm=summary_use_llm, - ) - section_summaries = build_section_summary_lookup(str(add_dir)) - except Exception as exc: - logger.warning(f"doc_nav enrichment failed (non-fatal): {exc}") - document_top_summary = load_nav_top_summary( - str(add_dir), str(source_file_name) - ) - if document_top_summary: - for chunk in chunks: - metadata = chunk.get("metadata") - if not isinstance(metadata, dict): - metadata = {} - chunk["metadata"] = metadata - metadata["document_top_summary"] = document_top_summary - - data_id = JobMetadataHelper.get_field(job_metadata, "data_id") - - lifecycle_service.update_progress( - job_id, progress=80, message="Generating ZIP package..." - ) - processing_completed_at = datetime.now(timezone.utc) - processing_timing_updates = { - "processing_completed_at": processing_completed_at.isoformat(), - "processing_duration_ms": max( - 0, - int( - ( - processing_completed_at - processing_started_at - ).total_seconds() - * 1000 - ), - ), - } - metadata_service.update_metadata(job_id, processing_timing_updates) - job_metadata.update(processing_timing_updates) - - zip_service = ZipResultService() - zip_file_path, checksum, statistics, zip_size = ( - zip_service.generate_zip_package( - job_id=job_id, - chunks=chunks, - add_dir=str(add_dir) if add_dir else "", - source_file_name=source_file_name, - data_id=data_id, - job_metadata=job_metadata, - parsed_df=parsed_contents_df, - temp_dir=task_workspace_dir, - ) - ) - del statistics - - checksum_value = ( - checksum.get("value", "") - if isinstance(checksum, dict) - else (checksum or "") - ) - - lifecycle_service.update_progress( - job_id, progress=90, message="Uploading results to S3..." - ) - - result_bundle = get_result_storage().upload( - job_id=job_id, - result_dir=str(add_dir) if add_dir else "", - zip_file_path=zip_file_path, - ) - result_s3_key = result_bundle.zip_key - - stored_count = 0 - - lifecycle_service.update_progress( - job_id, progress=100, message="Task complete!" - ) - - lifecycle_service.finalize_job_success( - job_id=job_id, - chunks=chunks, - result_s3_key=result_s3_key, - checksum=checksum_value, - zip_size=zip_size, - stored_count=stored_count, - delivery_mode="url", - section_summaries=section_summaries, - ) - - logger.info( - f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" - ) - - return { - "status": "success", - "job_id": job_id, - "add_dir": None, - "vectors_count": 0, - "contents_count": len(parsed_contents_df), - "stored_count": stored_count, - "delivery_mode": "url", - "result_s3_key": result_s3_key, - } - finally: - cleanup_task_workspace(task_workspace_dir) - - raise WorkerHandlingException( - user_message="We could not complete document processing", - internal_message=f"Parse workflow exited without a result for job_id={job_id}", - ) diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 8612bb376..c28db03dd 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -35,8 +35,8 @@ def _build_pending_file_job_metadata(source_file_name: str) -> dict[str, Any]: def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: import app.core.tasks.kb_tasks as kb_tasks + import app.services.document_ingestion.service as parse_job_service import app.services.document_parser.parse_service as parse_service - import app.services.workload.parse_job_service as parse_job_service from shared.core.database_sync import get_sync_engine from shared.services.redis.redis_sync_service import ( SyncJobInfoRedisService, From 38fdc8870b6ce4aa85d31477d02c55f86fb41848 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 01:30:30 +0800 Subject: [PATCH 13/40] refactor: consolidate job storage helpers --- .../app/services/jobs/result_projection.py | 9 +- .../services/document_ingestion/service.py | 11 +- .../services/document_ingestion/workspace.py | 41 +- .../services/document_parser/doc_parser.py | 2 +- .../services/document_parser/image_parser.py | 2 +- .../app/services/document_parser/md_parser.py | 2 +- .../document_parser/mineru_pdf_service.py | 6 +- .../services/document_parser/pptx_parser.py | 2 +- .../services/document_parser/table_parser.py | 2 +- .../services/document_parser/txt_parser.py | 2 +- .../services/storage/sync_storage_service.py | 140 ++--- .../contract/test_parse_task_contract.py | 51 +- .../shared/core/state_machine/service.py | 68 +-- .../core/state_machine/transition_payloads.py | 73 +++ .../shared/models/schemas/s3_file.py | 8 - .../services/storage/file_upload_service.py | 569 +----------------- .../services/storage/job_file_storage.py | 276 +++++++++ .../shared/services/storage/result_storage.py | 59 +- .../shared/services/webhook/dispatcher.py | 11 +- .../shared/utils/CommonHelper.py | 45 -- .../shared/utils/FileDownUpUtils.py | 253 -------- .../{CommonHelperSync.py => file_loading.py} | 16 +- .../shared/utils/zip_download.py | 79 +++ 23 files changed, 599 insertions(+), 1128 deletions(-) create mode 100644 packages/shared-python/shared/core/state_machine/transition_payloads.py delete mode 100644 packages/shared-python/shared/models/schemas/s3_file.py create mode 100644 packages/shared-python/shared/services/storage/job_file_storage.py delete mode 100644 packages/shared-python/shared/utils/CommonHelper.py delete mode 100644 packages/shared-python/shared/utils/FileDownUpUtils.py rename packages/shared-python/shared/utils/{CommonHelperSync.py => file_loading.py} (73%) create mode 100644 packages/shared-python/shared/utils/zip_download.py diff --git a/apps/api/app/services/jobs/result_projection.py b/apps/api/app/services/jobs/result_projection.py index 308f58116..faf7ab075 100644 --- a/apps/api/app/services/jobs/result_projection.py +++ b/apps/api/app/services/jobs/result_projection.py @@ -9,7 +9,7 @@ from shared.core.exceptions.domain_exceptions import JobOperationException from shared.models.schemas.job import JobResultResponse, StandardErrorObject from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.storage.file_upload_service import FileUploadService +from shared.services.storage.job_file_storage import JobFileStorage from shared.utils.error_details import normalize_error_details from shared.utils.utc_now import utc_now_naive @@ -118,9 +118,10 @@ async def _resolve_result_delivery( result_url_expires_at = job.created_at if job_result and job_result.result_s3_key: - upload_service = FileUploadService() - result_url_info = await upload_service.generate_download_url( - job_result.result_s3_key + result_storage = JobFileStorage() + result_url_info = result_storage.generate_download_url( + job_result.result_s3_key, + bucket=result_storage.results_bucket, ) result_url = result_url_info["download_url"] diff --git a/apps/worker/app/services/document_ingestion/service.py b/apps/worker/app/services/document_ingestion/service.py index 625464a41..5ce813f4d 100644 --- a/apps/worker/app/services/document_ingestion/service.py +++ b/apps/worker/app/services/document_ingestion/service.py @@ -20,10 +20,7 @@ download_s3_file_to_temp, ) from app.services.document_parser.stage_profiler import stage_timer -from app.services.storage.sync_storage_service import ( - generate_download_url, - verify_s3_file_exists, -) +from app.services.storage.sync_storage_service import verify_s3_file_exists from loguru import logger from sqlalchemy import select @@ -220,11 +217,7 @@ def _run_parse_job( filename = JobMetadataHelper.get_field(job_context.job_metadata, "source_file_name") file_ext = os.path.splitext(job_context.s3_key)[1].lower() if job_context.s3_key else "" - file_url = generate_download_url( - job_context.s3_key, - settings.S3_BUCKET_NAME, - )["download_url"] - local_temp_path = download_s3_file_to_temp(file_url, file_ext, input_dir) + local_temp_path = download_s3_file_to_temp(job_context.s3_key, file_ext, input_dir) logger.info(f"File downloaded: job_id={job_id}, local_path={local_temp_path}") from app.services.document_parser.internal_parse_name import ( diff --git a/apps/worker/app/services/document_ingestion/workspace.py b/apps/worker/app/services/document_ingestion/workspace.py index 80217df1f..2cb5e51d4 100644 --- a/apps/worker/app/services/document_ingestion/workspace.py +++ b/apps/worker/app/services/document_ingestion/workspace.py @@ -4,16 +4,15 @@ import shutil import tempfile -import requests from loguru import logger from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ( FileSystemException, - StorageServiceException, SystemSettingInvalidException, SystemSettingMissingException, ) +from shared.services.storage.job_file_storage import JobFileStorage def cleanup_temp_file(file_path: str | None) -> None: @@ -69,34 +68,12 @@ def create_task_workspace(job_id: str) -> str: ) from exc -def download_s3_file_to_temp(file_url: str, file_ext: str, temp_dir: str) -> str: +def download_s3_file_to_temp(s3_key: str, file_ext: str, temp_dir: str) -> str: """Download the source file from object storage into the task workspace.""" - local_temp_path: str | None = None - - try: - os.makedirs(temp_dir, exist_ok=True) - with tempfile.NamedTemporaryFile( - delete=False, - suffix=file_ext, - dir=temp_dir, - ) as temp_file: - local_temp_path = temp_file.name - with requests.get( - file_url, - timeout=120, - stream=True, - headers={"User-Agent": "Knowhere-Worker/1.0"}, - ) as response: - response.raise_for_status() - for chunk in response.iter_content(chunk_size=65536): - if chunk: - temp_file.write(chunk) - except requests.RequestException as exc: - cleanup_temp_file(local_temp_path) - raise StorageServiceException( - internal_message=f"Failed to download source file from object storage: {exc}", - operation="download_source_file", - original_exception=exc, - ) from exc - - return local_temp_path + storage = JobFileStorage() + return storage.download_to_temp( + s3_key, + suffix=file_ext, + temp_dir=temp_dir, + bucket=settings.S3_BUCKET_NAME, + ) diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index 8e50ab15e..5146e19df 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -39,7 +39,7 @@ from shared.core.exceptions.domain_exceptions import DocxParsingException from shared.core.exceptions.knowhere_exception import KnowhereException from shared.utils.chunk_refs import build_chunk_ref, has_chunk_ref -from shared.utils.CommonHelperSync import load_file_bytes +from shared.utils.file_loading import load_file_bytes from shared.utils.file_utils import path_handle from shared.utils.text_utils import tokenize2stw_remove diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/image_parser.py index 70c3c232a..80723ed4e 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/image_parser.py @@ -22,7 +22,7 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response from shared.utils.chunk_refs import build_chunk_ref -from shared.utils.CommonHelperSync import is_remote, load_file_bytes +from shared.utils.file_loading import is_remote, load_file_bytes from shared.utils.file_utils import path_handle from shared.utils.OpenAICompatibleClientSync import ( OpenAICompatibleClientSync, diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py index 33e7357af..80c1a99bd 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/md_parser.py @@ -264,7 +264,7 @@ def parse_md( relative_root=None, ): if md_lines is None and file_path is not None: - from shared.utils.CommonHelperSync import is_remote, load_file_bytes + from shared.utils.file_loading import is_remote, load_file_bytes if is_remote(file_path): file_bytes = load_file_bytes(file_path) diff --git a/apps/worker/app/services/document_parser/mineru_pdf_service.py b/apps/worker/app/services/document_parser/mineru_pdf_service.py index 24a2145cc..4e0c640ce 100644 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ b/apps/worker/app/services/document_parser/mineru_pdf_service.py @@ -20,8 +20,8 @@ UnavailableException, ) from shared.core.exceptions.knowhere_exception import KnowhereException -from shared.utils.CommonHelperSync import is_remote -from shared.utils.FileDownUpUtils import s3_download_extract_zip +from shared.utils.file_loading import is_remote +from shared.utils.zip_download import download_and_extract_zip MINERU_UPLOAD_TIMEOUT = ( settings.MINERU_UPLOAD_CONNECT_TIMEOUT, @@ -369,7 +369,7 @@ def poll_mineru_task( last_state = state if state == "done": - s3_download_extract_zip( + download_and_extract_zip( status["full_zip_url"], dest_dir=output_dir, keep_exts=(".md", ".jpg", ".jpeg", ".png", ".gif", ".json"), diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/pptx_parser.py index 408918655..84b8e02f7 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/pptx_parser.py @@ -29,7 +29,7 @@ FileSystemException, ) from shared.core.logging import LogEvent -from shared.utils.CommonHelperSync import load_file_bytes +from shared.utils.file_loading import load_file_bytes from shared.utils.file_utils import path_handle # ==================== LibreOffice conversion ==================== diff --git a/apps/worker/app/services/document_parser/table_parser.py b/apps/worker/app/services/document_parser/table_parser.py index 0916eab57..44b951baa 100755 --- a/apps/worker/app/services/document_parser/table_parser.py +++ b/apps/worker/app/services/document_parser/table_parser.py @@ -24,7 +24,7 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response from shared.utils.chunk_refs import build_chunk_ref -from shared.utils.CommonHelperSync import load_file_bytes +from shared.utils.file_loading import load_file_bytes from shared.utils.file_utils import path_handle from shared.utils.OpenAICompatibleClientSync import get_openai_client from shared.utils.text_utils import remove_duplicates_orderkept, tokenize2stw_remove diff --git a/apps/worker/app/services/document_parser/txt_parser.py b/apps/worker/app/services/document_parser/txt_parser.py index 3851817b5..e1d41956b 100755 --- a/apps/worker/app/services/document_parser/txt_parser.py +++ b/apps/worker/app/services/document_parser/txt_parser.py @@ -12,7 +12,7 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response from shared.utils.chunk_refs import CHUNK_REF_PATTERN -from shared.utils.CommonHelperSync import load_file_bytes +from shared.utils.file_loading import load_file_bytes from shared.utils.OpenAICompatibleClientSync import get_openai_client diff --git a/apps/worker/app/services/storage/sync_storage_service.py b/apps/worker/app/services/storage/sync_storage_service.py index 268c68b90..dc22139ba 100644 --- a/apps/worker/app/services/storage/sync_storage_service.py +++ b/apps/worker/app/services/storage/sync_storage_service.py @@ -1,102 +1,67 @@ -""" -Sync storage operations for worker tasks. -Provides S3 file operations and HTTP file downloads using sync adapters -that yield cooperatively under gevent. -""" +"""Sync adapter for shared Job file storage used by worker tasks.""" import os -import tempfile -from typing import Any, Dict, Optional +from typing import Any from loguru import logger from shared.core.config import settings -from shared.core.config.storage import get_cached_storage_adapter -from shared.core.exceptions.domain_exceptions import StorageServiceException -from shared.utils.pinned_outbound_http import download_pinned_outbound_file -from shared.utils.url_security import validate_http_url_and_resolve_ip +from shared.services.storage.job_file_storage import JobFileStorage -def get_storage_adapter(): - """Get the storage adapter for direct sync S3 operations.""" - return get_cached_storage_adapter() +def get_storage_adapter() -> JobFileStorage: + """Get the shared job file storage module for sync worker operations.""" + return JobFileStorage() -def verify_s3_file_exists(s3_key: str, bucket: Optional[str] = None) -> Dict[str, Any]: - """Verify S3 file exists using sync adapter calls.""" - adapter = get_storage_adapter() - bucket_name = bucket or settings.S3_BUCKET_NAME - try: - if not adapter.exists(s3_key, bucket_name): - return {"exists": False} - size = adapter.get_object_size(s3_key, bucket_name) - return {"exists": True, "size": size} - except Exception as e: - if "404" in str(e) or "not found" in str(e).lower(): - return {"exists": False} - raise StorageServiceException( - internal_message=f"S3 file verification failed: {e}", - operation="verify_s3_file_exists", - original_exception=e, - ) +def verify_s3_file_exists(s3_key: str, bucket: str | None = None) -> dict[str, Any]: + """Verify an uploaded source file exists.""" + storage = get_storage_adapter() + return storage.verify_exists( + s3_key, + bucket=bucket or settings.S3_BUCKET_NAME, + ) def generate_download_url( - s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600 -) -> Dict[str, Any]: - """Generate presigned download URL using sync adapter.""" - adapter = get_storage_adapter() - bucket_name = bucket or settings.S3_BUCKET_NAME - download_url = adapter.generate_presigned_url( - s3_key, expiration=expires_in, bucket=bucket_name, method="GET" + s3_key: str, bucket: str | None = None, expires_in: int = 3600 +) -> dict[str, Any]: + """Generate a presigned download URL for a stored object.""" + storage = get_storage_adapter() + return storage.generate_download_url( + s3_key, + bucket=bucket or settings.S3_BUCKET_NAME, + expires_in=expires_in, ) - return {"download_url": download_url, "expires_in": expires_in} -def upload_to_s3(local_file_path: str, s3_key: str, bucket: str): - """Upload file to S3 using sync adapter.""" - adapter = get_storage_adapter() - adapter.upload_file(local_file_path, s3_key, bucket) +def upload_to_s3(local_file_path: str, s3_key: str, bucket: str) -> None: + """Upload a local file using the shared job file storage rules.""" + storage = get_storage_adapter() + storage.upload_local_file(local_file_path, s3_key, bucket=bucket) def download_s3_object_to_temp( s3_key: str, suffix: str, temp_dir: str, - bucket: Optional[str] = None, + bucket: str | None = None, ) -> str: """Download an object-storage file into a task-local temp file.""" - adapter = get_storage_adapter() - bucket_name = bucket or settings.S3_BUCKET_NAME - local_temp_path: str | None = None - - try: - os.makedirs(temp_dir, exist_ok=True) - with tempfile.NamedTemporaryFile( - delete=False, - suffix=suffix, - dir=temp_dir, - ) as temp_file: - local_temp_path = temp_file.name - adapter.download_file(s3_key, local_temp_path, bucket_name) - return local_temp_path - except Exception as e: - if local_temp_path and os.path.exists(local_temp_path): - os.remove(local_temp_path) - raise StorageServiceException( - internal_message=( - f"Failed to download object-storage file to temp path: " - f"s3_key={s3_key}, temp_dir={temp_dir}, error={e}" - ), - operation="download_s3_object_to_temp", - original_exception=e, - ) from e + storage = get_storage_adapter() + return storage.download_to_temp( + s3_key, + suffix=suffix, + temp_dir=temp_dir, + bucket=bucket or settings.S3_BUCKET_NAME, + ) def upload_zip_result(job_id: str, zip_file_path: str) -> str: """Upload ZIP result file to S3 and cleanup temp file.""" - results_bucket = getattr(settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME) - s3_key = f"results/{job_id}.zip" + storage = get_storage_adapter() + results_bucket = storage.results_bucket + s3_key = storage.build_result_zip_key(job_id=job_id) upload_to_s3(zip_file_path, s3_key, results_bucket) logger.info(f"Result ZIP uploaded: job_id={job_id}, key={s3_key}") try: @@ -109,33 +74,6 @@ def upload_zip_result(job_id: str, zip_file_path: str) -> str: def download_file_from_url(file_url: str) -> str: """Download a URL file through SSRF validation and IP pinning.""" - temp_file_path = "" - try: - validation = validate_http_url_and_resolve_ip(file_url) - if not validation.is_valid or not validation.validated_ip: - raise StorageServiceException( - internal_message=f"Invalid URL: {validation.error_message}", - operation="download_from_url", - ) - - temp_dir = getattr(settings, "TMP_PATH", "/tmp") - os.makedirs(temp_dir, exist_ok=True) - download_result = download_pinned_outbound_file( - url=validation.url, - pinned_ip=validation.validated_ip, - timeout_seconds=300, - user_agent="Knowhere-FileDownloader/1.0", - temp_dir=temp_dir, - ) - temp_file_path = download_result.temp_file_path - return temp_file_path - except StorageServiceException: - raise - except Exception as e: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise StorageServiceException( - internal_message=f"Failed to download file: {e}", - operation="download_from_url", - original_exception=e, - ) + storage = get_storage_adapter() + temp_dir = getattr(settings, "TMP_PATH", "/tmp") + return storage.download_file_from_url(file_url, temp_dir=temp_dir) diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index c28db03dd..d93c3f863 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -175,23 +175,16 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: - return {"download_url": f"https://example.test/{storage_key}"} - monkeypatch.setattr( parse_job_service, "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr( - parse_job_service, - "generate_download_url", - fake_generate_download_url, - ) def fake_download_s3_file_to_temp( - file_url: str, file_ext: str, temp_dir: str + storage_key: str, file_ext: str, temp_dir: str ) -> str: + assert storage_key == s3_key assert file_ext == ".pdf" downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) @@ -776,12 +769,10 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: - return {"download_url": f"https://example.test/{storage_key}"} - def fake_download_s3_file_to_temp( - file_url: str, file_ext: str, temp_dir: str + storage_key: str, file_ext: str, temp_dir: str ) -> str: + assert storage_key == s3_key downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) return str(downloaded_path) @@ -870,11 +861,6 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr( - parse_job_service, - "generate_download_url", - fake_generate_download_url, - ) monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) @@ -1016,12 +1002,10 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: - return {"download_url": f"https://example.test/{storage_key}"} - def fake_download_s3_file_to_temp( - file_url: str, file_ext: str, temp_dir: str + storage_key: str, file_ext: str, temp_dir: str ) -> str: + assert storage_key in s3_keys.values() assert file_ext == ".pdf" downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) @@ -1073,11 +1057,6 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr( - parse_job_service, - "generate_download_url", - fake_generate_download_url, - ) monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr(parse_job_service.PageEstimator, "estimate", fake_estimate_page_count) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) @@ -1256,13 +1235,6 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr( - parse_job_service, - "generate_download_url", - lambda *_args, **_kwargs: (_ for _ in ()).throw( - AssertionError("terminal parse task should not request a download URL") - ), - ) monkeypatch.setattr( parse_service, "checkerboard_inject_parse", @@ -1354,23 +1326,16 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]: - return {"download_url": f"https://example.test/{storage_key}"} - monkeypatch.setattr( parse_job_service, "verify_s3_file_exists", fake_verify_s3_file_exists, ) - monkeypatch.setattr( - parse_job_service, - "generate_download_url", - fake_generate_download_url, - ) def fake_download_s3_file_to_temp( - file_url: str, file_ext: str, temp_dir: str + storage_key: str, file_ext: str, temp_dir: str ) -> str: + assert storage_key == s3_key downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) return str(downloaded_path) diff --git a/packages/shared-python/shared/core/state_machine/service.py b/packages/shared-python/shared/core/state_machine/service.py index ee5f421a4..366649acb 100644 --- a/packages/shared-python/shared/core/state_machine/service.py +++ b/packages/shared-python/shared/core/state_machine/service.py @@ -7,7 +7,6 @@ import asyncio import time -from datetime import datetime, timezone from typing import Any, Dict, Optional from loguru import logger @@ -19,18 +18,19 @@ JobStatus, is_valid_transition, ) +from shared.core.state_machine.transition_payloads import ( + build_failure_transition_metadata, + build_progress_cache_payload, + build_retry_transition, + serialize_transition_metadata, + utc_now_naive, +) from shared.models.database.job import Job from shared.models.database.job_state_audit_log import JobStateAuditLog from shared.services.redis import RedisServiceFactory -from shared.utils.error_details import normalize_error_details -from shared.utils.json_utils import make_json_safe from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder -def _utc_now_naive() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) - - class AsyncStateMachineService: """Async state machine service — used by the API (FastAPI + asyncpg).""" @@ -136,7 +136,14 @@ async def mark_failed( ) -> bool: """Mark a job as failed with error information.""" try: - normalized_details = normalize_error_details(error_details) + normalized_details, transition_metadata = ( + build_failure_transition_metadata( + error_message=error_message, + error_code=error_code, + error_details=error_details, + metadata=metadata, + ) + ) await self._update_job_error( db, job_id, @@ -145,12 +152,6 @@ async def mark_failed( normalized_details, ) - transition_metadata = (metadata or {}).copy() - transition_metadata["error_message"] = error_message - transition_metadata["error_code"] = error_code - if normalized_details: - transition_metadata["error_details"] = normalized_details - return await self.transition( db, job_id, @@ -208,17 +209,11 @@ async def handle_retry( logger.error(f"Job {job_id} has no status") return False - retry_target = ( - JobStatus.PENDING.value - if current_state == JobStatus.FAILED.value - else current_state + retry_target, retry_metadata = build_retry_transition( + current_state=current_state, + retry_metadata=retry_metadata, ) - retry_metadata = retry_metadata or {} - retry_metadata["retry_reason"] = "task_retry" - retry_metadata["retry_timestamp"] = str(int(time.time())) - retry_metadata["retry_count"] = retry_metadata.get("retry_count", 0) + 1 - # Always use full transition() for CAS protection — even same-state return await self.transition( db, @@ -291,7 +286,7 @@ async def _cas_update_state( .values( status=to_state, version=old_version + 1, - updated_at=_utc_now_naive(), + updated_at=utc_now_naive(), ) ) return result.rowcount > 0 @@ -307,13 +302,7 @@ async def _record_audit_log( operator_type: str, metadata: Optional[Dict[str, Any]], ) -> None: - serialized = None - if metadata: - try: - serialized = make_json_safe(metadata) - except Exception as e: - logger.warning(f"Metadata serialization failed: {e}") - serialized = {"error": "metadata_serialization_failed"} + serialized = serialize_transition_metadata(metadata) db.add( JobStateAuditLog( @@ -392,13 +381,20 @@ async def _update_redis_cache( ) progress_key = redis_key_builder.task_progress(job_id) - progress_data: Dict[str, Any] = { - "status": status, - "timestamp": str(int(time.time())), - } + progress_data: Dict[str, Any] = build_progress_cache_payload( + status=status, + metadata=None, + timestamp=int(time.time()), + ) if metadata: try: - progress_data.update(make_json_safe(metadata)) + progress_data.update( + build_progress_cache_payload( + status=status, + metadata=metadata, + timestamp=int(time.time()), + ) + ) except Exception as e: logger.warning(f"Metadata serialization skipped: {e}") diff --git a/packages/shared-python/shared/core/state_machine/transition_payloads.py b/packages/shared-python/shared/core/state_machine/transition_payloads.py new file mode 100644 index 000000000..c2a2cf58b --- /dev/null +++ b/packages/shared-python/shared/core/state_machine/transition_payloads.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +from shared.core.state_machine.states import JobStatus +from shared.utils.error_details import normalize_error_details +from shared.utils.json_utils import make_json_safe + + +def utc_now_naive() -> datetime: + return datetime.now(timezone.utc).replace(tzinfo=None) + + +def serialize_transition_metadata( + metadata: dict[str, Any] | None, +) -> dict[str, Any] | None: + if not metadata: + return None + + try: + return make_json_safe(metadata) + except Exception: + return {"error": "metadata_serialization_failed"} + + +def build_failure_transition_metadata( + *, + error_message: str, + error_code: str, + error_details: dict[str, Any] | None, + metadata: dict[str, Any] | None, +) -> tuple[dict[str, Any] | None, dict[str, Any]]: + normalized_details = normalize_error_details(error_details) + transition_metadata = (metadata or {}).copy() + transition_metadata["error_message"] = error_message + transition_metadata["error_code"] = error_code + if normalized_details: + transition_metadata["error_details"] = normalized_details + return normalized_details, transition_metadata + + +def build_retry_transition( + *, + current_state: str, + retry_metadata: dict[str, Any] | None, +) -> tuple[str, dict[str, Any]]: + retry_target = ( + JobStatus.PENDING.value + if current_state == JobStatus.FAILED.value + else current_state + ) + + resolved_metadata = retry_metadata or {} + resolved_metadata["retry_reason"] = "task_retry" + resolved_metadata["retry_timestamp"] = str(int(datetime.now(timezone.utc).timestamp())) + resolved_metadata["retry_count"] = resolved_metadata.get("retry_count", 0) + 1 + return retry_target, resolved_metadata + + +def build_progress_cache_payload( + *, + status: str, + metadata: dict[str, Any] | None, + timestamp: int, +) -> dict[str, Any]: + progress_data: dict[str, Any] = { + "status": status, + "timestamp": str(timestamp), + } + if metadata: + progress_data.update(make_json_safe(metadata)) + return progress_data diff --git a/packages/shared-python/shared/models/schemas/s3_file.py b/packages/shared-python/shared/models/schemas/s3_file.py deleted file mode 100644 index 4b3a0cbbc..000000000 --- a/packages/shared-python/shared/models/schemas/s3_file.py +++ /dev/null @@ -1,8 +0,0 @@ -from pydantic import BaseModel - - -class FliesDownload(BaseModel): - message: str - file_key: str - download_url: str - expires_in_seconds: int diff --git a/packages/shared-python/shared/services/storage/file_upload_service.py b/packages/shared-python/shared/services/storage/file_upload_service.py index c60f3c557..8b24793a4 100644 --- a/packages/shared-python/shared/services/storage/file_upload_service.py +++ b/packages/shared-python/shared/services/storage/file_upload_service.py @@ -1,147 +1,30 @@ -"""Storage upload service.""" +"""Async adapter for shared Job file storage.""" import asyncio -import json -import os -from typing import Any, Dict, Optional +from typing import Any, Optional from loguru import logger -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - KnowhereException, - StorageServiceException, -) -from shared.utils.pinned_outbound_http import ( - download_pinned_outbound_file_async, -) -from shared.utils.url_security import validate_http_url_and_resolve_ip_async +from shared.core.exceptions.domain_exceptions import StorageServiceException +from shared.services.storage.job_file_storage import JobFileStorage class FileUploadService: - """File upload service supporting S3, OSS, and MinIO.""" + """Async adapter over the shared Job file storage module.""" - def __init__(self): - self.adapter = settings.get_storage_adapter() - self.uploads_bucket = settings.S3_BUCKET_NAME - self.results_bucket = getattr( - settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME - ) - - async def handle_direct_upload(self, file_path: str, job_id: str) -> str: - """ - Handle a direct file upload. - - Args: - file_path: Local file path. - job_id: Job ID. - - Returns: - str: Storage key. - """ - try: - # Build the storage key. - file_extension = os.path.splitext(file_path)[1] - s3_key = f"uploads/{job_id}{file_extension}" - - # Upload the file. - await self._upload_to_s3(file_path, s3_key, self.uploads_bucket) - - logger.info(f"Direct file upload succeeded: {file_path} -> {s3_key}") - return s3_key - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Direct file upload failed: {e}") - raise StorageServiceException( - internal_message=f"Direct file upload failed: {str(e)}", - operation="direct_upload", - original_exception=e, - ) - - async def handle_url_upload(self, file_url: str, job_id: str) -> str: - """ - Handle a URL-based upload flow. - - Args: - file_url: File URL. - job_id: Job ID. - - Returns: - str: Storage key. - """ - try: - # Download the file into a temporary location first. - temp_file_path = await self._download_file_from_url(file_url) - - try: - # Build the storage key. - file_extension = os.path.splitext(file_url.split("?")[0])[1] - s3_key = f"uploads/{job_id}{file_extension}" - - # Upload the downloaded file. - await self._upload_to_s3(temp_file_path, s3_key, self.uploads_bucket) - - logger.info( - f"URL file download and upload succeeded: {file_url} -> {s3_key}" - ) - return s3_key - - finally: - # Clean up the temporary file. - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - - except KnowhereException: - raise - except Exception as e: - logger.error(f"URL file handling failed: {e}") - raise StorageServiceException( - internal_message=f"URL file handling failed: {str(e)}", - operation="url_upload", - original_exception=e, - ) + def __init__(self, *, storage: JobFileStorage | None = None) -> None: + self._storage = storage or JobFileStorage() async def generate_upload_url( self, job_id: str, file_extension: str = "" - ) -> Dict[str, Any]: - """ - Generate a presigned upload URL. - - Args: - job_id: Job ID. - file_extension: File extension. - - Returns: - Dict: Upload URL payload including the storage key. - """ + ) -> dict[str, Any]: try: - s3_key = f"uploads/{job_id}{file_extension}" - - # Infer a Content-Type from the file extension. - content_type = self.get_content_type(file_extension) - - # Use the job waiting expiry as the upload URL TTL. - upload_url = self.adapter.generate_presigned_url( - s3_key, - expiration=settings.JOB_WAITING_EXPIRE_SECONDS, - bucket=self.uploads_bucket, - method="PUT", - headers={"Content-Type": content_type}, + return await asyncio.to_thread( + self._storage.generate_upload_url, + job_id=job_id, + file_extension=file_extension, ) - logger.info(f"Generated presigned upload URL: {upload_url}") - - return { - "upload_url": upload_url, - "s3_key": s3_key, - "expires_in": settings.JOB_WAITING_EXPIRE_SECONDS, - "upload_headers": {"Content-Type": content_type}, - } - - except KnowhereException: - raise except Exception as e: logger.error(f"Failed to generate upload URL: {e}") raise StorageServiceException( @@ -152,29 +35,15 @@ async def generate_upload_url( async def generate_download_url( self, s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600 - ) -> Dict[str, Any]: - """ - Generate a presigned download URL. - - Args: - s3_key: Storage key. - bucket: Optional bucket name. - - Returns: - str: Download URL. - """ + ) -> dict[str, Any]: try: - bucket_name = bucket or self.results_bucket - - # Generate a one-hour presigned URL by default. - download_url = self.adapter.generate_presigned_url( - s3_key, expiration=expires_in, bucket=bucket_name, method="GET" + return await asyncio.to_thread( + self._storage.generate_download_url, + s3_key, + bucket=bucket or self._storage.results_bucket, + expires_in=expires_in, ) - return {"download_url": download_url, "expires_in": expires_in} - - except KnowhereException: - raise except Exception as e: logger.error(f"Failed to generate download URL: {e}") raise StorageServiceException( @@ -183,411 +52,19 @@ async def generate_download_url( original_exception=e, ) - async def get_file_info( + async def verify_s3_file_exists( self, s3_key: str, bucket: Optional[str] = None - ) -> Optional[Dict[str, Any]]: - """ - Get file information. - - Args: - s3_key: Storage key. - bucket: Optional bucket name. - - Returns: - Dict: File metadata. - """ - try: - bucket_name = bucket or self.results_bucket - - # Check existence and load the object size. - if not self.adapter.exists(s3_key, bucket_name): - return None - - size = self.adapter.get_object_size(s3_key, bucket_name) - return { - "size": size, - "content_type": None, # The adapter interface does not expose content_type yet. - "last_modified": None, - "etag": None, - } - - except Exception as e: - # Treat not-found responses as a missing object. - if "404" in str(e) or "not found" in str(e).lower(): - return None - logger.error(f"Failed to get file info: {e}") - raise StorageServiceException( - internal_message=f"Failed to get file info: {str(e)}", - operation="get_file_info", - original_exception=e, - ) - - async def upload_result_file( - self, local_file_path: str, job_id: str, file_extension: str = "" - ) -> str: - """ - Upload a result file. - - Args: - local_file_path: Local file path. - job_id: Job ID. - file_extension: File extension. - - Returns: - str: Storage key. - """ - try: - s3_key = f"results/{job_id}{file_extension}" - await self._upload_to_s3(local_file_path, s3_key, self.results_bucket) - - logger.info(f"Result file upload succeeded: {local_file_path} -> {s3_key}") - return s3_key - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Result file upload failed: {e}") - raise StorageServiceException( - internal_message=f"Result file upload failed: {str(e)}", - operation="upload_result_file", - original_exception=e, - ) - - async def upload_json_result( - self, - job_id: str, - result_data: Dict[str, Any], - *, - content_type: str = "application/json", - ) -> str: - """Upload a JSON result file; deprecated but kept for compatibility.""" + ) -> dict[str, Any]: try: - s3_key = f"results/{job_id}.json" - from io import BytesIO - - body = json.dumps(result_data, ensure_ascii=False).encode("utf-8") - self.adapter.upload_fileobj( - BytesIO(body), + return await asyncio.to_thread( + self._storage.verify_exists, s3_key, - bucket=self.results_bucket, - content_type=content_type, + bucket=bucket or self._storage.uploads_bucket, ) - logger.info(f"Result JSON upload succeeded: job_id={job_id}, key={s3_key}") - return s3_key - except KnowhereException: - raise except Exception as e: - logger.error(f"Failed to upload result JSON: {e}") - raise StorageServiceException( - internal_message=f"Failed to upload result JSON: {str(e)}", - operation="upload_json_result", - original_exception=e, - ) - - async def upload_zip_result( - self, - job_id: str, - zip_file_path: str, - ) -> str: - """Upload a ZIP result file.""" - try: - s3_key = f"results/{job_id}.zip" - await self._upload_to_s3(zip_file_path, s3_key, self.results_bucket) - logger.info(f"Result ZIP upload succeeded: job_id={job_id}, key={s3_key}") - - # Clean up the temporary ZIP after upload. - try: - if os.path.exists(zip_file_path): - os.remove(zip_file_path) - except Exception as e: - logger.warning(f"Failed to clean up temporary ZIP file: {e}") - - return s3_key - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to upload result ZIP: {e}") - raise StorageServiceException( - internal_message=f"Failed to upload result ZIP: {str(e)}", - operation="upload_zip_result", - original_exception=e, - ) - - def _ensure_bucket_exists(self, bucket_name: str) -> bool: - """ - Ensure the bucket is accessible. - - Args: - bucket_name: Bucket name. - - Returns: - bool: Whether the bucket check succeeded. - """ - try: - # In adapter mode, probe accessibility by listing objects. - adapter = settings.get_storage_adapter() - list(adapter.list_objects(prefix="", bucket=bucket_name)) - logger.debug(f"Bucket {bucket_name} is accessible") - return True - except Exception as e: - # The bucket is missing or inaccessible. - # For OSS, buckets should already exist; only accessibility is checked here. - logger.warning( - f"Bucket {bucket_name} may not exist or may be inaccessible: {e}" - ) - # In production, buckets should already be provisioned, so continue. - # Return False here instead if strict enforcement is ever needed. - return True - - async def _ensure_bucket_exists_async(self, bucket_name: str) -> bool: - """ - Asynchronously ensure the bucket is accessible. - - Args: - bucket_name: Bucket name. - - Returns: - bool: Whether the bucket check succeeded. - """ - - def _check_and_create(): - try: - # In adapter mode, probe accessibility by listing objects. - adapter = settings.get_storage_adapter() - list(adapter.list_objects(prefix="", bucket=bucket_name)) - logger.debug(f"Bucket {bucket_name} is accessible") - return True - except Exception as e: - # The bucket is missing or inaccessible. - logger.warning( - f"Bucket {bucket_name} may not exist or may be inaccessible: {e}" - ) - # In production, buckets should already be provisioned, so continue. - return True - - # Run the synchronous probe in a thread pool. - loop = asyncio.get_event_loop() - return await loop.run_in_executor(None, _check_and_create) - - async def _upload_to_s3(self, local_file_path: str, s3_key: str, bucket: str): - """Upload a file to storage.""" - # Ensure the bucket is accessible before uploading. - if not await self._ensure_bucket_exists_async(bucket): - raise StorageServiceException( - internal_message=f"Could not ensure bucket {bucket} exists", - operation="ensure_bucket", - ) - - def _upload(): - self.adapter.upload_file(local_file_path, s3_key, bucket) - - # Run the blocking upload in a thread pool. - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, _upload) - - async def download_from_s3(self, s3_key: str, bucket: Optional[str] = None) -> str: - """Download a file from storage into a local temporary directory.""" - import uuid - - if bucket is None: - bucket = settings.S3_BUCKET_NAME - - # Create the temporary destination directory. - temp_dir = getattr(settings, "TMP_PATH", "/tmp") - os.makedirs(temp_dir, exist_ok=True) - - # Generate a temporary filename while preserving the original extension. - file_extension = os.path.splitext(s3_key)[1] - temp_filename = f"temp_{uuid.uuid4().hex}{file_extension}" - temp_file_path = os.path.join(temp_dir, temp_filename) - - try: - # Use the adapter to download the file. - def _download(): - self.adapter.download_file(s3_key, temp_file_path, bucket) - - # Run the blocking download in the event loop executor. - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, _download) - - return temp_file_path - - except KnowhereException: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise - except Exception as e: - # Clean up the temporary file on failure. - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise StorageServiceException( - internal_message=f"Failed to download file from S3: {str(e)}", - operation="download_from_s3", - original_exception=e, - ) - - async def _download_file_from_url(self, file_url: str) -> str: - """Download a file from a URL into a temporary directory.""" - temp_file_path = "" - try: - validation = await validate_http_url_and_resolve_ip_async(file_url) - if not validation.is_valid or not validation.validated_ip: - raise StorageServiceException( - internal_message=f"Invalid URL: {validation.error_message}", - operation="download_from_url", - ) - - temp_dir = getattr(settings, "TMP_PATH", "/tmp") - os.makedirs(temp_dir, exist_ok=True) - download_result = await download_pinned_outbound_file_async( - url=validation.url, - pinned_ip=validation.validated_ip, - timeout_seconds=300, - user_agent="Knowhere-FileDownloader/1.0", - temp_dir=temp_dir, - ) - temp_file_path = download_result.temp_file_path - return temp_file_path - - except KnowhereException: - raise - except Exception as e: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise StorageServiceException( - internal_message=f"Failed to download file: {str(e)}", - operation="download_from_url", - original_exception=e, - ) - - async def verify_s3_file_exists( - self, s3_key: str, bucket: Optional[str] = None - ) -> Dict[str, Any]: - """ - Verify whether a file exists in storage. - - Args: - s3_key: Storage key. - bucket: Optional bucket name. - - Returns: - Dict: File info payload, or `{"exists": False}` when missing. - """ - try: - bucket_name = bucket or self.uploads_bucket - - # Use the adapter to check object existence. - exists = self.adapter.exists(s3_key, bucket_name) - if not exists: - return {"exists": False} - - size = self.adapter.get_object_size(s3_key, bucket_name) - return { - "exists": True, - "size": size, - "content_type": None, - "last_modified": None, - "etag": None, - } - - except Exception as e: - # Treat not-found responses as a missing object. - if "404" in str(e) or "not found" in str(e).lower(): - return {"exists": False} logger.error(f"Failed to verify file existence: {e}") raise StorageServiceException( internal_message=f"Failed to verify file existence: {str(e)}", operation="verify_s3_file_exists", original_exception=e, ) - - def get_content_type(self, file_extension: str) -> str: - """ - Return a Content-Type for a file extension. - - Args: - file_extension: File extension, such as `.pdf` or `.docx`. - - Returns: - str: Content-Type - """ - content_types = { - ".pdf": "application/pdf", - ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", - ".doc": "application/msword", - ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - ".xls": "application/vnd.ms-excel", - ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation", - ".ppt": "application/vnd.ms-powerpoint", - ".csv": "text/csv", - ".txt": "text/plain", - ".md": "text/markdown", - ".json": "application/json", - ".xml": "application/xml", - ".html": "text/html", - ".htm": "text/html", - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".gif": "image/gif", - ".bmp": "image/bmp", - ".tiff": "image/tiff", - ".svg": "image/svg+xml", - ".zip": "application/zip", - ".rar": "application/x-rar-compressed", - ".7z": "application/x-7z-compressed", - ".tar": "application/x-tar", - ".gz": "application/gzip", - } - return content_types.get(file_extension.lower(), "application/octet-stream") - - async def get_file_url( - self, s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600 - ) -> str: - """ - Get a file URL from a storage key. - - Args: - s3_key: Storage key. - bucket: Optional bucket name. - expires_in: URL TTL in seconds, defaulting to one hour. - - Returns: - str: File URL. - """ - try: - bucket_name = bucket or self.uploads_bucket - - # Generate a presigned GET URL. - file_url = self.adapter.generate_presigned_url( - s3_key, expiration=expires_in, bucket=bucket_name, method="GET" - ) - - logger.info(f"Generated file URL successfully: {s3_key} -> {file_url}") - return file_url - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to get file URL: {e}") - raise StorageServiceException( - internal_message=f"Failed to get file URL: {str(e)}", - operation="get_file_url", - original_exception=e, - ) - - def generate_s3_key( - self, job_id: str, file_extension: str = "", prefix: str = "uploads" - ) -> str: - """ - Generate a storage key. - - Args: - job_id: Job ID. - file_extension: File extension. - prefix: Key prefix such as `uploads` or `results`. - - Returns: - str: Storage key. - """ - return f"{prefix}/{job_id}{file_extension}" diff --git a/packages/shared-python/shared/services/storage/job_file_storage.py b/packages/shared-python/shared/services/storage/job_file_storage.py new file mode 100644 index 000000000..02777bc0a --- /dev/null +++ b/packages/shared-python/shared/services/storage/job_file_storage.py @@ -0,0 +1,276 @@ +from __future__ import annotations + +import os +import tempfile +from typing import Any, BinaryIO + +from shared.core.config import settings +from shared.core.config.storage import get_cached_storage_adapter +from shared.core.exceptions.domain_exceptions import StorageServiceException +from shared.services.storage.storage_adapter import StorageAdapter +from shared.utils.pinned_outbound_http import download_pinned_outbound_file +from shared.utils.url_security import validate_http_url_and_resolve_ip + + +class JobFileStorage: + """Own storage rules for Job source files and Job Result bundles.""" + + def __init__( + self, + *, + storage_adapter: StorageAdapter | None = None, + uploads_bucket: str | None = None, + results_bucket: str | None = None, + ) -> None: + self._storage_adapter = storage_adapter + self.uploads_bucket = uploads_bucket or settings.S3_BUCKET_NAME + self.results_bucket = results_bucket or getattr( + settings, + "S3_RESULTS_BUCKET", + settings.S3_BUCKET_NAME, + ) + + @property + def storage_adapter(self) -> StorageAdapter: + if self._storage_adapter is None: + self._storage_adapter = get_cached_storage_adapter() + return self._storage_adapter + + def build_upload_key(self, *, job_id: str, file_extension: str = "") -> str: + return f"uploads/{job_id}{file_extension}" + + def build_result_key(self, *, job_id: str, file_extension: str = "") -> str: + return f"results/{job_id}{file_extension}" + + def build_result_zip_key(self, *, job_id: str) -> str: + return self.build_result_key(job_id=job_id, file_extension=".zip") + + def build_result_raw_prefix(self, *, job_id: str) -> str: + return f"results/{job_id}/" + + def generate_upload_url( + self, + *, + job_id: str, + file_extension: str = "", + ) -> dict[str, Any]: + storage_key = self.build_upload_key( + job_id=job_id, + file_extension=file_extension, + ) + content_type = self.get_content_type(file_extension) + upload_url = self.storage_adapter.generate_presigned_url( + storage_key, + expiration=settings.JOB_WAITING_EXPIRE_SECONDS, + bucket=self.uploads_bucket, + method="PUT", + headers={"Content-Type": content_type}, + ) + return { + "upload_url": upload_url, + "s3_key": storage_key, + "expires_in": settings.JOB_WAITING_EXPIRE_SECONDS, + "upload_headers": {"Content-Type": content_type}, + } + + def generate_download_url( + self, + storage_key: str, + *, + bucket: str, + expires_in: int = 3600, + ) -> dict[str, Any]: + download_url = self.storage_adapter.generate_presigned_url( + storage_key, + expiration=expires_in, + bucket=bucket, + method="GET", + ) + return {"download_url": download_url, "expires_in": expires_in} + + def verify_exists( + self, + storage_key: str, + *, + bucket: str, + ) -> dict[str, Any]: + try: + if not self.storage_adapter.exists(storage_key, bucket): + return {"exists": False} + + size = self.storage_adapter.get_object_size(storage_key, bucket) + return { + "exists": True, + "size": size, + "content_type": None, + "last_modified": None, + "etag": None, + } + except Exception as exc: + if "404" in str(exc) or "not found" in str(exc).lower(): + return {"exists": False} + raise StorageServiceException( + internal_message=f"Storage file verification failed: {exc}", + operation="verify_exists", + original_exception=exc, + ) from exc + + def upload_local_file( + self, + local_file_path: str, + storage_key: str, + *, + bucket: str, + ) -> dict[str, Any]: + try: + return self.storage_adapter.upload_file(local_file_path, storage_key, bucket) + except Exception as exc: + raise StorageServiceException( + internal_message=f"Storage upload failed: {exc}", + operation="upload_local_file", + original_exception=exc, + ) from exc + + def upload_fileobj( + self, + file_obj: BinaryIO, + storage_key: str, + *, + bucket: str, + content_type: str | None = None, + ) -> dict[str, Any]: + try: + return self.storage_adapter.upload_fileobj( + file_obj, + storage_key, + bucket=bucket, + content_type=content_type, + ) + except Exception as exc: + raise StorageServiceException( + internal_message=f"Storage upload file object failed: {exc}", + operation="upload_fileobj", + original_exception=exc, + ) from exc + + def download_to_path( + self, + storage_key: str, + local_path: str, + *, + bucket: str, + ) -> str: + try: + return self.storage_adapter.download_file(storage_key, local_path, bucket) + except Exception as exc: + raise StorageServiceException( + internal_message=f"Storage download failed: {exc}", + operation="download_to_path", + original_exception=exc, + ) from exc + + def download_to_temp( + self, + storage_key: str, + *, + suffix: str, + temp_dir: str, + bucket: str, + ) -> str: + local_temp_path: str | None = None + + try: + os.makedirs(temp_dir, exist_ok=True) + with tempfile.NamedTemporaryFile( + delete=False, + suffix=suffix, + dir=temp_dir, + ) as temp_file: + local_temp_path = temp_file.name + + self.download_to_path( + storage_key, + local_temp_path, + bucket=bucket, + ) + return local_temp_path + except Exception as exc: + if local_temp_path and os.path.exists(local_temp_path): + os.remove(local_temp_path) + raise StorageServiceException( + internal_message=( + "Failed to download object-storage file to temp path: " + f"storage_key={storage_key}, temp_dir={temp_dir}, error={exc}" + ), + operation="download_to_temp", + original_exception=exc, + ) from exc + + def download_file_from_url( + self, + file_url: str, + *, + temp_dir: str | None = None, + ) -> str: + temp_file_path = "" + try: + validation = validate_http_url_and_resolve_ip(file_url) + if not validation.is_valid or not validation.validated_ip: + raise StorageServiceException( + internal_message=f"Invalid URL: {validation.error_message}", + operation="download_from_url", + ) + + effective_temp_dir = temp_dir or getattr(settings, "TMP_PATH", "/tmp") + os.makedirs(effective_temp_dir, exist_ok=True) + download_result = download_pinned_outbound_file( + url=validation.url, + pinned_ip=validation.validated_ip, + timeout_seconds=300, + user_agent="Knowhere-FileDownloader/1.0", + temp_dir=effective_temp_dir, + ) + temp_file_path = download_result.temp_file_path + return temp_file_path + except StorageServiceException: + raise + except Exception as exc: + if temp_file_path and os.path.exists(temp_file_path): + os.remove(temp_file_path) + raise StorageServiceException( + internal_message=f"Failed to download file: {exc}", + operation="download_from_url", + original_exception=exc, + ) from exc + + @staticmethod + def get_content_type(file_extension: str) -> str: + content_types = { + ".pdf": "application/pdf", + ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ".doc": "application/msword", + ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ".xls": "application/vnd.ms-excel", + ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation", + ".ppt": "application/vnd.ms-powerpoint", + ".csv": "text/csv", + ".txt": "text/plain", + ".md": "text/markdown", + ".json": "application/json", + ".xml": "application/xml", + ".html": "text/html", + ".htm": "text/html", + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".bmp": "image/bmp", + ".tiff": "image/tiff", + ".svg": "image/svg+xml", + ".zip": "application/zip", + ".rar": "application/x-rar-compressed", + ".7z": "application/x-7z-compressed", + ".tar": "application/x-tar", + ".gz": "application/gzip", + } + return content_types.get(file_extension.lower(), "application/octet-stream") diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 52d26783f..650019385 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -2,9 +2,13 @@ import os from dataclasses import dataclass +from collections.abc import Iterator from pathlib import Path from typing import Protocol +from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.storage.storage_adapter import StorageAdapter + _EXCLUDED_FILE_NAMES = {".DS_Store", "Thumbs.db"} _EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"} _CLIENT_ARTIFACT_DIRS = {"images", "tables"} @@ -29,32 +33,24 @@ def generate_artifact_url( def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: ... -class ResultS3: +class JobResultStorage: def __init__( - self, *, results_bucket: str | None = None, storage_adapter=None + self, + *, + results_bucket: str | None = None, + storage_adapter: StorageAdapter | None = None, ) -> None: - if results_bucket is None: - from shared.core.config import settings - - results_bucket = getattr( - settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME - ) - self.results_bucket = results_bucket - self._storage_adapter = storage_adapter - - @property - def storage_adapter(self): - if self._storage_adapter is None: - from shared.core.config.storage import get_cached_storage_adapter - - self._storage_adapter = get_cached_storage_adapter() - return self._storage_adapter + self._job_file_storage = JobFileStorage( + storage_adapter=storage_adapter, + results_bucket=results_bucket, + ) + self.results_bucket = self._job_file_storage.results_bucket def build_zip_key(self, *, job_id: str) -> str: - return f"results/{job_id}.zip" + return self._job_file_storage.build_result_zip_key(job_id=job_id) def build_raw_prefix(self, *, job_id: str) -> str: - return f"results/{job_id}/" + return self._job_file_storage.build_result_raw_prefix(job_id=job_id) def build_raw_key(self, *, job_id: str, relative_path: str) -> str: normalized = self._normalize_raw_relative_path(relative_path) @@ -82,15 +78,21 @@ def upload( if not zip_path.is_file(): raise ValueError(f"Result ZIP file does not exist: {zip_file_path}") zip_key = self.build_zip_key(job_id=job_id) - self.storage_adapter.upload_file(str(zip_path), zip_key, self.results_bucket) + self._job_file_storage.upload_local_file( + str(zip_path), + zip_key, + bucket=self.results_bucket, + ) self._cleanup_file(zip_path) raw_files: dict[str, str] = {} for file_path in self._iter_raw_files(result_path): relative_path = file_path.relative_to(result_path).as_posix() raw_key = self.build_raw_key(job_id=job_id, relative_path=relative_path) - self.storage_adapter.upload_file( - str(file_path), raw_key, self.results_bucket + self._job_file_storage.upload_local_file( + str(file_path), + raw_key, + bucket=self.results_bucket, ) raw_files[relative_path] = raw_key @@ -101,12 +103,11 @@ def upload( ) def generate_url(self, *, storage_key: str, expires_in: int = 3600) -> str | None: - return self.storage_adapter.generate_presigned_url( + return self._job_file_storage.generate_download_url( storage_key, - expiration=expires_in, bucket=self.results_bucket, - method="GET", - ) + expires_in=expires_in, + )["download_url"] def generate_artifact_url( self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 @@ -119,7 +120,7 @@ def generate_artifact_url( expires_in=expires_in, ) - def _iter_raw_files(self, result_dir: Path): + def _iter_raw_files(self, result_dir: Path) -> Iterator[Path]: for root, dir_names, file_names in os.walk(result_dir): dir_names[:] = [ dir_name @@ -160,4 +161,4 @@ def _cleanup_file(self, file_path: Path) -> None: def get_result_storage() -> ResultStorage: - return ResultS3() + return JobResultStorage() diff --git a/packages/shared-python/shared/services/webhook/dispatcher.py b/packages/shared-python/shared/services/webhook/dispatcher.py index e5ee4610d..4d32ed3f2 100644 --- a/packages/shared-python/shared/services/webhook/dispatcher.py +++ b/packages/shared-python/shared/services/webhook/dispatcher.py @@ -345,13 +345,12 @@ async def _enrich_payload(self, event: WebhookEvent) -> Dict[str, Any]: # Add result_url (fresh download link) if job_result.result_s3_key: - from shared.services.storage.file_upload_service import ( - FileUploadService, - ) + from shared.services.storage.job_file_storage import JobFileStorage - upload_service = FileUploadService() - url_info = await upload_service.generate_download_url( - job_result.result_s3_key + result_storage = JobFileStorage() + url_info = result_storage.generate_download_url( + job_result.result_s3_key, + bucket=result_storage.results_bucket, ) payload["result_url"] = url_info["download_url"] logger.debug( diff --git a/packages/shared-python/shared/utils/CommonHelper.py b/packages/shared-python/shared/utils/CommonHelper.py deleted file mode 100644 index 2790c7a6a..000000000 --- a/packages/shared-python/shared/utils/CommonHelper.py +++ /dev/null @@ -1,45 +0,0 @@ -from io import BytesIO -from pathlib import Path - -import httpx -import pandas as pd -from starlette.datastructures import UploadFile as StarletteUploadFile - -from shared.utils.FileDownUpUtils import s3_upload_file - - -def is_remote(path): - """Check whether a path is a remote URL.""" - if path is None: - return False - if not isinstance(path, str): - return False - return path.startswith("http://") or path.startswith("https://") - - -async def load_file_bytes(file_path, *, file_url="", timeout=None): - if isinstance(file_path, str) and is_remote(file_path): - # If file_path is already a full URL, use it directly. - url_to_use = file_path - if not isinstance(file_url, str): - file_url = file_url.geturl() - # Prefer file_url when provided; otherwise keep file_path. - if file_url: - url_to_use = file_url - async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client: - r = await client.get(url_to_use) # Fetch the resolved URL. - r.raise_for_status() - return r.content - else: - p = Path(file_path) - return p.read_bytes() - - -async def upload_dataframe_to_s3(df: pd.DataFrame, filename: str, prefix: str): - # Write the DataFrame into an in-memory BytesIO buffer. - buffer = BytesIO() - df.to_csv(buffer, index=False) - buffer.seek(0) # Reset the cursor to the buffer start. - - upload_file = StarletteUploadFile(file=buffer, filename=filename) - s3_upload_file(upload_file, prefix) diff --git a/packages/shared-python/shared/utils/FileDownUpUtils.py b/packages/shared-python/shared/utils/FileDownUpUtils.py deleted file mode 100644 index 0f124ebe8..000000000 --- a/packages/shared-python/shared/utils/FileDownUpUtils.py +++ /dev/null @@ -1,253 +0,0 @@ -import os -import uuid -import zipfile -from pathlib import Path -from typing import Optional, Union -from urllib.parse import urljoin - -import aiohttp -import requests -from botocore.exceptions import ClientError -from starlette.datastructures import UploadFile - -from shared.core.config import settings -from shared.core.config.storage import get_cached_storage_adapter -from shared.core.exceptions.domain_exceptions import ( - KnowhereException, - NotFoundException, - StorageServiceException, -) -from shared.models.schemas.s3_file import FliesDownload - - -def s3_upload_file(file: UploadFile, prefix: str): - """ - Upload a file object to S3 storage. - :param file: Input file such as ``abc15sa25ww.doc`` - :param prefix: Storage prefix such as ``upload/123`` - :return: Upload result payload - """ - if prefix and not prefix.endswith("/"): - prefix += "/" - object_key = f"{prefix}{file.filename}" - adapter = get_cached_storage_adapter() - try: - # ``upload_fileobj`` streams efficiently and avoids large in-memory copies. - adapter.upload_fileobj( - file.file, object_key, content_type="application/octet-stream" - ) - public_url = ( - f"{settings.S3_PRIVATE_DOMAIN}/{object_key}" - if settings.S3_PRIVATE_DOMAIN - else f"storage/{object_key}" - ) - content = { - "message": "File uploaded successfully", - "bucket": settings.S3_BUCKET_NAME, - "file_key": object_key, - "public_url_for_reference": public_url, - } - return content - - except KnowhereException: - raise - except Exception as e: - # Wrap storage upload failures in a domain exception. - raise StorageServiceException( - internal_message=f"Storage upload failed: {str(e)}", - operation="upload", - original_exception=e, - ) - - -def s3_download_extract_zip( - url: str, - dest_dir: Union[str, os.PathLike], - *, - filename: str = "parsed.zip", - headers: Optional[dict] = None, - timeout: int | None = None, - chunk_size: int | None = None, - keep_exts: tuple[str, ...] = (".md", ".json"), - exclude_patterns: tuple[str, ...] = (), - clean_empty_dirs: bool = True, -): - """ - Download and extract a zip file, keeping only specific file types. - - Args: - exclude_patterns: Tuple of filename patterns to exclude (e.g., ("content_list", "middle.json")) - """ - import fnmatch - - from shared.core.constants import APIConstants, ProcessingConstants - - # Use defaults when optional arguments are omitted. - if timeout is None: - timeout = APIConstants.S3_FILE_DOWNLOAD_TIMEOUT - if chunk_size is None: - chunk_size = ProcessingConstants.IMG_CHUNK_SIZE - - dest_dir = Path(dest_dir).expanduser().resolve() - dest_dir.mkdir(parents=True, exist_ok=True) - zip_path = dest_dir / filename - - # 1) Download to zip_path and extract. - with requests.get( - url, headers=headers or {}, timeout=timeout, stream=True, allow_redirects=True - ) as r: - r.raise_for_status() - with open(zip_path, "wb") as f: - for chunk in r.iter_content(chunk_size=chunk_size): - if chunk: - f.write(chunk) - - with zipfile.ZipFile(zip_path, "r") as zf: - zf.extractall(dest_dir) - - # 2) Remove files outside keep_exts or matching exclude_patterns. - kept_files = [] - for p in dest_dir.rglob("*"): - if p.is_file(): - # Check if file should be excluded by pattern - should_exclude = False - for pattern in exclude_patterns: - if pattern in p.name or fnmatch.fnmatch(p.name, pattern): - should_exclude = True - break - - if should_exclude: - p.unlink() - elif p.suffix.lower() in keep_exts: - kept_files.append(p) - else: - p.unlink() - - # 4) Remove empty directories when requested. - if clean_empty_dirs: - for d in sorted( - [p for p in dest_dir.rglob("*") if p.is_dir()], - key=lambda x: len(x.parts), - reverse=True, - ): - try: - next(d.iterdir()) - except StopIteration: - d.rmdir() - # 5) Delete the downloaded zip file. - zip_path.unlink(missing_ok=True) - - -def s3_get_download_url(file_key: str, expires_in: int = 3600): - """ - Get a file download URL from its storage key. - :param file_key: Full file path and name - :param expires_in: Desired URL lifetime - :return: Signed download payload - """ - s3_client = settings.get_s3_client() - try: - # Generate a pre-signed URL. - presigned_url = s3_client.generate_presigned_url( - "get_object", - Params={"Bucket": settings.S3_BUCKET_NAME, "Key": file_key}, - ExpiresIn=expires_in, # URL lifetime. - ) - fsdl = FliesDownload( - message="URL signed successfully", - file_key=file_key, - download_url=presigned_url, - expires_in_seconds=expires_in, - ) - return fsdl - - except ClientError as e: - # boto3 may still sign missing objects; the resulting URL can later 404. - raise NotFoundException( - resource="File", - resource_id=file_key, - internal_message=( - f"Could not generate the URL. Check whether the file is correct " - f"or the S3 configuration is valid: {str(e)}" - ), - ) - - -def get_url_file(path): - file_sig = s3_get_download_url(path, expires_in=3600) - # Assemble the final URL. - file_url = file_sig.download_url - response = requests.get(file_url, verify=True) - response.raise_for_status() # Ensure the request succeeds. - return response - - -def get_pub_fileurl(path): - """ - Build a public URL from a storage path. - :param path: - :return: Public URL - """ - base_url = settings.S3_PRIVATE_DOMAIN.rstrip("/") - clean_path = path.replace("\\", "/").strip() - full_url = urljoin(base_url + "/", clean_path) - return full_url - - -def s3_public_file_url(file_key: str) -> str: - permanent_url = f"{settings.S3_PRIVATE_DOMAIN}/{settings.S3_BUCKET_NAME}/{file_key}" - return permanent_url - - -async def download_and_upload_image( - img_url: str, prefix: str = "images/", temp_store_path=None -) -> dict: - """ - Download an image, rename it, upload it to S3, and clean up locally. - :param img_url: Image URL - :param prefix: S3 storage prefix - :return: Dict containing upload results and the new download reference - """ - # Generate a unique filename. - unique_filename = f"{uuid.uuid4()}.jpg" - # Temporary directory. - if temp_store_path is None: - temp_store_path = r"/Volumes/U/temp/output/" - local_file_path = Path(f"{settings.S3_TEMP_PATH or '/tmp'}{unique_filename}") - # Path(f"{temp_store_path}{unique_filename}") - try: - # Download the image asynchronously. - async with aiohttp.ClientSession() as session: - async with session.get(img_url) as response: - response.raise_for_status() - with open(local_file_path, "wb") as f: - f.write(await response.read()) - - # Create a temporary UploadFile wrapper. - from fastapi import UploadFile - - upload_file = UploadFile( - filename=unique_filename, file=open(local_file_path, "rb") - ) - - # Upload to S3. - result = s3_upload_file(upload_file, prefix) - - # Close the file handle and delete the local file. - upload_file.file.close() - os.remove(local_file_path) - return result - - except KnowhereException: - if local_file_path.exists(): - os.remove(local_file_path) - raise - except Exception as e: - # Always remove the local file on failure as well. - if local_file_path.exists(): - os.remove(local_file_path) - raise StorageServiceException( - internal_message=f"Failed to download and upload the image: {str(e)}", - operation="download_and_upload", - original_exception=e, - ) diff --git a/packages/shared-python/shared/utils/CommonHelperSync.py b/packages/shared-python/shared/utils/file_loading.py similarity index 73% rename from packages/shared-python/shared/utils/CommonHelperSync.py rename to packages/shared-python/shared/utils/file_loading.py index 924501ddb..522516cea 100644 --- a/packages/shared-python/shared/utils/CommonHelperSync.py +++ b/packages/shared-python/shared/utils/file_loading.py @@ -1,22 +1,24 @@ -"""Sync helpers for gevent worker code paths. - -Keep API async helpers in `CommonHelper.py`; worker should import this module. -""" +"""Sync file-loading helpers for worker parsing paths.""" from pathlib import Path -from typing import Optional +from urllib.parse import ParseResult import httpx -def is_remote(path): +def is_remote(path: object) -> bool: """Return True if `path` is an HTTP(S) URL.""" if path is None or not isinstance(path, str): return False return path.startswith("http://") or path.startswith("https://") -def load_file_bytes(file_path, *, file_url: str = "", timeout: Optional[float] = None): +def load_file_bytes( + file_path: str | Path, + *, + file_url: str | ParseResult = "", + timeout: float | None = None, +) -> bytes: """Load bytes from local path or remote URL synchronously.""" if isinstance(file_path, str) and is_remote(file_path): url_to_use = file_path diff --git a/packages/shared-python/shared/utils/zip_download.py b/packages/shared-python/shared/utils/zip_download.py new file mode 100644 index 000000000..630173c2f --- /dev/null +++ b/packages/shared-python/shared/utils/zip_download.py @@ -0,0 +1,79 @@ +"""Download-and-extract helpers for remote ZIP artifacts.""" + +import os +import zipfile +from pathlib import Path +from collections.abc import Mapping + +import requests + + +def download_and_extract_zip( + url: str, + dest_dir: str | os.PathLike[str], + *, + filename: str = "parsed.zip", + headers: Mapping[str, str] | None = None, + timeout: int | None = None, + chunk_size: int | None = None, + keep_exts: tuple[str, ...] = (".md", ".json"), + exclude_patterns: tuple[str, ...] = (), + clean_empty_dirs: bool = True, +) -> None: + """Download a ZIP file, extract it, and keep only the requested artifacts.""" + import fnmatch + + from shared.core.constants import APIConstants, ProcessingConstants + + if timeout is None: + timeout = APIConstants.S3_FILE_DOWNLOAD_TIMEOUT + if chunk_size is None: + chunk_size = ProcessingConstants.IMG_CHUNK_SIZE + + destination = Path(dest_dir).expanduser().resolve() + destination.mkdir(parents=True, exist_ok=True) + zip_path = destination / filename + + with requests.get( + url, + headers=headers or {}, + timeout=timeout, + stream=True, + allow_redirects=True, + ) as response: + response.raise_for_status() + with open(zip_path, "wb") as zip_file: + for chunk in response.iter_content(chunk_size=chunk_size): + if chunk: + zip_file.write(chunk) + + with zipfile.ZipFile(zip_path, "r") as extracted_zip: + extracted_zip.extractall(destination) + + for extracted_path in destination.rglob("*"): + if not extracted_path.is_file(): + continue + + should_exclude = False + for pattern in exclude_patterns: + if pattern in extracted_path.name or fnmatch.fnmatch(extracted_path.name, pattern): + should_exclude = True + break + + if should_exclude: + extracted_path.unlink() + elif extracted_path.suffix.lower() not in keep_exts: + extracted_path.unlink() + + if clean_empty_dirs: + for directory in sorted( + [path for path in destination.rglob("*") if path.is_dir()], + key=lambda path: len(path.parts), + reverse=True, + ): + try: + next(directory.iterdir()) + except StopIteration: + directory.rmdir() + + zip_path.unlink(missing_ok=True) From 559f86f6105984df74e020bb2635548ecbdc7efe Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 01:37:42 +0800 Subject: [PATCH 14/40] refactor: align sync state machine payloads --- .../test_stale_job_sweeper_contract.py | 86 ++++++++++++++++ .../shared/core/state_machine/service_sync.py | 98 ++++++++++++++----- 2 files changed, 157 insertions(+), 27 deletions(-) diff --git a/apps/worker/tests/contract/test_stale_job_sweeper_contract.py b/apps/worker/tests/contract/test_stale_job_sweeper_contract.py index 48cd82705..c5987ad8d 100644 --- a/apps/worker/tests/contract/test_stale_job_sweeper_contract.py +++ b/apps/worker/tests/contract/test_stale_job_sweeper_contract.py @@ -138,3 +138,89 @@ def test_should_skip_duplicate_beat_firing_with_the_real_periodic_redis_lock( "status": "skipped", "reason": "duplicate Beat firing", } + + +def test_should_record_retry_transition_through_sync_state_machine( + worker_contract_environment: None, +) -> None: + from shared.core.state_machine.service_sync import SyncStateMachineService + from shared.core.database_sync import get_sync_db_context + from shared.services.redis.redis_sync_service import SyncRedisServiceFactory + + job_id = f"job_retry_{uuid4().hex[:12]}" + user_id = f"worker-user-{uuid4().hex[:12]}" + _, engine = _load_worker_modules() + + with engine.begin() as connection: + insert_contract_user(connection, user_id=user_id) + insert_contract_job( + connection, + job_id=job_id, + user_id=user_id, + status="failed", + source_type="file", + webhook_enabled=False, + job_metadata=_build_file_job_metadata(), + error_code="TRANSIENT", + error_message="temporary failure", + ) + + redis_service = SyncRedisServiceFactory.get_service() + state_machine = SyncStateMachineService(redis_service=redis_service) + + with get_sync_db_context() as db: + did_retry = state_machine.handle_retry( + db, + job_id, + retry_metadata={"worker": "contract"}, + ) + + assert did_retry is True + + with engine.begin() as connection: + job_row = ( + connection.execute( + text( + """ + SELECT status + FROM jobs + WHERE job_id = :job_id + """ + ), + {"job_id": job_id}, + ) + .mappings() + .one() + ) + audit_log_row = ( + connection.execute( + text( + """ + SELECT from_state, to_state, transition_reason, operator_type, transition_metadata + FROM job_state_audit_logs + WHERE job_id = :job_id + ORDER BY id DESC + LIMIT 1 + """ + ), + {"job_id": job_id}, + ) + .mappings() + .one() + ) + + audit_metadata = dict(audit_log_row["transition_metadata"]) + progress = redis_service.hgetall(f"task:{job_id}:progress") + + assert job_row["status"] == "pending" + assert audit_log_row["from_state"] == "failed" + assert audit_log_row["to_state"] == "pending" + assert audit_log_row["transition_reason"] == "retry_transition" + assert audit_log_row["operator_type"] == "retry" + assert audit_metadata["worker"] == "contract" + assert audit_metadata["retry_reason"] == "task_retry" + assert audit_metadata["retry_count"] == 1 + assert audit_metadata["retry_timestamp"] + assert progress["status"] == "pending" + assert progress["worker"] == "contract" + assert progress["retry_count"] == 1 diff --git a/packages/shared-python/shared/core/state_machine/service_sync.py b/packages/shared-python/shared/core/state_machine/service_sync.py index 88e9aa3c6..175802459 100644 --- a/packages/shared-python/shared/core/state_machine/service_sync.py +++ b/packages/shared-python/shared/core/state_machine/service_sync.py @@ -7,7 +7,6 @@ """ import time -from datetime import datetime, timezone from typing import Any, Dict, Optional from loguru import logger @@ -18,18 +17,19 @@ JobStatus, is_valid_transition, ) +from shared.core.state_machine.transition_payloads import ( + build_failure_transition_metadata, + build_progress_cache_payload, + build_retry_transition, + serialize_transition_metadata, + utc_now_naive, +) from shared.models.database.job import Job from shared.models.database.job_state_audit_log import JobStateAuditLog from shared.services.redis.redis_sync_service import SyncRedisServiceFactory -from shared.utils.error_details import normalize_error_details -from shared.utils.json_utils import make_json_safe from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder -def _utc_now_naive() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) - - class SyncStateMachineService: """Sync state machine service — used by the Worker (gevent + psycopg2).""" @@ -122,7 +122,14 @@ def mark_failed( ) -> bool: """Mark a job as failed with error information.""" try: - normalized_details = normalize_error_details(error_details) + normalized_details, transition_metadata = ( + build_failure_transition_metadata( + error_message=error_message, + error_code=error_code, + error_details=error_details, + metadata=metadata, + ) + ) self._update_job_error( db, job_id, @@ -131,12 +138,6 @@ def mark_failed( normalized_details, ) - transition_metadata = (metadata or {}).copy() - transition_metadata["error_message"] = error_message - transition_metadata["error_code"] = error_code - if normalized_details: - transition_metadata["error_details"] = normalized_details - return self.transition( db, job_id, @@ -172,6 +173,48 @@ def mark_completed( logger.error(f"Failed to mark Job {job_id} as completed: {e}") return False + def handle_retry( + self, + db: Session, + job_id: str, + retry_metadata: Optional[Dict[str, Any]] = None, + operator_id: Optional[str] = None, + ) -> bool: + """Handle task retry — always goes through CAS-protected transition.""" + try: + job = self._get_job(db, job_id) + if not job: + logger.error(f"Job {job_id} does not exist") + return False + + current_state = job.status + if not current_state: + logger.error(f"Job {job_id} has no status") + return False + + retry_target, retry_metadata = build_retry_transition( + current_state=current_state, + retry_metadata=retry_metadata, + ) + + return self.transition( + db, + job_id, + retry_target, + "retry_transition", + operator_id, + "retry", + retry_metadata, + ) + except Exception as e: + logger.error(f"Job {job_id} retry failed: {e}") + try: + if db.is_active: + db.rollback() + except Exception as rollback_err: + logger.warning(f"Job {job_id} rollback failed: {rollback_err}") + return False + # ── Private helpers ───────────────────────────────────────────────── def _get_job(self, db: Session, job_id: str) -> Optional[Job]: @@ -199,7 +242,7 @@ def _cas_update_state( .values( status=to_state, version=old_version + 1, - updated_at=_utc_now_naive(), + updated_at=utc_now_naive(), ) ) return result.rowcount > 0 @@ -215,13 +258,7 @@ def _record_audit_log( operator_type: str, metadata: Optional[Dict[str, Any]], ) -> None: - serialized = None - if metadata: - try: - serialized = make_json_safe(metadata) - except Exception as e: - logger.warning(f"Metadata serialization failed: {e}") - serialized = {"error": "metadata_serialization_failed"} + serialized = serialize_transition_metadata(metadata) db.add( JobStateAuditLog( @@ -281,13 +318,20 @@ def _update_redis_cache( ) progress_key = redis_key_builder.task_progress(job_id) - progress_data: Dict[str, Any] = { - "status": status, - "timestamp": str(int(time.time())), - } + progress_data: Dict[str, Any] = build_progress_cache_payload( + status=status, + metadata=None, + timestamp=int(time.time()), + ) if metadata: try: - progress_data.update(make_json_safe(metadata)) + progress_data.update( + build_progress_cache_payload( + status=status, + metadata=metadata, + timestamp=int(time.time()), + ) + ) except Exception as e: logger.warning(f"Metadata serialization skipped: {e}") From 5109662cbfc81c520d5aff153b37b021af5bf270 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 01:48:29 +0800 Subject: [PATCH 15/40] refactor: centralize job result delivery --- .../app/services/jobs/result_projection.py | 39 ++-- .../test_webhook_recovery_contract.py | 171 +++++++++++++++++- .../shared/services/jobs/result_delivery.py | 76 ++++++++ .../shared/services/webhook/dispatcher.py | 24 +-- .../services/webhook/qstash_publisher.py | 15 +- 5 files changed, 271 insertions(+), 54 deletions(-) create mode 100644 packages/shared-python/shared/services/jobs/result_delivery.py diff --git a/apps/api/app/services/jobs/result_projection.py b/apps/api/app/services/jobs/result_projection.py index faf7ab075..c187e7225 100644 --- a/apps/api/app/services/jobs/result_projection.py +++ b/apps/api/app/services/jobs/result_projection.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone from typing import Any, Literal, Optional, cast from urllib.parse import urlparse @@ -9,9 +9,8 @@ from shared.core.exceptions.domain_exceptions import JobOperationException from shared.models.schemas.job import JobResultResponse, StandardErrorObject from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.jobs.result_delivery import JobResultDeliveryResolver from shared.utils.error_details import normalize_error_details -from shared.utils.utc_now import utc_now_naive JobStatusValue = Literal[ "pending", "waiting-file", "running", "converting", "done", "failed" @@ -112,27 +111,19 @@ def _resolve_duration_seconds(job: Any) -> float | None: async def _resolve_result_delivery( job: Any, ) -> tuple[dict[str, Any] | None, str | None, datetime]: - job_result = job.job_result - result_url = None - result = None - result_url_expires_at = job.created_at - - if job_result and job_result.result_s3_key: - result_storage = JobFileStorage() - result_url_info = result_storage.generate_download_url( - job_result.result_s3_key, - bucket=result_storage.results_bucket, - ) - result_url = result_url_info["download_url"] - - if job_result.inline_payload: - result = job_result.inline_payload - - if result_url: - expires_in = int(result_url_info.get("expires_in", 3600)) - result_url_expires_at = utc_now_naive() + timedelta(seconds=expires_in) - - return result, result_url, result_url_expires_at + default_expires_at = require_utc( + job.created_at, + field_name="created_at", + ) + delivery = JobResultDeliveryResolver().resolve( + job.job_result, + default_expires_at=default_expires_at, + ) + return ( + delivery.result, + delivery.result_url, + delivery.result_url_expires_at or default_expires_at, + ) async def build_job_result_response( diff --git a/apps/worker/tests/contract/test_webhook_recovery_contract.py b/apps/worker/tests/contract/test_webhook_recovery_contract.py index b30633b9f..569946a14 100644 --- a/apps/worker/tests/contract/test_webhook_recovery_contract.py +++ b/apps/worker/tests/contract/test_webhook_recovery_contract.py @@ -36,7 +36,9 @@ def _insert_webhook_event( created_at: datetime, updated_at: datetime | None = None, qstash_message_id: str | None = None, + payload: dict[str, Any] | None = None, ) -> None: + event_payload = payload or {"event": "job.failed", "job_id": job_id} connection.execute( text( """ @@ -69,7 +71,7 @@ def _insert_webhook_event( "id": event_id, "job_id": job_id, "target_url": target_url, - "payload": json.dumps({"event": "job.failed", "job_id": job_id}), + "payload": json.dumps(event_payload), "status": status, "attempts": attempts, "next_retry_at": None, @@ -80,6 +82,51 @@ def _insert_webhook_event( ) +def _insert_job_result( + connection: Connection, + *, + job_result_id: str, + job_id: str, + result_s3_key: str, + inline_payload: dict[str, Any], +) -> None: + timestamp = _utc_now() + connection.execute( + text( + """ + INSERT INTO job_results ( + id, + job_id, + delivery_mode, + inline_payload, + result_s3_key, + result_size, + created_at, + updated_at + ) VALUES ( + :id, + :job_id, + 'url', + CAST(:inline_payload AS JSON), + :result_s3_key, + :result_size, + :created_at, + :updated_at + ) + """ + ), + { + "id": job_result_id, + "job_id": job_id, + "inline_payload": json.dumps(inline_payload), + "result_s3_key": result_s3_key, + "result_size": 123, + "created_at": timestamp, + "updated_at": timestamp, + }, + ) + + def _load_worker_modules() -> tuple[Any, Any, Engine]: import app.core.tasks.webhook_tasks as webhook_tasks from shared.core.database_sync import get_sync_engine @@ -271,6 +318,128 @@ def publish(self, **kwargs: Any) -> SimpleNamespace: assert secrets_count_row["secrets_count"] == 1 +def test_should_publish_completed_webhook_with_result_delivery_payload( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + _, qstash_publisher, engine = _load_worker_modules() + from shared.services.storage.job_file_storage import JobFileStorage + + user_id = f"worker-user-{uuid4().hex[:12]}" + target_url = "https://hooks.contract.test/worker" + job_id = f"job_completed_{uuid4().hex[:12]}" + event_id = str(uuid4()) + result_s3_key = f"results/{job_id}.zip" + published_calls: list[dict[str, Any]] = [] + signed_url_calls: list[dict[str, Any]] = [] + + class FakeMessageClient: + def publish(self, **kwargs: Any) -> SimpleNamespace: + published_calls.append(kwargs) + return SimpleNamespace(message_id=f"msg_{event_id}") + + class FakeStorageAdapter: + def generate_presigned_url( + self, + key: str, + expiration: int = 3600, + bucket: str | None = None, + method: str = "GET", + headers: dict[str, str] | None = None, + ) -> str: + signed_url_calls.append( + { + "key": key, + "expiration": expiration, + "bucket": bucket, + "method": method, + "headers": headers, + } + ) + return f"signed://{bucket}/{key}?expires={expiration}" + + publisher = qstash_publisher.QStashWebhookPublisher() + monkeypatch.setattr( + qstash_publisher, + "validate_http_url_and_resolve_ip", + lambda *args, **kwargs: SimpleNamespace( + is_valid=True, + error_message=None, + validated_ip="93.184.216.34", + hostname="hooks.contract.test", + ), + ) + monkeypatch.setattr( + publisher, + "_get_client", + lambda: SimpleNamespace(message=FakeMessageClient()), + ) + monkeypatch.setattr( + qstash_publisher.JobResultDeliveryResolver, + "__init__", + lambda self: setattr( + self, + "_storage", + JobFileStorage(storage_adapter=FakeStorageAdapter()), + ), + ) + + now = _utc_now() + with engine.begin() as connection: + insert_contract_user(connection, user_id=user_id) + insert_contract_job( + connection, + job_id=job_id, + user_id=user_id, + status="done", + source_type="file", + webhook_url=target_url, + webhook_enabled=True, + job_metadata=_build_file_job_metadata(), + billing_status="charged", + ) + _insert_job_result( + connection, + job_result_id=str(uuid4()), + job_id=job_id, + result_s3_key=result_s3_key, + inline_payload={"checksum": "contract-checksum"}, + ) + _insert_webhook_event( + connection, + event_id=event_id, + job_id=job_id, + target_url=target_url, + status="pending", + attempts=0, + created_at=now, + payload={"event": "job.completed", "job_id": job_id}, + ) + + message_id = publisher.publish_event(event_id) + + assert message_id == f"msg_{event_id}" + assert len(published_calls) == 1 + assert signed_url_calls == [ + { + "key": result_s3_key, + "expiration": 3600, + "bucket": qstash_publisher.app_config.S3_RESULTS_BUCKET, + "method": "GET", + "headers": None, + } + ] + + published_payload = json.loads(published_calls[0]["body"]) + assert published_payload["event"] == "job.completed" + assert published_payload["job_id"] == job_id + assert published_payload["result"] == {"checksum": "contract-checksum"} + assert published_payload["result_url"] == ( + f"signed://{qstash_publisher.app_config.S3_RESULTS_BUCKET}/{result_s3_key}" + "?expires=3600" + ) + + def test_should_reconcile_stale_delivering_webhook_events_from_qstash_logs( worker_contract_environment: None, monkeypatch: MonkeyPatch, diff --git a/packages/shared-python/shared/services/jobs/result_delivery.py b/packages/shared-python/shared/services/jobs/result_delivery.py new file mode 100644 index 000000000..0a5c9b0d6 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/result_delivery.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any + +from shared.services.storage.job_file_storage import JobFileStorage +from shared.utils.utc_now import utc_now_naive + + +@dataclass(frozen=True) +class JobResultDelivery: + result: dict[str, Any] | None + result_url: str | None + result_url_expires_at: datetime | None + + +class JobResultDeliveryResolver: + """Resolve the public delivery fields exposed for a terminal Job Result.""" + + def __init__(self, *, storage: JobFileStorage | None = None) -> None: + self._storage = storage or JobFileStorage() + + def resolve( + self, + job_result: Any | None, + *, + default_expires_at: datetime | None = None, + ) -> JobResultDelivery: + result = None + result_url = None + result_url_expires_at = default_expires_at + + if not job_result: + return JobResultDelivery( + result=result, + result_url=result_url, + result_url_expires_at=result_url_expires_at, + ) + + inline_payload = getattr(job_result, "inline_payload", None) + if inline_payload: + result = inline_payload + + result_s3_key = getattr(job_result, "result_s3_key", None) + if result_s3_key: + url_info = self._storage.generate_download_url( + result_s3_key, + bucket=self._storage.results_bucket, + ) + result_url = url_info["download_url"] + expires_in = int(url_info.get("expires_in", 3600)) + result_url_expires_at = utc_now_naive() + timedelta(seconds=expires_in) + + return JobResultDelivery( + result=result, + result_url=result_url, + result_url_expires_at=result_url_expires_at, + ) + + def enrich_payload( + self, + payload: dict[str, Any], + *, + job_result: Any | None, + ) -> dict[str, Any]: + if payload.get("event") != "job.completed": + return payload + + delivery = self.resolve(job_result) + enriched = dict(payload) + if delivery.result_url: + enriched["result_url"] = delivery.result_url + if delivery.result: + enriched["result"] = delivery.result + return enriched diff --git a/packages/shared-python/shared/services/webhook/dispatcher.py b/packages/shared-python/shared/services/webhook/dispatcher.py index 4d32ed3f2..1c5890f4f 100644 --- a/packages/shared-python/shared/services/webhook/dispatcher.py +++ b/packages/shared-python/shared/services/webhook/dispatcher.py @@ -29,6 +29,7 @@ from shared.models.database.job import Job from shared.models.database.webhook import WebhookEvent, WebhookEventStatus from shared.models.database.webhook_log import WebhookLog +from shared.services.jobs.result_delivery import JobResultDeliveryResolver from shared.utils.pinned_outbound_http import ( send_pinned_outbound_request, ) @@ -341,25 +342,10 @@ async def _enrich_payload(self, event: WebhookEvent) -> Dict[str, Any]: ) return payload - job_result = job.job_result - - # Add result_url (fresh download link) - if job_result.result_s3_key: - from shared.services.storage.job_file_storage import JobFileStorage - - result_storage = JobFileStorage() - url_info = result_storage.generate_download_url( - job_result.result_s3_key, - bucket=result_storage.results_bucket, - ) - payload["result_url"] = url_info["download_url"] - logger.debug( - f"Enriched payload with result_url for job {event.job_id}" - ) - - # Add result (inline payload) - if job_result.inline_payload: - payload["result"] = job_result.inline_payload + payload = JobResultDeliveryResolver().enrich_payload( + payload, + job_result=job.job_result, + ) except Exception as e: logger.error(f"Failed to enrich payload for event {event.id}: {e}") diff --git a/packages/shared-python/shared/services/webhook/qstash_publisher.py b/packages/shared-python/shared/services/webhook/qstash_publisher.py index 8150100b2..03279d89e 100644 --- a/packages/shared-python/shared/services/webhook/qstash_publisher.py +++ b/packages/shared-python/shared/services/webhook/qstash_publisher.py @@ -23,6 +23,7 @@ from shared.core.config import app_config from shared.core.exceptions.domain_exceptions import QStashServiceException from shared.models.database.webhook import WebhookEventStatus +from shared.services.jobs.result_delivery import JobResultDeliveryResolver from shared.utils.url_security import ( validate_http_url_and_resolve_ip, ) @@ -272,16 +273,10 @@ def _enrich_payload(self, db: Any, event: Any) -> Dict[str, Any]: if not job or not job.job_result: return payload - job_result = job.job_result - if job_result.result_s3_key: - payload["result_url"] = app_config.get_storage_adapter().generate_presigned_url( - job_result.result_s3_key, - expiration=3600, - method="GET", - ) - - if job_result.inline_payload: - payload["result"] = job_result.inline_payload + payload = JobResultDeliveryResolver().enrich_payload( + payload, + job_result=job.job_result, + ) except Exception as exc: logger.error(f"Failed to enrich payload for event {event.id}: {exc}") From c452d98f44861e8b072f46ec82785e1c7d911eb1 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 02:28:51 +0800 Subject: [PATCH 16/40] refactor: centralize retrieval asset projection --- ...est_retrieval_asset_projection_contract.py | 124 ++++++++++++++++++ .../services/retrieval/agentic/evidence.py | 30 +---- .../shared/services/retrieval/assets.py | 89 ++++++++++++- .../services/retrieval/response_projection.py | 51 ++----- 4 files changed, 224 insertions(+), 70 deletions(-) create mode 100644 apps/api/tests/contract/test_retrieval_asset_projection_contract.py diff --git a/apps/api/tests/contract/test_retrieval_asset_projection_contract.py b/apps/api/tests/contract/test_retrieval_asset_projection_contract.py new file mode 100644 index 000000000..fd23c106e --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_asset_projection_contract.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from typing import Any + +import pytest + + +class FakeResultStorage: + def __init__(self) -> None: + self.should_fail: bool = False + self.generated_refs: list[tuple[str, str]] = [] + + def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: + if not artifact_ref: + return None + normalized = str(artifact_ref).strip().replace("\\", "/").lstrip("/") + if normalized.startswith("images/") or normalized.startswith("tables/"): + return normalized + return None + + def generate_artifact_url( + self, + *, + job_id: str, + artifact_ref: str, + expires_in: int = 3600, + ) -> str | None: + del expires_in + if self.should_fail: + raise RuntimeError("storage signing failed") + normalized = self.normalize_artifact_ref(artifact_ref) + if normalized is None: + return None + self.generated_refs.append((job_id, normalized)) + return f"https://assets.example.test/{job_id}/{normalized}" + + +@pytest.mark.asyncio +async def test_retrieval_asset_projection_should_attach_signed_urls_only_to_result_media_artifacts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from shared.services.retrieval.assets import enrich_rows_with_retrieval_asset_urls + + fake_storage = FakeResultStorage() + monkeypatch.setattr( + "shared.services.retrieval.assets.get_result_storage", + lambda: fake_storage, + ) + + rows: list[dict[str, Any]] = [ + { + "chunk_id": "image-chunk", + "chunk_type": "image", + "job_id": "job_123", + "file_path": "images/chart.png", + }, + { + "chunk_id": "text-chunk", + "chunk_type": "text", + "job_id": "job_123", + "file_path": "images/inline.png", + }, + { + "chunk_id": "external-chunk", + "chunk_type": "table", + "job_id": "job_123", + "file_path": "https://example.test/table.html", + }, + ] + + enriched_rows = await enrich_rows_with_retrieval_asset_urls( + rows, + log_context="contract projection", + ) + + assert enriched_rows[0]["asset_url"] == ( + "https://assets.example.test/job_123/images/chart.png" + ) + assert "asset_url" not in enriched_rows[1] + assert "asset_url" not in enriched_rows[2] + assert fake_storage.generated_refs == [("job_123", "images/chart.png")] + + +@pytest.mark.asyncio +async def test_retrieval_asset_projection_should_build_chunk_url_map_and_ignore_storage_failures( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from shared.services.retrieval.assets import build_retrieval_asset_url_map + + fake_storage = FakeResultStorage() + monkeypatch.setattr( + "shared.services.retrieval.assets.get_result_storage", + lambda: fake_storage, + ) + + url_map = await build_retrieval_asset_url_map( + [ + { + "chunk_id": "image-chunk", + "type": "image", + "job_id": "job_123", + "file_path": "images/chart.png", + }, + ], + log_context="contract map", + ) + + fake_storage.should_fail = True + failed_url_map = await build_retrieval_asset_url_map( + [ + { + "chunk_id": "table-chunk", + "type": "table", + "job_id": "job_123", + "file_path": "tables/data.html", + }, + ], + log_context="contract map", + ) + + assert url_map == { + "image-chunk": "https://assets.example.test/job_123/images/chart.png" + } + assert failed_url_map == {} diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence.py b/packages/shared-python/shared/services/retrieval/agentic/evidence.py index 1d88efef7..3d54fb47b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence.py @@ -9,10 +9,7 @@ from shared.models.database.document import RetrievalHitStat from shared.services.retrieval.agentic.budget import BudgetLedger from shared.services.retrieval.agentic.types import DocTreeNode -from shared.services.retrieval.assets import ( - generate_retrieval_asset_url, - is_client_result_artifact_ref, -) +from shared.services.retrieval.assets import build_retrieval_asset_url_map from shared.services.retrieval.hit_stats_service import compute_importance_score from shared.utils.token_estimate import estimate_tokens @@ -81,27 +78,10 @@ def collect_media_chunks_all( async def build_asset_url_map( media_chunks: list[dict[str, Any]], ) -> dict[str, str]: - url_map: dict[str, str] = {} - for chunk in media_chunks: - chunk_id = str(chunk.get("chunk_id") or "").strip() - file_path = chunk.get("file_path") or "" - job_id = chunk.get("job_id") or "" - if not chunk_id or not file_path or not job_id: - continue - if not is_client_result_artifact_ref(file_path): - continue - try: - url = await generate_retrieval_asset_url( - job_id=str(job_id), - artifact_ref=str(file_path), - ) - if url: - url_map[chunk_id] = url - except Exception as exc: - logger.warning( - f"Failed to generate asset URL for {chunk_id} (ignored): {exc}" - ) - return url_map + return await build_retrieval_asset_url_map( + media_chunks, + log_context="agentic evidence", + ) def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]: diff --git a/packages/shared-python/shared/services/retrieval/assets.py b/packages/shared-python/shared/services/retrieval/assets.py index 39badabea..54868e980 100644 --- a/packages/shared-python/shared/services/retrieval/assets.py +++ b/packages/shared-python/shared/services/retrieval/assets.py @@ -1,14 +1,93 @@ from __future__ import annotations +from typing import Any + +from loguru import logger + +from shared.services.retrieval.hydration import MEDIA_CHUNK_TYPES, normalize_chunk_type from shared.services.storage.result_storage import get_result_storage -async def generate_retrieval_asset_url(*, job_id: str, artifact_ref: str) -> str | None: - return get_result_storage().generate_artifact_url( - job_id=job_id, - artifact_ref=artifact_ref, +def _normalize_artifact_ref(asset_ref: object) -> str | None: + return get_result_storage().normalize_artifact_ref( + None if asset_ref is None else str(asset_ref) ) +def _is_retrieval_media_row(row: dict[str, Any]) -> bool: + raw_chunk_type = row.get("chunk_type") or row.get("type") + return normalize_chunk_type(raw_chunk_type) in MEDIA_CHUNK_TYPES + + +def _resolve_asset_request(row: dict[str, Any]) -> tuple[str, str] | None: + job_id = str(row.get("job_id") or "").strip() + if not job_id or not _is_retrieval_media_row(row): + return None + + artifact_ref = _normalize_artifact_ref(row.get("file_path")) + if artifact_ref is None: + return None + + return job_id, artifact_ref + + +async def _generate_retrieval_asset_url( + *, + row: dict[str, Any], + log_context: str, +) -> str | None: + request = _resolve_asset_request(row) + if request is None: + return None + + job_id, artifact_ref = request + try: + return get_result_storage().generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + ) + except Exception as exc: + logger.warning(f"Failed to generate {log_context} asset URL (ignored): {exc}") + return None + + +async def enrich_rows_with_retrieval_asset_urls( + rows: list[dict[str, Any]], + *, + log_context: str, +) -> list[dict[str, Any]]: + enriched_rows: list[dict[str, Any]] = [] + for row in rows: + enriched = dict(row) + asset_url = await _generate_retrieval_asset_url( + row=row, + log_context=log_context, + ) + if asset_url: + enriched["asset_url"] = asset_url + enriched_rows.append(enriched) + return enriched_rows + + +async def build_retrieval_asset_url_map( + rows: list[dict[str, Any]], + *, + log_context: str, +) -> dict[str, str]: + url_map: dict[str, str] = {} + for row in rows: + chunk_id = str(row.get("chunk_id") or "").strip() + if not chunk_id: + continue + + asset_url = await _generate_retrieval_asset_url( + row=row, + log_context=log_context, + ) + if asset_url: + url_map[chunk_id] = asset_url + return url_map + + def is_client_result_artifact_ref(asset_ref: str | None) -> bool: - return get_result_storage().normalize_artifact_ref(asset_ref) is not None + return _normalize_artifact_ref(asset_ref) is not None diff --git a/packages/shared-python/shared/services/retrieval/response_projection.py b/packages/shared-python/shared/services/retrieval/response_projection.py index 36e01dfbc..71c092267 100644 --- a/packages/shared-python/shared/services/retrieval/response_projection.py +++ b/packages/shared-python/shared/services/retrieval/response_projection.py @@ -2,15 +2,10 @@ from typing import Any -from loguru import logger - -from shared.services.retrieval.assets import generate_retrieval_asset_url, is_client_result_artifact_ref +from shared.services.retrieval.assets import enrich_rows_with_retrieval_asset_urls from shared.services.retrieval.hydration import ( - MEDIA_CHUNK_TYPES, PUBLIC_RESULT_FIELDS, PUBLIC_SOURCE_FIELDS, - is_media_chunk, - normalize_chunk_type, ) @@ -29,24 +24,10 @@ def to_public_source(row: dict[str, Any]) -> dict[str, Any]: async def enrich_referenced_chunks_with_asset_urls(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: - enriched_refs: list[dict[str, Any]] = [] - for ref in refs: - enriched = dict(ref) - chunk_type = normalize_chunk_type(ref.get('chunk_type')) - artifact_ref = ref.get('file_path', '') - job_id = ref.get('job_id', '') - if chunk_type in MEDIA_CHUNK_TYPES and job_id and is_client_result_artifact_ref(artifact_ref): - try: - asset_url = await generate_retrieval_asset_url( - job_id=str(job_id), - artifact_ref=str(artifact_ref), - ) - if asset_url: - enriched['asset_url'] = asset_url - except Exception as exc: - logger.warning(f'Failed to generate agentic asset URL (ignored): {exc}') - enriched_refs.append(enriched) - return enriched_refs + return await enrich_rows_with_retrieval_asset_urls( + refs, + log_context='agentic referenced chunk', + ) async def project_public_retrieval_response(response: dict[str, Any]) -> dict[str, Any]: @@ -62,25 +43,15 @@ async def project_public_retrieval_response(response: dict[str, Any]) -> dict[st if response.get('referenced_chunks') is not None: public_response['referenced_chunks'] = response['referenced_chunks'] + projected_rows = await enrich_rows_with_retrieval_asset_urls( + response.get('results', []), + log_context='retrieval result', + ) public_results: list[dict[str, Any]] = [] - for row in response.get('results', []): - artifact_ref = row.get('file_path') - asset_url = None - if is_media_chunk(row) and is_client_result_artifact_ref(artifact_ref) and row.get('job_id'): - try: - asset_url = await generate_retrieval_asset_url( - job_id=str(row['job_id']), - artifact_ref=str(artifact_ref), - ) - except Exception as exc: - logger.warning(f'Failed to generate retrieval asset URL (ignored): {exc}') - + for row in projected_rows: public_row: dict[str, Any] = {} for field in PUBLIC_RESULT_FIELDS: - if field == 'asset_url': - if asset_url: - public_row['asset_url'] = asset_url - elif field in row: + if field in row: public_row[field] = row[field] if 'source' in row: public_row['source'] = row['source'] From 96ae384511a242b0043609f313692258d3fbaa33 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 02:43:41 +0800 Subject: [PATCH 17/40] refactor: deepen job file storage interface --- .../services/document_ingestion/service.py | 4 +- .../services/document_ingestion/workspace.py | 3 +- .../document_parser/mineru_pdf_service.py | 13 +-- .../services/document_parser/pptx_parser.py | 4 +- .../services/storage/sync_storage_service.py | 79 ------------- .../services/workload/url_upload_service.py | 16 +-- .../contract/test_parse_task_contract.py | 43 +++---- .../contract/test_url_upload_contract.py | 19 ++-- .../test_worker_job_file_storage_contract.py | 105 ++++++++++++++++++ .../services/storage/job_file_storage.py | 40 +++++++ 10 files changed, 190 insertions(+), 136 deletions(-) delete mode 100644 apps/worker/app/services/storage/sync_storage_service.py create mode 100644 apps/worker/tests/contract/test_worker_job_file_storage_contract.py diff --git a/apps/worker/app/services/document_ingestion/service.py b/apps/worker/app/services/document_ingestion/service.py index 5ce813f4d..f6edb2878 100644 --- a/apps/worker/app/services/document_ingestion/service.py +++ b/apps/worker/app/services/document_ingestion/service.py @@ -20,7 +20,6 @@ download_s3_file_to_temp, ) from app.services.document_parser.stage_profiler import stage_timer -from app.services.storage.sync_storage_service import verify_s3_file_exists from loguru import logger from sqlalchemy import select @@ -43,6 +42,7 @@ SyncJobMetadataService, SyncRedisServiceFactory, ) +from shared.services.storage.job_file_storage import JobFileStorage from shared.services.storage.result_storage import get_result_storage from shared.services.storage.zip_result_service import ZipResultService @@ -164,7 +164,7 @@ def _load_parse_job_context( def _assert_source_file_within_size_limit(s3_key: str) -> None: - file_info = verify_s3_file_exists(s3_key) + file_info = JobFileStorage().verify_upload_exists(s3_key) if not file_info.get("exists"): raise NotFoundException( resource="S3File", diff --git a/apps/worker/app/services/document_ingestion/workspace.py b/apps/worker/app/services/document_ingestion/workspace.py index 2cb5e51d4..c6bd2ca85 100644 --- a/apps/worker/app/services/document_ingestion/workspace.py +++ b/apps/worker/app/services/document_ingestion/workspace.py @@ -71,9 +71,8 @@ def create_task_workspace(job_id: str) -> str: def download_s3_file_to_temp(s3_key: str, file_ext: str, temp_dir: str) -> str: """Download the source file from object storage into the task workspace.""" storage = JobFileStorage() - return storage.download_to_temp( + return storage.download_upload_to_temp( s3_key, suffix=file_ext, temp_dir=temp_dir, - bucket=settings.S3_BUCKET_NAME, ) diff --git a/apps/worker/app/services/document_parser/mineru_pdf_service.py b/apps/worker/app/services/document_parser/mineru_pdf_service.py index 4e0c640ce..d3cde60e4 100644 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ b/apps/worker/app/services/document_parser/mineru_pdf_service.py @@ -20,6 +20,7 @@ UnavailableException, ) from shared.core.exceptions.knowhere_exception import KnowhereException +from shared.services.storage.job_file_storage import JobFileStorage from shared.utils.file_loading import is_remote from shared.utils.zip_download import download_and_extract_zip @@ -120,10 +121,8 @@ def _inspect_mineru_source_s3_key(s3_key: Optional[str]) -> tuple[Optional[str], return None, False assert s3_key is not None - from app.services.storage.sync_storage_service import verify_s3_file_exists - try: - existing_file = verify_s3_file_exists(s3_key, settings.S3_BUCKET_NAME) + existing_file = JobFileStorage().verify_upload_exists(s3_key) except Exception as exc: _log_mineru_url_mode_storage_fallback( operation="verify_source_object", @@ -165,10 +164,8 @@ def resolve_mineru_source_s3_key( return None assert s3_key is not None - from app.services.storage.sync_storage_service import upload_to_s3 - try: - upload_to_s3(local_file_path, s3_key, settings.S3_BUCKET_NAME) + JobFileStorage().upload_source_file(local_file_path, s3_key) except Exception as exc: _log_mineru_url_mode_storage_fallback( operation="upload_source_object", @@ -729,9 +726,7 @@ def parse_via_full( if resolved_s3_key is not None: try: - from app.services.storage.sync_storage_service import generate_download_url - - presigned = generate_download_url( + presigned = JobFileStorage().generate_upload_download_url( resolved_s3_key, expires_in=settings.MINERU_URL_MODE_PRESIGN_EXPIRY ) presigned_url = presigned["download_url"] diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/pptx_parser.py index 84b8e02f7..8782be796 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/pptx_parser.py @@ -19,7 +19,6 @@ from app.services.document_parser.pptx_pdf_rendering import ( render_pdf_to_image_pdf as _render_pdf_to_image_pdf, ) -from app.services.storage.sync_storage_service import download_s3_object_to_temp from loguru import logger from markitdown import MarkItDown from pptx2md import ConversionConfig, convert @@ -29,6 +28,7 @@ FileSystemException, ) from shared.core.logging import LogEvent +from shared.services.storage.job_file_storage import JobFileStorage from shared.utils.file_loading import load_file_bytes from shared.utils.file_utils import path_handle @@ -331,7 +331,7 @@ def _parse_cached_rendered_pdf( logger.info( f"[parse_pptx] Reusing rendered PDF for MinerU URL mode: {rendered_pdf_s3_key}" ) - cached_rendered_pdf_path = download_s3_object_to_temp( + cached_rendered_pdf_path = JobFileStorage().download_upload_to_temp( cached_rendered_pdf_s3_key, suffix=".pdf", temp_dir=output_dir, diff --git a/apps/worker/app/services/storage/sync_storage_service.py b/apps/worker/app/services/storage/sync_storage_service.py deleted file mode 100644 index dc22139ba..000000000 --- a/apps/worker/app/services/storage/sync_storage_service.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Sync adapter for shared Job file storage used by worker tasks.""" - -import os -from typing import Any - -from loguru import logger - -from shared.core.config import settings -from shared.services.storage.job_file_storage import JobFileStorage - - -def get_storage_adapter() -> JobFileStorage: - """Get the shared job file storage module for sync worker operations.""" - return JobFileStorage() - - -def verify_s3_file_exists(s3_key: str, bucket: str | None = None) -> dict[str, Any]: - """Verify an uploaded source file exists.""" - storage = get_storage_adapter() - return storage.verify_exists( - s3_key, - bucket=bucket or settings.S3_BUCKET_NAME, - ) - - -def generate_download_url( - s3_key: str, bucket: str | None = None, expires_in: int = 3600 -) -> dict[str, Any]: - """Generate a presigned download URL for a stored object.""" - storage = get_storage_adapter() - return storage.generate_download_url( - s3_key, - bucket=bucket or settings.S3_BUCKET_NAME, - expires_in=expires_in, - ) - - -def upload_to_s3(local_file_path: str, s3_key: str, bucket: str) -> None: - """Upload a local file using the shared job file storage rules.""" - storage = get_storage_adapter() - storage.upload_local_file(local_file_path, s3_key, bucket=bucket) - - -def download_s3_object_to_temp( - s3_key: str, - suffix: str, - temp_dir: str, - bucket: str | None = None, -) -> str: - """Download an object-storage file into a task-local temp file.""" - storage = get_storage_adapter() - return storage.download_to_temp( - s3_key, - suffix=suffix, - temp_dir=temp_dir, - bucket=bucket or settings.S3_BUCKET_NAME, - ) - - -def upload_zip_result(job_id: str, zip_file_path: str) -> str: - """Upload ZIP result file to S3 and cleanup temp file.""" - storage = get_storage_adapter() - results_bucket = storage.results_bucket - s3_key = storage.build_result_zip_key(job_id=job_id) - upload_to_s3(zip_file_path, s3_key, results_bucket) - logger.info(f"Result ZIP uploaded: job_id={job_id}, key={s3_key}") - try: - if os.path.exists(zip_file_path): - os.remove(zip_file_path) - except Exception as e: - logger.warning(f"Failed to cleanup temp ZIP: {e}") - return s3_key - - -def download_file_from_url(file_url: str) -> str: - """Download a URL file through SSRF validation and IP pinning.""" - storage = get_storage_adapter() - temp_dir = getattr(settings, "TMP_PATH", "/tmp") - return storage.download_file_from_url(file_url, temp_dir=temp_dir) diff --git a/apps/worker/app/services/workload/url_upload_service.py b/apps/worker/app/services/workload/url_upload_service.py index 3618a31fd..bd5347a2b 100644 --- a/apps/worker/app/services/workload/url_upload_service.py +++ b/apps/worker/app/services/workload/url_upload_service.py @@ -3,11 +3,6 @@ import os from typing import Any -from app.services.storage.sync_storage_service import ( - download_file_from_url, - upload_to_s3, - verify_s3_file_exists, -) from loguru import logger from shared.core.config import settings @@ -22,6 +17,7 @@ SyncJobMetadataService, SyncRedisServiceFactory, ) +from shared.services.storage.job_file_storage import JobFileStorage from shared.utils.url_file_type import resolve_file_extension_sync @@ -78,8 +74,12 @@ def upload_url_file( lifecycle_service.update_progress( job_id, progress=10, message="Downloading file from URL..." ) + storage = JobFileStorage() try: - temp_file_path = download_file_from_url(source_url) + temp_file_path = storage.download_file_from_url( + source_url, + temp_dir=getattr(settings, "TMP_PATH", "/tmp"), + ) except Exception as exc: raise ValidationException( user_message="Failed to download file from URL", @@ -119,7 +119,7 @@ def upload_url_file( lifecycle_service.update_progress( job_id, progress=50, message="Uploading file to S3..." ) - upload_to_s3(temp_file_path, str(s3_key), settings.S3_BUCKET_NAME) + storage.upload_source_file(temp_file_path, str(s3_key)) logger.info(f"File uploaded to S3: {s3_key}") finally: @@ -130,7 +130,7 @@ def upload_url_file( lifecycle_service.update_progress( job_id, progress=80, message="Verifying upload result..." ) - file_info = verify_s3_file_exists(str(s3_key)) + file_info = storage.verify_upload_exists(str(s3_key)) if not file_info.get("exists"): raise StorageServiceException( user_message="We failed to verify your file upload", diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index d93c3f863..3c9dcc48b 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -84,6 +84,19 @@ def _save_worker_task_cache( return redis_service +def _patch_verify_upload_exists( + monkeypatch: MonkeyPatch, + file_info_for_storage_key: Any, +) -> None: + from shared.services.storage.job_file_storage import JobFileStorage + + monkeypatch.setattr( + JobFileStorage, + "verify_upload_exists", + lambda self, storage_key: file_info_for_storage_key(storage_key), + ) + + def _find_task_workspaces(root: Path, job_id: str) -> list[Path]: return sorted( path @@ -175,11 +188,7 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - monkeypatch.setattr( - parse_job_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) def fake_download_s3_file_to_temp( storage_key: str, file_ext: str, temp_dir: str @@ -856,11 +865,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: raw_files={}, ) - monkeypatch.setattr( - parse_job_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) @@ -1052,11 +1057,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: raw_files={}, ) - monkeypatch.setattr( - parse_job_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) monkeypatch.setattr(parse_job_service.PageEstimator, "estimate", fake_estimate_page_count) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) @@ -1230,11 +1231,7 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return {"exists": storage_key == s3_key, "size": 1024} - monkeypatch.setattr( - parse_job_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) monkeypatch.setattr( parse_service, "checkerboard_inject_parse", @@ -1326,11 +1323,7 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: "size": _SAMPLE_PDF_PATH.stat().st_size, } - monkeypatch.setattr( - parse_job_service, - "verify_s3_file_exists", - fake_verify_s3_file_exists, - ) + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) def fake_download_s3_file_to_temp( storage_key: str, file_ext: str, temp_dir: str diff --git a/apps/worker/tests/contract/test_url_upload_contract.py b/apps/worker/tests/contract/test_url_upload_contract.py index 0c3966620..f98f1a816 100644 --- a/apps/worker/tests/contract/test_url_upload_contract.py +++ b/apps/worker/tests/contract/test_url_upload_contract.py @@ -43,6 +43,7 @@ def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progres sync_job_info_service_cls, sync_redis_service_factory, ) = _load_upload_task_modules() + from shared.services.storage.job_file_storage import JobFileStorage user_id = f"worker-user-{uuid4().hex[:12]}" job_id = f"job_url_upload_{uuid4().hex[:12]}" @@ -58,21 +59,21 @@ def resolve_public_address( return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))] monkeypatch.setattr( - url_upload_service, + JobFileStorage, "download_file_from_url", - lambda _source_url: str(downloaded_path), + lambda self, _source_url, *, temp_dir=None: str(downloaded_path), ) monkeypatch.setattr( - url_upload_service, - "upload_to_s3", - lambda local_path, storage_key, bucket: uploaded_calls.append( - (local_path, storage_key, bucket) + JobFileStorage, + "upload_source_file", + lambda self, local_path, storage_key: uploaded_calls.append( + (local_path, storage_key, self.uploads_bucket) ), ) monkeypatch.setattr( - url_upload_service, - "verify_s3_file_exists", - lambda storage_key: {"exists": storage_key == s3_key, "size": 3}, + JobFileStorage, + "verify_upload_exists", + lambda self, storage_key: {"exists": storage_key == s3_key, "size": 3}, ) monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address) diff --git a/apps/worker/tests/contract/test_worker_job_file_storage_contract.py b/apps/worker/tests/contract/test_worker_job_file_storage_contract.py new file mode 100644 index 000000000..b41b0bd9d --- /dev/null +++ b/apps/worker/tests/contract/test_worker_job_file_storage_contract.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any, BinaryIO + + +class FakeStorageAdapter: + def __init__(self) -> None: + self.existing_keys: set[tuple[str, str]] = set() + self.object_sizes: dict[tuple[str, str], int] = {} + self.upload_calls: list[tuple[str, str, str]] = [] + self.download_calls: list[tuple[str, str, str]] = [] + self.presigned_calls: list[tuple[str, str, int, str]] = [] + + def generate_presigned_url( + self, + s3_key: str, + expiration: int = 3600, + bucket: str | None = None, + method: str = "GET", + headers: dict[str, str] | None = None, + ) -> str: + del headers + assert bucket is not None + self.presigned_calls.append((s3_key, bucket, expiration, method)) + return f"https://storage.example.test/{bucket}/{s3_key}" + + def exists(self, s3_key: str, bucket: str) -> bool: + return (s3_key, bucket) in self.existing_keys + + def get_object_size(self, s3_key: str, bucket: str) -> int: + return self.object_sizes[(s3_key, bucket)] + + def upload_file(self, file_path: str, s3_key: str, bucket: str) -> dict[str, Any]: + self.upload_calls.append((file_path, s3_key, bucket)) + self.existing_keys.add((s3_key, bucket)) + self.object_sizes[(s3_key, bucket)] = Path(file_path).stat().st_size + return {"bucket": bucket, "key": s3_key} + + def upload_fileobj( + self, + file_obj: BinaryIO, + s3_key: str, + bucket: str, + content_type: str | None = None, + ) -> dict[str, Any]: + del file_obj, content_type + return {"bucket": bucket, "key": s3_key} + + def download_file(self, s3_key: str, local_path: str, bucket: str) -> str: + self.download_calls.append((s3_key, local_path, bucket)) + Path(local_path).write_bytes(b"downloaded") + return local_path + + +def test_job_file_storage_should_hide_upload_bucket_rules_for_worker_source_files( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from shared.services.storage.job_file_storage import JobFileStorage + + del worker_contract_environment + + storage_adapter = FakeStorageAdapter() + storage = JobFileStorage( + storage_adapter=storage_adapter, + uploads_bucket="uploads-bucket", + results_bucket="results-bucket", + ) + source_path = tmp_path / "source.pdf" + source_path.write_bytes(b"pdf") + + storage.upload_source_file(str(source_path), "uploads/job_123.pdf") + file_info = storage.verify_upload_exists("uploads/job_123.pdf") + download_info = storage.generate_upload_download_url( + "uploads/job_123.pdf", + expires_in=60, + ) + downloaded_path = storage.download_upload_to_temp( + "uploads/job_123.pdf", + suffix=".pdf", + temp_dir=str(tmp_path), + ) + + assert file_info == { + "exists": True, + "size": 3, + "content_type": None, + "last_modified": None, + "etag": None, + } + assert download_info == { + "download_url": "https://storage.example.test/uploads-bucket/uploads/job_123.pdf", + "expires_in": 60, + } + assert Path(downloaded_path).read_bytes() == b"downloaded" + assert storage_adapter.upload_calls == [ + (str(source_path), "uploads/job_123.pdf", "uploads-bucket") + ] + assert storage_adapter.download_calls == [ + ("uploads/job_123.pdf", downloaded_path, "uploads-bucket") + ] + assert storage_adapter.presigned_calls == [ + ("uploads/job_123.pdf", "uploads-bucket", 60, "GET") + ] diff --git a/packages/shared-python/shared/services/storage/job_file_storage.py b/packages/shared-python/shared/services/storage/job_file_storage.py index 02777bc0a..cc46f6bd3 100644 --- a/packages/shared-python/shared/services/storage/job_file_storage.py +++ b/packages/shared-python/shared/services/storage/job_file_storage.py @@ -88,6 +88,18 @@ def generate_download_url( ) return {"download_url": download_url, "expires_in": expires_in} + def generate_upload_download_url( + self, + storage_key: str, + *, + expires_in: int = 3600, + ) -> dict[str, Any]: + return self.generate_download_url( + storage_key, + bucket=self.uploads_bucket, + expires_in=expires_in, + ) + def verify_exists( self, storage_key: str, @@ -115,6 +127,9 @@ def verify_exists( original_exception=exc, ) from exc + def verify_upload_exists(self, storage_key: str) -> dict[str, Any]: + return self.verify_exists(storage_key, bucket=self.uploads_bucket) + def upload_local_file( self, local_file_path: str, @@ -131,6 +146,17 @@ def upload_local_file( original_exception=exc, ) from exc + def upload_source_file( + self, + local_file_path: str, + storage_key: str, + ) -> dict[str, Any]: + return self.upload_local_file( + local_file_path, + storage_key, + bucket=self.uploads_bucket, + ) + def upload_fileobj( self, file_obj: BinaryIO, @@ -206,6 +232,20 @@ def download_to_temp( original_exception=exc, ) from exc + def download_upload_to_temp( + self, + storage_key: str, + *, + suffix: str, + temp_dir: str, + ) -> str: + return self.download_to_temp( + storage_key, + suffix=suffix, + temp_dir=temp_dir, + bucket=self.uploads_bucket, + ) + def download_file_from_url( self, file_url: str, From 3c543627ca4b513a9c5271f2db6de5e11c47d978 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 02:48:29 +0800 Subject: [PATCH 18/40] refactor: remove api state machine facade --- apps/api/app/repositories/job_repository.py | 4 +- .../confirmation_service.py | 4 +- .../s3_events/upload_event_service.py | 6 +- .../app/services/state_machine/__init__.py | 21 --- .../api/app/services/state_machine/manager.py | 128 ------------------ 5 files changed, 7 insertions(+), 156 deletions(-) delete mode 100644 apps/api/app/services/state_machine/__init__.py delete mode 100644 apps/api/app/services/state_machine/manager.py diff --git a/apps/api/app/repositories/job_repository.py b/apps/api/app/repositories/job_repository.py index fd9deafe3..c3750352b 100644 --- a/apps/api/app/repositories/job_repository.py +++ b/apps/api/app/repositories/job_repository.py @@ -3,13 +3,13 @@ from datetime import datetime from typing import Any, Dict, Optional, Sequence -from app.services.state_machine import JobStateMachine from loguru import logger from sqlalchemy import and_, desc, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload +from shared.core.state_machine.service import AsyncStateMachineService from shared.models.database.job import Job from shared.models.database.job_state_history import JobStateHistory @@ -18,7 +18,7 @@ class JobRepository: """Repository for Job persistence operations.""" def __init__(self): - self.state_machine = JobStateMachine() + self.state_machine = AsyncStateMachineService() async def create_job( self, diff --git a/apps/api/app/services/document_ingestion/confirmation_service.py b/apps/api/app/services/document_ingestion/confirmation_service.py index 970520d61..73198721d 100644 --- a/apps/api/app/services/document_ingestion/confirmation_service.py +++ b/apps/api/app/services/document_ingestion/confirmation_service.py @@ -3,7 +3,6 @@ from app.repositories.job_repository import JobRepository from app.services.jobs import check_job_permission from app.services.knowledge.kb_orchestrator import KBOrchestrator -from app.services.state_machine import JobStateMachine from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession @@ -13,6 +12,7 @@ PermissionDeniedException, ValidationException, ) +from shared.core.state_machine.service import AsyncStateMachineService from shared.core.state_machine.states import JobStatus from shared.services.storage.file_upload_service import FileUploadService @@ -94,7 +94,7 @@ async def _transition_job_to_uploaded( job_id: str, trigger: str = "manual_upload_completed", ) -> None: - state_machine = JobStateMachine() + state_machine = AsyncStateMachineService() await state_machine.transition( db, job_id, diff --git a/apps/api/app/services/s3_events/upload_event_service.py b/apps/api/app/services/s3_events/upload_event_service.py index df124f61f..e30dd8fde 100644 --- a/apps/api/app/services/s3_events/upload_event_service.py +++ b/apps/api/app/services/s3_events/upload_event_service.py @@ -5,10 +5,10 @@ from app.repositories.job_repository import JobRepository from app.services.knowledge.kb_orchestrator import KBOrchestrator -from app.services.state_machine import JobStateMachine from loguru import logger from shared.core.database import get_db_context +from shared.core.state_machine.service import AsyncStateMachineService from shared.core.state_machine.states import JobStatus from shared.models.schemas.s3_event import S3Event @@ -55,7 +55,7 @@ async def process_upload_events(s3_event: S3Event) -> None: if is_job_expired(job.updated_at, settings.JOB_WAITING_EXPIRE_SECONDS): logger.warning(f"Job {job_id} upload expired, marking failed") - state_machine = JobStateMachine() + state_machine = AsyncStateMachineService() await state_machine.mark_failed( db, job_id, @@ -64,7 +64,7 @@ async def process_upload_events(s3_event: S3Event) -> None: ) continue - state_machine = JobStateMachine() + state_machine = AsyncStateMachineService() await state_machine.transition( db, job_id, diff --git a/apps/api/app/services/state_machine/__init__.py b/apps/api/app/services/state_machine/__init__.py deleted file mode 100644 index 2fc629999..000000000 --- a/apps/api/app/services/state_machine/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -""" -API-side compatibility layer for the shared state machine. - -The canonical implementation now lives under ``shared.core.state_machine``. -This package keeps stable imports for API callers while avoiding duplicate -state machine logic inside ``apps/api``. -""" - -from shared.core.state_machine.service import AsyncStateMachineService -from shared.core.state_machine.states import JobStatus - -from .manager import JobStateMachine - -StateMachineService = AsyncStateMachineService - -__all__ = [ - "AsyncStateMachineService", - "JobStateMachine", - "JobStatus", - "StateMachineService", -] diff --git a/apps/api/app/services/state_machine/manager.py b/apps/api/app/services/state_machine/manager.py deleted file mode 100644 index f23804115..000000000 --- a/apps/api/app/services/state_machine/manager.py +++ /dev/null @@ -1,128 +0,0 @@ -""" -API-side state machine facade. - -Core transition logic lives in ``shared.core.state_machine.service``. This -module keeps the ``JobStateMachine`` entry point that API code already uses. -""" - -from typing import Any, Dict, Optional - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.state_machine.service import AsyncStateMachineService -from shared.services.redis import RedisServiceFactory - - -class JobStateMachine: - """Compatibility facade over ``AsyncStateMachineService``.""" - - def __init__(self, redis_service: Optional[Any] = None) -> None: - self.redis = redis_service or RedisServiceFactory.get_service() - self.state_machine = AsyncStateMachineService(self.redis) - - async def transition( - self, - db: AsyncSession, - job_id: str, - to_state: str, - transition_reason: str = "normal_transition", - operator_id: Optional[str] = None, - operator_type: str = "system", - metadata: Optional[Dict[str, Any]] = None, - auto_commit: bool = True, - ) -> bool: - """Execute a CAS-protected state transition.""" - try: - return await self.state_machine.transition( - db=db, - job_id=job_id, - to_state=to_state, - transition_reason=transition_reason, - operator_id=operator_id, - operator_type=operator_type, - metadata=metadata, - auto_commit=auto_commit, - ) - except Exception as err: - logger.error(f"Job {job_id} transition failed: {err}") - return False - - async def mark_failed( - self, - db: AsyncSession, - job_id: str, - error_message: str, - error_code: str = "UNKNOWN", - error_details: Optional[Dict[str, Any]] = None, - operator_id: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - auto_commit: bool = True, - ) -> bool: - """Mark a job as failed.""" - try: - return await self.state_machine.mark_failed( - db=db, - job_id=job_id, - error_message=error_message, - error_code=error_code, - error_details=error_details, - operator_id=operator_id, - metadata=metadata, - auto_commit=auto_commit, - ) - except Exception as err: - logger.error(f"Failed to mark Job {job_id} as failed: {err}") - return False - - async def mark_completed( - self, - db: AsyncSession, - job_id: str, - result_metadata: Optional[Dict[str, Any]] = None, - operator_id: Optional[str] = None, - auto_commit: bool = True, - ) -> bool: - """Mark a job as completed.""" - try: - return await self.state_machine.mark_completed( - db=db, - job_id=job_id, - result_metadata=result_metadata, - operator_id=operator_id, - auto_commit=auto_commit, - ) - except Exception as err: - logger.error(f"Failed to mark Job {job_id} as completed: {err}") - return False - - async def handle_retry( - self, - db: AsyncSession, - job_id: str, - retry_metadata: Optional[Dict[str, Any]] = None, - operator_id: Optional[str] = None, - ) -> bool: - """Retry a job through the shared state machine.""" - try: - return await self.state_machine.handle_retry( - db=db, - job_id=job_id, - retry_metadata=retry_metadata, - operator_id=operator_id, - ) - except Exception as err: - logger.error(f"Failed to retry Job {job_id}: {err}") - return False - - async def get_current_state( - self, - db: AsyncSession, - job_id: str, - ) -> Optional[str]: - """Read the current state through the shared service.""" - try: - return await self.state_machine.get_current_state(db=db, job_id=job_id) - except Exception as err: - logger.error(f"Failed to read Job {job_id} state: {err}") - return None From 11d624e9c3c78a3c02f18fbd5c9f911b8855a829 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 13:18:55 +0800 Subject: [PATCH 19/40] refactor deepen workflow modules --- .../billing/billing_command_workflow.py | 103 +++ .../services/billing/billing_read_model.py | 288 +++++++++ .../billing/billing_workflow_service.py | 352 ++--------- apps/api/app/services/jobs/job_read_model.py | 209 ++++++ apps/api/app/services/jobs/read_service.py | 167 +---- .../document_ingestion/processing_run.py | 583 +++++++++++++++++ .../services/document_ingestion/service.py | 593 +----------------- .../shared/services/retrieval/app_service.py | 510 ++------------- .../services/retrieval/execution_plan.py | 550 ++++++++++++++++ 9 files changed, 1865 insertions(+), 1490 deletions(-) create mode 100644 apps/api/app/services/billing/billing_command_workflow.py create mode 100644 apps/api/app/services/billing/billing_read_model.py create mode 100644 apps/api/app/services/jobs/job_read_model.py create mode 100644 apps/worker/app/services/document_ingestion/processing_run.py create mode 100644 packages/shared-python/shared/services/retrieval/execution_plan.py diff --git a/apps/api/app/services/billing/billing_command_workflow.py b/apps/api/app/services/billing/billing_command_workflow.py new file mode 100644 index 000000000..bd3d62073 --- /dev/null +++ b/apps/api/app/services/billing/billing_command_workflow.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +from app.services.billing.stripe_purchase_service import StripePurchaseService +from app.services.billing.stripe_webhook_service import StripeWebhookService +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.billing import MicroDollar +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import StripeServiceException +from shared.models.database.user import User +from shared.models.schemas.billing import ( + BuyCreditsPackageRequest, + BuyCreditsRequest, + CheckoutSessionResponse, + PaymentIntentResponse, +) + + +class BillingCommandWorkflow: + async def buy_credits( + self, + *, + request: BuyCreditsRequest, + user_id: str, + ) -> PaymentIntentResponse: + stripe_purchase_service = StripePurchaseService() + try: + amount_cny = request.credits_amount * 0.02 + amount_cents = int(amount_cny * 100) + payment_intent = await stripe_purchase_service.create_payment_intent( + user_id=user_id, + amount=amount_cents, + credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, + currency="cny", + ) + + return PaymentIntentResponse( + client_secret=payment_intent["client_secret"], + payment_intent_id=payment_intent["payment_intent_id"], + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to buy credits: {str(exc)}" + ) + + async def buy_credits_package( + self, + db: AsyncSession, + *, + request: BuyCreditsPackageRequest, + user_id: str, + ) -> CheckoutSessionResponse: + stripe_purchase_service = StripePurchaseService() + try: + result = await db.execute(select(User.email).where(User.id == user_id)) + user_email = result.scalar_one_or_none() + + frontend_url = settings.FRONTEND_URL + success_url = f"{frontend_url}/billing?success=true&type=credits_package" + cancel_url = f"{frontend_url}/billing?canceled=true" + + checkout_url = await stripe_purchase_service.create_credits_package_checkout_session( + db=db, + user_id=user_id, + price_id=request.price_id, + success_url=success_url, + cancel_url=cancel_url, + quantity=request.quantity, + email=user_email, + ) + + return CheckoutSessionResponse(checkout_url=checkout_url, session_id="") + except Exception as exc: + raise StripeServiceException( + internal_message=( + "Failed to create credits package purchase: " + f"{str(exc)}" + ) + ) + + async def handle_stripe_webhook( + self, + db: AsyncSession, + *, + payload: bytes, + stripe_signature: str | None, + ) -> dict[str, object]: + stripe_webhook_service = StripeWebhookService() + try: + if not stripe_signature: + raise StripeServiceException( + internal_message="Missing stripe-signature header" + ) + return await stripe_webhook_service.handle_webhook( + db, + payload=payload, + sig_header=stripe_signature, + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to handle webhook: {str(exc)}" + ) diff --git a/apps/api/app/services/billing/billing_read_model.py b/apps/api/app/services/billing/billing_read_model.py new file mode 100644 index 000000000..ccc17fe09 --- /dev/null +++ b/apps/api/app/services/billing/billing_read_model.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +from typing import Optional + +from app.services.billing.price_config_service import PriceConfigService +from pydantic import BaseModel +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.billing import MicroDollar +from shared.core.exceptions.domain_exceptions import StripeServiceException +from shared.models.database.credits_transaction import CreditsTransaction +from shared.models.database.job import Job +from shared.models.database.stripe_price_config import StripePriceConfig +from shared.models.schemas.billing import ( + CreditsBalanceResponse, + TransactionHistoryResponse, + UsageStatsResponse, +) +from shared.services.billing import CreditsService + + +class ParseUsageResponse(BaseModel): + request_total: int + mom_growth: float + credits_used: float + estimated_amount: Optional[float] + success_rate: float + avg_processing_time: float + + +class BillingReadModel: + def __init__( + self, + *, + price_config_service: PriceConfigService | None = None, + credits_service: CreditsService | None = None, + ) -> None: + self._price_config_service = price_config_service or PriceConfigService() + self._credits_service = credits_service or CreditsService() + + async def get_credits_balance( + self, + db: AsyncSession, + *, + user_id: str, + ) -> CreditsBalanceResponse: + try: + await self._credits_service.ensure_user_initialized(db, user_id) + await db.commit() + + balance_micro_dollar = await self._credits_service.get_balance(db, user_id) + return CreditsBalanceResponse( + credits_balance=MicroDollar(balance_micro_dollar).to_credit() + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get credits balance: {str(exc)}" + ) + + async def get_usage_stats( + self, + db: AsyncSession, + *, + user_id: str, + period: str, + ) -> UsageStatsResponse: + try: + stats = await self._credits_service.get_usage_stats(db, user_id, period) + return UsageStatsResponse( + period=stats["period"], + total_credits_used=MicroDollar(stats["total_used"]).to_credit(), + api_calls_count=stats["transaction_count"], + success_rate=95.0, + average_response_time=stats.get("avg_response_time", 0), + top_endpoints=[], + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get usage statistics: {str(exc)}" + ) + + async def get_parse_usage_overview( + self, + db: AsyncSession, + *, + user_id: str, + ) -> ParseUsageResponse: + try: + total_micro_credits_used = await self._load_total_parse_micro_credits_used( + db, + user_id=user_id, + ) + success_rate, avg_processing_time = await self._load_parse_job_usage_stats( + db, + user_id=user_id, + ) + estimated_amount = await self._estimate_parse_usage_amount( + db, + total_micro_credits_used=total_micro_credits_used, + ) + + return ParseUsageResponse( + request_total=0, + mom_growth=0.0, + credits_used=MicroDollar(total_micro_credits_used).to_credit() or 0, + estimated_amount=estimated_amount, + success_rate=round(success_rate, 2), + avg_processing_time=avg_processing_time, + ) + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get parse usage overview: {str(exc)}" + ) + + async def get_transaction_history( + self, + db: AsyncSession, + *, + user_id: str, + limit: int, + ) -> list[TransactionHistoryResponse]: + try: + transactions = await self._credits_service.get_transaction_history( + db, + user_id, + limit, + ) + return [ + TransactionHistoryResponse( + id=transaction.id, + credits_amount=MicroDollar(transaction.credits_amount).to_credit(), + transaction_type=transaction.transaction_type, + description=transaction.description, + created_at=transaction.created_at, + ) + for transaction in transactions + ] + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get transaction history: {str(exc)}" + ) + + async def get_price_configs( + self, + db: AsyncSession, + *, + product_type: str | None, + ) -> dict[str, list[dict[str, object]]]: + try: + if product_type == "subscription": + configs = await self._price_config_service.repository.get_all_active(db) + return { + "subscriptions": [ + _subscription_config_payload(config) + for config in configs + if config.product_type == "subscription" + ], + "credits_packages": [], + } + + if product_type == "credits_package": + credits_configs = await self._price_config_service.get_all_credits_packages( + db + ) + return { + "subscriptions": [], + "credits_packages": [ + _credits_package_config_payload(config) + for config in credits_configs + ], + } + + configs = await self._price_config_service.repository.get_all_active(db) + return { + "subscriptions": [ + _subscription_config_payload(config) + for config in configs + if config.product_type == "subscription" + ], + "credits_packages": [ + _credits_package_config_payload(config) + for config in configs + if config.product_type == "credits_package" + ], + } + except Exception as exc: + raise StripeServiceException( + internal_message=f"Failed to get price configurations: {str(exc)}" + ) + + async def _load_total_parse_micro_credits_used( + self, + db: AsyncSession, + *, + user_id: str, + ) -> int: + credits_row = await db.execute( + select(func.coalesce(func.sum(CreditsTransaction.credits_amount), 0)) + .where(CreditsTransaction.user_id == user_id) + .where(CreditsTransaction.transaction_type.in_(["usage", "refund"])) + ) + return int(abs(credits_row.scalar_one() or 0)) + + async def _load_parse_job_usage_stats( + self, + db: AsyncSession, + *, + user_id: str, + ) -> tuple[float, float]: + job_row = await db.execute( + select( + func.count().filter(Job.status == "done").label("done_cnt"), + func.count() + .filter(Job.status.in_(["done", "failed"])) + .label("terminal_cnt"), + func.avg(func.extract("epoch", Job.updated_at - Job.created_at)) + .filter(Job.status.in_(["done", "failed"])) + .label("avg_secs"), + ).where(Job.user_id == user_id) + ) + job_stats = job_row.first() or (0, 0, 0.0) + done_count = getattr(job_stats, "done_cnt", 0) or 0 + terminal_count = getattr(job_stats, "terminal_cnt", 0) or 0 + success_rate = ( + done_count / terminal_count * 100 if terminal_count > 0 else 0.0 + ) + avg_processing_time = round( + float(getattr(job_stats, "avg_secs", 0.0) or 0.0), + 2, + ) + return success_rate, avg_processing_time + + async def _estimate_parse_usage_amount( + self, + db: AsyncSession, + *, + total_micro_credits_used: int, + ) -> float | None: + price_row = await db.execute( + select(StripePriceConfig) + .where(StripePriceConfig.product_type == "credits_package") + .where(StripePriceConfig.is_active.is_(True)) + .order_by(StripePriceConfig.created_at) + .limit(1) + ) + price_cfg = price_row.scalar_one_or_none() + if not price_cfg or not price_cfg.credits_amount or price_cfg.credits_amount <= 0: + return None + + return round( + price_cfg.amount_cents + * total_micro_credits_used + / (100 * price_cfg.credits_amount), + 4, + ) + + +def _subscription_config_payload(config: StripePriceConfig) -> dict[str, object]: + metadata = config.extra_metadata or {} + return { + "id": config.plan_id, + "plan_id": config.plan_id, + "price_id": config.price_id, + "name": metadata.get("display_name", config.plan_id.upper()), + "description": metadata.get("description", ""), + "features": metadata.get("features", []), + "popular": metadata.get("frontend_config", {}).get("popular", False), + "amount_cents": config.amount_cents, + "currency": config.currency, + "metadata": metadata, + } + + +def _credits_package_config_payload(config: StripePriceConfig) -> dict[str, object]: + metadata = config.extra_metadata or {} + credit_amount = MicroDollar(config.credits_amount).to_credit() + return { + "id": config.plan_id, + "plan_id": config.plan_id, + "price_id": config.price_id, + "name": metadata.get("display_name", f"{credit_amount} Credits"), + "description": metadata.get("description", ""), + "credits_amount": credit_amount, + "amount_cents": config.amount_cents, + "currency": config.currency, + "metadata": metadata, + } diff --git a/apps/api/app/services/billing/billing_workflow_service.py b/apps/api/app/services/billing/billing_workflow_service.py index bf55ae6f7..91cd17df4 100644 --- a/apps/api/app/services/billing/billing_workflow_service.py +++ b/apps/api/app/services/billing/billing_workflow_service.py @@ -1,21 +1,12 @@ from __future__ import annotations -from typing import Optional - -from app.services.billing.price_config_service import PriceConfigService -from app.services.billing.stripe_purchase_service import StripePurchaseService -from app.services.billing.stripe_webhook_service import StripeWebhookService -from pydantic import BaseModel -from sqlalchemy import func, select +import app.services.billing.billing_command_workflow as billing_command_workflow +from app.services.billing.billing_command_workflow import BillingCommandWorkflow +from app.services.billing.billing_command_workflow import StripePurchaseService +from app.services.billing.billing_command_workflow import StripeWebhookService +from app.services.billing.billing_read_model import BillingReadModel, ParseUsageResponse from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.billing import MicroDollar -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import StripeServiceException -from shared.models.database.credits_transaction import CreditsTransaction -from shared.models.database.job import Job -from shared.models.database.stripe_price_config import StripePriceConfig -from shared.models.database.user import User from shared.models.schemas.billing import ( BuyCreditsPackageRequest, BuyCreditsRequest, @@ -25,27 +16,24 @@ TransactionHistoryResponse, UsageStatsResponse, ) -from shared.services.billing import CreditsService - -class ParseUsageResponse(BaseModel): - request_total: int - mom_growth: float - credits_used: float - estimated_amount: Optional[float] - success_rate: float - avg_processing_time: float +__all__ = [ + "BillingWorkflowService", + "ParseUsageResponse", + "StripePurchaseService", + "StripeWebhookService", +] class BillingWorkflowService: def __init__( self, *, - price_config_service: PriceConfigService | None = None, - credits_service: CreditsService | None = None, + command_workflow: BillingCommandWorkflow | None = None, + read_model: BillingReadModel | None = None, ) -> None: - self._price_config_service = price_config_service or PriceConfigService() - self._credits_service = credits_service or CreditsService() + self._command_workflow = command_workflow or BillingCommandWorkflow() + self._read_model = read_model or BillingReadModel() async def buy_credits( self, @@ -53,25 +41,11 @@ async def buy_credits( request: BuyCreditsRequest, user_id: str, ) -> PaymentIntentResponse: - stripe_purchase_service = self._create_stripe_purchase_service() - try: - amount_cny = request.credits_amount * 0.02 - amount_cents = int(amount_cny * 100) - payment_intent = await stripe_purchase_service.create_payment_intent( - user_id=user_id, - amount=amount_cents, - credits_amount=MicroDollar.from_dollars(request.credits_amount).amount, - currency="cny", - ) - - return PaymentIntentResponse( - client_secret=payment_intent["client_secret"], - payment_intent_id=payment_intent["payment_intent_id"], - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to buy credits: {str(exc)}" - ) + _sync_legacy_adapter_overrides() + return await self._command_workflow.buy_credits( + request=request, + user_id=user_id, + ) async def get_credits_balance( self, @@ -79,18 +53,7 @@ async def get_credits_balance( *, user_id: str, ) -> CreditsBalanceResponse: - try: - await self._credits_service.ensure_user_initialized(db, user_id) - await db.commit() - - balance_micro_dollar = await self._credits_service.get_balance(db, user_id) - return CreditsBalanceResponse( - credits_balance=MicroDollar(balance_micro_dollar).to_credit() - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get credits balance: {str(exc)}" - ) + return await self._read_model.get_credits_balance(db, user_id=user_id) async def get_usage_stats( self, @@ -99,20 +62,11 @@ async def get_usage_stats( user_id: str, period: str, ) -> UsageStatsResponse: - try: - stats = await self._credits_service.get_usage_stats(db, user_id, period) - return UsageStatsResponse( - period=stats["period"], - total_credits_used=MicroDollar(stats["total_used"]).to_credit(), - api_calls_count=stats["transaction_count"], - success_rate=95.0, - average_response_time=stats.get("avg_response_time", 0), - top_endpoints=[], - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get usage statistics: {str(exc)}" - ) + return await self._read_model.get_usage_stats( + db, + user_id=user_id, + period=period, + ) async def get_parse_usage_overview( self, @@ -120,32 +74,7 @@ async def get_parse_usage_overview( *, user_id: str, ) -> ParseUsageResponse: - try: - total_micro_credits_used = await self._load_total_parse_micro_credits_used( - db, - user_id=user_id, - ) - success_rate, avg_processing_time = await self._load_parse_job_usage_stats( - db, - user_id=user_id, - ) - estimated_amount = await self._estimate_parse_usage_amount( - db, - total_micro_credits_used=total_micro_credits_used, - ) - - return ParseUsageResponse( - request_total=0, - mom_growth=0.0, - credits_used=MicroDollar(total_micro_credits_used).to_credit() or 0, - estimated_amount=estimated_amount, - success_rate=round(success_rate, 2), - avg_processing_time=avg_processing_time, - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get parse usage overview: {str(exc)}" - ) + return await self._read_model.get_parse_usage_overview(db, user_id=user_id) async def get_transaction_history( self, @@ -154,26 +83,11 @@ async def get_transaction_history( user_id: str, limit: int, ) -> list[TransactionHistoryResponse]: - try: - transactions = await self._credits_service.get_transaction_history( - db, - user_id, - limit, - ) - return [ - TransactionHistoryResponse( - id=transaction.id, - credits_amount=MicroDollar(transaction.credits_amount).to_credit(), - transaction_type=transaction.transaction_type, - description=transaction.description, - created_at=transaction.created_at, - ) - for transaction in transactions - ] - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get transaction history: {str(exc)}" - ) + return await self._read_model.get_transaction_history( + db, + user_id=user_id, + limit=limit, + ) async def get_price_configs( self, @@ -181,47 +95,10 @@ async def get_price_configs( *, product_type: str | None, ) -> dict[str, list[dict[str, object]]]: - try: - if product_type == "subscription": - configs = await self._price_config_service.repository.get_all_active(db) - return { - "subscriptions": [ - _subscription_config_payload(config) - for config in configs - if config.product_type == "subscription" - ], - "credits_packages": [], - } - - if product_type == "credits_package": - credits_configs = await self._price_config_service.get_all_credits_packages( - db - ) - return { - "subscriptions": [], - "credits_packages": [ - _credits_package_config_payload(config) - for config in credits_configs - ], - } - - configs = await self._price_config_service.repository.get_all_active(db) - return { - "subscriptions": [ - _subscription_config_payload(config) - for config in configs - if config.product_type == "subscription" - ], - "credits_packages": [ - _credits_package_config_payload(config) - for config in configs - if config.product_type == "credits_package" - ], - } - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to get price configurations: {str(exc)}" - ) + return await self._read_model.get_price_configs( + db, + product_type=product_type, + ) async def buy_credits_package( self, @@ -230,33 +107,12 @@ async def buy_credits_package( request: BuyCreditsPackageRequest, user_id: str, ) -> CheckoutSessionResponse: - stripe_purchase_service = self._create_stripe_purchase_service() - try: - result = await db.execute(select(User.email).where(User.id == user_id)) - user_email = result.scalar_one_or_none() - - frontend_url = settings.FRONTEND_URL - success_url = f"{frontend_url}/billing?success=true&type=credits_package" - cancel_url = f"{frontend_url}/billing?canceled=true" - - checkout_url = await stripe_purchase_service.create_credits_package_checkout_session( - db=db, - user_id=user_id, - price_id=request.price_id, - success_url=success_url, - cancel_url=cancel_url, - quantity=request.quantity, - email=user_email, - ) - - return CheckoutSessionResponse(checkout_url=checkout_url, session_id="") - except Exception as exc: - raise StripeServiceException( - internal_message=( - "Failed to create credits package purchase: " - f"{str(exc)}" - ) - ) + _sync_legacy_adapter_overrides() + return await self._command_workflow.buy_credits_package( + db, + request=request, + user_id=user_id, + ) async def handle_stripe_webhook( self, @@ -265,122 +121,14 @@ async def handle_stripe_webhook( payload: bytes, stripe_signature: str | None, ) -> dict[str, object]: - stripe_webhook_service = self._create_stripe_webhook_service() - try: - if not stripe_signature: - raise StripeServiceException( - internal_message="Missing stripe-signature header" - ) - return await stripe_webhook_service.handle_webhook( - db, - payload=payload, - sig_header=stripe_signature, - ) - except Exception as exc: - raise StripeServiceException( - internal_message=f"Failed to handle webhook: {str(exc)}" - ) - - def _create_stripe_purchase_service(self) -> StripePurchaseService: - return StripePurchaseService() - - def _create_stripe_webhook_service(self) -> StripeWebhookService: - return StripeWebhookService() - - async def _load_total_parse_micro_credits_used( - self, - db: AsyncSession, - *, - user_id: str, - ) -> int: - credits_row = await db.execute( - select(func.coalesce(func.sum(CreditsTransaction.credits_amount), 0)) - .where(CreditsTransaction.user_id == user_id) - .where(CreditsTransaction.transaction_type.in_(["usage", "refund"])) - ) - return int(abs(credits_row.scalar_one() or 0)) - - async def _load_parse_job_usage_stats( - self, - db: AsyncSession, - *, - user_id: str, - ) -> tuple[float, float]: - job_row = await db.execute( - select( - func.count().filter(Job.status == "done").label("done_cnt"), - func.count() - .filter(Job.status.in_(["done", "failed"])) - .label("terminal_cnt"), - func.avg(func.extract("epoch", Job.updated_at - Job.created_at)) - .filter(Job.status.in_(["done", "failed"])) - .label("avg_secs"), - ).where(Job.user_id == user_id) - ) - job_stats = job_row.first() or (0, 0, 0.0) - done_count = getattr(job_stats, "done_cnt", 0) or 0 - terminal_count = getattr(job_stats, "terminal_cnt", 0) or 0 - success_rate = ( - done_count / terminal_count * 100 if terminal_count > 0 else 0.0 - ) - avg_processing_time = round( - float(getattr(job_stats, "avg_secs", 0.0) or 0.0), - 2, + _sync_legacy_adapter_overrides() + return await self._command_workflow.handle_stripe_webhook( + db, + payload=payload, + stripe_signature=stripe_signature, ) - return success_rate, avg_processing_time - - async def _estimate_parse_usage_amount( - self, - db: AsyncSession, - *, - total_micro_credits_used: int, - ) -> float | None: - price_row = await db.execute( - select(StripePriceConfig) - .where(StripePriceConfig.product_type == "credits_package") - .where(StripePriceConfig.is_active.is_(True)) - .order_by(StripePriceConfig.created_at) - .limit(1) - ) - price_cfg = price_row.scalar_one_or_none() - if not price_cfg or not price_cfg.credits_amount or price_cfg.credits_amount <= 0: - return None - - return round( - price_cfg.amount_cents - * total_micro_credits_used - / (100 * price_cfg.credits_amount), - 4, - ) - - -def _subscription_config_payload(config: StripePriceConfig) -> dict[str, object]: - metadata = config.extra_metadata or {} - return { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": metadata.get("display_name", config.plan_id.upper()), - "description": metadata.get("description", ""), - "features": metadata.get("features", []), - "popular": metadata.get("frontend_config", {}).get("popular", False), - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": metadata, - } -def _credits_package_config_payload(config: StripePriceConfig) -> dict[str, object]: - metadata = config.extra_metadata or {} - credit_amount = MicroDollar(config.credits_amount).to_credit() - return { - "id": config.plan_id, - "plan_id": config.plan_id, - "price_id": config.price_id, - "name": metadata.get("display_name", f"{credit_amount} Credits"), - "description": metadata.get("description", ""), - "credits_amount": credit_amount, - "amount_cents": config.amount_cents, - "currency": config.currency, - "metadata": metadata, - } +def _sync_legacy_adapter_overrides() -> None: + billing_command_workflow.StripePurchaseService = StripePurchaseService + billing_command_workflow.StripeWebhookService = StripeWebhookService diff --git a/apps/api/app/services/jobs/job_read_model.py b/apps/api/app/services/jobs/job_read_model.py new file mode 100644 index 000000000..1d0a4ab39 --- /dev/null +++ b/apps/api/app/services/jobs/job_read_model.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import math +from datetime import datetime, timedelta, timezone +from typing import Optional + +from app.repositories.job_repository import JobRepository +from app.services.jobs.result_projection import ( + build_job_result_response, + to_job_status_value, +) +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + JobOperationException, + NotFoundException, + PermissionDeniedException, + ValidationException, +) +from shared.models.schemas.job import JobList, JobResultResponse +from shared.services.redis import RedisServiceFactory +from shared.utils.utc_now import utc_now_naive + + +class JobReadModel: + async def list_jobs_for_user( + self, + db: AsyncSession, + *, + user_id: str, + page: int, + page_size: int, + job_status: Optional[str], + job_type: Optional[str], + recent_days: Optional[int], + start_time: Optional[datetime], + end_time: Optional[datetime], + ) -> JobList: + return await list_jobs_for_user( + db, + user_id=user_id, + page=page, + page_size=page_size, + job_status=job_status, + job_type=job_type, + recent_days=recent_days, + start_time=start_time, + end_time=end_time, + ) + + async def get_job_result_for_user( + self, + db: AsyncSession, + *, + job_id: str, + user_id: str, + ) -> JobResultResponse: + return await get_job_result_for_user(db, job_id=job_id, user_id=user_id) + + +def check_job_permission(job, user_id: str, job_id: str) -> None: + if not job: + raise NotFoundException( + resource="Job", resource_id=job_id, internal_message="Job not found" + ) + + if str(job.user_id) != user_id: + raise PermissionDeniedException( + user_message="You don't have permission to access this job", + ) + + +def normalize_naive_utc_filter_datetime(dt: Optional[datetime]) -> Optional[datetime]: + if dt is None: + return None + if dt.tzinfo is None or dt.utcoffset() is None: + return dt + return dt.astimezone(timezone.utc).replace(tzinfo=None) + + +async def list_jobs_for_user( + db: AsyncSession, + *, + user_id: str, + page: int, + page_size: int, + job_status: Optional[str], + job_type: Optional[str], + recent_days: Optional[int], + start_time: Optional[datetime], + end_time: Optional[datetime], +) -> JobList: + try: + job_repo = JobRepository() + + if recent_days not in (None, 1, 7, 30): + raise ValidationException( + user_message="recent_days only supports 1, 7, or 30", + violations=[{"field": "recent_days", "description": "Invalid value"}], + ) + + created_after: Optional[datetime] = None + if recent_days: + created_after = utc_now_naive() - timedelta(days=recent_days) + + normalized_start_time = normalize_naive_utc_filter_datetime(start_time) + normalized_end_time = normalize_naive_utc_filter_datetime(end_time) + + if ( + normalized_start_time + and normalized_end_time + and normalized_start_time > normalized_end_time + ): + raise ValidationException( + user_message="start_time cannot be later than end_time", + violations=[ + {"field": "start_time", "description": "Must be before end_time"} + ], + ) + + if normalized_start_time: + created_after = normalized_start_time + created_before = normalized_end_time + + total_count = await job_repo.count_jobs_by_user( + db=db, + user_id=user_id, + created_after=created_after, + created_before=created_before, + job_type=job_type, + job_status=job_status, + ) + jobs = await job_repo.get_jobs_by_user( + db=db, + user_id=user_id, + limit=page_size, + offset=(page - 1) * page_size, + created_after=created_after, + created_before=created_before, + job_type=job_type, + job_status=job_status, + ) + + redis_service = RedisServiceFactory.get_service() + job_responses = [] + for job in jobs: + job_metadata = await job_repo.get_job_metadata( + db, job.job_id, redis_service + ) + job_responses.append( + await build_job_result_response( + job=job, + job_metadata=job_metadata, + progress=None, + ) + ) + + total_pages = math.ceil(total_count / page_size) if total_count > 0 else 0 + return JobList( + jobs=job_responses, + total=total_count, + page=page, + page_size=page_size, + total_pages=total_pages, + ) + + except ValidationException: + raise + except Exception as exc: + logger.error(f"Failed to list jobs: {exc}") + raise JobOperationException( + internal_message=f"Failed to get job list: {str(exc)}" + ) + + +async def get_job_result_for_user( + db: AsyncSession, + *, + job_id: str, + user_id: str, +) -> JobResultResponse: + try: + job_repo = JobRepository() + job = await job_repo.get_job_by_id(db, job_id) + check_job_permission(job, user_id, job_id) + assert job is not None + + progress = None + if to_job_status_value(job.status) == "running": + progress = {"total_pages": 10, "processed_pages": 5} + + redis_service = RedisServiceFactory.get_service() + job_metadata = await job_repo.get_job_metadata(db, job_id, redis_service) + return await build_job_result_response( + job=job, + job_metadata=job_metadata, + progress=progress, + ) + + except NotFoundException: + raise + except PermissionDeniedException: + raise + except Exception as exc: + logger.error(f"Failed to get job result: {exc}") + raise JobOperationException( + internal_message=f"Failed to get job result: {str(exc)}" + ) diff --git a/apps/api/app/services/jobs/read_service.py b/apps/api/app/services/jobs/read_service.py index 70deacf5d..c055dd2e5 100644 --- a/apps/api/app/services/jobs/read_service.py +++ b/apps/api/app/services/jobs/read_service.py @@ -1,46 +1,19 @@ from __future__ import annotations -import math -from datetime import datetime, timedelta, timezone +from datetime import datetime from typing import Optional -from app.repositories.job_repository import JobRepository -from app.services.jobs.result_projection import ( - build_job_result_response, - to_job_status_value, -) -from loguru import logger +from app.services.jobs.job_read_model import JobReadModel +from app.services.jobs.job_read_model import check_job_permission from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.exceptions.domain_exceptions import ( - JobOperationException, - NotFoundException, - PermissionDeniedException, - ValidationException, -) from shared.models.schemas.job import JobList, JobResultResponse -from shared.services.redis import RedisServiceFactory -from shared.utils.utc_now import utc_now_naive - -def check_job_permission(job, user_id: str, job_id: str) -> None: - if not job: - raise NotFoundException( - resource="Job", resource_id=job_id, internal_message="Job not found" - ) - - if str(job.user_id) != user_id: - raise PermissionDeniedException( - user_message="You don't have permission to access this job", - ) - - -def normalize_naive_utc_filter_datetime(dt: Optional[datetime]) -> Optional[datetime]: - if dt is None: - return None - if dt.tzinfo is None or dt.utcoffset() is None: - return dt - return dt.astimezone(timezone.utc).replace(tzinfo=None) +__all__ = [ + "check_job_permission", + "get_job_result_for_user", + "list_jobs_for_user", +] async def list_jobs_for_user( @@ -55,87 +28,17 @@ async def list_jobs_for_user( start_time: Optional[datetime], end_time: Optional[datetime], ) -> JobList: - try: - job_repo = JobRepository() - - if recent_days not in (None, 1, 7, 30): - raise ValidationException( - user_message="recent_days only supports 1, 7, or 30", - violations=[{"field": "recent_days", "description": "Invalid value"}], - ) - - created_after: Optional[datetime] = None - if recent_days: - created_after = utc_now_naive() - timedelta(days=recent_days) - - normalized_start_time = normalize_naive_utc_filter_datetime(start_time) - normalized_end_time = normalize_naive_utc_filter_datetime(end_time) - - if ( - normalized_start_time - and normalized_end_time - and normalized_start_time > normalized_end_time - ): - raise ValidationException( - user_message="start_time cannot be later than end_time", - violations=[ - {"field": "start_time", "description": "Must be before end_time"} - ], - ) - - if normalized_start_time: - created_after = normalized_start_time - created_before = normalized_end_time - - total_count = await job_repo.count_jobs_by_user( - db=db, - user_id=user_id, - created_after=created_after, - created_before=created_before, - job_type=job_type, - job_status=job_status, - ) - jobs = await job_repo.get_jobs_by_user( - db=db, - user_id=user_id, - limit=page_size, - offset=(page - 1) * page_size, - created_after=created_after, - created_before=created_before, - job_type=job_type, - job_status=job_status, - ) - - redis_service = RedisServiceFactory.get_service() - job_responses = [] - for job in jobs: - job_metadata = await job_repo.get_job_metadata( - db, job.job_id, redis_service - ) - job_responses.append( - await build_job_result_response( - job=job, - job_metadata=job_metadata, - progress=None, - ) - ) - - total_pages = math.ceil(total_count / page_size) if total_count > 0 else 0 - return JobList( - jobs=job_responses, - total=total_count, - page=page, - page_size=page_size, - total_pages=total_pages, - ) - - except ValidationException: - raise - except Exception as exc: - logger.error(f"Failed to list jobs: {exc}") - raise JobOperationException( - internal_message=f"Failed to get job list: {str(exc)}" - ) + return await JobReadModel().list_jobs_for_user( + db, + user_id=user_id, + page=page, + page_size=page_size, + job_status=job_status, + job_type=job_type, + recent_days=recent_days, + start_time=start_time, + end_time=end_time, + ) async def get_job_result_for_user( @@ -144,30 +47,8 @@ async def get_job_result_for_user( job_id: str, user_id: str, ) -> JobResultResponse: - try: - job_repo = JobRepository() - job = await job_repo.get_job_by_id(db, job_id) - check_job_permission(job, user_id, job_id) - assert job is not None - - progress = None - if to_job_status_value(job.status) == "running": - progress = {"total_pages": 10, "processed_pages": 5} - - redis_service = RedisServiceFactory.get_service() - job_metadata = await job_repo.get_job_metadata(db, job_id, redis_service) - return await build_job_result_response( - job=job, - job_metadata=job_metadata, - progress=progress, - ) - - except NotFoundException: - raise - except PermissionDeniedException: - raise - except Exception as exc: - logger.error(f"Failed to get job result: {exc}") - raise JobOperationException( - internal_message=f"Failed to get job result: {str(exc)}" - ) + return await JobReadModel().get_job_result_for_user( + db, + job_id=job_id, + user_id=user_id, + ) diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py new file mode 100644 index 000000000..f09484adb --- /dev/null +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -0,0 +1,583 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any + +import pandas as pd +from app.services.connect_builder.summary_builder import ( + build_section_summary_lookup, + enrich_doc_nav_summaries, + ensure_doc_nav_json, + load_nav_top_summary, +) +from app.services.document_ingestion.job_state_gate import mark_job_running +from app.services.document_ingestion.page_estimator import PageEstimator +from app.services.document_ingestion.workspace import ( + cleanup_task_workspace, + create_task_workspace, + download_s3_file_to_temp, +) +from app.services.document_parser.stage_profiler import stage_timer +from loguru import logger +from sqlalchemy import select + +from shared.core.config import settings +from shared.core.database_sync import get_sync_db_context +from shared.core.exceptions.domain_exceptions import ( + InsufficientCreditsException, + NotFoundException, + ValidationException, + WorkerHandlingException, +) +from shared.models.database.job import Job +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.billing.work_billing_service import WorkBillingService +from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks +from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service +from shared.services.redis.distributed_lock import RedisJobLock +from shared.services.redis.redis_sync_service import ( + SyncJobInfoRedisService, + SyncJobMetadataService, + SyncRedisServiceFactory, +) +from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.storage.result_storage import get_result_storage +from shared.services.storage.zip_result_service import ZipResultService + + +@dataclass(frozen=True) +class _ParseJobContext: + job_metadata: dict[str, object] + job_user_id: str | None + metadata_service: SyncJobMetadataService + redis_service: Any + s3_key: str + + +@dataclass(frozen=True) +class _ParseJobBillingSnapshot: + billing_amount_micro_dollars: int + billing_credits: float + billing_status: str + + +class DocumentProcessingRun: + """Run worker-side Document Ingestion for an uploaded file Job.""" + + def execute(self, job_id: str, user_id: str | None) -> dict[str, object]: + logger.info(f"Parse started: job_id={job_id}, user_id={user_id}") + lifecycle_service = get_sync_job_lifecycle_service() + + redis_service = SyncRedisServiceFactory.get_service() + job_context = _load_parse_job_context(job_id, user_id, redis_service) + _assert_source_file_within_size_limit(job_context.s3_key) + + should_process = mark_job_running(job_id, job_context.redis_service) + if not should_process: + logger.warning(f"Skipping parse_task for inactive job: job_id={job_id}") + return { + "status": "skipped", + "job_id": job_id, + "reason": "job_already_terminal", + } + + with RedisJobLock(job_context.redis_service, job_id): + task_workspace_dir, input_dir, output_dir = _prepare_task_workspace(job_id) + try: + return _run_parse_job( + job_id=job_id, + job_context=job_context, + lifecycle_service=lifecycle_service, + input_dir=input_dir, + output_dir=output_dir, + task_workspace_dir=task_workspace_dir, + ) + finally: + cleanup_task_workspace(task_workspace_dir) + + raise WorkerHandlingException( + user_message="We could not complete document processing", + internal_message=( + f"Parse workflow exited without a result for job_id={job_id}" + ), + ) + + +def _load_parse_job_context( + job_id: str, + requested_user_id: str | None, + redis_service: Any, +) -> _ParseJobContext: + job_info_service = SyncJobInfoRedisService(redis_service) + job_info = job_info_service.get_job_info(job_id) + + if not job_info: + logger.warning( + f"JobInfo not found in Redis for job_id={job_id}; falling back to database" + ) + with get_sync_db_context() as fallback_db: + job_row = fallback_db.execute( + select(Job).where(Job.job_id == job_id) + ).scalar_one_or_none() + + if not job_row or not job_row.s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id=job_id, + internal_message="job info not found in Redis or database", + ) + + s3_key: str = job_row.s3_key + job_user_id: str | None = ( + str(job_row.user_id) if job_row.user_id else requested_user_id + ) + logger.info(f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}") + else: + raw_s3_key = job_info.get("s3_key") + if not isinstance(raw_s3_key, str) or not raw_s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id="s3_key", + internal_message="Missing s3_key in job_info", + ) + + s3_key = raw_s3_key + raw_job_user_id = job_info.get("user_id") + job_user_id = ( + raw_job_user_id if isinstance(raw_job_user_id, str) else requested_user_id + ) + + metadata_service = SyncJobMetadataService(redis_service) + raw_job_metadata = metadata_service.get_metadata(job_id) + if not isinstance(raw_job_metadata, dict) or not raw_job_metadata: + raise NotFoundException( + resource="JobMetadata", + resource_id=job_id, + internal_message=f"Job metadata not found for job_id={job_id}", + ) + + return _ParseJobContext( + job_metadata=dict(raw_job_metadata), + job_user_id=job_user_id, + metadata_service=metadata_service, + redis_service=redis_service, + s3_key=s3_key, + ) + + +def _assert_source_file_within_size_limit(s3_key: str) -> None: + file_info = JobFileStorage().verify_upload_exists(s3_key) + if not file_info.get("exists"): + raise NotFoundException( + resource="S3File", + resource_id=s3_key, + internal_message=f"S3 file not found: {s3_key}", + ) + + logger.info(f"S3 file verified: {s3_key}") + + file_size = file_info.get("size", 0) + file_extension = os.path.splitext(s3_key)[1].lower() + if file_size > settings.MAX_FILE_SIZE: + limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) + raise ValidationException( + user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", + violations=[ + { + "field": "file_size", + "description": ( + f"Size {file_size} bytes exceeds limit of " + f"{settings.MAX_FILE_SIZE} bytes" + ), + } + ], + ) + + +def _prepare_task_workspace(job_id: str) -> tuple[str, str, str]: + task_workspace_dir = create_task_workspace(job_id) + input_dir = os.path.join(task_workspace_dir, "input") + output_dir = os.path.join(task_workspace_dir, "output") + os.makedirs(input_dir, exist_ok=True) + os.makedirs(output_dir, exist_ok=True) + logger.info( + f"Task workspace ready: job_id={job_id}, workspace={task_workspace_dir}" + ) + return task_workspace_dir, input_dir, output_dir + + +def _run_parse_job( + *, + job_id: str, + job_context: _ParseJobContext, + lifecycle_service: Any, + input_dir: str, + output_dir: str, + task_workspace_dir: str, +) -> dict[str, object]: + lifecycle_service.update_progress(job_id, progress=10, message="Parsing document...") + + filename = JobMetadataHelper.get_field(job_context.job_metadata, "source_file_name") + file_ext = os.path.splitext(job_context.s3_key)[1].lower() if job_context.s3_key else "" + local_temp_path = download_s3_file_to_temp(job_context.s3_key, file_ext, input_dir) + logger.info(f"File downloaded: job_id={job_id}, local_path={local_temp_path}") + + from app.services.document_parser.internal_parse_name import ( + prepare_internal_parse_input, + ) + from app.services.document_parser import parse_service + + prepared_parse_input = prepare_internal_parse_input( + local_temp_path, + filename, + fallback_ext=file_ext, + prefer_fallback_ext=True, + ) + internal_parse_name = prepared_parse_input.internal_filename + local_temp_path = prepared_parse_input.file_path + logger.info( + f"File prepared for parsing: job_id={job_id}, " + f"internal_filename={internal_parse_name}, local_path={local_temp_path}" + ) + + page_count = PageEstimator.estimate(local_temp_path) + logger.info(f"Workload estimation: job_id={job_id}, page_count={page_count}") + + processing_started_at = datetime.now(timezone.utc) + billing_snapshot = _charge_parse_job_pages( + job_id=job_id, + filename=filename, + job_user_id=job_context.job_user_id, + page_count=page_count, + ) + _record_processing_start( + job_id=job_id, + job_context=job_context, + billing_snapshot=billing_snapshot, + page_count=page_count, + processing_started_at=processing_started_at, + ) + + doc_type = JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "doc_type", + "auto", + ) + logger.info( + f"Start parse: job_id={job_id}, filename={filename}, " + f"internal_filename={internal_parse_name}, type={doc_type}" + ) + + with stage_timer( + "worker.parse.document", + job_id=job_id, + filename=filename, + doc_type=doc_type, + ): + add_dir, parsed_contents_df = parse_service.checkerboard_inject_parse( + file_full_path=local_temp_path, + filename=filename, + output_dir=output_dir, + job_id=job_id, + internal_output_filename=internal_parse_name, + kb_dir=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "kb_dir", + "Default_Root", + ), + doc_type=doc_type, + smart_title_parse=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "smart_title_parse", + True, + ), + summary_image=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_image", + True, + ), + summary_table=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_table", + True, + ), + summary_txt=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_txt", + True, + ), + add_frag_desc=JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "add_frag_desc", + "", + ), + s3_key=job_context.s3_key, + ) + + logger.info( + "File parsing completed: " + f"job_id={job_id}, add_dir={add_dir}, " + f"chunks={len(parsed_contents_df) if parsed_contents_df is not None else 0}" + ) + + if parsed_contents_df is None: + raise WorkerHandlingException( + user_message="We could not extract content from your file", + internal_message="File parsing failed, no content returned from parser", + ) + + if parsed_contents_df.empty: + logger.warning( + f"No content returned from file parsing: job_id={job_id}, filename={filename}" + ) + + lifecycle_service.update_progress( + job_id, + progress=30, + message="Parse completed, preparing chunks...", + ) + chunks = dataframe_to_chunks(parsed_contents_df) + + lifecycle_service.update_progress( + job_id, + progress=70, + message="Chunks ready, generating zip...", + ) + logger.info(f"Chunks prepared: job_id={job_id}, count={len(chunks)}") + + return _finalize_parse_job_success( + add_dir=add_dir, + chunks=chunks, + job_context=job_context, + job_id=job_id, + lifecycle_service=lifecycle_service, + parsed_contents_df=parsed_contents_df, + processing_started_at=processing_started_at, + task_workspace_dir=task_workspace_dir, + ) + + +def _charge_parse_job_pages( + *, + job_id: str, + filename: str | None, + job_user_id: str | None, + page_count: int, +) -> _ParseJobBillingSnapshot: + if not job_user_id: + raise NotFoundException( + resource="JobInfo", + resource_id="user_id", + internal_message=f"Missing user_id in job info for job_id={job_id}", + ) + + billing_service = WorkBillingService() + billing_filename = filename or "" + billing_status = "skipped" + billing_amount_micro_dollars = 0 + billing_credits = 0.0 + + with get_sync_db_context() as db: + job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()) + job = job_result.scalar_one_or_none() + + if job and getattr(job, "billing_status", "") == "charged": + logger.info(f"Job already charged: {job_id}") + billing_status = "charged" + billing_amount_micro_dollars = int(job.credits_charged or 0) + billing_credits = billing_amount_micro_dollars / 1_000_000 + else: + try: + billing_result = billing_service.charge_for_pages( + session=db, + user_id=job_user_id, + page_count=page_count, + filename=billing_filename, + ) + except InsufficientCreditsException: + logger.warning(f"Billing failed: job_id={job_id}, user_id={job_user_id}") + billing_amount = billing_service.estimate_page_charge( + page_count=page_count + ) + if job: + job.page_count = page_count + job.credits_charged = billing_amount.amount_micro_dollars + job.billing_status = "billing_failed" + db.commit() + + raise InsufficientCreditsException( + user_message=( + "Insufficient credits to process this document " + f"({page_count} pages required, cost: " + f"{billing_amount.credits})." + ), + required_credits=billing_amount.credits, + internal_message=( + f"job_id={job_id}, user_id={job_user_id}, " + f"page_count={page_count}" + ), + ) + + billing_status = billing_result.billing_status + billing_amount_micro_dollars = billing_result.amount_micro_dollars + billing_credits = billing_result.credits + if job: + job.page_count = page_count + job.credits_charged = billing_amount_micro_dollars + job.billing_status = billing_status + + return _ParseJobBillingSnapshot( + billing_amount_micro_dollars=billing_amount_micro_dollars, + billing_credits=billing_credits, + billing_status=billing_status, + ) + + +def _record_processing_start( + *, + job_id: str, + job_context: _ParseJobContext, + billing_snapshot: _ParseJobBillingSnapshot, + page_count: int, + processing_started_at: datetime, +) -> None: + metadata_updates = { + "page_count": page_count, + "billing_status": billing_snapshot.billing_status, + "billing_amount_micro_dollars": billing_snapshot.billing_amount_micro_dollars, + "billing_credits": billing_snapshot.billing_credits, + "processing_started_at": processing_started_at.isoformat(), + } + job_context.metadata_service.update_metadata(job_id, metadata_updates) + job_context.job_metadata.update(metadata_updates) + + +def _finalize_parse_job_success( + *, + add_dir: str, + chunks: list[dict[str, Any]], + job_context: _ParseJobContext, + job_id: str, + lifecycle_service: Any, + parsed_contents_df: pd.DataFrame, + processing_started_at: datetime, + task_workspace_dir: str, +) -> dict[str, object]: + source_file_name = JobMetadataHelper.get_field( + job_context.job_metadata, + "source_file_name", + ) or JobMetadataHelper.get_field(job_context.job_metadata, "source_url") + if isinstance(source_file_name, str) and "/" in source_file_name: + source_file_name = os.path.basename(source_file_name) + + document_top_summary = "" + section_summaries: dict[str, str] = {} + if add_dir and source_file_name: + if "path" in parsed_contents_df.columns: + ensure_doc_nav_json( + str(add_dir), + chunks, + source_file_name=str(source_file_name), + ) + try: + kb_dir_for_enrich = os.path.dirname(str(add_dir)) + summary_use_llm = JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "summary_use_llm", + False, + ) + enrich_doc_nav_summaries( + kb_dir_for_enrich, + source_file=str(source_file_name), + use_llm=summary_use_llm, + ) + section_summaries = build_section_summary_lookup(str(add_dir)) + except Exception as exc: + logger.warning(f"doc_nav enrichment failed (non-fatal): {exc}") + document_top_summary = load_nav_top_summary(str(add_dir), str(source_file_name)) + + if document_top_summary: + for chunk in chunks: + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + chunk["metadata"] = metadata + metadata["document_top_summary"] = document_top_summary + + lifecycle_service.update_progress( + job_id, + progress=80, + message="Generating ZIP package...", + ) + processing_completed_at = datetime.now(timezone.utc) + processing_timing_updates = { + "processing_completed_at": processing_completed_at.isoformat(), + "processing_duration_ms": max( + 0, + int((processing_completed_at - processing_started_at).total_seconds() * 1000), + ), + } + job_context.metadata_service.update_metadata(job_id, processing_timing_updates) + job_context.job_metadata.update(processing_timing_updates) + + data_id = JobMetadataHelper.get_field(job_context.job_metadata, "data_id") + zip_service = ZipResultService() + zip_file_path, checksum, statistics, zip_size = zip_service.generate_zip_package( + job_id=job_id, + chunks=chunks, + add_dir=str(add_dir) if add_dir else "", + source_file_name=source_file_name, + data_id=data_id, + job_metadata=job_context.job_metadata, + parsed_df=parsed_contents_df, + temp_dir=task_workspace_dir, + ) + del statistics + + checksum_value = ( + checksum.get("value", "") + if isinstance(checksum, dict) + else (checksum or "") + ) + + lifecycle_service.update_progress( + job_id, + progress=90, + message="Uploading results to S3...", + ) + result_bundle = get_result_storage().upload( + job_id=job_id, + result_dir=str(add_dir) if add_dir else "", + zip_file_path=zip_file_path, + ) + result_s3_key = result_bundle.zip_key + stored_count = 0 + + lifecycle_service.update_progress(job_id, progress=100, message="Task complete!") + lifecycle_service.finalize_job_success( + job_id=job_id, + chunks=chunks, + result_s3_key=result_s3_key, + checksum=checksum_value, + zip_size=zip_size, + stored_count=stored_count, + delivery_mode="url", + section_summaries=section_summaries, + ) + + logger.info( + f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" + ) + + return { + "status": "success", + "job_id": job_id, + "add_dir": None, + "vectors_count": 0, + "contents_count": len(parsed_contents_df), + "stored_count": stored_count, + "delivery_mode": "url", + "result_s3_key": result_s3_key, + } diff --git a/apps/worker/app/services/document_ingestion/service.py b/apps/worker/app/services/document_ingestion/service.py index f6edb2878..cabc0e451 100644 --- a/apps/worker/app/services/document_ingestion/service.py +++ b/apps/worker/app/services/document_ingestion/service.py @@ -1,579 +1,30 @@ from __future__ import annotations -import os -from dataclasses import dataclass -from datetime import datetime, timezone -from typing import Any - -import pandas as pd -from app.services.connect_builder.summary_builder import ( - build_section_summary_lookup, - enrich_doc_nav_summaries, - ensure_doc_nav_json, - load_nav_top_summary, -) -from app.services.document_ingestion.job_state_gate import mark_job_running -from app.services.document_ingestion.page_estimator import PageEstimator -from app.services.document_ingestion.workspace import ( - cleanup_task_workspace, - create_task_workspace, - download_s3_file_to_temp, -) -from app.services.document_parser.stage_profiler import stage_timer -from loguru import logger -from sqlalchemy import select - -from shared.core.config import settings -from shared.core.database_sync import get_sync_db_context -from shared.core.exceptions.domain_exceptions import ( - InsufficientCreditsException, - NotFoundException, - ValidationException, - WorkerHandlingException, -) -from shared.models.database.job import Job -from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.billing.work_billing_service import WorkBillingService -from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks -from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service -from shared.services.redis.distributed_lock import RedisJobLock -from shared.services.redis.redis_sync_service import ( - SyncJobInfoRedisService, - SyncJobMetadataService, - SyncRedisServiceFactory, -) -from shared.services.storage.job_file_storage import JobFileStorage -from shared.services.storage.result_storage import get_result_storage -from shared.services.storage.zip_result_service import ZipResultService - - -@dataclass(frozen=True) -class _ParseJobContext: - job_metadata: dict[str, object] - job_user_id: str | None - metadata_service: SyncJobMetadataService - redis_service: Any - s3_key: str - - -@dataclass(frozen=True) -class _ParseJobBillingSnapshot: - billing_amount_micro_dollars: int - billing_credits: float - billing_status: str +import app.services.document_ingestion.processing_run as processing_run +from app.services.document_ingestion.processing_run import DocumentProcessingRun +from app.services.document_ingestion.processing_run import PageEstimator +from app.services.document_ingestion.processing_run import cleanup_task_workspace +from app.services.document_ingestion.processing_run import download_s3_file_to_temp +from app.services.document_ingestion.processing_run import get_result_storage +from app.services.document_ingestion.processing_run import settings + +__all__ = [ + "PageEstimator", + "cleanup_task_workspace", + "download_s3_file_to_temp", + "get_result_storage", + "parse_uploaded_file_job", + "settings", +] def parse_uploaded_file_job(job_id: str, user_id: str | None) -> dict[str, object]: """Run worker-side Document Ingestion for an uploaded file Job.""" - logger.info(f"Parse started: job_id={job_id}, user_id={user_id}") - lifecycle_service = get_sync_job_lifecycle_service() - - redis_service = SyncRedisServiceFactory.get_service() - job_context = _load_parse_job_context(job_id, user_id, redis_service) - _assert_source_file_within_size_limit(job_context.s3_key) - - should_process = mark_job_running(job_id, job_context.redis_service) - if not should_process: - logger.warning(f"Skipping parse_task for inactive job: job_id={job_id}") - return { - "status": "skipped", - "job_id": job_id, - "reason": "job_already_terminal", - } - - with RedisJobLock(job_context.redis_service, job_id): - task_workspace_dir, input_dir, output_dir = _prepare_task_workspace(job_id) - try: - return _run_parse_job( - job_id=job_id, - job_context=job_context, - lifecycle_service=lifecycle_service, - input_dir=input_dir, - output_dir=output_dir, - task_workspace_dir=task_workspace_dir, - ) - finally: - cleanup_task_workspace(task_workspace_dir) - - raise WorkerHandlingException( - user_message="We could not complete document processing", - internal_message=f"Parse workflow exited without a result for job_id={job_id}", - ) - - -def _load_parse_job_context( - job_id: str, - requested_user_id: str | None, - redis_service: Any, -) -> _ParseJobContext: - job_info_service = SyncJobInfoRedisService(redis_service) - job_info = job_info_service.get_job_info(job_id) - - if not job_info: - logger.warning( - f"JobInfo not found in Redis for job_id={job_id}; falling back to database" - ) - with get_sync_db_context() as fallback_db: - job_row = fallback_db.execute( - select(Job).where(Job.job_id == job_id) - ).scalar_one_or_none() - - if not job_row or not job_row.s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id=job_id, - internal_message="job info not found in Redis or database", - ) - - s3_key: str = job_row.s3_key - job_user_id: str | None = ( - str(job_row.user_id) if job_row.user_id else requested_user_id - ) - logger.info(f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}") - else: - raw_s3_key = job_info.get("s3_key") - if not isinstance(raw_s3_key, str) or not raw_s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id="s3_key", - internal_message="Missing s3_key in job_info", - ) - - s3_key = raw_s3_key - raw_job_user_id = job_info.get("user_id") - job_user_id = ( - raw_job_user_id if isinstance(raw_job_user_id, str) else requested_user_id - ) - - metadata_service = SyncJobMetadataService(redis_service) - raw_job_metadata = metadata_service.get_metadata(job_id) - if not isinstance(raw_job_metadata, dict) or not raw_job_metadata: - raise NotFoundException( - resource="JobMetadata", - resource_id=job_id, - internal_message=f"Job metadata not found for job_id={job_id}", - ) - - return _ParseJobContext( - job_metadata=dict(raw_job_metadata), - job_user_id=job_user_id, - metadata_service=metadata_service, - redis_service=redis_service, - s3_key=s3_key, - ) - - -def _assert_source_file_within_size_limit(s3_key: str) -> None: - file_info = JobFileStorage().verify_upload_exists(s3_key) - if not file_info.get("exists"): - raise NotFoundException( - resource="S3File", - resource_id=s3_key, - internal_message=f"S3 file not found: {s3_key}", - ) - - logger.info(f"S3 file verified: {s3_key}") - - file_size = file_info.get("size", 0) - file_extension = os.path.splitext(s3_key)[1].lower() - if file_size > settings.MAX_FILE_SIZE: - limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) - raise ValidationException( - user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", - violations=[ - { - "field": "file_size", - "description": ( - f"Size {file_size} bytes exceeds limit of " - f"{settings.MAX_FILE_SIZE} bytes" - ), - } - ], - ) - - -def _prepare_task_workspace(job_id: str) -> tuple[str, str, str]: - task_workspace_dir = create_task_workspace(job_id) - input_dir = os.path.join(task_workspace_dir, "input") - output_dir = os.path.join(task_workspace_dir, "output") - os.makedirs(input_dir, exist_ok=True) - os.makedirs(output_dir, exist_ok=True) - logger.info( - f"Task workspace ready: job_id={job_id}, workspace={task_workspace_dir}" - ) - return task_workspace_dir, input_dir, output_dir - - -def _run_parse_job( - *, - job_id: str, - job_context: _ParseJobContext, - lifecycle_service: Any, - input_dir: str, - output_dir: str, - task_workspace_dir: str, -) -> dict[str, object]: - lifecycle_service.update_progress(job_id, progress=10, message="Parsing document...") - - filename = JobMetadataHelper.get_field(job_context.job_metadata, "source_file_name") - file_ext = os.path.splitext(job_context.s3_key)[1].lower() if job_context.s3_key else "" - local_temp_path = download_s3_file_to_temp(job_context.s3_key, file_ext, input_dir) - logger.info(f"File downloaded: job_id={job_id}, local_path={local_temp_path}") - - from app.services.document_parser.internal_parse_name import ( - prepare_internal_parse_input, - ) - from app.services.document_parser import parse_service - - prepared_parse_input = prepare_internal_parse_input( - local_temp_path, - filename, - fallback_ext=file_ext, - prefer_fallback_ext=True, - ) - internal_parse_name = prepared_parse_input.internal_filename - local_temp_path = prepared_parse_input.file_path - logger.info( - f"File prepared for parsing: job_id={job_id}, " - f"internal_filename={internal_parse_name}, local_path={local_temp_path}" - ) - - page_count = PageEstimator.estimate(local_temp_path) - logger.info(f"Workload estimation: job_id={job_id}, page_count={page_count}") - - processing_started_at = datetime.now(timezone.utc) - billing_snapshot = _charge_parse_job_pages( - job_id=job_id, - filename=filename, - job_user_id=job_context.job_user_id, - page_count=page_count, - ) - _record_processing_start( - job_id=job_id, - job_context=job_context, - billing_snapshot=billing_snapshot, - page_count=page_count, - processing_started_at=processing_started_at, - ) - - doc_type = JobMetadataHelper.get_parsing_param( - job_context.job_metadata, - "doc_type", - "auto", - ) - logger.info( - f"Start parse: job_id={job_id}, filename={filename}, " - f"internal_filename={internal_parse_name}, type={doc_type}" - ) - - with stage_timer( - "worker.parse.document", - job_id=job_id, - filename=filename, - doc_type=doc_type, - ): - add_dir, parsed_contents_df = parse_service.checkerboard_inject_parse( - file_full_path=local_temp_path, - filename=filename, - output_dir=output_dir, - job_id=job_id, - internal_output_filename=internal_parse_name, - kb_dir=JobMetadataHelper.get_parsing_param( - job_context.job_metadata, - "kb_dir", - "Default_Root", - ), - doc_type=doc_type, - smart_title_parse=JobMetadataHelper.get_parsing_param( - job_context.job_metadata, - "smart_title_parse", - True, - ), - summary_image=JobMetadataHelper.get_parsing_param( - job_context.job_metadata, - "summary_image", - True, - ), - summary_table=JobMetadataHelper.get_parsing_param( - job_context.job_metadata, - "summary_table", - True, - ), - summary_txt=JobMetadataHelper.get_parsing_param( - job_context.job_metadata, - "summary_txt", - True, - ), - add_frag_desc=JobMetadataHelper.get_parsing_param( - job_context.job_metadata, - "add_frag_desc", - "", - ), - s3_key=job_context.s3_key, - ) - - logger.info( - "File parsing completed: " - f"job_id={job_id}, add_dir={add_dir}, " - f"chunks={len(parsed_contents_df) if parsed_contents_df is not None else 0}" - ) - - if parsed_contents_df is None: - raise WorkerHandlingException( - user_message="We could not extract content from your file", - internal_message="File parsing failed, no content returned from parser", - ) - - if parsed_contents_df.empty: - logger.warning( - f"No content returned from file parsing: job_id={job_id}, filename={filename}" - ) - - lifecycle_service.update_progress( - job_id, - progress=30, - message="Parse completed, preparing chunks...", - ) - chunks = dataframe_to_chunks(parsed_contents_df) - - lifecycle_service.update_progress( - job_id, - progress=70, - message="Chunks ready, generating zip...", - ) - logger.info(f"Chunks prepared: job_id={job_id}, count={len(chunks)}") - - return _finalize_parse_job_success( - add_dir=add_dir, - chunks=chunks, - job_context=job_context, - job_id=job_id, - lifecycle_service=lifecycle_service, - parsed_contents_df=parsed_contents_df, - processing_started_at=processing_started_at, - task_workspace_dir=task_workspace_dir, - ) - - -def _charge_parse_job_pages( - *, - job_id: str, - filename: str | None, - job_user_id: str | None, - page_count: int, -) -> _ParseJobBillingSnapshot: - if not job_user_id: - raise NotFoundException( - resource="JobInfo", - resource_id="user_id", - internal_message=f"Missing user_id in job info for job_id={job_id}", - ) - - billing_service = WorkBillingService() - billing_filename = filename or "" - billing_status = "skipped" - billing_amount_micro_dollars = 0 - billing_credits = 0.0 - - with get_sync_db_context() as db: - job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()) - job = job_result.scalar_one_or_none() - - if job and getattr(job, "billing_status", "") == "charged": - logger.info(f"Job already charged: {job_id}") - billing_status = "charged" - billing_amount_micro_dollars = int(job.credits_charged or 0) - billing_credits = billing_amount_micro_dollars / 1_000_000 - else: - try: - billing_result = billing_service.charge_for_pages( - session=db, - user_id=job_user_id, - page_count=page_count, - filename=billing_filename, - ) - except InsufficientCreditsException: - logger.warning(f"Billing failed: job_id={job_id}, user_id={job_user_id}") - billing_amount = billing_service.estimate_page_charge( - page_count=page_count - ) - if job: - job.page_count = page_count - job.credits_charged = billing_amount.amount_micro_dollars - job.billing_status = "billing_failed" - db.commit() - - raise InsufficientCreditsException( - user_message=( - "Insufficient credits to process this document " - f"({page_count} pages required, cost: " - f"{billing_amount.credits})." - ), - required_credits=billing_amount.credits, - internal_message=( - f"job_id={job_id}, user_id={job_user_id}, " - f"page_count={page_count}" - ), - ) - - billing_status = billing_result.billing_status - billing_amount_micro_dollars = billing_result.amount_micro_dollars - billing_credits = billing_result.credits - if job: - job.page_count = page_count - job.credits_charged = billing_amount_micro_dollars - job.billing_status = billing_status - - return _ParseJobBillingSnapshot( - billing_amount_micro_dollars=billing_amount_micro_dollars, - billing_credits=billing_credits, - billing_status=billing_status, - ) - - -def _record_processing_start( - *, - job_id: str, - job_context: _ParseJobContext, - billing_snapshot: _ParseJobBillingSnapshot, - page_count: int, - processing_started_at: datetime, -) -> None: - metadata_updates = { - "page_count": page_count, - "billing_status": billing_snapshot.billing_status, - "billing_amount_micro_dollars": billing_snapshot.billing_amount_micro_dollars, - "billing_credits": billing_snapshot.billing_credits, - "processing_started_at": processing_started_at.isoformat(), - } - job_context.metadata_service.update_metadata(job_id, metadata_updates) - job_context.job_metadata.update(metadata_updates) - - -def _finalize_parse_job_success( - *, - add_dir: str, - chunks: list[dict[str, Any]], - job_context: _ParseJobContext, - job_id: str, - lifecycle_service: Any, - parsed_contents_df: pd.DataFrame, - processing_started_at: datetime, - task_workspace_dir: str, -) -> dict[str, object]: - source_file_name = JobMetadataHelper.get_field( - job_context.job_metadata, - "source_file_name", - ) or JobMetadataHelper.get_field(job_context.job_metadata, "source_url") - if isinstance(source_file_name, str) and "/" in source_file_name: - source_file_name = os.path.basename(source_file_name) - - document_top_summary = "" - section_summaries: dict[str, str] = {} - if add_dir and source_file_name: - if "path" in parsed_contents_df.columns: - ensure_doc_nav_json( - str(add_dir), - chunks, - source_file_name=str(source_file_name), - ) - try: - kb_dir_for_enrich = os.path.dirname(str(add_dir)) - summary_use_llm = JobMetadataHelper.get_parsing_param( - job_context.job_metadata, - "summary_use_llm", - False, - ) - enrich_doc_nav_summaries( - kb_dir_for_enrich, - source_file=str(source_file_name), - use_llm=summary_use_llm, - ) - section_summaries = build_section_summary_lookup(str(add_dir)) - except Exception as exc: - logger.warning(f"doc_nav enrichment failed (non-fatal): {exc}") - document_top_summary = load_nav_top_summary(str(add_dir), str(source_file_name)) - - if document_top_summary: - for chunk in chunks: - metadata = chunk.get("metadata") - if not isinstance(metadata, dict): - metadata = {} - chunk["metadata"] = metadata - metadata["document_top_summary"] = document_top_summary - - lifecycle_service.update_progress( - job_id, - progress=80, - message="Generating ZIP package...", - ) - processing_completed_at = datetime.now(timezone.utc) - processing_timing_updates = { - "processing_completed_at": processing_completed_at.isoformat(), - "processing_duration_ms": max( - 0, - int((processing_completed_at - processing_started_at).total_seconds() * 1000), - ), - } - job_context.metadata_service.update_metadata(job_id, processing_timing_updates) - job_context.job_metadata.update(processing_timing_updates) - - data_id = JobMetadataHelper.get_field(job_context.job_metadata, "data_id") - zip_service = ZipResultService() - zip_file_path, checksum, statistics, zip_size = zip_service.generate_zip_package( - job_id=job_id, - chunks=chunks, - add_dir=str(add_dir) if add_dir else "", - source_file_name=source_file_name, - data_id=data_id, - job_metadata=job_context.job_metadata, - parsed_df=parsed_contents_df, - temp_dir=task_workspace_dir, - ) - del statistics - - checksum_value = ( - checksum.get("value", "") - if isinstance(checksum, dict) - else (checksum or "") - ) - - lifecycle_service.update_progress( - job_id, - progress=90, - message="Uploading results to S3...", - ) - result_bundle = get_result_storage().upload( - job_id=job_id, - result_dir=str(add_dir) if add_dir else "", - zip_file_path=zip_file_path, - ) - result_s3_key = result_bundle.zip_key - stored_count = 0 - - lifecycle_service.update_progress(job_id, progress=100, message="Task complete!") - lifecycle_service.finalize_job_success( - job_id=job_id, - chunks=chunks, - result_s3_key=result_s3_key, - checksum=checksum_value, - zip_size=zip_size, - stored_count=stored_count, - delivery_mode="url", - section_summaries=section_summaries, - ) + _sync_legacy_adapter_overrides() + return DocumentProcessingRun().execute(job_id, user_id) - logger.info( - f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" - ) - return { - "status": "success", - "job_id": job_id, - "add_dir": None, - "vectors_count": 0, - "contents_count": len(parsed_contents_df), - "stored_count": stored_count, - "delivery_mode": "url", - "result_s3_key": result_s3_key, - } +def _sync_legacy_adapter_overrides() -> None: + processing_run.download_s3_file_to_temp = download_s3_file_to_temp + processing_run.get_result_storage = get_result_storage + processing_run.cleanup_task_workspace = cleanup_task_workspace diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 76b70c8c5..91959a9e4 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -1,70 +1,23 @@ from __future__ import annotations -import os -import time from typing import Any -from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.cache_service import get_cached_retrieval_query_result, set_cached_retrieval_query_result -from shared.services.retrieval.channels import path_channel, content_channel, term_channel -from shared.services.retrieval.graph_service import GraphQueryService -from shared.services.retrieval.hit_stats_recorder import schedule_retrieval_hit_stats_update -from shared.services.retrieval.hydration import ( - assemble_retrieval_results, - hydrate_referenced_chunk_rows, -) -from shared.services.retrieval.response_projection import ( - attach_citation, - enrich_referenced_chunks_with_asset_urls, - project_public_retrieval_response, -) -from shared.services.retrieval.scoring import ( - get_row_path, - merge_channels_rrf, - merge_same_section_rows, - normalize_row_scores, -) -from shared.services.retrieval.ranking import rank_retrieval_candidates -from shared.services.retrieval.scoped_corpus import count_scoped_chunks, load_all_scoped_chunks -from shared.services.retrieval.settings import ( - CHANNEL_WEIGHT_CONTENT as _CHANNEL_WEIGHT_CONTENT, - CHANNEL_WEIGHT_PATH as _CHANNEL_WEIGHT_PATH, - CHANNEL_WEIGHT_TERM as _CHANNEL_WEIGHT_TERM, - INTERNAL_RECALL_K_MULTIPLIER as _INTERNAL_RECALL_K_MULTIPLIER, - resolve_allowed_chunk_types as _resolve_allowed_chunk_types, -) +import shared.services.retrieval.execution_plan as execution_plan +from shared.services.retrieval.channels import content_channel, path_channel, term_channel +from shared.services.retrieval.execution_plan import RetrievalExecutionPlan +from shared.services.retrieval.execution_plan import list_graph_routed_chunks +from shared.services.retrieval.scoring import merge_channels_rrf - -async def list_graph_routed_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - service = GraphQueryService() - entry_document_ids = await service.find_entry_documents( - db, - user_id=user_id, - namespace=namespace, - query=query, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - return await service.collect_candidate_chunks( - db, - user_id=user_id, - namespace=namespace, - entry_document_ids=entry_document_ids, - query=query, - top_k=top_k * _INTERNAL_RECALL_K_MULTIPLIER, - exclude_sections=exclude_sections, - ) +__all__ = [ + "content_channel", + "list_graph_routed_chunks", + "merge_channels_rrf", + "path_channel", + "run_retrieval_query", + "term_channel", +] async def run_retrieval_query( @@ -78,7 +31,7 @@ async def run_retrieval_query( exclude_sections: list[dict[str, str]], data_type: int = 1, signal_paths: list[str] | None = None, - filter_mode: str = 'delete', + filter_mode: str = "delete", channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, rerank: bool = False, @@ -86,421 +39,30 @@ async def run_retrieval_query( internal_recall_k: int | None = None, use_agentic: bool | None = None, ) -> dict[str, Any]: - """Checkerboard retrieval: 3 independent channels -> RRF -> agent/graph union -> assembly.""" - t_start = time.monotonic() - query = query.strip() - logger.info('\n' + '█' * 70) - logger.info(' 🚀 RETRIEVAL PIPELINE START') - logger.info(f' query="{query}"') - logger.info(f' user={user_id} ns={namespace} top_k={top_k} data_type={data_type}') - logger.info(f' exclude_docs={exclude_document_ids} exclude_secs={len(exclude_sections)}') - logger.info('█' * 70) - - if not query: - logger.info(' ⛔ Empty query filtered, skipping retrieval pipeline') - return { - "namespace": namespace, - "query": query, - "router_used": "empty_query_filtered", - "results": [], - } - - allowed_chunk_types = _resolve_allowed_chunk_types(data_type) - effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER - logger.info(f' allowed_chunk_types={allowed_chunk_types} effective_recall_k={effective_recall_k} signal_paths={signal_paths} filter_mode={filter_mode} rerank={rerank} threshold={threshold}') - - cache_extra = dict( - data_type=data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - rerank=rerank, - threshold=threshold, - internal_recall_k=internal_recall_k, - # Always True: agentic mode now always routes through workflow - decomposition_enabled=True, - ) - - cache_version: int | None = None - try: - cache_version, cached = await get_cached_retrieval_query_result( - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - **cache_extra, - ) - if cached: - logger.info(f'retrieval: cache_hit=True version={cache_version}') - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, - namespace=namespace, - results=cached.get("results", []), - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - return await project_public_retrieval_response(cached) - except Exception as e: - logger.warning(f"Failed to read retrieval cache (ignored): {e}") - - logger.debug(f' 📦 Cache miss (version={cache_version}), running full pipeline') - - # ── Small KB optimization ── - try: - total_chunk_count = await count_scoped_chunks( - db, user_id=user_id, namespace=namespace, - exclude_document_ids=exclude_document_ids, - allowed_chunk_types=allowed_chunk_types, - ) - except Exception as e: - logger.warning(f"Failed to count scoped chunks, skipping small KB optimization: {e}") - total_chunk_count = top_k + 1 - logger.info(f'\n 📊 Total chunks in scope: {total_chunk_count}') - if total_chunk_count <= top_k: - logger.info(f' ⚡ Small KB optimization: {total_chunk_count} chunks <= top_k={top_k}, returning all') - all_rows = await load_all_scoped_chunks( - db, user_id=user_id, namespace=namespace, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths or [], - filter_mode=filter_mode, - ) - logger.info(f' small_kb load: loaded={len(all_rows)} rows after signal/exclude filters') - assembled_rows = await assemble_retrieval_results( - db=db, rows=all_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - ) - results = [attach_citation(row) for row in assembled_rows] - response = { - "namespace": namespace, "query": query, - "router_used": "small_kb_all", "results": results, - } - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, namespace=namespace, version=cache_version, - query=query, top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") - try: - schedule_retrieval_hit_stats_update(user_id=user_id, namespace=namespace, results=results) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info(f' ✅ Small KB: {len(results)} results in {elapsed_total}ms') - return await project_public_retrieval_response(response) - - # ══ Route: agentic (unified workflow) vs legacy ══ - if use_agentic is not None: - _agentic_enabled = use_agentic - else: - _agentic_enabled = os.environ.get('RETRIEVAL_AGENTIC_ENABLED', 'true') == 'true' - if _agentic_enabled: - # ── Unified agentic path via WorkflowOrchestrator ── - # Simple queries: planner returns a single-step plan (no decomposition). - # Complex queries: planner returns a multi-step plan with synthesize. - # Both go through the same code path. - from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator - - workflow = WorkflowOrchestrator() - workflow_result = await workflow.run( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - ) - - enriched_refs = await enrich_referenced_chunks_with_asset_urls( - workflow_result.referenced_chunks, - ) - - workflow_result_rows = await hydrate_referenced_chunk_rows( - db=db, - user_id=user_id, - namespace=namespace, - refs=enriched_refs, - ) - scoped_reference_keys = { - ( - str(row.get('document_id') or '').strip(), - str(row.get('chunk_id') or '').strip(), - ) - for row in workflow_result_rows - } - enriched_refs = [ - ref for ref in enriched_refs - if ( - str(ref.get('document_id') or '').strip(), - str(ref.get('chunk_id') or '').strip(), - ) in scoped_reference_keys - ] - assembled_workflow_rows = await assemble_retrieval_results( - db=db, - rows=workflow_result_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - ) - response = workflow_result.to_api_response() - # Override referenced_chunks with enriched versions - response['referenced_chunks'] = enriched_refs - response['results'] = [attach_citation(row) for row in assembled_workflow_rows] - - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, namespace=namespace, version=cache_version, - query=query, top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") - - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, namespace=namespace, - results=enriched_refs, - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info( - f'\n{"█" * 70}\n' - f' ✅ AGENTIC RETRIEVAL COMPLETE: ' - f'{len(enriched_refs)} chunks | ' - f'answer={len(workflow_result.answer_text)} chars | ' - f'router={workflow_result.router_used} | {elapsed_total}ms\n' - f'{"█" * 70}' - ) - - return await project_public_retrieval_response(response) - - else: - - # ── LEGACY path (existing code, unchanged) ── - - # ── Channel execution ── - active_channels = set(channels) if channels else {'path', 'content', 'term'} - logger.info(f'\n 📡 PHASE 1: Bottom-Layer Discovery (channels={sorted(active_channels)})') - logger.info(f' effective_recall_k={effective_recall_k}') - - path_rows: list[dict[str, Any]] = [] - content_rows: list[dict[str, Any]] = [] - term_rows: list[dict[str, Any]] = [] - - if 'path' in active_channels: - t_ch = time.monotonic() - path_rows = await path_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 path_channel: {len(path_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(path_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}') - if len(path_rows) > 5: - logger.info(f' ... and {len(path_rows) - 5} more') - - if 'content' in active_channels: - t_ch = time.monotonic() - content_rows = await content_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 content_channel: {len(content_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(content_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} content={str(r.get("content",""))[:80]}') - if len(content_rows) > 5: - logger.info(f' ... and {len(content_rows) - 5} more') - - if 'term' in active_channels: - t_ch = time.monotonic() - term_rows = await term_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 term_channel: {len(term_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(term_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}') - if len(term_rows) > 5: - logger.info(f' ... and {len(term_rows) - 5} more') - - # ── RRF fusion with configurable weights ── - default_weights = { - 'path': _CHANNEL_WEIGHT_PATH, - 'content': _CHANNEL_WEIGHT_CONTENT, - 'term': _CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM)) - - fused_rows = merge_channels_rrf(channel_lists, weight_list, effective_recall_k) if channel_lists else [] - logger.info(f'\n 🔀 RRF Fusion: {len(fused_rows)} rows from {len(channel_lists)} channels (weights={dict(zip(["path","content","term"][:len(weight_list)], weight_list))})') - for i, r in enumerate(fused_rows[:5]): - logger.info(f' [{i}] rrf_score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")}') - if len(fused_rows) > 5: - logger.info(f' ... and {len(fused_rows) - 5} more') - - # ── Section merging ── - pre_merge = len(fused_rows) - fused_rows = merge_same_section_rows(fused_rows) - if len(fused_rows) != pre_merge: - logger.info(f'retrieval: section_merge={pre_merge}->{len(fused_rows)}') - - # ── Threshold filtering ── - if threshold > 0.0 and fused_rows: - pre_count = len(fused_rows) - fused_rows = [row for row in fused_rows if row.get('score', 0.0) >= threshold] - logger.info(f'retrieval: threshold_filter={pre_count}->{len(fused_rows)} (threshold={threshold})') - - if fused_rows: - normalize_row_scores( - fused_rows, - source_field='score', - target_field='discovery_score', - default=0.5, - ) - - # ── Legacy graph routing ── - logger.info('\n 🧭 PHASE 2: Legacy Graph Routing') - router_used = 'discovery_only' - agent_rows: list[dict[str, Any]] = [] - - try: - agent_rows = await list_graph_routed_chunks( - db, user_id=user_id, namespace=namespace, query=query, - top_k=top_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if agent_rows: - router_used = 'discovery+graph' - logger.info(f' 📊 Graph routing: {len(agent_rows)} rows') - except Exception as exc: - logger.error(f' ❌ Graph routing failed (ignored): {exc}') - agent_rows = [] - - if agent_rows: - normalize_row_scores( - agent_rows, - source_field='score', - target_field='agent_score', - default=0.5, - ) - - ranked_rows = await rank_retrieval_candidates( - db, - user_id=user_id, - namespace=namespace, - discovery_rows=fused_rows, - routed_rows=agent_rows, - top_k=top_k, - ) - if ranked_rows: - logger.info(f'\n 🧮 Unified candidate ranking: {len(ranked_rows)} rows') - for i, row in enumerate(ranked_rows[:10]): - logger.info( - ' ' - f'[{i}] evidence={row.get("evidence_score", 0.0):.4f} ' - f'discovery={row.get("discovery_score", 0.0):.4f} ' - f'agent={row.get("agent_score", 0.0):.4f} ' - f'path={get_row_path(row)}' - ) - - assembled_rows = await assemble_retrieval_results( - db=db, - rows=ranked_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - ) - results = [attach_citation(row) for row in assembled_rows] - - response = { + _sync_legacy_adapter_overrides() + request: dict[str, Any] = { + "db": db, + "user_id": user_id, "namespace": namespace, "query": query, - "router_used": router_used, - "results": results, + "top_k": top_k, + "exclude_document_ids": exclude_document_ids, + "exclude_sections": exclude_sections, + "data_type": data_type, + "signal_paths": signal_paths, + "filter_mode": filter_mode, + "channels": channels, + "channel_weights": channel_weights, + "rerank": rerank, + "threshold": threshold, + "internal_recall_k": internal_recall_k, + "use_agentic": use_agentic, } + return await RetrievalExecutionPlan(request).execute() - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, - namespace=namespace, - version=cache_version, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, - **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") - - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, - namespace=namespace, - results=results, - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info(f'\n{"█" * 70}') - logger.info(f' ✅ RETRIEVAL COMPLETE: {len(results)} results | router={router_used} | {elapsed_total}ms') - for i, r in enumerate(results[:10]): - src = r.get('source', {}) - logger.info( - f' [{i+1}] type={r.get("chunk_type","?")} score={r.get("score",0):.4f}' - f' path={src.get("section_path","")}' - f' file={src.get("source_file_name","")}' - ) - if len(results) > 10: - logger.info(f' ... and {len(results) - 10} more') - logger.info(f'{"█" * 70}') - return await project_public_retrieval_response(response) +def _sync_legacy_adapter_overrides() -> None: + execution_plan.path_channel = path_channel + execution_plan.content_channel = content_channel + execution_plan.term_channel = term_channel + execution_plan.list_graph_routed_chunks = list_graph_routed_chunks diff --git a/packages/shared-python/shared/services/retrieval/execution_plan.py b/packages/shared-python/shared/services/retrieval/execution_plan.py new file mode 100644 index 000000000..f3c580d76 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution_plan.py @@ -0,0 +1,550 @@ +from __future__ import annotations + +import os +import time +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.cache_service import get_cached_retrieval_query_result, set_cached_retrieval_query_result +from shared.services.retrieval.channels import path_channel, content_channel, term_channel +from shared.services.retrieval.graph_service import GraphQueryService +from shared.services.retrieval.hit_stats_recorder import schedule_retrieval_hit_stats_update +from shared.services.retrieval.hydration import ( + assemble_retrieval_results, + hydrate_referenced_chunk_rows, +) +from shared.services.retrieval.response_projection import ( + attach_citation, + enrich_referenced_chunks_with_asset_urls, + project_public_retrieval_response, +) +from shared.services.retrieval.scoring import ( + get_row_path, + merge_channels_rrf, + merge_same_section_rows, + normalize_row_scores, +) +from shared.services.retrieval.ranking import rank_retrieval_candidates +from shared.services.retrieval.scoped_corpus import count_scoped_chunks, load_all_scoped_chunks +from shared.services.retrieval.settings import ( + CHANNEL_WEIGHT_CONTENT as _CHANNEL_WEIGHT_CONTENT, + CHANNEL_WEIGHT_PATH as _CHANNEL_WEIGHT_PATH, + CHANNEL_WEIGHT_TERM as _CHANNEL_WEIGHT_TERM, + INTERNAL_RECALL_K_MULTIPLIER as _INTERNAL_RECALL_K_MULTIPLIER, + resolve_allowed_chunk_types as _resolve_allowed_chunk_types, +) + + +async def list_graph_routed_chunks( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + service = GraphQueryService() + entry_document_ids = await service.find_entry_documents( + db, + user_id=user_id, + namespace=namespace, + query=query, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + return await service.collect_candidate_chunks( + db, + user_id=user_id, + namespace=namespace, + entry_document_ids=entry_document_ids, + query=query, + top_k=top_k * _INTERNAL_RECALL_K_MULTIPLIER, + exclude_sections=exclude_sections, + ) + + +async def run_retrieval_query( + *, + db: AsyncSession, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int = 1, + signal_paths: list[str] | None = None, + filter_mode: str = 'delete', + channels: list[str] | None = None, + channel_weights: dict[str, float] | None = None, + rerank: bool = False, + threshold: float = 0.0, + internal_recall_k: int | None = None, + use_agentic: bool | None = None, +) -> dict[str, Any]: + """Checkerboard retrieval: 3 independent channels -> RRF -> agent/graph union -> assembly.""" + return await RetrievalExecutionPlan( + { + "db": db, + "user_id": user_id, + "namespace": namespace, + "query": query, + "top_k": top_k, + "exclude_document_ids": exclude_document_ids, + "exclude_sections": exclude_sections, + "data_type": data_type, + "signal_paths": signal_paths, + "filter_mode": filter_mode, + "channels": channels, + "channel_weights": channel_weights, + "rerank": rerank, + "threshold": threshold, + "internal_recall_k": internal_recall_k, + "use_agentic": use_agentic, + } + ).execute() + + +class RetrievalExecutionPlan: + def __init__(self, request: dict[str, Any]) -> None: + self.request = request + + async def execute(self) -> dict[str, Any]: + db = self.request["db"] + user_id = self.request["user_id"] + namespace = self.request["namespace"] + query = self.request["query"] + top_k = self.request["top_k"] + exclude_document_ids = self.request["exclude_document_ids"] + exclude_sections = self.request["exclude_sections"] + data_type = self.request["data_type"] + signal_paths = self.request["signal_paths"] + filter_mode = self.request["filter_mode"] + channels = self.request["channels"] + channel_weights = self.request["channel_weights"] + rerank = self.request["rerank"] + threshold = self.request["threshold"] + internal_recall_k = self.request["internal_recall_k"] + use_agentic = self.request["use_agentic"] + + t_start = time.monotonic() + query = query.strip() + logger.info('\n' + '█' * 70) + logger.info(' 🚀 RETRIEVAL PIPELINE START') + logger.info(f' query="{query}"') + logger.info(f' user={user_id} ns={namespace} top_k={top_k} data_type={data_type}') + logger.info(f' exclude_docs={exclude_document_ids} exclude_secs={len(exclude_sections)}') + logger.info('█' * 70) + + if not query: + logger.info(' ⛔ Empty query filtered, skipping retrieval pipeline') + return { + "namespace": namespace, + "query": query, + "router_used": "empty_query_filtered", + "results": [], + } + + allowed_chunk_types = _resolve_allowed_chunk_types(data_type) + effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER + logger.info(f' allowed_chunk_types={allowed_chunk_types} effective_recall_k={effective_recall_k} signal_paths={signal_paths} filter_mode={filter_mode} rerank={rerank} threshold={threshold}') + + cache_extra = dict( + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + rerank=rerank, + threshold=threshold, + internal_recall_k=internal_recall_k, + # Always True: agentic mode now always routes through workflow + decomposition_enabled=True, + ) + + cache_version: int | None = None + try: + cache_version, cached = await get_cached_retrieval_query_result( + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + **cache_extra, + ) + if cached: + logger.info(f'retrieval: cache_hit=True version={cache_version}') + try: + schedule_retrieval_hit_stats_update( + user_id=user_id, + namespace=namespace, + results=cached.get("results", []), + ) + except Exception as e: + logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") + return await project_public_retrieval_response(cached) + except Exception as e: + logger.warning(f"Failed to read retrieval cache (ignored): {e}") + + logger.debug(f' 📦 Cache miss (version={cache_version}), running full pipeline') + + # ── Small KB optimization ── + try: + total_chunk_count = await count_scoped_chunks( + db, user_id=user_id, namespace=namespace, + exclude_document_ids=exclude_document_ids, + allowed_chunk_types=allowed_chunk_types, + ) + except Exception as e: + logger.warning(f"Failed to count scoped chunks, skipping small KB optimization: {e}") + total_chunk_count = top_k + 1 + logger.info(f'\n 📊 Total chunks in scope: {total_chunk_count}') + if total_chunk_count <= top_k: + logger.info(f' ⚡ Small KB optimization: {total_chunk_count} chunks <= top_k={top_k}, returning all') + all_rows = await load_all_scoped_chunks( + db, user_id=user_id, namespace=namespace, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths or [], + filter_mode=filter_mode, + ) + logger.info(f' small_kb load: loaded={len(all_rows)} rows after signal/exclude filters') + assembled_rows = await assemble_retrieval_results( + db=db, rows=all_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + ) + results = [attach_citation(row) for row in assembled_rows] + response = { + "namespace": namespace, "query": query, + "router_used": "small_kb_all", "results": results, + } + if cache_version is not None: + try: + await set_cached_retrieval_query_result( + user_id=user_id, namespace=namespace, version=cache_version, + query=query, top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + response=response, **cache_extra, + ) + except Exception as e: + logger.warning(f"Failed to write retrieval cache (ignored): {e}") + try: + schedule_retrieval_hit_stats_update(user_id=user_id, namespace=namespace, results=results) + except Exception as e: + logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") + elapsed_total = round((time.monotonic() - t_start) * 1000) + logger.info(f' ✅ Small KB: {len(results)} results in {elapsed_total}ms') + return await project_public_retrieval_response(response) + + # ══ Route: agentic (unified workflow) vs legacy ══ + if use_agentic is not None: + _agentic_enabled = use_agentic + else: + _agentic_enabled = os.environ.get('RETRIEVAL_AGENTIC_ENABLED', 'true') == 'true' + if _agentic_enabled: + # ── Unified agentic path via WorkflowOrchestrator ── + # Simple queries: planner returns a single-step plan (no decomposition). + # Complex queries: planner returns a multi-step plan with synthesize. + # Both go through the same code path. + from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator + + workflow = WorkflowOrchestrator() + workflow_result = await workflow.run( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + ) + + enriched_refs = await enrich_referenced_chunks_with_asset_urls( + workflow_result.referenced_chunks, + ) + + workflow_result_rows = await hydrate_referenced_chunk_rows( + db=db, + user_id=user_id, + namespace=namespace, + refs=enriched_refs, + ) + scoped_reference_keys = { + ( + str(row.get('document_id') or '').strip(), + str(row.get('chunk_id') or '').strip(), + ) + for row in workflow_result_rows + } + enriched_refs = [ + ref for ref in enriched_refs + if ( + str(ref.get('document_id') or '').strip(), + str(ref.get('chunk_id') or '').strip(), + ) in scoped_reference_keys + ] + assembled_workflow_rows = await assemble_retrieval_results( + db=db, + rows=workflow_result_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + ) + response = workflow_result.to_api_response() + # Override referenced_chunks with enriched versions + response['referenced_chunks'] = enriched_refs + response['results'] = [attach_citation(row) for row in assembled_workflow_rows] + + if cache_version is not None: + try: + await set_cached_retrieval_query_result( + user_id=user_id, namespace=namespace, version=cache_version, + query=query, top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + response=response, **cache_extra, + ) + except Exception as e: + logger.warning(f"Failed to write retrieval cache (ignored): {e}") + + try: + schedule_retrieval_hit_stats_update( + user_id=user_id, namespace=namespace, + results=enriched_refs, + ) + except Exception as e: + logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") + + elapsed_total = round((time.monotonic() - t_start) * 1000) + logger.info( + f'\n{"█" * 70}\n' + f' ✅ AGENTIC RETRIEVAL COMPLETE: ' + f'{len(enriched_refs)} chunks | ' + f'answer={len(workflow_result.answer_text)} chars | ' + f'router={workflow_result.router_used} | {elapsed_total}ms\n' + f'{"█" * 70}' + ) + + return await project_public_retrieval_response(response) + + else: + + # ── LEGACY path (existing code, unchanged) ── + + # ── Channel execution ── + active_channels = set(channels) if channels else {'path', 'content', 'term'} + logger.info(f'\n 📡 PHASE 1: Bottom-Layer Discovery (channels={sorted(active_channels)})') + logger.info(f' effective_recall_k={effective_recall_k}') + + path_rows: list[dict[str, Any]] = [] + content_rows: list[dict[str, Any]] = [] + term_rows: list[dict[str, Any]] = [] + + if 'path' in active_channels: + t_ch = time.monotonic() + path_rows = await path_channel( + db, user_id=user_id, namespace=namespace, query=query, + top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, filter_mode=filter_mode, + ) + elapsed_ch = round((time.monotonic() - t_ch) * 1000) + logger.info(f'\n 📡 path_channel: {len(path_rows)} rows in {elapsed_ch}ms') + for i, r in enumerate(path_rows[:5]): + logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}') + if len(path_rows) > 5: + logger.info(f' ... and {len(path_rows) - 5} more') + + if 'content' in active_channels: + t_ch = time.monotonic() + content_rows = await content_channel( + db, user_id=user_id, namespace=namespace, query=query, + top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, filter_mode=filter_mode, + ) + elapsed_ch = round((time.monotonic() - t_ch) * 1000) + logger.info(f'\n 📡 content_channel: {len(content_rows)} rows in {elapsed_ch}ms') + for i, r in enumerate(content_rows[:5]): + logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} content={str(r.get("content",""))[:80]}') + if len(content_rows) > 5: + logger.info(f' ... and {len(content_rows) - 5} more') + + if 'term' in active_channels: + t_ch = time.monotonic() + term_rows = await term_channel( + db, user_id=user_id, namespace=namespace, query=query, + top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, + signal_paths=signal_paths, filter_mode=filter_mode, + ) + elapsed_ch = round((time.monotonic() - t_ch) * 1000) + logger.info(f'\n 📡 term_channel: {len(term_rows)} rows in {elapsed_ch}ms') + for i, r in enumerate(term_rows[:5]): + logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}') + if len(term_rows) > 5: + logger.info(f' ... and {len(term_rows) - 5} more') + + # ── RRF fusion with configurable weights ── + default_weights = { + 'path': _CHANNEL_WEIGHT_PATH, + 'content': _CHANNEL_WEIGHT_CONTENT, + 'term': _CHANNEL_WEIGHT_TERM, + } + effective_weights = {**default_weights, **(channel_weights or {})} + + channel_lists: list[list[dict[str, Any]]] = [] + weight_list: list[float] = [] + + if path_rows: + channel_lists.append(path_rows) + weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH)) + if content_rows: + channel_lists.append(content_rows) + weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT)) + if term_rows: + channel_lists.append(term_rows) + weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM)) + + fused_rows = merge_channels_rrf(channel_lists, weight_list, effective_recall_k) if channel_lists else [] + logger.info(f'\n 🔀 RRF Fusion: {len(fused_rows)} rows from {len(channel_lists)} channels (weights={dict(zip(["path","content","term"][:len(weight_list)], weight_list))})') + for i, r in enumerate(fused_rows[:5]): + logger.info(f' [{i}] rrf_score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")}') + if len(fused_rows) > 5: + logger.info(f' ... and {len(fused_rows) - 5} more') + + # ── Section merging ── + pre_merge = len(fused_rows) + fused_rows = merge_same_section_rows(fused_rows) + if len(fused_rows) != pre_merge: + logger.info(f'retrieval: section_merge={pre_merge}->{len(fused_rows)}') + + # ── Threshold filtering ── + if threshold > 0.0 and fused_rows: + pre_count = len(fused_rows) + fused_rows = [row for row in fused_rows if row.get('score', 0.0) >= threshold] + logger.info(f'retrieval: threshold_filter={pre_count}->{len(fused_rows)} (threshold={threshold})') + + if fused_rows: + normalize_row_scores( + fused_rows, + source_field='score', + target_field='discovery_score', + default=0.5, + ) + + # ── Legacy graph routing ── + logger.info('\n 🧭 PHASE 2: Legacy Graph Routing') + router_used = 'discovery_only' + agent_rows: list[dict[str, Any]] = [] + + try: + agent_rows = await list_graph_routed_chunks( + db, user_id=user_id, namespace=namespace, query=query, + top_k=top_k, exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + if agent_rows: + router_used = 'discovery+graph' + logger.info(f' 📊 Graph routing: {len(agent_rows)} rows') + except Exception as exc: + logger.error(f' ❌ Graph routing failed (ignored): {exc}') + agent_rows = [] + + if agent_rows: + normalize_row_scores( + agent_rows, + source_field='score', + target_field='agent_score', + default=0.5, + ) + + ranked_rows = await rank_retrieval_candidates( + db, + user_id=user_id, + namespace=namespace, + discovery_rows=fused_rows, + routed_rows=agent_rows, + top_k=top_k, + ) + if ranked_rows: + logger.info(f'\n 🧮 Unified candidate ranking: {len(ranked_rows)} rows') + for i, row in enumerate(ranked_rows[:10]): + logger.info( + ' ' + f'[{i}] evidence={row.get("evidence_score", 0.0):.4f} ' + f'discovery={row.get("discovery_score", 0.0):.4f} ' + f'agent={row.get("agent_score", 0.0):.4f} ' + f'path={get_row_path(row)}' + ) + + assembled_rows = await assemble_retrieval_results( + db=db, + rows=ranked_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, + ) + results = [attach_citation(row) for row in assembled_rows] + + response = { + "namespace": namespace, + "query": query, + "router_used": router_used, + "results": results, + } + + if cache_version is not None: + try: + await set_cached_retrieval_query_result( + user_id=user_id, + namespace=namespace, + version=cache_version, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + response=response, + **cache_extra, + ) + except Exception as e: + logger.warning(f"Failed to write retrieval cache (ignored): {e}") + + try: + schedule_retrieval_hit_stats_update( + user_id=user_id, + namespace=namespace, + results=results, + ) + except Exception as e: + logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") + + elapsed_total = round((time.monotonic() - t_start) * 1000) + logger.info(f'\n{"█" * 70}') + logger.info(f' ✅ RETRIEVAL COMPLETE: {len(results)} results | router={router_used} | {elapsed_total}ms') + for i, r in enumerate(results[:10]): + src = r.get('source', {}) + logger.info( + f' [{i+1}] type={r.get("chunk_type","?")} score={r.get("score",0):.4f}' + f' path={src.get("section_path","")}' + f' file={src.get("source_file_name","")}' + ) + if len(results) > 10: + logger.info(f' ... and {len(results) - 10} more') + logger.info(f'{"█" * 70}') + + return await project_public_retrieval_response(response) From d0c039c16fe15622776b32d0edca110518556a1e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 13:50:30 +0800 Subject: [PATCH 20/40] refactor deepen retrieval and ingestion modules --- .../billing/billing_workflow_service.py | 18 +- .../document_ingestion/creation_service.py | 17 +- .../services/document_ingestion/service.py | 9 +- .../app/services/jobs/result_projection.py | 15 +- .../tests/contract/test_billing_contract.py | 6 +- .../tests/contract/test_retrieval_contract.py | 6 +- .../document_ingestion/processing_billing.py | 118 ++++ .../document_ingestion/processing_context.py | 121 ++++ .../document_ingestion/processing_run.py | 248 +------ .../services/document_ingestion/service.py | 22 +- .../contract/test_parse_task_contract.py | 41 +- .../shared/models/schemas/job_metadata.py | 76 ++ .../shared/services/job_lifecycle_sync.py | 3 +- .../shared/services/retrieval/app_service.py | 20 +- .../services/retrieval/execution_plan.py | 664 ++++++------------ .../services/retrieval/execution_routes.py | 527 ++++++++++++++ 16 files changed, 1140 insertions(+), 771 deletions(-) create mode 100644 apps/worker/app/services/document_ingestion/processing_billing.py create mode 100644 apps/worker/app/services/document_ingestion/processing_context.py create mode 100644 packages/shared-python/shared/services/retrieval/execution_routes.py diff --git a/apps/api/app/services/billing/billing_workflow_service.py b/apps/api/app/services/billing/billing_workflow_service.py index 91cd17df4..1c87939f5 100644 --- a/apps/api/app/services/billing/billing_workflow_service.py +++ b/apps/api/app/services/billing/billing_workflow_service.py @@ -1,9 +1,6 @@ from __future__ import annotations -import app.services.billing.billing_command_workflow as billing_command_workflow from app.services.billing.billing_command_workflow import BillingCommandWorkflow -from app.services.billing.billing_command_workflow import StripePurchaseService -from app.services.billing.billing_command_workflow import StripeWebhookService from app.services.billing.billing_read_model import BillingReadModel, ParseUsageResponse from sqlalchemy.ext.asyncio import AsyncSession @@ -17,12 +14,7 @@ UsageStatsResponse, ) -__all__ = [ - "BillingWorkflowService", - "ParseUsageResponse", - "StripePurchaseService", - "StripeWebhookService", -] +__all__ = ["BillingWorkflowService", "ParseUsageResponse"] class BillingWorkflowService: @@ -41,7 +33,6 @@ async def buy_credits( request: BuyCreditsRequest, user_id: str, ) -> PaymentIntentResponse: - _sync_legacy_adapter_overrides() return await self._command_workflow.buy_credits( request=request, user_id=user_id, @@ -107,7 +98,6 @@ async def buy_credits_package( request: BuyCreditsPackageRequest, user_id: str, ) -> CheckoutSessionResponse: - _sync_legacy_adapter_overrides() return await self._command_workflow.buy_credits_package( db, request=request, @@ -121,14 +111,8 @@ async def handle_stripe_webhook( payload: bytes, stripe_signature: str | None, ) -> dict[str, object]: - _sync_legacy_adapter_overrides() return await self._command_workflow.handle_stripe_webhook( db, payload=payload, stripe_signature=stripe_signature, ) - - -def _sync_legacy_adapter_overrides() -> None: - billing_command_workflow.StripePurchaseService = StripePurchaseService - billing_command_workflow.StripeWebhookService = StripeWebhookService diff --git a/apps/api/app/services/document_ingestion/creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py index 1e418360c..a76977b65 100644 --- a/apps/api/app/services/document_ingestion/creation_service.py +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -25,6 +25,7 @@ from shared.core.state_machine.states import JobStatus from shared.models.database.job import Job from shared.models.schemas.job import JobCreate, JobResponse +from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.redis import JobInfoRedisService, RedisServiceFactory from shared.services.redis.job_metadata_service import JobMetadataService from shared.services.storage.file_upload_service import FileUploadService @@ -151,8 +152,10 @@ async def _create_file_job( assert payload.file_name is not None file_extension = os.path.splitext(payload.file_name)[1] s3_key = f"uploads/{job_id}{file_extension}" - scope.job_metadata["source_file_name"] = payload.file_name - scope.job_metadata["source_type"] = "file" + JobMetadataHelper.set_file_source( + scope.job_metadata, + source_file_name=payload.file_name, + ) job = await self._create_waiting_job( db, @@ -224,12 +227,10 @@ async def _create_url_job( file_extension=file_extension, ) s3_key = f"uploads/{job_id}{file_extension}" - scope.job_metadata.update( - { - "source_file_name": source_file_name, - "source_url": payload.source_url, - "source_type": "url", - } + JobMetadataHelper.set_url_source( + scope.job_metadata, + source_file_name=source_file_name, + source_url=payload.source_url, ) job = await self._create_waiting_job( diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index 022ccc340..4573f8051 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -204,7 +204,7 @@ async def _resolve_scope( current_user: CurrentUser, ) -> ResolvedDocumentIngestionScope: job_metadata = cast(JobMetadata, JobMetadataHelper.create_from_request(payload)) - requested_document_id = cast(str | None, job_metadata.get("document_id")) + requested_document_id = JobMetadataHelper.get_document_id(job_metadata) if requested_document_id: active_job = await find_active_job_for_document( db, @@ -239,8 +239,11 @@ async def _resolve_scope( active_job_id=active_job.job_id, ) - job_metadata["document_id"] = effective_document_id - job_metadata["namespace"] = effective_namespace + JobMetadataHelper.set_document_scope( + job_metadata, + document_id=effective_document_id, + namespace=effective_namespace, + ) return ResolvedDocumentIngestionScope( job_metadata=job_metadata, document_id=effective_document_id, diff --git a/apps/api/app/services/jobs/result_projection.py b/apps/api/app/services/jobs/result_projection.py index c187e7225..0a00ce2d2 100644 --- a/apps/api/app/services/jobs/result_projection.py +++ b/apps/api/app/services/jobs/result_projection.py @@ -66,12 +66,7 @@ def to_job_status_value(status: str) -> JobStatusValue: def _resolve_original_request(job_metadata: Optional[dict[str, Any]]) -> dict[str, Any]: - original_request = ( - job_metadata.get("original_request") - if isinstance(job_metadata, dict) - else {} - ) - return original_request if isinstance(original_request, dict) else {} + return JobMetadataHelper.get_original_request(job_metadata) def _resolve_source_file_name(original_request: dict[str, Any]) -> str | None: @@ -97,8 +92,8 @@ def _resolve_parsing_params( original_request: dict[str, Any], ) -> dict[str, Any]: parsing_params = original_request.get("parsing_params") or {} - if not parsing_params and isinstance(job_metadata, dict): - parsing_params = job_metadata.get("parsing_params") or {} + if not parsing_params: + parsing_params = JobMetadataHelper.get_parsing_params_dict(job_metadata) return parsing_params if isinstance(parsing_params, dict) else {} @@ -139,11 +134,11 @@ async def build_job_result_response( return JobResultResponse( job_id=job.job_id, - namespace=JobMetadataHelper.get_field(job_metadata, "namespace"), + namespace=JobMetadataHelper.get_namespace(job_metadata), document_id=resolve_public_document_id(job), status=to_job_status_value(job.status), source_type=job.source_type, - data_id=JobMetadataHelper.get_field(job_metadata, "data_id"), + data_id=JobMetadataHelper.get_data_id(job_metadata), created_at=require_utc(job.created_at, field_name="created_at"), progress=progress, error=build_error_response(job, job_metadata), diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index de7e3cc04..6f5eb2be6 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -453,7 +453,7 @@ async def create_credits_package_checkout_session( async with developer_api_client_factory() as api_client: billing_service_module = importlib.import_module( - "app.services.billing.billing_workflow_service" + "app.services.billing.billing_command_workflow" ) monkeypatch.setattr( billing_service_module, @@ -498,7 +498,7 @@ async def create_payment_intent( async with developer_api_client_factory() as api_client: billing_service_module = importlib.import_module( - "app.services.billing.billing_workflow_service" + "app.services.billing.billing_command_workflow" ) monkeypatch.setattr( billing_service_module, @@ -543,7 +543,7 @@ async def handle_webhook( async with developer_api_client_factory() as api_client: billing_service_module = importlib.import_module( - "app.services.billing.billing_workflow_service" + "app.services.billing.billing_command_workflow" ) monkeypatch.setattr( billing_service_module, diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index cade4fbf7..7bad9363a 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -405,15 +405,15 @@ async def fake_graph_routing(*_args: object, **_kwargs: object) -> list[dict[str return [] monkeypatch.setattr( - "shared.services.retrieval.app_service.path_channel", + "shared.services.retrieval.execution_routes.path_channel", fake_path_channel, ) monkeypatch.setattr( - "shared.services.retrieval.app_service.content_channel", + "shared.services.retrieval.execution_routes.content_channel", fake_content_channel, ) monkeypatch.setattr( - "shared.services.retrieval.app_service.list_graph_routed_chunks", + "shared.services.retrieval.execution_routes.list_graph_routed_chunks", fake_graph_routing, ) diff --git a/apps/worker/app/services/document_ingestion/processing_billing.py b/apps/worker/app/services/document_ingestion/processing_billing.py new file mode 100644 index 000000000..12aaa1159 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/processing_billing.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime + +from app.services.document_ingestion.processing_context import ParseJobContext +from loguru import logger +from sqlalchemy import select + +from shared.core.database_sync import get_sync_db_context +from shared.core.exceptions.domain_exceptions import ( + InsufficientCreditsException, + NotFoundException, +) +from shared.models.database.job import Job +from shared.services.billing.work_billing_service import WorkBillingService + + +@dataclass(frozen=True) +class ParseJobBillingSnapshot: + billing_amount_micro_dollars: int + billing_credits: float + billing_status: str + + +def charge_parse_job_pages( + *, + job_id: str, + filename: str | None, + job_user_id: str | None, + page_count: int, +) -> ParseJobBillingSnapshot: + if not job_user_id: + raise NotFoundException( + resource="JobInfo", + resource_id="user_id", + internal_message=f"Missing user_id in job info for job_id={job_id}", + ) + + billing_service = WorkBillingService() + billing_filename = filename or "" + billing_status = "skipped" + billing_amount_micro_dollars = 0 + billing_credits = 0.0 + + with get_sync_db_context() as db: + job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()) + job = job_result.scalar_one_or_none() + + if job and getattr(job, "billing_status", "") == "charged": + logger.info(f"Job already charged: {job_id}") + billing_status = "charged" + billing_amount_micro_dollars = int(job.credits_charged or 0) + billing_credits = billing_amount_micro_dollars / 1_000_000 + else: + try: + billing_result = billing_service.charge_for_pages( + session=db, + user_id=job_user_id, + page_count=page_count, + filename=billing_filename, + ) + except InsufficientCreditsException: + logger.warning(f"Billing failed: job_id={job_id}, user_id={job_user_id}") + billing_amount = billing_service.estimate_page_charge( + page_count=page_count + ) + if job: + job.page_count = page_count + job.credits_charged = billing_amount.amount_micro_dollars + job.billing_status = "billing_failed" + db.commit() + + raise InsufficientCreditsException( + user_message=( + "Insufficient credits to process this document " + f"({page_count} pages required, cost: " + f"{billing_amount.credits})." + ), + required_credits=billing_amount.credits, + internal_message=( + f"job_id={job_id}, user_id={job_user_id}, " + f"page_count={page_count}" + ), + ) + + billing_status = billing_result.billing_status + billing_amount_micro_dollars = billing_result.amount_micro_dollars + billing_credits = billing_result.credits + if job: + job.page_count = page_count + job.credits_charged = billing_amount_micro_dollars + job.billing_status = billing_status + + return ParseJobBillingSnapshot( + billing_amount_micro_dollars=billing_amount_micro_dollars, + billing_credits=billing_credits, + billing_status=billing_status, + ) + + +def record_processing_start( + *, + job_id: str, + job_context: ParseJobContext, + billing_snapshot: ParseJobBillingSnapshot, + page_count: int, + processing_started_at: datetime, +) -> None: + metadata_updates = { + "page_count": page_count, + "billing_status": billing_snapshot.billing_status, + "billing_amount_micro_dollars": billing_snapshot.billing_amount_micro_dollars, + "billing_credits": billing_snapshot.billing_credits, + "processing_started_at": processing_started_at.isoformat(), + } + job_context.metadata_service.update_metadata(job_id, metadata_updates) + job_context.job_metadata.update(metadata_updates) diff --git a/apps/worker/app/services/document_ingestion/processing_context.py b/apps/worker/app/services/document_ingestion/processing_context.py new file mode 100644 index 000000000..51a40cd43 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/processing_context.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any + +from loguru import logger +from sqlalchemy import select + +from shared.core.config import settings +from shared.core.database_sync import get_sync_db_context +from shared.core.exceptions.domain_exceptions import ( + NotFoundException, + ValidationException, +) +from shared.models.database.job import Job +from shared.services.redis.redis_sync_service import ( + SyncJobInfoRedisService, + SyncJobMetadataService, +) +from shared.services.storage.job_file_storage import JobFileStorage + + +@dataclass(frozen=True) +class ParseJobContext: + job_metadata: dict[str, object] + job_user_id: str | None + metadata_service: SyncJobMetadataService + redis_service: Any + s3_key: str + + +def load_parse_job_context( + job_id: str, + requested_user_id: str | None, + redis_service: Any, +) -> ParseJobContext: + job_info_service = SyncJobInfoRedisService(redis_service) + job_info = job_info_service.get_job_info(job_id) + + if not job_info: + logger.warning( + f"JobInfo not found in Redis for job_id={job_id}; falling back to database" + ) + with get_sync_db_context() as fallback_db: + job_row = fallback_db.execute( + select(Job).where(Job.job_id == job_id) + ).scalar_one_or_none() + + if not job_row or not job_row.s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id=job_id, + internal_message="job info not found in Redis or database", + ) + + s3_key: str = job_row.s3_key + job_user_id: str | None = ( + str(job_row.user_id) if job_row.user_id else requested_user_id + ) + logger.info(f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}") + else: + raw_s3_key = job_info.get("s3_key") + if not isinstance(raw_s3_key, str) or not raw_s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id="s3_key", + internal_message="Missing s3_key in job_info", + ) + + s3_key = raw_s3_key + raw_job_user_id = job_info.get("user_id") + job_user_id = ( + raw_job_user_id if isinstance(raw_job_user_id, str) else requested_user_id + ) + + metadata_service = SyncJobMetadataService(redis_service) + raw_job_metadata = metadata_service.get_metadata(job_id) + if not isinstance(raw_job_metadata, dict) or not raw_job_metadata: + raise NotFoundException( + resource="JobMetadata", + resource_id=job_id, + internal_message=f"Job metadata not found for job_id={job_id}", + ) + + return ParseJobContext( + job_metadata=dict(raw_job_metadata), + job_user_id=job_user_id, + metadata_service=metadata_service, + redis_service=redis_service, + s3_key=s3_key, + ) + + +def assert_source_file_within_size_limit(s3_key: str) -> None: + file_info = JobFileStorage().verify_upload_exists(s3_key) + if not file_info.get("exists"): + raise NotFoundException( + resource="S3File", + resource_id=s3_key, + internal_message=f"S3 file not found: {s3_key}", + ) + + logger.info(f"S3 file verified: {s3_key}") + + file_size = file_info.get("size", 0) + file_extension = os.path.splitext(s3_key)[1].lower() + if file_size > settings.MAX_FILE_SIZE: + limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) + raise ValidationException( + user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", + violations=[ + { + "field": "file_size", + "description": ( + f"Size {file_size} bytes exceeds limit of " + f"{settings.MAX_FILE_SIZE} bytes" + ), + } + ], + ) diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index f09484adb..89dc4d765 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -1,7 +1,6 @@ from __future__ import annotations import os -from dataclasses import dataclass from datetime import datetime, timezone from typing import Any @@ -14,6 +13,15 @@ ) from app.services.document_ingestion.job_state_gate import mark_job_running from app.services.document_ingestion.page_estimator import PageEstimator +from app.services.document_ingestion.processing_billing import ( + charge_parse_job_pages, + record_processing_start, +) +from app.services.document_ingestion.processing_context import ( + ParseJobContext, + assert_source_file_within_size_limit, + load_parse_job_context, +) from app.services.document_ingestion.workspace import ( cleanup_task_workspace, create_task_workspace, @@ -21,48 +29,21 @@ ) from app.services.document_parser.stage_profiler import stage_timer from loguru import logger -from sqlalchemy import select -from shared.core.config import settings -from shared.core.database_sync import get_sync_db_context from shared.core.exceptions.domain_exceptions import ( - InsufficientCreditsException, - NotFoundException, - ValidationException, WorkerHandlingException, ) -from shared.models.database.job import Job from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.billing.work_billing_service import WorkBillingService from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service from shared.services.redis.distributed_lock import RedisJobLock from shared.services.redis.redis_sync_service import ( - SyncJobInfoRedisService, - SyncJobMetadataService, SyncRedisServiceFactory, ) -from shared.services.storage.job_file_storage import JobFileStorage from shared.services.storage.result_storage import get_result_storage from shared.services.storage.zip_result_service import ZipResultService -@dataclass(frozen=True) -class _ParseJobContext: - job_metadata: dict[str, object] - job_user_id: str | None - metadata_service: SyncJobMetadataService - redis_service: Any - s3_key: str - - -@dataclass(frozen=True) -class _ParseJobBillingSnapshot: - billing_amount_micro_dollars: int - billing_credits: float - billing_status: str - - class DocumentProcessingRun: """Run worker-side Document Ingestion for an uploaded file Job.""" @@ -71,8 +52,8 @@ def execute(self, job_id: str, user_id: str | None) -> dict[str, object]: lifecycle_service = get_sync_job_lifecycle_service() redis_service = SyncRedisServiceFactory.get_service() - job_context = _load_parse_job_context(job_id, user_id, redis_service) - _assert_source_file_within_size_limit(job_context.s3_key) + job_context = load_parse_job_context(job_id, user_id, redis_service) + assert_source_file_within_size_limit(job_context.s3_key) should_process = mark_job_running(job_id, job_context.redis_service) if not should_process: @@ -105,97 +86,6 @@ def execute(self, job_id: str, user_id: str | None) -> dict[str, object]: ) -def _load_parse_job_context( - job_id: str, - requested_user_id: str | None, - redis_service: Any, -) -> _ParseJobContext: - job_info_service = SyncJobInfoRedisService(redis_service) - job_info = job_info_service.get_job_info(job_id) - - if not job_info: - logger.warning( - f"JobInfo not found in Redis for job_id={job_id}; falling back to database" - ) - with get_sync_db_context() as fallback_db: - job_row = fallback_db.execute( - select(Job).where(Job.job_id == job_id) - ).scalar_one_or_none() - - if not job_row or not job_row.s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id=job_id, - internal_message="job info not found in Redis or database", - ) - - s3_key: str = job_row.s3_key - job_user_id: str | None = ( - str(job_row.user_id) if job_row.user_id else requested_user_id - ) - logger.info(f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}") - else: - raw_s3_key = job_info.get("s3_key") - if not isinstance(raw_s3_key, str) or not raw_s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id="s3_key", - internal_message="Missing s3_key in job_info", - ) - - s3_key = raw_s3_key - raw_job_user_id = job_info.get("user_id") - job_user_id = ( - raw_job_user_id if isinstance(raw_job_user_id, str) else requested_user_id - ) - - metadata_service = SyncJobMetadataService(redis_service) - raw_job_metadata = metadata_service.get_metadata(job_id) - if not isinstance(raw_job_metadata, dict) or not raw_job_metadata: - raise NotFoundException( - resource="JobMetadata", - resource_id=job_id, - internal_message=f"Job metadata not found for job_id={job_id}", - ) - - return _ParseJobContext( - job_metadata=dict(raw_job_metadata), - job_user_id=job_user_id, - metadata_service=metadata_service, - redis_service=redis_service, - s3_key=s3_key, - ) - - -def _assert_source_file_within_size_limit(s3_key: str) -> None: - file_info = JobFileStorage().verify_upload_exists(s3_key) - if not file_info.get("exists"): - raise NotFoundException( - resource="S3File", - resource_id=s3_key, - internal_message=f"S3 file not found: {s3_key}", - ) - - logger.info(f"S3 file verified: {s3_key}") - - file_size = file_info.get("size", 0) - file_extension = os.path.splitext(s3_key)[1].lower() - if file_size > settings.MAX_FILE_SIZE: - limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) - raise ValidationException( - user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", - violations=[ - { - "field": "file_size", - "description": ( - f"Size {file_size} bytes exceeds limit of " - f"{settings.MAX_FILE_SIZE} bytes" - ), - } - ], - ) - - def _prepare_task_workspace(job_id: str) -> tuple[str, str, str]: task_workspace_dir = create_task_workspace(job_id) input_dir = os.path.join(task_workspace_dir, "input") @@ -211,7 +101,7 @@ def _prepare_task_workspace(job_id: str) -> tuple[str, str, str]: def _run_parse_job( *, job_id: str, - job_context: _ParseJobContext, + job_context: ParseJobContext, lifecycle_service: Any, input_dir: str, output_dir: str, @@ -219,7 +109,9 @@ def _run_parse_job( ) -> dict[str, object]: lifecycle_service.update_progress(job_id, progress=10, message="Parsing document...") - filename = JobMetadataHelper.get_field(job_context.job_metadata, "source_file_name") + filename = JobMetadataHelper.get_source_file_name( + job_context.job_metadata, + ) or os.path.basename(job_context.s3_key) file_ext = os.path.splitext(job_context.s3_key)[1].lower() if job_context.s3_key else "" local_temp_path = download_s3_file_to_temp(job_context.s3_key, file_ext, input_dir) logger.info(f"File downloaded: job_id={job_id}, local_path={local_temp_path}") @@ -246,13 +138,13 @@ def _run_parse_job( logger.info(f"Workload estimation: job_id={job_id}, page_count={page_count}") processing_started_at = datetime.now(timezone.utc) - billing_snapshot = _charge_parse_job_pages( + billing_snapshot = charge_parse_job_pages( job_id=job_id, filename=filename, job_user_id=job_context.job_user_id, page_count=page_count, ) - _record_processing_start( + record_processing_start( job_id=job_id, job_context=job_context, billing_snapshot=billing_snapshot, @@ -359,116 +251,22 @@ def _run_parse_job( ) -def _charge_parse_job_pages( - *, - job_id: str, - filename: str | None, - job_user_id: str | None, - page_count: int, -) -> _ParseJobBillingSnapshot: - if not job_user_id: - raise NotFoundException( - resource="JobInfo", - resource_id="user_id", - internal_message=f"Missing user_id in job info for job_id={job_id}", - ) - - billing_service = WorkBillingService() - billing_filename = filename or "" - billing_status = "skipped" - billing_amount_micro_dollars = 0 - billing_credits = 0.0 - - with get_sync_db_context() as db: - job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()) - job = job_result.scalar_one_or_none() - - if job and getattr(job, "billing_status", "") == "charged": - logger.info(f"Job already charged: {job_id}") - billing_status = "charged" - billing_amount_micro_dollars = int(job.credits_charged or 0) - billing_credits = billing_amount_micro_dollars / 1_000_000 - else: - try: - billing_result = billing_service.charge_for_pages( - session=db, - user_id=job_user_id, - page_count=page_count, - filename=billing_filename, - ) - except InsufficientCreditsException: - logger.warning(f"Billing failed: job_id={job_id}, user_id={job_user_id}") - billing_amount = billing_service.estimate_page_charge( - page_count=page_count - ) - if job: - job.page_count = page_count - job.credits_charged = billing_amount.amount_micro_dollars - job.billing_status = "billing_failed" - db.commit() - - raise InsufficientCreditsException( - user_message=( - "Insufficient credits to process this document " - f"({page_count} pages required, cost: " - f"{billing_amount.credits})." - ), - required_credits=billing_amount.credits, - internal_message=( - f"job_id={job_id}, user_id={job_user_id}, " - f"page_count={page_count}" - ), - ) - - billing_status = billing_result.billing_status - billing_amount_micro_dollars = billing_result.amount_micro_dollars - billing_credits = billing_result.credits - if job: - job.page_count = page_count - job.credits_charged = billing_amount_micro_dollars - job.billing_status = billing_status - - return _ParseJobBillingSnapshot( - billing_amount_micro_dollars=billing_amount_micro_dollars, - billing_credits=billing_credits, - billing_status=billing_status, - ) - - -def _record_processing_start( - *, - job_id: str, - job_context: _ParseJobContext, - billing_snapshot: _ParseJobBillingSnapshot, - page_count: int, - processing_started_at: datetime, -) -> None: - metadata_updates = { - "page_count": page_count, - "billing_status": billing_snapshot.billing_status, - "billing_amount_micro_dollars": billing_snapshot.billing_amount_micro_dollars, - "billing_credits": billing_snapshot.billing_credits, - "processing_started_at": processing_started_at.isoformat(), - } - job_context.metadata_service.update_metadata(job_id, metadata_updates) - job_context.job_metadata.update(metadata_updates) - - def _finalize_parse_job_success( *, add_dir: str, chunks: list[dict[str, Any]], - job_context: _ParseJobContext, + job_context: ParseJobContext, job_id: str, lifecycle_service: Any, parsed_contents_df: pd.DataFrame, processing_started_at: datetime, task_workspace_dir: str, ) -> dict[str, object]: - source_file_name = JobMetadataHelper.get_field( + source_file_name = JobMetadataHelper.get_source_file_name( job_context.job_metadata, - "source_file_name", - ) or JobMetadataHelper.get_field(job_context.job_metadata, "source_url") + ) or JobMetadataHelper.get_source_url(job_context.job_metadata) + if not source_file_name: + source_file_name = os.path.basename(job_context.s3_key) if isinstance(source_file_name, str) and "/" in source_file_name: source_file_name = os.path.basename(source_file_name) @@ -522,7 +320,7 @@ def _finalize_parse_job_success( job_context.metadata_service.update_metadata(job_id, processing_timing_updates) job_context.job_metadata.update(processing_timing_updates) - data_id = JobMetadataHelper.get_field(job_context.job_metadata, "data_id") + data_id = JobMetadataHelper.get_data_id(job_context.job_metadata) zip_service = ZipResultService() zip_file_path, checksum, statistics, zip_size = zip_service.generate_zip_package( job_id=job_id, diff --git a/apps/worker/app/services/document_ingestion/service.py b/apps/worker/app/services/document_ingestion/service.py index cabc0e451..292bb155f 100644 --- a/apps/worker/app/services/document_ingestion/service.py +++ b/apps/worker/app/services/document_ingestion/service.py @@ -1,30 +1,10 @@ from __future__ import annotations -import app.services.document_ingestion.processing_run as processing_run from app.services.document_ingestion.processing_run import DocumentProcessingRun -from app.services.document_ingestion.processing_run import PageEstimator -from app.services.document_ingestion.processing_run import cleanup_task_workspace -from app.services.document_ingestion.processing_run import download_s3_file_to_temp -from app.services.document_ingestion.processing_run import get_result_storage -from app.services.document_ingestion.processing_run import settings -__all__ = [ - "PageEstimator", - "cleanup_task_workspace", - "download_s3_file_to_temp", - "get_result_storage", - "parse_uploaded_file_job", - "settings", -] +__all__ = ["parse_uploaded_file_job"] def parse_uploaded_file_job(job_id: str, user_id: str | None) -> dict[str, object]: """Run worker-side Document Ingestion for an uploaded file Job.""" - _sync_legacy_adapter_overrides() return DocumentProcessingRun().execute(job_id, user_id) - - -def _sync_legacy_adapter_overrides() -> None: - processing_run.download_s3_file_to_temp = download_s3_file_to_temp - processing_run.get_result_storage = get_result_storage - processing_run.cleanup_task_workspace = cleanup_task_workspace diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 3c9dcc48b..b4c3aae49 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -35,7 +35,7 @@ def _build_pending_file_job_metadata(source_file_name: str) -> dict[str, Any]: def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: import app.core.tasks.kb_tasks as kb_tasks - import app.services.document_ingestion.service as parse_job_service + import app.services.document_ingestion.processing_run as parse_job_service import app.services.document_parser.parse_service as parse_service from shared.core.database_sync import get_sync_engine from shared.services.redis.redis_sync_service import ( @@ -55,6 +55,12 @@ def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: ) +def _load_worker_settings() -> Any: + from shared.core.config import settings + + return settings + + def _save_worker_task_cache( *, job_id: str, @@ -143,6 +149,7 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state( sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-user-{uuid4().hex[:12]}" job_id: str = f"job_parse_success_{uuid4().hex[:12]}" @@ -179,8 +186,8 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state( ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(parse_job_service.settings, "BILLING_ENABLED", billing_enabled) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "BILLING_ENABLED", billing_enabled) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { @@ -338,8 +345,8 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: }, }, ] - expected_credits_charged = 3 * int(parse_job_service.settings.MICRO_DOLLARS_PER_PAGE) - expected_initial_balance = int(parse_job_service.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + expected_credits_charged = 3 * int(settings.MICRO_DOLLARS_PER_PAGE) + expected_initial_balance = int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 assert result == { "status": "success", @@ -583,6 +590,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-user-{uuid4().hex[:12]}" existing_job_id: str = f"job_existing_{uuid4().hex[:12]}" @@ -765,8 +773,8 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(parse_job_service.settings, "BILLING_ENABLED", False) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "BILLING_ENABLED", False) def fake_cleanup_task_workspace(workspace_dir: str | None) -> bool: captured_artifacts["workspace_dir"] = workspace_dir @@ -958,6 +966,7 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks( sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-concurrent-user-{uuid4().hex[:12]}" job_ids: list[str] = [f"job_cb_{index}_{uuid4().hex[:12]}" for index in range(2)] @@ -998,8 +1007,8 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks( ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) - monkeypatch.setattr(parse_job_service.settings, "BILLING_ENABLED", True) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "BILLING_ENABLED", True) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { @@ -1069,9 +1078,9 @@ def run_parse_task(job_id: str) -> dict[str, Any]: with ThreadPoolExecutor(max_workers=len(job_ids)) as executor: results = list(executor.map(run_parse_task, job_ids)) - expected_credits_charged = int(parse_job_service.settings.MICRO_DOLLARS_PER_PAGE) + expected_credits_charged = int(settings.MICRO_DOLLARS_PER_PAGE) expected_initial_balance = ( - int(parse_job_service.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 ) with engine.begin() as connection: @@ -1195,6 +1204,7 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-user-{uuid4().hex[:12]}" job_id: str = f"job_parse_skipped_{uuid4().hex[:12]}" @@ -1227,7 +1237,7 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return {"exists": storage_key == s3_key, "size": 1024} @@ -1284,6 +1294,7 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi sync_job_metadata_service_cls, sync_redis_service_factory, ) = _load_parse_task_modules() + settings = _load_worker_settings() user_id: str = f"worker-user-{uuid4().hex[:12]}" job_id: str = f"job_parse_failure_{uuid4().hex[:12]}" @@ -1316,7 +1327,7 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi ) _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) - monkeypatch.setattr(parse_job_service.settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { "exists": storage_key == s3_key, @@ -1355,8 +1366,8 @@ def fake_download_s3_file_to_temp( assert result.status == "FAILURE" assert _find_task_workspaces(tmp_path, job_id) == [] - expected_credits_charged = 3 * int(parse_job_service.settings.MICRO_DOLLARS_PER_PAGE) - expected_initial_balance = int(parse_job_service.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 + expected_credits_charged = 3 * int(settings.MICRO_DOLLARS_PER_PAGE) + expected_initial_balance = int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000 with engine.begin() as connection: job_row = ( diff --git a/packages/shared-python/shared/models/schemas/job_metadata.py b/packages/shared-python/shared/models/schemas/job_metadata.py index 79cfbff81..a4d2c7d96 100644 --- a/packages/shared-python/shared/models/schemas/job_metadata.py +++ b/packages/shared-python/shared/models/schemas/job_metadata.py @@ -52,6 +52,35 @@ def create_from_request(request, **kwargs) -> Dict[str, Any]: metadata.update(kwargs) return metadata + @staticmethod + def set_document_scope( + metadata: Dict[str, Any], + *, + document_id: str, + namespace: str, + ) -> None: + """Store the effective retrieval document scope.""" + metadata["document_id"] = document_id + metadata["namespace"] = namespace + + @staticmethod + def set_file_source(metadata: Dict[str, Any], *, source_file_name: str) -> None: + """Store source metadata for direct file uploads.""" + metadata["source_file_name"] = source_file_name + metadata["source_type"] = "file" + + @staticmethod + def set_url_source( + metadata: Dict[str, Any], + *, + source_file_name: str, + source_url: str, + ) -> None: + """Store source metadata for URL ingestion.""" + metadata["source_file_name"] = source_file_name + metadata["source_url"] = source_url + metadata["source_type"] = "url" + @staticmethod def get_field( metadata: Optional[Dict[str, Any]], field: str, default: Any = None @@ -61,6 +90,53 @@ def get_field( return default return metadata.get(field, default) + @staticmethod + def get_string_field( + metadata: Optional[Dict[str, Any]], field: str, default: str | None = None + ) -> str | None: + """Read a string field from metadata.""" + value = JobMetadataHelper.get_field(metadata, field, default) + return value if isinstance(value, str) else default + + @staticmethod + def get_original_request(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Return the stored creation request payload.""" + original_request = JobMetadataHelper.get_field(metadata, "original_request", {}) + return original_request if isinstance(original_request, dict) else {} + + @staticmethod + def get_parsing_params_dict(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Return stored parsing parameters as a dictionary.""" + parsing_params = JobMetadataHelper.get_field(metadata, "parsing_params", {}) + return parsing_params if isinstance(parsing_params, dict) else {} + + @staticmethod + def get_namespace( + metadata: Optional[Dict[str, Any]], default: str | None = None + ) -> str | None: + """Return the retrieval namespace stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "namespace", default) + + @staticmethod + def get_document_id(metadata: Optional[Dict[str, Any]]) -> str | None: + """Return the retrieval document id stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "document_id") + + @staticmethod + def get_data_id(metadata: Optional[Dict[str, Any]]) -> str | None: + """Return the user-defined data id stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "data_id") + + @staticmethod + def get_source_file_name(metadata: Optional[Dict[str, Any]]) -> str | None: + """Return the source file name stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "source_file_name") + + @staticmethod + def get_source_url(metadata: Optional[Dict[str, Any]]) -> str | None: + """Return the source URL stored in metadata.""" + return JobMetadataHelper.get_string_field(metadata, "source_url") + @staticmethod def get_parsing_param( metadata: Optional[Dict[str, Any]], param: str, default: Any = None diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index cd9dcc8ef..1c8b70f0a 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -25,6 +25,7 @@ from shared.models.database.document import DocumentSection from shared.models.database.job_result import JobChunk, JobResult from shared.models.database.webhook import WebhookEvent, WebhookEventStatus +from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.billing.credits_sync_service import SyncCreditsService from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, @@ -374,7 +375,7 @@ def _build_retrieval_cache_invalidation( namespaces: list[str] = [] metadata = job.job_metadata or {} - new_namespace = metadata.get("namespace") or "default" + new_namespace = JobMetadataHelper.get_namespace(metadata, "default") or "default" namespaces.append(new_namespace) if previous_document_scope and previous_document_scope.get("namespace"): diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 91959a9e4..13c60625e 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -4,20 +4,10 @@ from sqlalchemy.ext.asyncio import AsyncSession -import shared.services.retrieval.execution_plan as execution_plan -from shared.services.retrieval.channels import content_channel, path_channel, term_channel from shared.services.retrieval.execution_plan import RetrievalExecutionPlan -from shared.services.retrieval.execution_plan import list_graph_routed_chunks from shared.services.retrieval.scoring import merge_channels_rrf -__all__ = [ - "content_channel", - "list_graph_routed_chunks", - "merge_channels_rrf", - "path_channel", - "run_retrieval_query", - "term_channel", -] +__all__ = ["merge_channels_rrf", "run_retrieval_query"] async def run_retrieval_query( @@ -39,7 +29,6 @@ async def run_retrieval_query( internal_recall_k: int | None = None, use_agentic: bool | None = None, ) -> dict[str, Any]: - _sync_legacy_adapter_overrides() request: dict[str, Any] = { "db": db, "user_id": user_id, @@ -59,10 +48,3 @@ async def run_retrieval_query( "use_agentic": use_agentic, } return await RetrievalExecutionPlan(request).execute() - - -def _sync_legacy_adapter_overrides() -> None: - execution_plan.path_channel = path_channel - execution_plan.content_channel = content_channel - execution_plan.term_channel = term_channel - execution_plan.list_graph_routed_chunks = list_graph_routed_chunks diff --git a/packages/shared-python/shared/services/retrieval/execution_plan.py b/packages/shared-python/shared/services/retrieval/execution_plan.py index f3c580d76..96e0e15c7 100644 --- a/packages/shared-python/shared/services/retrieval/execution_plan.py +++ b/packages/shared-python/shared/services/retrieval/execution_plan.py @@ -1,72 +1,31 @@ from __future__ import annotations -import os import time from typing import Any from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.cache_service import get_cached_retrieval_query_result, set_cached_retrieval_query_result -from shared.services.retrieval.channels import path_channel, content_channel, term_channel -from shared.services.retrieval.graph_service import GraphQueryService -from shared.services.retrieval.hit_stats_recorder import schedule_retrieval_hit_stats_update -from shared.services.retrieval.hydration import ( - assemble_retrieval_results, - hydrate_referenced_chunk_rows, +from shared.services.retrieval.cache_service import ( + get_cached_retrieval_query_result, + set_cached_retrieval_query_result, +) +from shared.services.retrieval.execution_routes import ( + RetrievalRouteContext, + run_retrieval_route, +) +from shared.services.retrieval.hit_stats_recorder import ( + schedule_retrieval_hit_stats_update, ) from shared.services.retrieval.response_projection import ( - attach_citation, - enrich_referenced_chunks_with_asset_urls, project_public_retrieval_response, ) -from shared.services.retrieval.scoring import ( - get_row_path, - merge_channels_rrf, - merge_same_section_rows, - normalize_row_scores, -) -from shared.services.retrieval.ranking import rank_retrieval_candidates -from shared.services.retrieval.scoped_corpus import count_scoped_chunks, load_all_scoped_chunks from shared.services.retrieval.settings import ( - CHANNEL_WEIGHT_CONTENT as _CHANNEL_WEIGHT_CONTENT, - CHANNEL_WEIGHT_PATH as _CHANNEL_WEIGHT_PATH, - CHANNEL_WEIGHT_TERM as _CHANNEL_WEIGHT_TERM, - INTERNAL_RECALL_K_MULTIPLIER as _INTERNAL_RECALL_K_MULTIPLIER, - resolve_allowed_chunk_types as _resolve_allowed_chunk_types, + INTERNAL_RECALL_K_MULTIPLIER, + resolve_allowed_chunk_types, ) -async def list_graph_routed_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - service = GraphQueryService() - entry_document_ids = await service.find_entry_documents( - db, - user_id=user_id, - namespace=namespace, - query=query, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - return await service.collect_candidate_chunks( - db, - user_id=user_id, - namespace=namespace, - entry_document_ids=entry_document_ids, - query=query, - top_k=top_k * _INTERNAL_RECALL_K_MULTIPLIER, - exclude_sections=exclude_sections, - ) - - async def run_retrieval_query( *, db: AsyncSession, @@ -78,7 +37,7 @@ async def run_retrieval_query( exclude_sections: list[dict[str, str]], data_type: int = 1, signal_paths: list[str] | None = None, - filter_mode: str = 'delete', + filter_mode: str = "delete", channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, rerank: bool = False, @@ -86,7 +45,7 @@ async def run_retrieval_query( internal_recall_k: int | None = None, use_agentic: bool | None = None, ) -> dict[str, Any]: - """Checkerboard retrieval: 3 independent channels -> RRF -> agent/graph union -> assembly.""" + """Run retrieval through the plan module.""" return await RetrievalExecutionPlan( { "db": db, @@ -114,34 +73,36 @@ def __init__(self, request: dict[str, Any]) -> None: self.request = request async def execute(self) -> dict[str, Any]: - db = self.request["db"] - user_id = self.request["user_id"] - namespace = self.request["namespace"] - query = self.request["query"] - top_k = self.request["top_k"] - exclude_document_ids = self.request["exclude_document_ids"] - exclude_sections = self.request["exclude_sections"] - data_type = self.request["data_type"] - signal_paths = self.request["signal_paths"] - filter_mode = self.request["filter_mode"] - channels = self.request["channels"] - channel_weights = self.request["channel_weights"] - rerank = self.request["rerank"] - threshold = self.request["threshold"] - internal_recall_k = self.request["internal_recall_k"] - use_agentic = self.request["use_agentic"] - - t_start = time.monotonic() - query = query.strip() - logger.info('\n' + '█' * 70) - logger.info(' 🚀 RETRIEVAL PIPELINE START') - logger.info(f' query="{query}"') - logger.info(f' user={user_id} ns={namespace} top_k={top_k} data_type={data_type}') - logger.info(f' exclude_docs={exclude_document_ids} exclude_secs={len(exclude_sections)}') - logger.info('█' * 70) + db: AsyncSession = self.request["db"] + user_id: str = self.request["user_id"] + namespace: str = self.request["namespace"] + query: str = str(self.request["query"]).strip() + top_k: int = self.request["top_k"] + exclude_document_ids: list[str] = self.request["exclude_document_ids"] + exclude_sections: list[dict[str, str]] = self.request["exclude_sections"] + data_type: int = self.request["data_type"] + signal_paths: list[str] | None = self.request["signal_paths"] + filter_mode: str = self.request["filter_mode"] + channels: list[str] | None = self.request["channels"] + channel_weights: dict[str, float] | None = self.request["channel_weights"] + rerank: bool = self.request["rerank"] + threshold: float = self.request["threshold"] + internal_recall_k: int | None = self.request["internal_recall_k"] + use_agentic: bool | None = self.request["use_agentic"] + + start_time = time.monotonic() + _log_retrieval_start( + query=query, + user_id=user_id, + namespace=namespace, + top_k=top_k, + data_type=data_type, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) if not query: - logger.info(' ⛔ Empty query filtered, skipping retrieval pipeline') + logger.info(" ⛔ Empty query filtered, skipping retrieval pipeline") return { "namespace": namespace, "query": query, @@ -149,402 +110,213 @@ async def execute(self) -> dict[str, Any]: "results": [], } - allowed_chunk_types = _resolve_allowed_chunk_types(data_type) - effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER - logger.info(f' allowed_chunk_types={allowed_chunk_types} effective_recall_k={effective_recall_k} signal_paths={signal_paths} filter_mode={filter_mode} rerank={rerank} threshold={threshold}') + allowed_chunk_types: set[str] | None = resolve_allowed_chunk_types(data_type) + effective_recall_k = ( + internal_recall_k + if internal_recall_k is not None + else top_k * INTERNAL_RECALL_K_MULTIPLIER + ) + logger.info( + f" allowed_chunk_types={allowed_chunk_types} " + f"effective_recall_k={effective_recall_k} " + f"signal_paths={signal_paths} filter_mode={filter_mode} " + f"rerank={rerank} threshold={threshold}" + ) + + cache_extra = { + "data_type": data_type, + "signal_paths": signal_paths, + "filter_mode": filter_mode, + "channels": channels, + "channel_weights": channel_weights, + "rerank": rerank, + "threshold": threshold, + "internal_recall_k": internal_recall_k, + "decomposition_enabled": True, + } + cache_version, cached_response = await _read_cached_response( + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + cache_extra=cache_extra, + ) + if cached_response is not None: + return cached_response - cache_extra = dict( + logger.debug(f" 📦 Cache miss (version={cache_version}), running full pipeline") + + route_context = RetrievalRouteContext( + db=db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + allowed_chunk_types=allowed_chunk_types, data_type=data_type, signal_paths=signal_paths, filter_mode=filter_mode, channels=channels, channel_weights=channel_weights, - rerank=rerank, threshold=threshold, - internal_recall_k=internal_recall_k, - # Always True: agentic mode now always routes through workflow - decomposition_enabled=True, + effective_recall_k=effective_recall_k, + use_agentic=use_agentic, ) + outcome = await run_retrieval_route(route_context) - cache_version: int | None = None - try: - cache_version, cached = await get_cached_retrieval_query_result( + if cache_version is not None: + await _write_cached_response( user_id=user_id, namespace=namespace, + version=cache_version, query=query, top_k=top_k, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, - **cache_extra, + response=outcome.response, + cache_extra=cache_extra, ) - if cached: - logger.info(f'retrieval: cache_hit=True version={cache_version}') - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, - namespace=namespace, - results=cached.get("results", []), - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - return await project_public_retrieval_response(cached) - except Exception as e: - logger.warning(f"Failed to read retrieval cache (ignored): {e}") - - logger.debug(f' 📦 Cache miss (version={cache_version}), running full pipeline') - - # ── Small KB optimization ── - try: - total_chunk_count = await count_scoped_chunks( - db, user_id=user_id, namespace=namespace, - exclude_document_ids=exclude_document_ids, - allowed_chunk_types=allowed_chunk_types, - ) - except Exception as e: - logger.warning(f"Failed to count scoped chunks, skipping small KB optimization: {e}") - total_chunk_count = top_k + 1 - logger.info(f'\n 📊 Total chunks in scope: {total_chunk_count}') - if total_chunk_count <= top_k: - logger.info(f' ⚡ Small KB optimization: {total_chunk_count} chunks <= top_k={top_k}, returning all') - all_rows = await load_all_scoped_chunks( - db, user_id=user_id, namespace=namespace, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths or [], - filter_mode=filter_mode, - ) - logger.info(f' small_kb load: loaded={len(all_rows)} rows after signal/exclude filters') - assembled_rows = await assemble_retrieval_results( - db=db, rows=all_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - ) - results = [attach_citation(row) for row in assembled_rows] - response = { - "namespace": namespace, "query": query, - "router_used": "small_kb_all", "results": results, - } - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, namespace=namespace, version=cache_version, - query=query, top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") - try: - schedule_retrieval_hit_stats_update(user_id=user_id, namespace=namespace, results=results) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info(f' ✅ Small KB: {len(results)} results in {elapsed_total}ms') - return await project_public_retrieval_response(response) - # ══ Route: agentic (unified workflow) vs legacy ══ - if use_agentic is not None: - _agentic_enabled = use_agentic - else: - _agentic_enabled = os.environ.get('RETRIEVAL_AGENTIC_ENABLED', 'true') == 'true' - if _agentic_enabled: - # ── Unified agentic path via WorkflowOrchestrator ── - # Simple queries: planner returns a single-step plan (no decomposition). - # Complex queries: planner returns a multi-step plan with synthesize. - # Both go through the same code path. - from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator - - workflow = WorkflowOrchestrator() - workflow_result = await workflow.run( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - ) + _schedule_hit_stats_update( + user_id=user_id, + namespace=namespace, + results=outcome.hit_stats_results, + ) + _log_retrieval_complete( + outcome=outcome.response, + label=outcome.completion_label, + count=outcome.completion_count, + detail=outcome.completion_detail, + elapsed_ms=round((time.monotonic() - start_time) * 1000), + ) + return await project_public_retrieval_response(outcome.response) - enriched_refs = await enrich_referenced_chunks_with_asset_urls( - workflow_result.referenced_chunks, - ) - workflow_result_rows = await hydrate_referenced_chunk_rows( - db=db, +async def _read_cached_response( + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + cache_extra: dict[str, Any], +) -> tuple[int | None, dict[str, Any] | None]: + cache_version: int | None = None + try: + cache_version, cached = await get_cached_retrieval_query_result( + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + **cache_extra, + ) + if cached: + logger.info(f"retrieval: cache_hit=True version={cache_version}") + _schedule_hit_stats_update( user_id=user_id, namespace=namespace, - refs=enriched_refs, - ) - scoped_reference_keys = { - ( - str(row.get('document_id') or '').strip(), - str(row.get('chunk_id') or '').strip(), - ) - for row in workflow_result_rows - } - enriched_refs = [ - ref for ref in enriched_refs - if ( - str(ref.get('document_id') or '').strip(), - str(ref.get('chunk_id') or '').strip(), - ) in scoped_reference_keys - ] - assembled_workflow_rows = await assemble_retrieval_results( - db=db, - rows=workflow_result_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - ) - response = workflow_result.to_api_response() - # Override referenced_chunks with enriched versions - response['referenced_chunks'] = enriched_refs - response['results'] = [attach_citation(row) for row in assembled_workflow_rows] - - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, namespace=namespace, version=cache_version, - query=query, top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") - - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, namespace=namespace, - results=enriched_refs, - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info( - f'\n{"█" * 70}\n' - f' ✅ AGENTIC RETRIEVAL COMPLETE: ' - f'{len(enriched_refs)} chunks | ' - f'answer={len(workflow_result.answer_text)} chars | ' - f'router={workflow_result.router_used} | {elapsed_total}ms\n' - f'{"█" * 70}' + results=cached.get("results", []), ) + return cache_version, await project_public_retrieval_response(cached) + except Exception as exc: + logger.warning(f"Failed to read retrieval cache (ignored): {exc}") + return cache_version, None - return await project_public_retrieval_response(response) - - else: - - # ── LEGACY path (existing code, unchanged) ── - - # ── Channel execution ── - active_channels = set(channels) if channels else {'path', 'content', 'term'} - logger.info(f'\n 📡 PHASE 1: Bottom-Layer Discovery (channels={sorted(active_channels)})') - logger.info(f' effective_recall_k={effective_recall_k}') - - path_rows: list[dict[str, Any]] = [] - content_rows: list[dict[str, Any]] = [] - term_rows: list[dict[str, Any]] = [] - - if 'path' in active_channels: - t_ch = time.monotonic() - path_rows = await path_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 path_channel: {len(path_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(path_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}') - if len(path_rows) > 5: - logger.info(f' ... and {len(path_rows) - 5} more') - if 'content' in active_channels: - t_ch = time.monotonic() - content_rows = await content_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 content_channel: {len(content_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(content_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} content={str(r.get("content",""))[:80]}') - if len(content_rows) > 5: - logger.info(f' ... and {len(content_rows) - 5} more') - - if 'term' in active_channels: - t_ch = time.monotonic() - term_rows = await term_channel( - db, user_id=user_id, namespace=namespace, query=query, - top_k=effective_recall_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, filter_mode=filter_mode, - ) - elapsed_ch = round((time.monotonic() - t_ch) * 1000) - logger.info(f'\n 📡 term_channel: {len(term_rows)} rows in {elapsed_ch}ms') - for i, r in enumerate(term_rows[:5]): - logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}') - if len(term_rows) > 5: - logger.info(f' ... and {len(term_rows) - 5} more') - - # ── RRF fusion with configurable weights ── - default_weights = { - 'path': _CHANNEL_WEIGHT_PATH, - 'content': _CHANNEL_WEIGHT_CONTENT, - 'term': _CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM)) - - fused_rows = merge_channels_rrf(channel_lists, weight_list, effective_recall_k) if channel_lists else [] - logger.info(f'\n 🔀 RRF Fusion: {len(fused_rows)} rows from {len(channel_lists)} channels (weights={dict(zip(["path","content","term"][:len(weight_list)], weight_list))})') - for i, r in enumerate(fused_rows[:5]): - logger.info(f' [{i}] rrf_score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")}') - if len(fused_rows) > 5: - logger.info(f' ... and {len(fused_rows) - 5} more') - - # ── Section merging ── - pre_merge = len(fused_rows) - fused_rows = merge_same_section_rows(fused_rows) - if len(fused_rows) != pre_merge: - logger.info(f'retrieval: section_merge={pre_merge}->{len(fused_rows)}') - - # ── Threshold filtering ── - if threshold > 0.0 and fused_rows: - pre_count = len(fused_rows) - fused_rows = [row for row in fused_rows if row.get('score', 0.0) >= threshold] - logger.info(f'retrieval: threshold_filter={pre_count}->{len(fused_rows)} (threshold={threshold})') - - if fused_rows: - normalize_row_scores( - fused_rows, - source_field='score', - target_field='discovery_score', - default=0.5, - ) - - # ── Legacy graph routing ── - logger.info('\n 🧭 PHASE 2: Legacy Graph Routing') - router_used = 'discovery_only' - agent_rows: list[dict[str, Any]] = [] - - try: - agent_rows = await list_graph_routed_chunks( - db, user_id=user_id, namespace=namespace, query=query, - top_k=top_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if agent_rows: - router_used = 'discovery+graph' - logger.info(f' 📊 Graph routing: {len(agent_rows)} rows') - except Exception as exc: - logger.error(f' ❌ Graph routing failed (ignored): {exc}') - agent_rows = [] - - if agent_rows: - normalize_row_scores( - agent_rows, - source_field='score', - target_field='agent_score', - default=0.5, - ) - - ranked_rows = await rank_retrieval_candidates( - db, - user_id=user_id, - namespace=namespace, - discovery_rows=fused_rows, - routed_rows=agent_rows, - top_k=top_k, - ) - if ranked_rows: - logger.info(f'\n 🧮 Unified candidate ranking: {len(ranked_rows)} rows') - for i, row in enumerate(ranked_rows[:10]): - logger.info( - ' ' - f'[{i}] evidence={row.get("evidence_score", 0.0):.4f} ' - f'discovery={row.get("discovery_score", 0.0):.4f} ' - f'agent={row.get("agent_score", 0.0):.4f} ' - f'path={get_row_path(row)}' - ) - - assembled_rows = await assemble_retrieval_results( - db=db, - rows=ranked_rows, +async def _write_cached_response( + *, + user_id: str, + namespace: str, + version: int, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + response: dict[str, Any], + cache_extra: dict[str, Any], +) -> None: + try: + await set_cached_retrieval_query_result( + user_id=user_id, + namespace=namespace, + version=version, + query=query, + top_k=top_k, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, + response=response, + **cache_extra, ) - results = [attach_citation(row) for row in assembled_rows] + except Exception as exc: + logger.warning(f"Failed to write retrieval cache (ignored): {exc}") - response = { - "namespace": namespace, - "query": query, - "router_used": router_used, - "results": results, - } - if cache_version is not None: - try: - await set_cached_retrieval_query_result( - user_id=user_id, - namespace=namespace, - version=cache_version, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - response=response, - **cache_extra, - ) - except Exception as e: - logger.warning(f"Failed to write retrieval cache (ignored): {e}") +def _schedule_hit_stats_update( + *, + user_id: str, + namespace: str, + results: list[dict[str, Any]], +) -> None: + try: + schedule_retrieval_hit_stats_update( + user_id=user_id, + namespace=namespace, + results=results, + ) + except Exception as exc: + logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {exc}") - try: - schedule_retrieval_hit_stats_update( - user_id=user_id, - namespace=namespace, - results=results, - ) - except Exception as e: - logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") - elapsed_total = round((time.monotonic() - t_start) * 1000) - logger.info(f'\n{"█" * 70}') - logger.info(f' ✅ RETRIEVAL COMPLETE: {len(results)} results | router={router_used} | {elapsed_total}ms') - for i, r in enumerate(results[:10]): - src = r.get('source', {}) +def _log_retrieval_start( + *, + query: str, + user_id: str, + namespace: str, + top_k: int, + data_type: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> None: + logger.info("\n" + "█" * 70) + logger.info(" 🚀 RETRIEVAL PIPELINE START") + logger.info(f' query="{query}"') + logger.info( + f" user={user_id} ns={namespace} top_k={top_k} data_type={data_type}" + ) + logger.info( + f" exclude_docs={exclude_document_ids} " + f"exclude_secs={len(exclude_sections)}" + ) + logger.info("█" * 70) + + +def _log_retrieval_complete( + *, + outcome: dict[str, Any], + label: str, + count: int, + detail: str, + elapsed_ms: int, +) -> None: + logger.info(f'\n{"█" * 70}') + logger.info(f" ✅ {label} COMPLETE: {count} {detail} | {elapsed_ms}ms") + results = outcome.get("results", []) + if isinstance(results, list): + for index, result in enumerate(results[:10]): + source = result.get("source", {}) logger.info( - f' [{i+1}] type={r.get("chunk_type","?")} score={r.get("score",0):.4f}' - f' path={src.get("section_path","")}' - f' file={src.get("source_file_name","")}' + f" [{index + 1}] type={result.get('chunk_type', '?')} " + f"score={result.get('score', 0):.4f}" + f" path={source.get('section_path', '')}" + f" file={source.get('source_file_name', '')}" ) if len(results) > 10: - logger.info(f' ... and {len(results) - 10} more') - logger.info(f'{"█" * 70}') - - return await project_public_retrieval_response(response) + logger.info(f" ... and {len(results) - 10} more") + logger.info(f'{"█" * 70}') diff --git a/packages/shared-python/shared/services/retrieval/execution_routes.py b/packages/shared-python/shared/services/retrieval/execution_routes.py new file mode 100644 index 000000000..14dcc619f --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution_routes.py @@ -0,0 +1,527 @@ +from __future__ import annotations + +import os +import time +from dataclasses import dataclass +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.channels import content_channel, path_channel, term_channel +from shared.services.retrieval.graph_service import GraphQueryService +from shared.services.retrieval.hydration import ( + assemble_retrieval_results, + hydrate_referenced_chunk_rows, +) +from shared.services.retrieval.ranking import rank_retrieval_candidates +from shared.services.retrieval.response_projection import ( + attach_citation, + enrich_referenced_chunks_with_asset_urls, +) +from shared.services.retrieval.scoped_corpus import ( + count_scoped_chunks, + load_all_scoped_chunks, +) +from shared.services.retrieval.scoring import ( + get_row_path, + merge_channels_rrf, + merge_same_section_rows, + normalize_row_scores, +) +from shared.services.retrieval.settings import ( + CHANNEL_WEIGHT_CONTENT, + CHANNEL_WEIGHT_PATH, + CHANNEL_WEIGHT_TERM, + INTERNAL_RECALL_K_MULTIPLIER, +) + + +@dataclass(frozen=True) +class RetrievalRouteContext: + db: AsyncSession + user_id: str + namespace: str + query: str + top_k: int + exclude_document_ids: list[str] + exclude_sections: list[dict[str, str]] + allowed_chunk_types: set[str] | None + data_type: int + signal_paths: list[str] | None + filter_mode: str + channels: list[str] | None + channel_weights: dict[str, float] | None + threshold: float + effective_recall_k: int + use_agentic: bool | None + + +@dataclass(frozen=True) +class RetrievalRouteOutcome: + response: dict[str, Any] + hit_stats_results: list[dict[str, Any]] + completion_label: str + completion_count: int + completion_detail: str + + +async def list_graph_routed_chunks( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + service = GraphQueryService() + entry_document_ids = await service.find_entry_documents( + db, + user_id=user_id, + namespace=namespace, + query=query, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + return await service.collect_candidate_chunks( + db, + user_id=user_id, + namespace=namespace, + entry_document_ids=entry_document_ids, + query=query, + top_k=top_k * INTERNAL_RECALL_K_MULTIPLIER, + exclude_sections=exclude_sections, + ) + + +async def run_retrieval_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome: + small_kb_outcome = await _try_run_small_kb_route(context) + if small_kb_outcome is not None: + return small_kb_outcome + + if _should_use_agentic_route(context.use_agentic): + return await _run_agentic_route(context) + + return await _run_legacy_route(context) + + +async def _try_run_small_kb_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome | None: + try: + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + ) + except Exception as exc: + logger.warning( + f"Failed to count scoped chunks, skipping small KB optimization: {exc}" + ) + total_chunk_count = context.top_k + 1 + + logger.info(f"\n Total chunks in scope: {total_chunk_count}") + if total_chunk_count > context.top_k: + return None + + logger.info( + f" Small KB optimization: {total_chunk_count} chunks " + f"<= top_k={context.top_k}, returning all" + ) + all_rows = await load_all_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths or [], + filter_mode=context.filter_mode, + ) + logger.info( + f" small_kb load: loaded={len(all_rows)} rows after signal/exclude filters" + ) + assembled_rows = await assemble_retrieval_results( + db=context.db, + rows=all_rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) + results = [attach_citation(row) for row in assembled_rows] + response = { + "namespace": context.namespace, + "query": context.query, + "router_used": "small_kb_all", + "results": results, + } + return RetrievalRouteOutcome( + response=response, + hit_stats_results=results, + completion_label="Small KB", + completion_count=len(results), + completion_detail="results", + ) + + +def _should_use_agentic_route(use_agentic: bool | None) -> bool: + if use_agentic is not None: + return use_agentic + return os.environ.get("RETRIEVAL_AGENTIC_ENABLED", "true") == "true" + + +async def _run_agentic_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome: + from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator + + workflow = WorkflowOrchestrator() + workflow_result = await workflow.run( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.top_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + data_type=context.data_type, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + channels=context.channels, + channel_weights=context.channel_weights, + ) + + enriched_refs = await enrich_referenced_chunks_with_asset_urls( + workflow_result.referenced_chunks, + ) + + workflow_result_rows = await hydrate_referenced_chunk_rows( + db=context.db, + user_id=context.user_id, + namespace=context.namespace, + refs=enriched_refs, + ) + scoped_reference_keys = { + ( + str(row.get("document_id") or "").strip(), + str(row.get("chunk_id") or "").strip(), + ) + for row in workflow_result_rows + } + enriched_refs = [ + ref + for ref in enriched_refs + if ( + str(ref.get("document_id") or "").strip(), + str(ref.get("chunk_id") or "").strip(), + ) + in scoped_reference_keys + ] + assembled_workflow_rows = await assemble_retrieval_results( + db=context.db, + rows=workflow_result_rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) + response = workflow_result.to_api_response() + response["referenced_chunks"] = enriched_refs + response["results"] = [attach_citation(row) for row in assembled_workflow_rows] + + completion_detail = ( + f"chunks | answer={len(workflow_result.answer_text)} chars | " + f"router={workflow_result.router_used}" + ) + return RetrievalRouteOutcome( + response=response, + hit_stats_results=enriched_refs, + completion_label="AGENTIC RETRIEVAL", + completion_count=len(enriched_refs), + completion_detail=completion_detail, + ) + + +async def _run_legacy_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome: + active_channels = set(context.channels) if context.channels else { + "path", + "content", + "term", + } + logger.info( + f"\n PHASE 1: Bottom-Layer Discovery " + f"(channels={sorted(active_channels)})" + ) + logger.info(f" effective_recall_k={context.effective_recall_k}") + + path_rows = await _load_path_rows(context, active_channels) + content_rows = await _load_content_rows(context, active_channels) + term_rows = await _load_term_rows(context, active_channels) + + fused_rows = _fuse_legacy_rows( + context=context, + path_rows=path_rows, + content_rows=content_rows, + term_rows=term_rows, + ) + router_used = "discovery_only" + agent_rows: list[dict[str, Any]] = [] + + logger.info("\n PHASE 2: Legacy Graph Routing") + try: + agent_rows = await list_graph_routed_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.top_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + ) + if agent_rows: + router_used = "discovery+graph" + logger.info(f" Graph routing: {len(agent_rows)} rows") + except Exception as exc: + logger.error(f" Graph routing failed (ignored): {exc}") + agent_rows = [] + + if agent_rows: + normalize_row_scores( + agent_rows, + source_field="score", + target_field="agent_score", + default=0.5, + ) + + ranked_rows = await rank_retrieval_candidates( + context.db, + user_id=context.user_id, + namespace=context.namespace, + discovery_rows=fused_rows, + routed_rows=agent_rows, + top_k=context.top_k, + ) + if ranked_rows: + logger.info(f"\n Unified candidate ranking: {len(ranked_rows)} rows") + for index, row in enumerate(ranked_rows[:10]): + logger.info( + " " + f"[{index}] evidence={row.get('evidence_score', 0.0):.4f} " + f"discovery={row.get('discovery_score', 0.0):.4f} " + f"agent={row.get('agent_score', 0.0):.4f} " + f"path={get_row_path(row)}" + ) + + assembled_rows = await assemble_retrieval_results( + db=context.db, + rows=ranked_rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) + results = [attach_citation(row) for row in assembled_rows] + response = { + "namespace": context.namespace, + "query": context.query, + "router_used": router_used, + "results": results, + } + return RetrievalRouteOutcome( + response=response, + hit_stats_results=results, + completion_label="RETRIEVAL", + completion_count=len(results), + completion_detail=f"results | router={router_used}", + ) + + +async def _load_path_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "path" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await path_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n path_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"type={row.get('chunk_type', '?')}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +async def _load_content_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "content" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await content_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n content_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"content={str(row.get('content', ''))[:80]}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +async def _load_term_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "term" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await term_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n term_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"type={row.get('chunk_type', '?')}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +def _fuse_legacy_rows( + *, + context: RetrievalRouteContext, + path_rows: list[dict[str, Any]], + content_rows: list[dict[str, Any]], + term_rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + default_weights = { + "path": CHANNEL_WEIGHT_PATH, + "content": CHANNEL_WEIGHT_CONTENT, + "term": CHANNEL_WEIGHT_TERM, + } + effective_weights = {**default_weights, **(context.channel_weights or {})} + + channel_lists: list[list[dict[str, Any]]] = [] + weight_list: list[float] = [] + + if path_rows: + channel_lists.append(path_rows) + weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) + if content_rows: + channel_lists.append(content_rows) + weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) + if term_rows: + channel_lists.append(term_rows) + weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) + + if channel_lists: + fused_rows = merge_channels_rrf( + channel_lists, + weight_list, + context.effective_recall_k, + ) + else: + fused_rows = [] + logger.info( + f"\n RRF Fusion: {len(fused_rows)} rows from " + f"{len(channel_lists)} channels " + f"(weights={dict(zip(['path', 'content', 'term'][:len(weight_list)], weight_list))})" + ) + for index, row in enumerate(fused_rows[:5]): + logger.info( + f" [{index}] rrf_score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')}" + ) + if len(fused_rows) > 5: + logger.info(f" ... and {len(fused_rows) - 5} more") + + pre_merge = len(fused_rows) + fused_rows = merge_same_section_rows(fused_rows) + if len(fused_rows) != pre_merge: + logger.info(f"retrieval: section_merge={pre_merge}->{len(fused_rows)}") + + if context.channel_weights is not None: + logger.debug(f"retrieval: channel_weights={context.channel_weights}") + + fused_rows = _filter_rows_by_threshold(fused_rows, context) + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + + return fused_rows + + +def _filter_rows_by_threshold( + rows: list[dict[str, Any]], + context: RetrievalRouteContext, +) -> list[dict[str, Any]]: + if context.threshold <= 0.0 or not rows: + return rows + + pre_count = len(rows) + filtered_rows = [ + row for row in rows if row.get("score", 0.0) >= context.threshold + ] + logger.info( + f"retrieval: threshold_filter={pre_count}->{len(filtered_rows)} " + f"(threshold={context.threshold})" + ) + return filtered_rows From 94bf0878fb84f795e0e4d3a444300726aa50a1ec Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 14:21:40 +0800 Subject: [PATCH 21/40] refactor deepen job lifecycle and upload handoff --- .../confirmation_service.py | 67 +--- .../document_ingestion/handoff_service.py | 72 ++++ .../s3_events/upload_event_service.py | 43 +- .../contract/test_job_creation_contract.py | 18 +- .../tests/contract/test_s3_event_contract.py | 12 +- .../services/workload/url_upload_context.py | 42 ++ .../services/workload/url_upload_service.py | 122 ++---- .../services/workload/url_upload_transfer.py | 106 +++++ .../contract/test_url_upload_contract.py | 5 +- .../shared/services/job_failure_sync.py | 96 +++++ .../shared/services/job_lifecycle_sync.py | 379 +----------------- .../shared/services/job_publication_sync.py | 161 ++++++++ .../shared/services/job_result_sync.py | 76 ++++ .../shared/services/job_success_sync.py | 114 ++++++ .../services/job_webhook_outbox_sync.py | 87 ++++ 15 files changed, 835 insertions(+), 565 deletions(-) create mode 100644 apps/api/app/services/document_ingestion/handoff_service.py create mode 100644 apps/worker/app/services/workload/url_upload_context.py create mode 100644 apps/worker/app/services/workload/url_upload_transfer.py create mode 100644 packages/shared-python/shared/services/job_failure_sync.py create mode 100644 packages/shared-python/shared/services/job_publication_sync.py create mode 100644 packages/shared-python/shared/services/job_result_sync.py create mode 100644 packages/shared-python/shared/services/job_success_sync.py create mode 100644 packages/shared-python/shared/services/job_webhook_outbox_sync.py diff --git a/apps/api/app/services/document_ingestion/confirmation_service.py b/apps/api/app/services/document_ingestion/confirmation_service.py index 73198721d..cc513d22b 100644 --- a/apps/api/app/services/document_ingestion/confirmation_service.py +++ b/apps/api/app/services/document_ingestion/confirmation_service.py @@ -1,8 +1,10 @@ from __future__ import annotations from app.repositories.job_repository import JobRepository +from app.services.document_ingestion.handoff_service import ( + DocumentIngestionHandoffService, +) from app.services.jobs import check_job_permission -from app.services.knowledge.kb_orchestrator import KBOrchestrator from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession @@ -12,12 +14,9 @@ PermissionDeniedException, ValidationException, ) -from shared.core.state_machine.service import AsyncStateMachineService from shared.core.state_machine.states import JobStatus from shared.services.storage.file_upload_service import FileUploadService -_JOB_TYPE_KB_MANAGEMENT = "kb_management" - class DocumentIngestionConfirmationService: def __init__( @@ -25,9 +24,11 @@ def __init__( *, job_repository: JobRepository | None = None, file_upload_service: FileUploadService | None = None, + handoff_service: DocumentIngestionHandoffService | None = None, ) -> None: self._job_repository = job_repository or JobRepository() self._file_upload_service = file_upload_service or FileUploadService() + self._handoff_service = handoff_service or DocumentIngestionHandoffService() async def confirm_upload( self, @@ -66,13 +67,11 @@ async def confirm_upload( ], ) - await _transition_job_to_uploaded(db, job_id=job_id) - await _start_job_workflow( + await self._handoff_service.start_uploaded_file_workflow( db=db, - job_id=job_id, - job_type=job.job_type, - source_type="file", + job=job, user_id=user_id, + trigger="manual_upload_completed", ) return {"message": "File upload confirmed; processing started"} except NotFoundException: @@ -86,53 +85,3 @@ async def confirm_upload( raise JobOperationException( internal_message=f"Failed to confirm upload: {str(exc)}" ) - - -async def _transition_job_to_uploaded( - db: AsyncSession, - *, - job_id: str, - trigger: str = "manual_upload_completed", -) -> None: - state_machine = AsyncStateMachineService() - await state_machine.transition( - db, - job_id, - JobStatus.PENDING.value, - trigger, - None, - "system", - ) - - -async def _start_job_workflow( - db: AsyncSession, - *, - job_id: str, - job_type: str, - source_type: str, - user_id: str, - file_path: str | None = None, - file_url: str | None = None, -) -> None: - if job_type == _JOB_TYPE_KB_MANAGEMENT: - orchestrator = KBOrchestrator() - await orchestrator.start_workflow( - db=db, - job_id=job_id, - source_type=source_type, - file_path=file_path, - file_url=file_url, - user_id=user_id, - ) - return - - raise ValidationException( - user_message="Unsupported job type", - violations=[ - { - "field": "job_type", - "description": f"Job type '{job_type}' is not supported", - } - ], - ) diff --git a/apps/api/app/services/document_ingestion/handoff_service.py b/apps/api/app/services/document_ingestion/handoff_service.py new file mode 100644 index 000000000..9a14d92bb --- /dev/null +++ b/apps/api/app/services/document_ingestion/handoff_service.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from typing import Any + +from app.services.knowledge.kb_orchestrator import KBOrchestrator +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ValidationException +from shared.core.state_machine.service import AsyncStateMachineService +from shared.core.state_machine.states import JobStatus + +_JOB_TYPE_KB_MANAGEMENT = "kb_management" + + +class DocumentIngestionHandoffService: + """Advance uploaded Document Ingestion Jobs into worker parsing.""" + + def __init__( + self, + *, + state_machine: AsyncStateMachineService | None = None, + orchestrator: KBOrchestrator | None = None, + ) -> None: + self._state_machine = state_machine or AsyncStateMachineService() + self._orchestrator = orchestrator or KBOrchestrator() + + async def start_uploaded_file_workflow( + self, + db: AsyncSession, + *, + job: Any, + user_id: str, + trigger: str, + ) -> None: + await self._state_machine.transition( + db, + job.job_id, + JobStatus.PENDING.value, + trigger, + None, + "system", + ) + + if job.job_type != _JOB_TYPE_KB_MANAGEMENT: + raise ValidationException( + user_message="Unsupported job type", + violations=[ + { + "field": "job_type", + "description": ( + f"Job type '{job.job_type}' is not supported" + ), + } + ], + ) + + await self._orchestrator.start_workflow( + db=db, + job_id=job.job_id, + source_type="file", + file_path=None, + file_url=None, + user_id=user_id, + ) + + async def mark_upload_expired(self, db: AsyncSession, *, job: Any) -> None: + await self._state_machine.mark_failed( + db, + job.job_id, + "Upload expired: file was not uploaded within the allowed time window", + error_code="UPLOAD_EXPIRED", + ) diff --git a/apps/api/app/services/s3_events/upload_event_service.py b/apps/api/app/services/s3_events/upload_event_service.py index e30dd8fde..8b41ca310 100644 --- a/apps/api/app/services/s3_events/upload_event_service.py +++ b/apps/api/app/services/s3_events/upload_event_service.py @@ -4,12 +4,12 @@ import os from app.repositories.job_repository import JobRepository -from app.services.knowledge.kb_orchestrator import KBOrchestrator +from app.services.document_ingestion.handoff_service import ( + DocumentIngestionHandoffService, +) from loguru import logger from shared.core.database import get_db_context -from shared.core.state_machine.service import AsyncStateMachineService -from shared.core.state_machine.states import JobStatus from shared.models.schemas.s3_event import S3Event @@ -25,6 +25,7 @@ async def process_upload_events(s3_event: S3Event) -> None: try: upload_events = s3_event.get_upload_events() job_repo = JobRepository() + handoff_service = DocumentIngestionHandoffService() for event in upload_events: s3_key = event.object_key or event.s3.get("object", {}).get("key") @@ -55,40 +56,16 @@ async def process_upload_events(s3_event: S3Event) -> None: if is_job_expired(job.updated_at, settings.JOB_WAITING_EXPIRE_SECONDS): logger.warning(f"Job {job_id} upload expired, marking failed") - state_machine = AsyncStateMachineService() - await state_machine.mark_failed( - db, - job_id, - "Upload expired: file was not uploaded within the allowed time window", - error_code="UPLOAD_EXPIRED", - ) + await handoff_service.mark_upload_expired(db, job=job) continue - state_machine = AsyncStateMachineService() - await state_machine.transition( - db, - job_id, - JobStatus.PENDING.value, - "s3_upload_completed", - None, - "system", + await handoff_service.start_uploaded_file_workflow( + db=db, + job=job, + user_id=str(job.user_id), + trigger="s3_upload_completed", ) - if job.job_type == "kb_management": - orchestrator = KBOrchestrator() - await orchestrator.start_workflow( - db=db, - job_id=job_id, - source_type="file", - file_path=None, - file_url=None, - user_id=str(job.user_id), - ) - else: - logger.warning( - f"Unsupported job type for upload event: {job.job_type}, job_id={job_id}" - ) - logger.info(f"Triggered processing for job {job_id}") except Exception as exc: diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index 32c52ff51..9a9d47926 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -980,19 +980,18 @@ async def _fake_verify_s3_file_exists( assert bucket is None return {"exists": True, "s3_key": s3_key} - async def _fake_start_workflow_for_job( + async def _fake_start_workflow( + self: object, db: object, job_id: str, - job_type: str, source_type: str, + file_path: str | None, + file_url: str | None, user_id: str, - file_path: str | None = None, - file_url: str | None = None, ) -> None: started_workflows.append( { "job_id": job_id, - "job_type": job_type, "source_type": source_type, "user_id": user_id, "file_path": file_path, @@ -1002,7 +1001,7 @@ async def _fake_start_workflow_for_job( ) async with developer_api_client_factory() as api_client: - import app.services.document_ingestion.confirmation_service as document_ingestion_confirmation_service + import app.services.knowledge.kb_orchestrator as kb_orchestrator_module import shared.services.storage.file_upload_service as file_upload_service_module monkeypatch.setattr( @@ -1011,9 +1010,9 @@ async def _fake_start_workflow_for_job( _fake_verify_s3_file_exists, ) monkeypatch.setattr( - document_ingestion_confirmation_service, - "_start_job_workflow", - _fake_start_workflow_for_job, + kb_orchestrator_module.KBOrchestrator, + "start_workflow", + _fake_start_workflow, ) create_response = await api_client.post("/api/v1/jobs", json=payload) @@ -1035,7 +1034,6 @@ async def _fake_start_workflow_for_job( assert started_workflows == [ { "job_id": job_id, - "job_type": "kb_management", "source_type": "file", "user_id": "local-dev-user", "file_path": None, diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index ac5bf3a27..b8e8561de 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -93,11 +93,11 @@ async def start_workflow( async with api_client_factory() as api_client: user_id, job_id = await _insert_waiting_file_job() - upload_event_service = importlib.import_module( - "app.services.s3_events.upload_event_service" + handoff_service = importlib.import_module( + "app.services.document_ingestion.handoff_service" ) monkeypatch.setattr( - upload_event_service, + handoff_service, "KBOrchestrator", FakeKBOrchestrator, ) @@ -145,11 +145,11 @@ async def start_workflow( async with api_client_factory() as api_client: _, job_id = await _insert_waiting_file_job() - upload_event_service = importlib.import_module( - "app.services.s3_events.upload_event_service" + handoff_service = importlib.import_module( + "app.services.document_ingestion.handoff_service" ) monkeypatch.setattr( - upload_event_service, + handoff_service, "KBOrchestrator", FakeKBOrchestrator, ) diff --git a/apps/worker/app/services/workload/url_upload_context.py b/apps/worker/app/services/workload/url_upload_context.py new file mode 100644 index 000000000..af7fb807b --- /dev/null +++ b/apps/worker/app/services/workload/url_upload_context.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.services.redis.redis_sync_service import ( + SyncJobInfoRedisService, + SyncJobMetadataService, +) + + +@dataclass(frozen=True) +class UrlUploadContext: + s3_key: str + + +def load_url_upload_context(job_id: str, redis_service: Any) -> UrlUploadContext: + job_info_service = SyncJobInfoRedisService(redis_service) + job_info = job_info_service.get_job_info(job_id) + + if job_info: + raw_s3_key = job_info.get("s3_key") + else: + metadata_service = SyncJobMetadataService(redis_service) + job_metadata = metadata_service.get_metadata(job_id) + if not job_metadata: + raise NotFoundException( + resource="JobInfo", + resource_id=job_id, + internal_message="Job info not found in Redis or Metadata", + ) + raw_s3_key = job_metadata.get("s3_key") + + if not raw_s3_key: + raise NotFoundException( + resource="JobInfo", + resource_id="s3_key", + internal_message=f"Missing s3_key in Redis job info for job_id={job_id}", + ) + + return UrlUploadContext(s3_key=str(raw_s3_key)) diff --git a/apps/worker/app/services/workload/url_upload_service.py b/apps/worker/app/services/workload/url_upload_service.py index bd5347a2b..cad1e2744 100644 --- a/apps/worker/app/services/workload/url_upload_service.py +++ b/apps/worker/app/services/workload/url_upload_service.py @@ -1,24 +1,20 @@ from __future__ import annotations -import os from typing import Any from loguru import logger -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ( - NotFoundException, - StorageServiceException, - ValidationException, +from app.services.workload.url_upload_context import load_url_upload_context +from app.services.workload.url_upload_transfer import ( + assert_temp_file_within_size_limit, + cleanup_temp_file, + download_source_url_to_temp, + resolve_supported_url_extension, + upload_temp_file_to_source_storage, + verify_source_upload, ) from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service -from shared.services.redis.redis_sync_service import ( - SyncJobInfoRedisService, - SyncJobMetadataService, - SyncRedisServiceFactory, -) -from shared.services.storage.job_file_storage import JobFileStorage -from shared.utils.url_file_type import resolve_file_extension_sync +from shared.services.redis.redis_sync_service import SyncRedisServiceFactory def upload_url_file( @@ -31,111 +27,42 @@ def upload_url_file( lifecycle_service = get_sync_job_lifecycle_service() redis_service = SyncRedisServiceFactory.get_service() - job_info_service = SyncJobInfoRedisService(redis_service) - job_info = job_info_service.get_job_info(job_id) - - if not job_info: - metadata_service = SyncJobMetadataService(redis_service) - job_metadata = metadata_service.get_metadata(job_id) - if job_metadata: - s3_key = job_metadata.get("s3_key") - else: - raise NotFoundException( - resource="JobInfo", - resource_id=job_id, - internal_message="Job info not found in Redis or Metadata", - ) - else: - s3_key = job_info.get("s3_key") - - if not s3_key: - raise NotFoundException( - resource="JobInfo", - resource_id="s3_key", - internal_message=f"Missing s3_key in Redis job info for job_id={job_id}", - ) + upload_context = load_url_upload_context(job_id, redis_service) lifecycle_service.update_progress( job_id, progress=3, message="Validating URL file type..." ) - file_extension = resolve_file_extension_sync(source_url) - if not file_extension: - supported_formats = ", ".join(sorted(settings.get_supported_extensions())) - raise ValidationException( - user_message="Unsupported file type", - violations=[ - { - "field": "file_extension", - "description": f"Must be one of: {supported_formats}", - } - ], - ) + file_extension = resolve_supported_url_extension(source_url) lifecycle_service.update_progress( job_id, progress=10, message="Downloading file from URL..." ) - storage = JobFileStorage() - try: - temp_file_path = storage.download_file_from_url( - source_url, - temp_dir=getattr(settings, "TMP_PATH", "/tmp"), - ) - except Exception as exc: - raise ValidationException( - user_message="Failed to download file from URL", - violations=[ - { - "field": "source_url", - "description": "Could not download file from the provided URL", - } - ], - internal_message=( - f"Failed to download file from URL: {source_url}, error: {exc}" - ), - ) + temp_file_path = download_source_url_to_temp(source_url) try: lifecycle_service.update_progress( job_id, progress=30, message="Validating file size..." ) - file_size = os.path.getsize(temp_file_path) - if file_size > settings.MAX_FILE_SIZE: - limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) - raise ValidationException( - user_message=( - f"File size exceeds limit (max {limit_mb}MB for {file_extension})" - ), - violations=[ - { - "field": "file_size", - "description": ( - f"Size {file_size} bytes exceeds limit of " - f"{settings.MAX_FILE_SIZE} bytes" - ), - } - ], - ) + assert_temp_file_within_size_limit( + temp_file_path=temp_file_path, + file_extension=file_extension, + ) lifecycle_service.update_progress( job_id, progress=50, message="Uploading file to S3..." ) - storage.upload_source_file(temp_file_path, str(s3_key)) - logger.info(f"File uploaded to S3: {s3_key}") + upload_temp_file_to_source_storage( + temp_file_path=temp_file_path, + s3_key=upload_context.s3_key, + ) finally: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - logger.debug(f"Temp file cleaned up: {temp_file_path}") + cleanup_temp_file(temp_file_path) lifecycle_service.update_progress( job_id, progress=80, message="Verifying upload result..." ) - file_info = storage.verify_upload_exists(str(s3_key)) - if not file_info.get("exists"): - raise StorageServiceException( - user_message="We failed to verify your file upload", - internal_message=f"S3 file verification failed for {s3_key}", - ) + file_info = verify_source_upload(upload_context.s3_key) lifecycle_service.update_progress( job_id, @@ -143,12 +70,13 @@ def upload_url_file( message="URL file upload complete, waiting for processing...", ) logger.info( - f"URL file upload complete, waiting for S3 webhook: {job_id} -> {s3_key}" + "URL file upload complete, waiting for S3 webhook: " + f"{job_id} -> {upload_context.s3_key}" ) return { "status": "success", "job_id": job_id, - "s3_key": s3_key, + "s3_key": upload_context.s3_key, "file_size": file_info.get("size"), } diff --git a/apps/worker/app/services/workload/url_upload_transfer.py b/apps/worker/app/services/workload/url_upload_transfer.py new file mode 100644 index 000000000..25f4faf13 --- /dev/null +++ b/apps/worker/app/services/workload/url_upload_transfer.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import os + +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + StorageServiceException, + ValidationException, +) +from shared.services.storage.job_file_storage import JobFileStorage +from shared.utils.url_file_type import resolve_file_extension_sync + + +def resolve_supported_url_extension(source_url: str) -> str: + file_extension = resolve_file_extension_sync(source_url) + if file_extension: + return file_extension + + supported_formats = ", ".join(sorted(settings.get_supported_extensions())) + raise ValidationException( + user_message="Unsupported file type", + violations=[ + { + "field": "file_extension", + "description": f"Must be one of: {supported_formats}", + } + ], + ) + + +def download_source_url_to_temp(source_url: str) -> str: + storage = JobFileStorage() + try: + return storage.download_file_from_url( + source_url, + temp_dir=getattr(settings, "TMP_PATH", "/tmp"), + ) + except Exception as exc: + raise ValidationException( + user_message="Failed to download file from URL", + violations=[ + { + "field": "source_url", + "description": "Could not download file from the provided URL", + } + ], + internal_message=( + f"Failed to download file from URL: {source_url}, error: {exc}" + ), + ) + + +def assert_temp_file_within_size_limit( + *, + temp_file_path: str, + file_extension: str, +) -> int: + file_size = os.path.getsize(temp_file_path) + if file_size <= settings.MAX_FILE_SIZE: + return file_size + + limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024) + raise ValidationException( + user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})", + violations=[ + { + "field": "file_size", + "description": ( + f"Size {file_size} bytes exceeds limit of " + f"{settings.MAX_FILE_SIZE} bytes" + ), + } + ], + ) + + +def upload_temp_file_to_source_storage( + *, + temp_file_path: str, + s3_key: str, +) -> None: + storage = JobFileStorage() + storage.upload_source_file(temp_file_path, s3_key) + logger.info(f"File uploaded to S3: {s3_key}") + + +def verify_source_upload(s3_key: str) -> dict[str, object]: + storage = JobFileStorage() + file_info = storage.verify_upload_exists(s3_key) + if file_info.get("exists"): + return dict(file_info) + + raise StorageServiceException( + user_message="We failed to verify your file upload", + internal_message=f"S3 file verification failed for {s3_key}", + ) + + +def cleanup_temp_file(temp_file_path: str | None) -> None: + if not temp_file_path: + return + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + logger.debug(f"Temp file cleaned up: {temp_file_path}") diff --git a/apps/worker/tests/contract/test_url_upload_contract.py b/apps/worker/tests/contract/test_url_upload_contract.py index f98f1a816..cdf4a3de4 100644 --- a/apps/worker/tests/contract/test_url_upload_contract.py +++ b/apps/worker/tests/contract/test_url_upload_contract.py @@ -38,11 +38,12 @@ def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progres ) -> None: ( kb_tasks, - url_upload_service, + _url_upload_service, engine, sync_job_info_service_cls, sync_redis_service_factory, ) = _load_upload_task_modules() + from shared.core.config import settings from shared.services.storage.job_file_storage import JobFileStorage user_id = f"worker-user-{uuid4().hex[:12]}" @@ -124,7 +125,7 @@ def resolve_public_address( "file_size": 3, } assert uploaded_calls == [ - (str(downloaded_path), s3_key, url_upload_service.settings.S3_BUCKET_NAME), + (str(downloaded_path), s3_key, settings.S3_BUCKET_NAME), ] assert os.path.exists(downloaded_path) is False diff --git a/packages/shared-python/shared/services/job_failure_sync.py b/packages/shared-python/shared/services/job_failure_sync.py new file mode 100644 index 000000000..0888b1c27 --- /dev/null +++ b/packages/shared-python/shared/services/job_failure_sync.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.orm import Session + +from shared.core.response import build_standard_error_response +from shared.core.state_machine.service_sync import SyncStateMachineService +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent +from shared.services.billing.credits_sync_service import SyncCreditsService +from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox +from shared.utils.error_details import normalize_error_details + + +class SyncJobFailureFinalizer: + """Finalize failed Jobs inside the lifecycle transaction.""" + + def __init__( + self, + *, + state_machine: SyncStateMachineService | None = None, + webhook_outbox: SyncJobWebhookOutbox | None = None, + credits_service: SyncCreditsService | None = None, + ) -> None: + self._state_machine = state_machine or SyncStateMachineService() + self._webhook_outbox = webhook_outbox or SyncJobWebhookOutbox() + self._credits_service = credits_service or SyncCreditsService() + + def finalize( + self, + db: Session, + *, + job_id: str, + error_message: str, + error_code: str, + error_details: dict[str, Any] | None, + should_refund: bool, + ) -> tuple[bool, WebhookEvent | None]: + transition_ok = self._state_machine.mark_failed( + db, + job_id, + error_message, + error_code=error_code, + error_details=error_details, + ) + if not transition_ok: + logger.error(f"Job {job_id} mark_failed transition failed") + return False, None + + if should_refund: + self._try_refund_credits(db, job_id) + + normalized_error_details = normalize_error_details(error_details) + webhook_event = self._webhook_outbox.create_event( + db, + job_id=job_id, + event_type="job.failed", + extra_payload={ + "error": build_standard_error_response( + code=error_code, + message=error_message, + request_id=job_id, + details=normalized_error_details, + ), + }, + ) + return True, webhook_event + + def enqueue_webhook_after_commit(self, webhook_event: WebhookEvent | None) -> None: + self._webhook_outbox.enqueue_after_commit(webhook_event) + + def _try_refund_credits(self, db: Session, job_id: str) -> None: + try: + result = db.execute(select(Job).where(Job.job_id == job_id)) + job = result.scalar_one_or_none() + if not job: + return + + amount = getattr(job, "credits_charged", 0) or 0 + billing_status = getattr(job, "billing_status", "") + if amount <= 0 or billing_status != "charged": + return + + self._credits_service.refund_job_credits( + session=db, + user_id=str(job.user_id), + amount=amount, + job_id=job_id, + ) + job.billing_status = "refunded" + logger.info(f"Refunded {amount} credits for job {job_id}") + except Exception as exc: + logger.error(f"Credit refund failed for job {job_id}: {exc}") diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index 1c8b70f0a..9d79b4ce0 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -10,35 +10,19 @@ from __future__ import annotations import time -from datetime import datetime, timezone from typing import Any, Dict, List, Optional -from uuid import uuid4 from loguru import logger -from sqlalchemy import delete, select, update -from sqlalchemy.orm import Session from shared.core.database_sync import get_sync_db_context -from shared.core.response import build_standard_error_response -from shared.core.state_machine.service_sync import SyncStateMachineService -from shared.models.database.job import Job -from shared.models.database.document import DocumentSection -from shared.models.database.job_result import JobChunk, JobResult -from shared.models.database.webhook import WebhookEvent, WebhookEventStatus -from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.billing.credits_sync_service import SyncCreditsService +from shared.services.job_failure_sync import SyncJobFailureFinalizer +from shared.services.job_success_sync import SyncJobSuccessFinalizer from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, ) -from shared.services.retrieval.publication_service import RetrievalPublicationService -from shared.utils.error_details import normalize_error_details from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder -def _utc_now_naive() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) - - class SyncJobLifecycleService: """Manages job lifecycle transitions in the worker process (sync/gevent). @@ -46,8 +30,8 @@ class SyncJobLifecycleService: """ def __init__(self) -> None: - self._state_machine = SyncStateMachineService() - self._retrieval_publication = RetrievalPublicationService() + self._success_finalizer = SyncJobSuccessFinalizer() + self._failure_finalizer = SyncJobFailureFinalizer() # ── Public API ────────────────────────────────────────────────────── @@ -76,88 +60,27 @@ def finalize_job_success( with get_sync_db_context() as db: try: - inline_payload = {"checksum": checksum} - job_result = self._upsert_job_result( - db, - job_id, - delivery_mode, - inline_payload=inline_payload, - result_s3_key=result_s3_key, - result_size=zip_size, - ) - - normalized_chunks = chunks or [] - self._replace_chunks(db, job_result.id, normalized_chunks) - previous_document_scope = ( - self._retrieval_publication.get_existing_document_scope( - db, - job_id=job_id, - ) - ) - published_document_state = ( - self._retrieval_publication.publish_document_state( - db, - job_id=job_id, - job_result_id=job_result.id, - chunks=normalized_chunks, - ) - ) - if published_document_state is not None and not published_document_state.get("skipped_all_duplicate"): - # Backfill DocumentSection.summary from enriched doc_nav data - if section_summaries: - self._backfill_section_summaries( - db, - document_id=published_document_state.get("document_id", ""), - job_result_id=job_result.id, - section_summaries=section_summaries, - ) - self._retrieval_publication.publish_document_graph( - db, - job_id=job_id, - job_result_id=job_result.id, - ) - cache_invalidation = self._build_retrieval_cache_invalidation( + finalization = self._success_finalizer.finalize( db, job_id=job_id, - published_document_state=published_document_state, - previous_document_scope=previous_document_scope, - ) - - transition_ok = self._state_machine.mark_completed( - db, - job_id, - result_metadata={ - "storage_completed": True, - "stored_count": stored_count, - "delivery_mode": delivery_mode, - }, + result_s3_key=result_s3_key, + checksum=checksum, + zip_size=zip_size, + chunks=chunks or [], + stored_count=stored_count, + delivery_mode=delivery_mode, + section_summaries=section_summaries, ) - if not transition_ok: - logger.error(f"Job {job_id} mark_completed transition failed") + if finalization.response.get("status") != "success": db.rollback() - return { - "status": "failed", - "job_id": job_id, - "reason": "state_transition_failed", - } - - webhook_event = self._maybe_create_webhook_event( - db, - job_id, - event_type="job.completed", - ) + return finalization.response db.commit() logger.info(f"Job {job_id} success transaction committed") - self._post_commit_invalidate_retrieval_cache(cache_invalidation) - self._post_commit_enqueue_webhook(webhook_event) + self._success_finalizer.run_post_commit_actions(finalization) - return { - "status": "success", - "job_id": job_id, - "stored_count": stored_count, - } + return finalization.response except Exception as exc: logger.error(f"Failed to finalize job success {job_id}: {exc}") @@ -185,40 +108,22 @@ def finalize_job_failure( with get_sync_db_context() as db: try: - transition_ok = self._state_machine.mark_failed( + transition_ok, webhook_event = self._failure_finalizer.finalize( db, - job_id, - error_message, + job_id=job_id, + error_message=error_message, error_code=error_code, error_details=error_details, + should_refund=should_refund, ) if not transition_ok: - logger.error(f"Job {job_id} mark_failed transition failed") db.rollback() return False - if should_refund: - self._try_refund_credits(db, job_id) - - normalized_error_details = normalize_error_details(error_details) - webhook_event = self._maybe_create_webhook_event( - db, - job_id, - event_type="job.failed", - extra_payload={ - "error": build_standard_error_response( - code=error_code, - message=error_message, - request_id=job_id, - details=normalized_error_details, - ), - }, - ) - db.commit() logger.info(f"Job {job_id} failure transaction committed") - self._post_commit_enqueue_webhook(webhook_event) + self._failure_finalizer.enqueue_webhook_after_commit(webhook_event) return True @@ -260,248 +165,6 @@ def update_progress( logger.warning(f"Failed to update progress for job {job_id}: {exc}") return False - # ── Private helpers ───────────────────────────────────────────────── - - def _backfill_section_summaries( - self, - db: Session, - *, - document_id: str, - job_result_id: str, - section_summaries: Dict[str, str], - ) -> None: - """Populate DocumentSection.summary from enriched doc_nav data. - - Runs UPDATE statements within the existing transaction so no extra - commit is needed. Overwrites any existing summary value since the - enriched doc_nav data is the authoritative source. - """ - if not document_id or not section_summaries: - return - try: - for path, summary in section_summaries.items(): - if not path or not summary: - continue - db.execute( - update(DocumentSection) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.section_path == path) - .values(summary=summary) - ) - db.flush() - logger.debug( - f"Backfilled section summaries: document_id={document_id}, " - f"count={len(section_summaries)}" - ) - except Exception as exc: - logger.warning(f"Section summary backfill failed (non-fatal): {exc}") - - def _upsert_job_result( - self, - db: Session, - job_id: str, - delivery_mode: str, - *, - inline_payload: Optional[Dict[str, Any]] = None, - result_s3_key: Optional[str] = None, - result_size: Optional[int] = None, - ) -> JobResult: - """Create or update JobResult row.""" - result = db.execute(select(JobResult).where(JobResult.job_id == job_id)) - existing = result.scalar_one_or_none() - - if existing: - existing.delivery_mode = delivery_mode - existing.inline_payload = inline_payload - existing.result_s3_key = result_s3_key - existing.result_size = result_size - db.flush() - return existing - - job_result = JobResult( - job_id=job_id, - delivery_mode=delivery_mode, - document_metadata={}, - inline_payload=inline_payload, - result_s3_key=result_s3_key, - result_size=result_size, - ) - db.add(job_result) - db.flush() - return job_result - - def _replace_chunks( - self, - db: Session, - job_result_id: str, - chunks: List[Dict[str, Any]], - ) -> Optional[Dict[str, Any]]: - """Delete existing chunks and insert new ones.""" - db.execute(delete(JobChunk).where(JobChunk.job_result_id == job_result_id)) - - if not chunks: - db.flush() - return - - chunk_models = [] - for index, chunk in enumerate(chunks): - chunk_identifier = chunk.get("chunk_id") or str(uuid4()) - chunk_models.append( - JobChunk( - job_result_id=job_result_id, - chunk_id=chunk_identifier, - chunk_type=chunk.get("type", "paragraph"), - text=chunk.get("text"), - path=chunk.get("metadata", {}).get("path"), - chunk_metadata=chunk.get("metadata"), - sort_order=chunk.get("order", index), - ) - ) - db.add_all(chunk_models) - db.flush() - - def _build_retrieval_cache_invalidation( - self, - db: Session, - *, - job_id: str, - published_document_state: Optional[Dict[str, str]], - previous_document_scope: Optional[Dict[str, str]], - ) -> Optional[Dict[str, Any]]: - job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() - if not job: - return None - - namespaces: list[str] = [] - metadata = job.job_metadata or {} - new_namespace = JobMetadataHelper.get_namespace(metadata, "default") or "default" - namespaces.append(new_namespace) - - if previous_document_scope and previous_document_scope.get("namespace"): - namespaces.append(previous_document_scope["namespace"]) - if published_document_state and published_document_state.get("namespace"): - namespaces.append(published_document_state["namespace"]) - - return {"user_id": str(job.user_id), "namespaces": namespaces, "job_id": job_id} - - def _post_commit_invalidate_retrieval_cache( - self, cache_invalidation: Optional[Dict[str, Any]] - ) -> None: - if not cache_invalidation: - return - try: - redis_service = SyncRedisServiceFactory.get_service() - user_id = cache_invalidation["user_id"] - seen: set[str] = set() - for namespace in cache_invalidation["namespaces"]: - if not namespace or namespace in seen: - continue - seen.add(namespace) - redis_service.incr(f"retrieval:version:{user_id}:{namespace}") - except Exception as exc: - logger.warning( - f"Failed to invalidate retrieval cache after publication (ignored): job_id={cache_invalidation.get('job_id')}, error={exc}" - ) - - def _maybe_create_webhook_event( - self, - db: Session, - job_id: str, - event_type: str, - extra_payload: Optional[Dict[str, Any]] = None, - ) -> Optional[WebhookEvent]: - """Create a WebhookEvent if the job has webhooks enabled.""" - result = db.execute(select(Job).where(Job.job_id == job_id)) - job = result.scalar_one_or_none() - - if not job: - logger.warning(f"Job not found for webhook check: {job_id}") - return None - - webhook_url = getattr(job, "webhook_url", None) - if not job.webhook_enabled or not webhook_url: - return None - - status = "completed" if event_type == "job.completed" else "failed" - timestamp_key = f"{status}_at" - payload: Dict[str, Any] = { - "event": event_type, - "job_id": job_id, - "status": status, - timestamp_key: _utc_now_naive().isoformat(), - } - if extra_payload: - payload.update(extra_payload) - - event = WebhookEvent( - job_id=job_id, - target_url=webhook_url, - payload=payload, - status=WebhookEventStatus.PENDING, - attempts=0, - ) - db.add(event) - db.flush() - logger.info(f"WebhookEvent created: event_id={event.id}, job_id={job_id}") - return event - - def _try_refund_credits(self, db: Session, job_id: str) -> None: - """Attempt to refund credits for a failed job.""" - try: - result = db.execute(select(Job).where(Job.job_id == job_id)) - job = result.scalar_one_or_none() - if not job: - return - - amount = getattr(job, "credits_charged", 0) or 0 - billing_status = getattr(job, "billing_status", "") - if amount <= 0 or billing_status != "charged": - return - - credits_service = SyncCreditsService() - credits_service.refund_job_credits( - session=db, - user_id=str(job.user_id), - amount=amount, - job_id=job_id, - ) - job.billing_status = "refunded" - logger.info(f"Refunded {amount} credits for job {job_id}") - except Exception as exc: - logger.error(f"Credit refund failed for job {job_id}: {exc}") - - def _post_commit_enqueue_webhook( - self, - webhook_event: Optional[WebhookEvent], - ) -> None: - """Publish a persisted webhook via QStash after commit (best-effort).""" - if not webhook_event: - return - - try: - from shared.services.webhook.qstash_publisher import ( - get_qstash_webhook_publisher, - ) - - publisher = get_qstash_webhook_publisher() - message_id = publisher.publish_event(webhook_event.id) - if not message_id: - logger.warning( - f"Webhook publish failed after commit: event_id={webhook_event.id}" - ) - return - logger.info( - f"Webhook published after commit: event_id={webhook_event.id}, " - f"message_id={message_id}" - ) - except Exception as exc: - logger.error( - f"Failed to publish webhook after commit (event persisted): " - f"event_id={webhook_event.id}, error={exc}" - ) - - # Module-level singleton _lifecycle_service: Optional[SyncJobLifecycleService] = None diff --git a/packages/shared-python/shared/services/job_publication_sync.py b/packages/shared-python/shared/services/job_publication_sync.py new file mode 100644 index 000000000..b6b2d69a1 --- /dev/null +++ b/packages/shared-python/shared/services/job_publication_sync.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from loguru import logger +from sqlalchemy import select, update +from sqlalchemy.orm import Session + +from shared.models.database.document import DocumentSection +from shared.models.database.job import Job +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.redis.redis_sync_service import SyncRedisServiceFactory +from shared.services.retrieval.publication_service import RetrievalPublicationService + + +@dataclass(frozen=True) +class JobPublicationOutcome: + published_document_state: dict[str, str] | None + cache_invalidation: dict[str, Any] | None + + +class SyncJobPublicationFinalizer: + """Publish terminal parse results and invalidate retrieval cache after commit.""" + + def __init__( + self, + *, + retrieval_publication: RetrievalPublicationService | None = None, + ) -> None: + self._retrieval_publication = ( + retrieval_publication or RetrievalPublicationService() + ) + + def publish_result( + self, + db: Session, + *, + job_id: str, + job_result_id: str, + chunks: list[dict[str, Any]], + section_summaries: dict[str, str] | None, + ) -> JobPublicationOutcome: + previous_document_scope = self._retrieval_publication.get_existing_document_scope( + db, + job_id=job_id, + ) + published_document_state = self._retrieval_publication.publish_document_state( + db, + job_id=job_id, + job_result_id=job_result_id, + chunks=chunks, + ) + if _should_publish_document_graph(published_document_state): + assert published_document_state is not None + if section_summaries: + self._backfill_section_summaries( + db, + document_id=published_document_state.get("document_id", ""), + job_result_id=job_result_id, + section_summaries=section_summaries, + ) + self._retrieval_publication.publish_document_graph( + db, + job_id=job_id, + job_result_id=job_result_id, + ) + + cache_invalidation = self._build_cache_invalidation( + db, + job_id=job_id, + published_document_state=published_document_state, + previous_document_scope=previous_document_scope, + ) + return JobPublicationOutcome( + published_document_state=published_document_state, + cache_invalidation=cache_invalidation, + ) + + def invalidate_cache_after_commit( + self, + cache_invalidation: dict[str, Any] | None, + ) -> None: + if not cache_invalidation: + return + + try: + redis_service = SyncRedisServiceFactory.get_service() + user_id = cache_invalidation["user_id"] + seen: set[str] = set() + for namespace in cache_invalidation["namespaces"]: + if not namespace or namespace in seen: + continue + seen.add(namespace) + redis_service.incr(f"retrieval:version:{user_id}:{namespace}") + except Exception as exc: + logger.warning( + "Failed to invalidate retrieval cache after publication " + f"(ignored): job_id={cache_invalidation.get('job_id')}, error={exc}" + ) + + def _backfill_section_summaries( + self, + db: Session, + *, + document_id: str, + job_result_id: str, + section_summaries: dict[str, str], + ) -> None: + if not document_id or not section_summaries: + return + + try: + for path, summary in section_summaries.items(): + if not path or not summary: + continue + db.execute( + update(DocumentSection) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + .where(DocumentSection.section_path == path) + .values(summary=summary) + ) + db.flush() + logger.debug( + f"Backfilled section summaries: document_id={document_id}, " + f"count={len(section_summaries)}" + ) + except Exception as exc: + logger.warning(f"Section summary backfill failed (non-fatal): {exc}") + + def _build_cache_invalidation( + self, + db: Session, + *, + job_id: str, + published_document_state: dict[str, str] | None, + previous_document_scope: dict[str, str] | None, + ) -> dict[str, Any] | None: + job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() + if not job: + return None + + metadata = job.job_metadata or {} + namespaces = [ + JobMetadataHelper.get_namespace(metadata, "default") or "default", + ] + if previous_document_scope and previous_document_scope.get("namespace"): + namespaces.append(previous_document_scope["namespace"]) + if published_document_state and published_document_state.get("namespace"): + namespaces.append(published_document_state["namespace"]) + + return {"user_id": str(job.user_id), "namespaces": namespaces, "job_id": job_id} + + +def _should_publish_document_graph( + published_document_state: dict[str, str] | None, +) -> bool: + return published_document_state is not None and not published_document_state.get( + "skipped_all_duplicate" + ) diff --git a/packages/shared-python/shared/services/job_result_sync.py b/packages/shared-python/shared/services/job_result_sync.py new file mode 100644 index 000000000..2c7b2f5aa --- /dev/null +++ b/packages/shared-python/shared/services/job_result_sync.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from typing import Any +from uuid import uuid4 + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from shared.models.database.job_result import JobChunk, JobResult + + +class SyncJobResultWriter: + """Persist terminal Job Result artifacts inside an existing transaction.""" + + def upsert_job_result( + self, + db: Session, + job_id: str, + delivery_mode: str, + *, + inline_payload: dict[str, Any] | None = None, + result_s3_key: str | None = None, + result_size: int | None = None, + ) -> JobResult: + result = db.execute(select(JobResult).where(JobResult.job_id == job_id)) + existing = result.scalar_one_or_none() + + if existing: + existing.delivery_mode = delivery_mode + existing.inline_payload = inline_payload + existing.result_s3_key = result_s3_key + existing.result_size = result_size + db.flush() + return existing + + job_result = JobResult( + job_id=job_id, + delivery_mode=delivery_mode, + document_metadata={}, + inline_payload=inline_payload, + result_s3_key=result_s3_key, + result_size=result_size, + ) + db.add(job_result) + db.flush() + return job_result + + def replace_chunks( + self, + db: Session, + job_result_id: str, + chunks: list[dict[str, Any]], + ) -> None: + db.execute(delete(JobChunk).where(JobChunk.job_result_id == job_result_id)) + + if not chunks: + db.flush() + return + + chunk_models = [] + for index, chunk in enumerate(chunks): + chunk_identifier = chunk.get("chunk_id") or str(uuid4()) + metadata = chunk.get("metadata") + chunk_models.append( + JobChunk( + job_result_id=job_result_id, + chunk_id=chunk_identifier, + chunk_type=chunk.get("type", "paragraph"), + text=chunk.get("text"), + path=metadata.get("path") if isinstance(metadata, dict) else None, + chunk_metadata=metadata, + sort_order=chunk.get("order", index), + ) + ) + db.add_all(chunk_models) + db.flush() diff --git a/packages/shared-python/shared/services/job_success_sync.py b/packages/shared-python/shared/services/job_success_sync.py new file mode 100644 index 000000000..110e46df3 --- /dev/null +++ b/packages/shared-python/shared/services/job_success_sync.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from loguru import logger +from sqlalchemy.orm import Session + +from shared.core.state_machine.service_sync import SyncStateMachineService +from shared.models.database.webhook import WebhookEvent +from shared.services.job_publication_sync import SyncJobPublicationFinalizer +from shared.services.job_result_sync import SyncJobResultWriter +from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox + + +@dataclass(frozen=True) +class JobSuccessFinalization: + response: dict[str, Any] + cache_invalidation: dict[str, Any] | None + webhook_event: WebhookEvent | None + + +class SyncJobSuccessFinalizer: + """Finalize successful Jobs inside the lifecycle transaction.""" + + def __init__( + self, + *, + state_machine: SyncStateMachineService | None = None, + result_writer: SyncJobResultWriter | None = None, + publication_finalizer: SyncJobPublicationFinalizer | None = None, + webhook_outbox: SyncJobWebhookOutbox | None = None, + ) -> None: + self._state_machine = state_machine or SyncStateMachineService() + self._result_writer = result_writer or SyncJobResultWriter() + self._publication_finalizer = ( + publication_finalizer or SyncJobPublicationFinalizer() + ) + self._webhook_outbox = webhook_outbox or SyncJobWebhookOutbox() + + def finalize( + self, + db: Session, + *, + job_id: str, + result_s3_key: str, + checksum: str, + zip_size: int, + chunks: list[dict[str, Any]], + stored_count: int, + delivery_mode: str, + section_summaries: dict[str, str] | None, + ) -> JobSuccessFinalization: + job_result = self._result_writer.upsert_job_result( + db, + job_id, + delivery_mode, + inline_payload={"checksum": checksum}, + result_s3_key=result_s3_key, + result_size=zip_size, + ) + self._result_writer.replace_chunks(db, job_result.id, 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, + ) + + transition_ok = self._state_machine.mark_completed( + db, + job_id, + result_metadata={ + "storage_completed": True, + "stored_count": stored_count, + "delivery_mode": delivery_mode, + }, + ) + if not transition_ok: + logger.error(f"Job {job_id} mark_completed transition failed") + return JobSuccessFinalization( + response={ + "status": "failed", + "job_id": job_id, + "reason": "state_transition_failed", + }, + cache_invalidation=None, + webhook_event=None, + ) + + webhook_event = self._webhook_outbox.create_event( + db, + job_id=job_id, + event_type="job.completed", + ) + return JobSuccessFinalization( + response={ + "status": "success", + "job_id": job_id, + "stored_count": stored_count, + }, + cache_invalidation=publication_outcome.cache_invalidation, + webhook_event=webhook_event, + ) + + def run_post_commit_actions( + self, + finalization: JobSuccessFinalization, + ) -> None: + self._publication_finalizer.invalidate_cache_after_commit( + finalization.cache_invalidation, + ) + self._webhook_outbox.enqueue_after_commit(finalization.webhook_event) diff --git a/packages/shared-python/shared/services/job_webhook_outbox_sync.py b/packages/shared-python/shared/services/job_webhook_outbox_sync.py new file mode 100644 index 000000000..d699a26d5 --- /dev/null +++ b/packages/shared-python/shared/services/job_webhook_outbox_sync.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.orm import Session + +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent, WebhookEventStatus + + +class SyncJobWebhookOutbox: + """Create webhook events in-transaction and publish them after commit.""" + + def create_event( + self, + db: Session, + *, + job_id: str, + event_type: str, + extra_payload: dict[str, Any] | None = None, + ) -> WebhookEvent | None: + result = db.execute(select(Job).where(Job.job_id == job_id)) + job = result.scalar_one_or_none() + + if not job: + logger.warning(f"Job not found for webhook check: {job_id}") + return None + + webhook_url = getattr(job, "webhook_url", None) + if not job.webhook_enabled or not webhook_url: + return None + + status = "completed" if event_type == "job.completed" else "failed" + timestamp_key = f"{status}_at" + payload: dict[str, Any] = { + "event": event_type, + "job_id": job_id, + "status": status, + timestamp_key: _utc_now_naive().isoformat(), + } + if extra_payload: + payload.update(extra_payload) + + event = WebhookEvent( + job_id=job_id, + target_url=webhook_url, + payload=payload, + status=WebhookEventStatus.PENDING, + attempts=0, + ) + db.add(event) + db.flush() + logger.info(f"WebhookEvent created: event_id={event.id}, job_id={job_id}") + return event + + def enqueue_after_commit(self, webhook_event: WebhookEvent | None) -> None: + if not webhook_event: + return + + try: + from shared.services.webhook.qstash_publisher import ( + get_qstash_webhook_publisher, + ) + + publisher = get_qstash_webhook_publisher() + message_id = publisher.publish_event(webhook_event.id) + if not message_id: + logger.warning( + f"Webhook publish failed after commit: event_id={webhook_event.id}" + ) + return + logger.info( + f"Webhook published after commit: event_id={webhook_event.id}, " + f"message_id={message_id}" + ) + except Exception as exc: + logger.error( + "Failed to publish webhook after commit (event persisted): " + f"event_id={webhook_event.id}, error={exc}" + ) + + +def _utc_now_naive() -> datetime: + return datetime.now(timezone.utc).replace(tzinfo=None) From 0d66fe6b8202516957eb837ea29b3bae07e147ec Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 14:48:49 +0800 Subject: [PATCH 22/40] refactor deepen workflow and demo projections --- apps/api/app/services/demo_source_catalog.py | 254 ++------------ .../app/services/demo_source_projection.py | 246 ++++++++++++++ .../tests/contract/test_retrieval_contract.py | 4 +- .../retrieval/workflow/orchestrator.py | 313 ++---------------- .../retrieval/workflow/plan_service.py | 59 ++++ .../workflow/reference_projection.py | 40 +++ .../retrieval/workflow/runtime_config.py | 33 ++ .../retrieval/workflow/step_runner.py | 192 +++++++++++ 8 files changed, 630 insertions(+), 511 deletions(-) create mode 100644 apps/api/app/services/demo_source_projection.py create mode 100644 packages/shared-python/shared/services/retrieval/workflow/plan_service.py create mode 100644 packages/shared-python/shared/services/retrieval/workflow/reference_projection.py create mode 100644 packages/shared-python/shared/services/retrieval/workflow/runtime_config.py create mode 100644 packages/shared-python/shared/services/retrieval/workflow/step_runner.py diff --git a/apps/api/app/services/demo_source_catalog.py b/apps/api/app/services/demo_source_catalog.py index 9a7ecef40..ec6328af7 100644 --- a/apps/api/app/services/demo_source_catalog.py +++ b/apps/api/app/services/demo_source_catalog.py @@ -1,4 +1,4 @@ -"""Canonical Demo Source catalog and payload shaping.""" +"""Canonical Demo Source catalog and file access.""" from __future__ import annotations @@ -8,7 +8,8 @@ from functools import lru_cache from pathlib import Path from typing import Any -from urllib.parse import quote + +from app.services.demo_source_projection import DemoSourceProjection @dataclass(frozen=True) @@ -129,10 +130,16 @@ class DemoSourceDefinition: class DemoSourceCatalog: + def __init__(self, *, projection: DemoSourceProjection | None = None) -> None: + self._projection = projection or DemoSourceProjection() + def get_catalog(self) -> dict[str, Any]: return { "sources": [ - self._source_catalog_payload(source) + self._projection.source_catalog_payload( + source=source, + chunks=_load_source_chunks(source), + ) for source in _DEMO_SOURCE_DEFINITIONS ], } @@ -157,8 +164,12 @@ def list_chunks( "title": source.title, "mime_type": source.mime_type, "chunks": [ - _chunk_payload(source=source, chunk=chunk) - for chunk in page_chunks + self._projection.chunk_payload( + source=source, + chunk=chunk, + sort_order=start + index, + ) + for index, chunk in enumerate(page_chunks) ], "pagination": { "page": page, @@ -178,12 +189,21 @@ def get_chunk( if source is None: return None - for chunk in _load_source_chunks(source): - if demo_chunk_id in {_canonical_chunk_id(source, chunk), chunk["chunk_id"]}: + chunks = _load_source_chunks(source) + for sort_order, chunk in enumerate(chunks): + if self._projection.matches_chunk_id( + source=source, + chunk=chunk, + demo_chunk_id=demo_chunk_id, + ): return { "demo_source_id": source.demo_source_id, "canonical_document_id": source.canonical_document_id, - "chunk": _chunk_payload(source=source, chunk=chunk), + "chunk": self._projection.chunk_payload( + source=source, + chunk=chunk, + sort_order=sort_order, + ), } return None @@ -237,95 +257,10 @@ def source_directory(self, source: DemoSourceDefinition) -> Path: return _DATA_ROOT / source.asset_directory def publication_chunks(self, source: DemoSourceDefinition) -> list[dict[str, Any]]: - return [ - _publication_chunk(source=source, chunk=chunk) - for chunk in _load_source_chunks(source) - ] - - def _source_catalog_payload(self, source: DemoSourceDefinition) -> dict[str, Any]: - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "title": source.title, - "mime_type": source.mime_type, - "size_bytes": source.size_bytes, - "status": "ready", - "chunk_count": source.chunk_count, - "original_file": { - "url": f"/api/v1/demo/sources/{source.demo_source_id}/original", - "mime_type": source.mime_type, - "size_bytes": source.size_bytes, - "can_download": False, - }, - "examples": [ - self._example_payload(source=source, example=example) - for example in source.examples - ], - } - - def _example_payload( - self, - *, - source: DemoSourceDefinition, - example: DemoExampleDefinition, - ) -> dict[str, Any]: - return { - "id": example.id, - "question": example.question, - "answer": example.answer, - "citations": [ - _citation_payload(source=source, citation=citation) - for citation in example.citations - ], - } - - -def _publication_chunk( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> dict[str, Any]: - materialized_chunk = dict(chunk) - metadata = _metadata(materialized_chunk) - raw_path = _first_string(metadata.get("path"), materialized_chunk.get("path")) - publication_path = _publication_path(source=source, raw_path=raw_path) - file_path = _first_string( - metadata.get("file_path"), - metadata.get("filePath"), - materialized_chunk.get("file_path"), - materialized_chunk.get("path") if _is_media_chunk(materialized_chunk) else None, - ) - - metadata["path"] = publication_path - if file_path: - metadata["file_path"] = file_path - materialized_chunk["file_path"] = file_path - materialized_chunk["path"] = publication_path - materialized_chunk["metadata"] = metadata - return materialized_chunk - - -def _publication_path( - *, - source: DemoSourceDefinition, - raw_path: str | None, -) -> str: - prefix = f"Default_Root/{source.title}" - raw = str(raw_path or "").strip() - if not raw: - return prefix - - if "-->" in raw: - sections = [part.strip() for part in raw.split("-->")[1:] if part.strip()] - return "/".join([prefix, *sections]) if sections else prefix - - if raw.startswith("images/") or raw.startswith("tables/"): - return f"{prefix}/Assets/{raw}" - - parts = [part.strip() for part in raw.split("/") if part.strip()] - if len(parts) >= 2 and parts[0] == "Default_Root": - return raw - return prefix + return self._projection.publication_chunks( + source=source, + chunks=_load_source_chunks(source), + ) def _normalize_asset_path(asset_path: str) -> Path | None: @@ -338,131 +273,6 @@ def _normalize_asset_path(asset_path: str) -> Path | None: return Path(*parts) -def _citation_payload( - *, - source: DemoSourceDefinition, - citation: DemoCitationDefinition, -) -> dict[str, Any]: - chunk = _resolve_citation_chunk(source=source, citation=citation) - return { - "demo_source_id": source.demo_source_id, - "canonical_document_id": source.canonical_document_id, - "canonical_chunk_id": _canonical_chunk_id(source, chunk), - "chunk_id": chunk["chunk_id"], - "chunk_type": _normalize_chunk_type(chunk.get("type")), - "content": citation.content, - "description": citation.description, - "source": { - "document_id": source.canonical_document_id, - "source_file_name": source.title, - "section_path": citation.section_path, - }, - } - - -def _resolve_citation_chunk( - *, - source: DemoSourceDefinition, - citation: DemoCitationDefinition, -) -> dict[str, Any]: - chunks = _load_source_chunks(source) - normalized_content = _normalize_text(citation.content) - if normalized_content: - for chunk in chunks: - if normalized_content in _normalize_text(str(chunk.get("content") or "")): - return chunk - - for chunk in chunks: - if str(chunk.get("path") or "") == citation.section_path: - return chunk - - raise ValueError( - "Demo citation does not resolve to a canonical chunk: " - f"demo_source_id={source.demo_source_id}, section_path={citation.section_path}" - ) - - -def _chunk_payload( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> dict[str, Any]: - metadata = _metadata(chunk) - file_path = _first_string( - metadata.get("file_path"), - metadata.get("filePath"), - chunk.get("file_path"), - chunk.get("path") if _is_media_chunk(chunk) else None, - ) - return { - "id": _canonical_chunk_id(source, chunk), - "chunk_id": chunk["chunk_id"], - "chunk_type": _normalize_chunk_type(chunk.get("type")), - "content": str(chunk.get("content") or ""), - "section_path": str(chunk.get("path") or "") or None, - "source_chunk_path": str(chunk.get("path") or "") or None, - "file_path": file_path, - "sort_order": _sort_order(source=source, chunk=chunk), - "metadata": metadata, - "asset_url": _asset_url(source=source, file_path=file_path), - "created_at": None, - } - - -def _sort_order( - *, - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> int: - try: - return _load_source_chunks(source).index(chunk) - except ValueError: - return 0 - - -def _metadata(chunk: dict[str, Any]) -> dict[str, Any]: - metadata = chunk.get("metadata") - return dict(metadata) if isinstance(metadata, dict) else {} - - -def _is_media_chunk(chunk: dict[str, Any]) -> bool: - return _normalize_chunk_type(chunk.get("type")) in {"image", "table"} - - -def _canonical_chunk_id( - source: DemoSourceDefinition, - chunk: dict[str, Any], -) -> str: - return f"{source.demo_source_id}:{chunk['chunk_id']}" - - -def _asset_url( - *, - source: DemoSourceDefinition, - file_path: str | None, -) -> str | None: - if not file_path: - return None - encoded_path = quote(file_path, safe="/") - return f"/api/v1/demo/sources/{source.demo_source_id}/assets/{encoded_path}" - - -def _normalize_chunk_type(value: object) -> str: - raw = str(value or "").strip().split("\n", 1)[0].lower() - return raw if raw in {"text", "image", "table"} else "text" - - -def _first_string(*values: object) -> str | None: - for value in values: - if isinstance(value, str) and value.strip(): - return value.strip() - return None - - -def _normalize_text(value: str) -> str: - return " ".join(value.lower().split()) - - @lru_cache(maxsize=8) def _load_source_chunks(source: DemoSourceDefinition) -> tuple[dict[str, Any], ...]: chunks_path = (_DATA_ROOT / source.asset_directory) / "chunks.json" diff --git a/apps/api/app/services/demo_source_projection.py b/apps/api/app/services/demo_source_projection.py new file mode 100644 index 000000000..483e52de8 --- /dev/null +++ b/apps/api/app/services/demo_source_projection.py @@ -0,0 +1,246 @@ +"""Projection logic for canonical Demo Source data.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any +from urllib.parse import quote + +if TYPE_CHECKING: + from app.services.demo_source_catalog import ( + DemoCitationDefinition, + DemoExampleDefinition, + DemoSourceDefinition, + ) + + +class DemoSourceProjection: + def source_catalog_payload( + self, + *, + source: DemoSourceDefinition, + chunks: tuple[dict[str, Any], ...], + ) -> dict[str, Any]: + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "title": source.title, + "mime_type": source.mime_type, + "size_bytes": source.size_bytes, + "status": "ready", + "chunk_count": source.chunk_count, + "original_file": { + "url": f"/api/v1/demo/sources/{source.demo_source_id}/original", + "mime_type": source.mime_type, + "size_bytes": source.size_bytes, + "can_download": False, + }, + "examples": [ + self._example_payload(source=source, example=example, chunks=chunks) + for example in source.examples + ], + } + + def chunk_payload( + self, + *, + source: DemoSourceDefinition, + chunk: dict[str, Any], + sort_order: int, + ) -> dict[str, Any]: + metadata = _metadata(chunk) + file_path = _first_string( + metadata.get("file_path"), + metadata.get("filePath"), + chunk.get("file_path"), + chunk.get("path") if _is_media_chunk(chunk) else None, + ) + return { + "id": self.canonical_chunk_id(source=source, chunk=chunk), + "chunk_id": chunk["chunk_id"], + "chunk_type": _normalize_chunk_type(chunk.get("type")), + "content": str(chunk.get("content") or ""), + "section_path": str(chunk.get("path") or "") or None, + "source_chunk_path": str(chunk.get("path") or "") or None, + "file_path": file_path, + "sort_order": sort_order, + "metadata": metadata, + "asset_url": _asset_url(source=source, file_path=file_path), + "created_at": None, + } + + def publication_chunks( + self, + *, + source: DemoSourceDefinition, + chunks: tuple[dict[str, Any], ...], + ) -> list[dict[str, Any]]: + return [ + _publication_chunk(source=source, chunk=chunk) + for chunk in chunks + ] + + def canonical_chunk_id( + self, + *, + source: DemoSourceDefinition, + chunk: dict[str, Any], + ) -> str: + return f"{source.demo_source_id}:{chunk['chunk_id']}" + + def matches_chunk_id( + self, + *, + source: DemoSourceDefinition, + chunk: dict[str, Any], + demo_chunk_id: str, + ) -> bool: + return demo_chunk_id in { + self.canonical_chunk_id(source=source, chunk=chunk), + chunk["chunk_id"], + } + + def _example_payload( + self, + *, + source: DemoSourceDefinition, + example: DemoExampleDefinition, + chunks: tuple[dict[str, Any], ...], + ) -> dict[str, Any]: + return { + "id": example.id, + "question": example.question, + "answer": example.answer, + "citations": [ + self._citation_payload(source=source, citation=citation, chunks=chunks) + for citation in example.citations + ], + } + + def _citation_payload( + self, + *, + source: DemoSourceDefinition, + citation: DemoCitationDefinition, + chunks: tuple[dict[str, Any], ...], + ) -> dict[str, Any]: + chunk = _resolve_citation_chunk(source=source, citation=citation, chunks=chunks) + return { + "demo_source_id": source.demo_source_id, + "canonical_document_id": source.canonical_document_id, + "canonical_chunk_id": self.canonical_chunk_id(source=source, chunk=chunk), + "chunk_id": chunk["chunk_id"], + "chunk_type": _normalize_chunk_type(chunk.get("type")), + "content": citation.content, + "description": citation.description, + "source": { + "document_id": source.canonical_document_id, + "source_file_name": source.title, + "section_path": citation.section_path, + }, + } + + +def _publication_chunk( + *, + source: DemoSourceDefinition, + chunk: dict[str, Any], +) -> dict[str, Any]: + materialized_chunk = dict(chunk) + metadata = _metadata(materialized_chunk) + raw_path = _first_string(metadata.get("path"), materialized_chunk.get("path")) + publication_path = _publication_path(source=source, raw_path=raw_path) + file_path = _first_string( + metadata.get("file_path"), + metadata.get("filePath"), + materialized_chunk.get("file_path"), + materialized_chunk.get("path") if _is_media_chunk(materialized_chunk) else None, + ) + + metadata["path"] = publication_path + if file_path: + metadata["file_path"] = file_path + materialized_chunk["file_path"] = file_path + materialized_chunk["path"] = publication_path + materialized_chunk["metadata"] = metadata + return materialized_chunk + + +def _publication_path( + *, + source: DemoSourceDefinition, + raw_path: str | None, +) -> str: + prefix = f"Default_Root/{source.title}" + raw = str(raw_path or "").strip() + if not raw: + return prefix + + if "-->" in raw: + sections = [part.strip() for part in raw.split("-->")[1:] if part.strip()] + return "/".join([prefix, *sections]) if sections else prefix + + if raw.startswith("images/") or raw.startswith("tables/"): + return f"{prefix}/Assets/{raw}" + + parts = [part.strip() for part in raw.split("/") if part.strip()] + if len(parts) >= 2 and parts[0] == "Default_Root": + return raw + return prefix + + +def _resolve_citation_chunk( + *, + source: DemoSourceDefinition, + citation: DemoCitationDefinition, + chunks: tuple[dict[str, Any], ...], +) -> dict[str, Any]: + normalized_content = _normalize_text(citation.content) + if normalized_content: + for chunk in chunks: + if normalized_content in _normalize_text(str(chunk.get("content") or "")): + return chunk + + for chunk in chunks: + if str(chunk.get("path") or "") == citation.section_path: + return chunk + + raise ValueError( + "Demo citation does not resolve to a canonical chunk: " + f"demo_source_id={source.demo_source_id}, section_path={citation.section_path}" + ) + + +def _metadata(chunk: dict[str, Any]) -> dict[str, Any]: + metadata = chunk.get("metadata") + return dict(metadata) if isinstance(metadata, dict) else {} + + +def _is_media_chunk(chunk: dict[str, Any]) -> bool: + return _normalize_chunk_type(chunk.get("type")) in {"image", "table"} + + +def _asset_url( + *, + source: DemoSourceDefinition, + file_path: str | None, +) -> str | None: + if not file_path: + return None + encoded_path = quote(file_path, safe="/") + return f"/api/v1/demo/sources/{source.demo_source_id}/assets/{encoded_path}" + + +def _normalize_chunk_type(value: object) -> str: + raw = str(value or "").strip().split("\n", 1)[0].lower() + return raw if raw in {"text", "image", "table"} else "text" + + +def _first_string(*values: object) -> str | None: + for value in values: + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _normalize_text(value: str) -> str: + return " ".join(value.lower().split()) diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 7bad9363a..3ccf8b5d0 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -702,7 +702,7 @@ async def fake_retrieval_run( fake_plan, ) monkeypatch.setattr( - "shared.services.retrieval.workflow.orchestrator.RetrievalAgent.run", + "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", fake_retrieval_run, ) @@ -813,7 +813,7 @@ async def fake_retrieval_run( fake_plan, ) monkeypatch.setattr( - "shared.services.retrieval.workflow.orchestrator.RetrievalAgent.run", + "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", fake_retrieval_run, ) diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index d63735be5..3fd425521 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -2,30 +2,26 @@ from __future__ import annotations import asyncio -import os import time from collections.abc import Callable from contextlib import AbstractAsyncContextManager -from typing import Any from uuid import uuid4 from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.agentic.budget import BudgetLedger -from shared.services.retrieval.agentic.orchestrator import RetrievalAgent, _load_budget_inventory -from shared.services.retrieval.agentic.types import AgenticResult -from shared.services.retrieval.cache_service import ( - get_cached_workflow_plan, - set_cached_workflow_plan, -) +from shared.services.retrieval.agentic.orchestrator import _load_budget_inventory from shared.services.retrieval.llm_adapter import ( create_retrieval_llm_fn, create_retrieval_planner_fn, ) -from shared.services.retrieval.workflow.planner import QueryPlanner -from shared.services.retrieval.workflow.synthesizer import compose_final_answer, synthesize_step -from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult, WorkflowResult +from shared.services.retrieval.workflow.plan_service import WorkflowPlanService +from shared.services.retrieval.workflow.reference_projection import WorkflowReferenceProjection +from shared.services.retrieval.workflow.runtime_config import WorkflowRuntimeConfig +from shared.services.retrieval.workflow.step_runner import WorkflowStepRunner +from shared.services.retrieval.workflow.synthesizer import compose_final_answer +from shared.services.retrieval.workflow.types import StepResult, WorkflowResult from shared.services.retrieval.workflow.wallet import BudgetWallet DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] @@ -64,18 +60,14 @@ async def run( llm_fn=None, ) -> WorkflowResult: t0 = time.monotonic() + config = WorkflowRuntimeConfig.from_env() llm_fn = llm_fn or create_retrieval_llm_fn() planner_llm = create_retrieval_planner_fn(thinking=True) - planner_budget = _env_int('RETRIEVAL_PLANNER_THINKING_BUDGET', 4000) - wallet_total = _env_int('RETRIEVAL_WALLET_TOTAL_BUDGET', 200000) - per_retrieve = _env_int('RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET', 40000) - per_synthesize = _env_int('RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET', 6000) - max_steps = _env_int('RETRIEVAL_DECOMPOSITION_MAX_STEPS', 5) planner_ledger = BudgetLedger( - total=planner_budget, + total=config.planner_budget, planning_ratio=0.0, - bootstrap=planner_budget, + bootstrap=config.planner_budget, per_doc_min_share=0, ) total_chunks, total_docs, _chunks_count_by_doc = await _load_budget_inventory( @@ -86,32 +78,36 @@ async def run( ) planner_ledger.total_chunks = total_chunks planner_ledger.total_docs = total_docs - plan = await self._load_or_plan( + plan = await WorkflowPlanService().load_or_create( user_id=user_id, namespace=namespace, query=query, planner_llm=planner_llm, planner_ledger=planner_ledger, - max_steps=max_steps, - wallet_total=wallet_total, - per_retrieve=per_retrieve, + max_steps=config.max_steps, + wallet_total=config.wallet_total_budget, + per_retrieve=config.per_retrieve_step_budget, kb_total_docs=total_docs, kb_total_chunks=total_chunks, ) wallet = BudgetWallet( - total=wallet_total, - per_retrieve_step_default=per_retrieve, - per_synthesize_step_default=per_synthesize, + total=config.wallet_total_budget, + per_retrieve_step_default=config.per_retrieve_step_budget, + per_synthesize_step_default=config.per_synthesize_step_budget, ) ledgers = await wallet.allocate(plan) results_by_id: dict[str, StepResult] = {} - sem = asyncio.Semaphore(_env_int('RETRIEVAL_WORKFLOW_PARALLEL_MAX', 3)) + sem = asyncio.Semaphore(config.parallel_max) + step_runner = WorkflowStepRunner( + db_factory=self._get_db_factory(), + parent_run_id=self.parent_run_id, + ) for batch in plan.topological_batches(): await asyncio.gather( *[ - self._run_step( + step_runner.run_step( step=step, ledger=ledgers[step.id], results_by_id=results_by_id, @@ -136,10 +132,11 @@ async def run( answer_text = compose_final_answer(plan, results_by_id) ordered_results = [results_by_id[step.id] for step in plan.steps if step.id in results_by_id] - referenced_chunks = _dedupe_references( + reference_projection = WorkflowReferenceProjection() + referenced_chunks = reference_projection.dedupe( ref for step_result in ordered_results for ref in step_result.referenced_chunks ) - api_results = _references_to_results(referenced_chunks) + api_results = reference_projection.to_api_results(referenced_chunks) elapsed_ms = int((time.monotonic() - t0) * 1000) logger.info( 'workflow retrieval DONE: steps={} refs={} answer_chars={} elapsed={}ms', @@ -162,261 +159,3 @@ async def run( planner_snapshot=planner_ledger.snapshot(), parent_run_id=self.parent_run_id, ) - - async def _load_or_plan( - self, - *, - user_id: str, - namespace: str, - query: str, - planner_llm, - planner_ledger: BudgetLedger, - max_steps: int, - wallet_total: int, - per_retrieve: int, - kb_total_docs: int, - kb_total_chunks: int, - ) -> QueryPlan: - try: - cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query) - if cached: - return QueryPlan.from_dict(cached, original_query=query) - except Exception as exc: - logger.warning(f'workflow plan cache read failed (ignored): {exc}') - - planner = QueryPlanner( - llm_fn=planner_llm, - planner_ledger=planner_ledger, - max_steps=max_steps, - total_budget=wallet_total, - per_step_budget=per_retrieve, - ) - plan = await planner.plan( - query=query, - kb_total_docs=kb_total_docs, - kb_total_chunks=kb_total_chunks, - ) - try: - await set_cached_workflow_plan( - user_id=user_id, - namespace=namespace, - query=query, - plan=plan.to_dict(), - ) - except Exception as exc: - logger.warning(f'workflow plan cache write failed (ignored): {exc}') - return plan - - async def _run_step( - self, - *, - step: PlannedStep, - ledger: BudgetLedger, - results_by_id: dict[str, StepResult], - semaphore: asyncio.Semaphore, - user_id: str, - namespace: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - data_type: int, - signal_paths: list[str] | None, - filter_mode: str, - channels: list[str] | None, - channel_weights: dict[str, float] | None, - llm_fn, - ) -> None: - async with semaphore: - if step.step_kind == 'synthesize': - await self._run_synthesize_step(step, ledger, results_by_id, llm_fn) - return - await self._run_retrieve_step( - step=step, - ledger=ledger, - results_by_id=results_by_id, - user_id=user_id, - namespace=namespace, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - llm_fn=llm_fn, - ) - - async def _run_retrieve_step( - self, - *, - step: PlannedStep, - ledger: BudgetLedger, - results_by_id: dict[str, StepResult], - user_id: str, - namespace: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - data_type: int, - signal_paths: list[str] | None, - filter_mode: str, - channels: list[str] | None, - channel_weights: dict[str, float] | None, - llm_fn, - ) -> None: - try: - # AsyncSession is not safe for concurrent use. Workflow steps may - # run in the same topological batch, so each retrieve step opens an - # isolated session and leaves the parent session untouched. - db_factory = self._get_db_factory() - async with db_factory() as step_db: - agentic_result = await RetrievalAgent().run( - step_db, - user_id=user_id, - namespace=namespace, - query=step.sub_query, - top_k=top_k, - llm_fn=llm_fn, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - ledger=ledger, - parent_run_id=self.parent_run_id, - workflow_step_id=step.id, - ) - results_by_id[step.id] = _step_result_from_agentic(step, agentic_result) - except Exception as exc: - logger.exception(f'workflow retrieve step failed: step_id={step.id}') - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status='error', - error=str(exc), - budget_snapshot=ledger.snapshot(), - ) - - async def _run_synthesize_step( - self, - step: PlannedStep, - ledger: BudgetLedger, - results_by_id: dict[str, StepResult], - llm_fn, - ) -> None: - if llm_fn is None: - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status='skipped', - answer_text='', - error='llm unavailable for synthesis', - budget_snapshot=ledger.snapshot(), - ) - return - prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id} - try: - answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger) - refs = _dedupe_references( - ref for result in prior.values() for ref in result.referenced_chunks - ) - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status='done', - answer_text=answer, - referenced_chunks=refs, - budget_snapshot=ledger.snapshot(), - ) - except Exception as exc: - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status='budget_stop' if 'budget' in str(exc).lower() else 'error', - answer_text='(budget exhausted)' if 'budget' in str(exc).lower() else '', - error=str(exc), - budget_snapshot=ledger.snapshot(), - ) - - -def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult: - if result.answer_text: - status = 'done' - elif result.failure_reason: - status = 'not_found' - elif 'budget' in (result.stop_reason or ''): - status = 'budget_stop' - else: - status = 'done' - return StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status=status, # type: ignore[arg-type] - answer_text=result.answer_text, - evidence_text=result.evidence_text, - referenced_chunks=result.referenced_chunks, - budget_snapshot=result.budget_snapshot, - router_used=result.router_used, - stop_reason=result.stop_reason, - failure_reason=result.failure_reason, - ) - - -def _dedupe_references(refs) -> list[dict[str, Any]]: - seen: set[str] = set() - out: list[dict[str, Any]] = [] - for ref in refs: - document_id = str(ref.get('document_id') or '').strip() - chunk_id = str(ref.get('chunk_id') or '').strip() - section_path = str(ref.get('section_path') or '').strip() - file_path = str(ref.get('file_path') or '').strip() - key = ( - f'{document_id}:{chunk_id}:{section_path}:{file_path}' - if document_id and chunk_id - else str(ref) - ) - if key in seen: - continue - seen.add(key) - out.append(dict(ref)) - return out - - -def _references_to_results(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: - return [ - { - 'chunk_id': ref.get('chunk_id'), - 'document_id': ref.get('document_id'), - 'chunk_type': ref.get('chunk_type'), - 'source': { - 'document_id': ref.get('document_id'), - 'section_path': ref.get('section_path'), - }, - } - for ref in refs - ] - - -def _env_int(name: str, default: int) -> int: - try: - return int(os.environ.get(name, str(default))) - except (TypeError, ValueError): - return default diff --git a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py new file mode 100644 index 000000000..a5e3a0d43 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py @@ -0,0 +1,59 @@ +"""Plan loading and creation for decomposed retrieval workflows.""" +from __future__ import annotations + +from loguru import logger + +from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.cache_service import ( + get_cached_workflow_plan, + set_cached_workflow_plan, +) +from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.workflow.planner import QueryPlanner +from shared.services.retrieval.workflow.types import QueryPlan + + +class WorkflowPlanService: + async def load_or_create( + self, + *, + user_id: str, + namespace: str, + query: str, + planner_llm: LLMFn | None, + planner_ledger: BudgetLedger, + max_steps: int, + wallet_total: int, + per_retrieve: int, + kb_total_docs: int, + kb_total_chunks: int, + ) -> QueryPlan: + try: + cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query) + if cached: + return QueryPlan.from_dict(cached, original_query=query) + except Exception as exc: + logger.warning(f"workflow plan cache read failed (ignored): {exc}") + + planner = QueryPlanner( + llm_fn=planner_llm, + planner_ledger=planner_ledger, + max_steps=max_steps, + total_budget=wallet_total, + per_step_budget=per_retrieve, + ) + plan = await planner.plan( + query=query, + kb_total_docs=kb_total_docs, + kb_total_chunks=kb_total_chunks, + ) + try: + await set_cached_workflow_plan( + user_id=user_id, + namespace=namespace, + query=query, + plan=plan.to_dict(), + ) + except Exception as exc: + logger.warning(f"workflow plan cache write failed (ignored): {exc}") + return plan diff --git a/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py b/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py new file mode 100644 index 000000000..992085c00 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py @@ -0,0 +1,40 @@ +"""Reference projection for decomposed retrieval workflows.""" +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + + +class WorkflowReferenceProjection: + def dedupe(self, refs: Iterable[dict[str, Any]]) -> list[dict[str, Any]]: + seen: set[str] = set() + out: list[dict[str, Any]] = [] + for ref in refs: + document_id = str(ref.get("document_id") or "").strip() + chunk_id = str(ref.get("chunk_id") or "").strip() + section_path = str(ref.get("section_path") or "").strip() + file_path = str(ref.get("file_path") or "").strip() + key = ( + f"{document_id}:{chunk_id}:{section_path}:{file_path}" + if document_id and chunk_id + else str(ref) + ) + if key in seen: + continue + seen.add(key) + out.append(dict(ref)) + return out + + def to_api_results(self, refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + { + "chunk_id": ref.get("chunk_id"), + "document_id": ref.get("document_id"), + "chunk_type": ref.get("chunk_type"), + "source": { + "document_id": ref.get("document_id"), + "section_path": ref.get("section_path"), + }, + } + for ref in refs + ] diff --git a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py new file mode 100644 index 000000000..bcb6ab262 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py @@ -0,0 +1,33 @@ +"""Runtime configuration for decomposed retrieval workflows.""" +from __future__ import annotations + +import os +from dataclasses import dataclass + + +@dataclass(frozen=True) +class WorkflowRuntimeConfig: + planner_budget: int = 4000 + wallet_total_budget: int = 200000 + per_retrieve_step_budget: int = 40000 + per_synthesize_step_budget: int = 6000 + max_steps: int = 5 + parallel_max: int = 3 + + @classmethod + def from_env(cls) -> "WorkflowRuntimeConfig": + return cls( + planner_budget=_env_int("RETRIEVAL_PLANNER_THINKING_BUDGET", 4000), + wallet_total_budget=_env_int("RETRIEVAL_WALLET_TOTAL_BUDGET", 200000), + per_retrieve_step_budget=_env_int("RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET", 40000), + per_synthesize_step_budget=_env_int("RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET", 6000), + max_steps=_env_int("RETRIEVAL_DECOMPOSITION_MAX_STEPS", 5), + parallel_max=_env_int("RETRIEVAL_WORKFLOW_PARALLEL_MAX", 3), + ) + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, str(default))) + except (TypeError, ValueError): + return default diff --git a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py new file mode 100644 index 000000000..8a6be3d9d --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py @@ -0,0 +1,192 @@ +"""Step execution for decomposed retrieval workflows.""" +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.orchestrator import RetrievalAgent +from shared.services.retrieval.agentic.types import AgenticResult +from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.workflow.reference_projection import WorkflowReferenceProjection +from shared.services.retrieval.workflow.synthesizer import synthesize_step +from shared.services.retrieval.workflow.types import PlannedStep, StepResult, StepStatus + +DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] + + +class WorkflowStepRunner: + def __init__(self, *, db_factory: DbSessionFactory, parent_run_id: str) -> None: + self._db_factory = db_factory + self._parent_run_id = parent_run_id + self._references = WorkflowReferenceProjection() + + async def run_step( + self, + *, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + semaphore: asyncio.Semaphore, + user_id: str, + namespace: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int, + signal_paths: list[str] | None, + filter_mode: str, + channels: list[str] | None, + channel_weights: dict[str, float] | None, + llm_fn: LLMFn | None, + ) -> None: + async with semaphore: + if step.step_kind == "synthesize": + await self._run_synthesize_step(step, ledger, results_by_id, llm_fn) + return + await self._run_retrieve_step( + step=step, + ledger=ledger, + results_by_id=results_by_id, + user_id=user_id, + namespace=namespace, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + llm_fn=llm_fn, + ) + + async def _run_retrieve_step( + self, + *, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + user_id: str, + namespace: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int, + signal_paths: list[str] | None, + filter_mode: str, + channels: list[str] | None, + channel_weights: dict[str, float] | None, + llm_fn: LLMFn | None, + ) -> None: + try: + async with self._db_factory() as step_db: + agentic_result = await RetrievalAgent().run( + step_db, + user_id=user_id, + namespace=namespace, + query=step.sub_query, + top_k=top_k, + llm_fn=llm_fn, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + ledger=ledger, + parent_run_id=self._parent_run_id, + workflow_step_id=step.id, + ) + results_by_id[step.id] = _step_result_from_agentic(step, agentic_result) + except Exception as exc: + logger.exception(f"workflow retrieve step failed: step_id={step.id}") + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="error", + error=str(exc), + budget_snapshot=ledger.snapshot(), + ) + + async def _run_synthesize_step( + self, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + llm_fn: LLMFn | None, + ) -> None: + if llm_fn is None: + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="skipped", + answer_text="", + error="llm unavailable for synthesis", + budget_snapshot=ledger.snapshot(), + ) + return + prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id} + try: + answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger) + refs = self._references.dedupe(ref for result in prior.values() for ref in result.referenced_chunks) + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="done", + answer_text=answer, + referenced_chunks=refs, + budget_snapshot=ledger.snapshot(), + ) + except Exception as exc: + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status="budget_stop" if "budget" in str(exc).lower() else "error", + answer_text="(budget exhausted)" if "budget" in str(exc).lower() else "", + error=str(exc), + budget_snapshot=ledger.snapshot(), + ) + + +def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult: + if result.answer_text: + status: StepStatus = "done" + elif result.failure_reason: + status = "not_found" + elif "budget" in (result.stop_reason or ""): + status = "budget_stop" + else: + status = "done" + return StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status=status, + answer_text=result.answer_text, + evidence_text=result.evidence_text, + referenced_chunks=result.referenced_chunks, + budget_snapshot=result.budget_snapshot, + router_used=result.router_used, + stop_reason=result.stop_reason, + failure_reason=result.failure_reason, + ) From 556ebd2b0400c02b87affba22cbac0481a3e96b9 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:08:57 +0800 Subject: [PATCH 23/40] refactor deepen agentic and zip modules --- .../app/services/connect_builder/builder.py | 527 ------ .../services/connect_builder/graph_builder.py | 1538 ----------------- .../connect_builder/summary_builder.py | 9 +- .../retrieval/agentic/discovery_phase.py | 252 +++ .../retrieval/agentic/document_navigation.py | 414 +++++ .../retrieval/agentic/orchestrator.py | 666 +------ .../services/retrieval/agentic/runtime.py | 146 ++ .../services/retrieval/graph_service.py | 16 +- .../services/storage/zip_result_schema.py | 345 ++++ .../services/storage/zip_result_service.py | 426 +---- 10 files changed, 1239 insertions(+), 3100 deletions(-) delete mode 100644 apps/worker/app/services/connect_builder/builder.py delete mode 100644 apps/worker/app/services/connect_builder/graph_builder.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/document_navigation.py create mode 100644 packages/shared-python/shared/services/retrieval/agentic/runtime.py create mode 100644 packages/shared-python/shared/services/storage/zip_result_schema.py diff --git a/apps/worker/app/services/connect_builder/builder.py b/apps/worker/app/services/connect_builder/builder.py deleted file mode 100644 index b788c7ff1..000000000 --- a/apps/worker/app/services/connect_builder/builder.py +++ /dev/null @@ -1,527 +0,0 @@ -""" -ConnectTo Builder — KB-level post-processor for inter-chunk relationships. - -This module discovers relationships between chunks across different files -within a knowledge base, populating the `connectto` column in the DataFrame. -""" - -import json -import re -from collections import defaultdict -from difflib import SequenceMatcher -from typing import Any, Dict, List, Optional, Tuple - -from loguru import logger - -from shared.utils.chunk_refs import CHUNK_REF_RE - -# ─── Relation Type Registry (extensible, not hard-coded) ────────────────────── - -RELATION_REGISTRY: Dict[str, Dict[str, Any]] = { - "related": { - "description": "Chunks share common concepts or topics", - "requires_llm": False, - }, - # TODO: LLM-classified relation types — uncomment when classify_relation() is implemented - # "contradicts": { - # "description": "Chunks describe opposing or contradictory facts", - # "requires_llm": True, - # }, - # "causal": { - # "description": "Chunks have a cause-effect relationship", - # "requires_llm": True, - # }, - # "extends": { - # "description": "One chunk extends or improves upon the other", - # "requires_llm": True, - # }, - # "supports": { - # "description": "One chunk provides evidence supporting the other", - # "requires_llm": True, - # }, - # "same_method": { - # "description": "Chunks discuss the same methodology or technique", - # "requires_llm": True, - # }, - # "same_data": { - # "description": "Chunks reference the same dataset", - # "requires_llm": True, - # }, -} - - -# ─── Default Configuration ──────────────────────────────────────────────────── - -DEFAULT_CONFIG: Dict[str, Any] = { - # Minimum number of shared keywords to consider a connection - "min_keyword_overlap": 3, - # Weight multiplier for keyword score (linear) - "keyword_score_weight": 1.0, - # Maximum connections per chunk (top-N by score) - "max_connections_per_chunk": 10, - # Minimum score threshold to create a connection - "min_score_threshold": 0.8, - # Only connect chunks from different files (skip intra-file) - "cross_file_only": True, - # Maximum character overlap ratio to allow (filter near-duplicates) - # Pairs with SequenceMatcher.ratio() >= this are considered duplicates - "max_content_overlap": 0.8, -} - - -# ─── Keyword Normalization ──────────────────────────────────────────────────── - -# TODO: Synonym dictionary for advanced normalization (e.g. "RL" ↔ "reinforcement learning") -_SYNONYM_MAP: Dict[str, str] = {} - - -def _normalize_keyword(keyword: str) -> str: - """ - Normalize a keyword for matching: - - lowercase - - strip whitespace - - collapse multiple spaces - - apply synonym mapping (TODO) - - Args: - keyword: Raw keyword string. - - Returns: - Normalized keyword string. - """ - kw = keyword.lower().strip() - kw = re.sub(r"\s+", " ", kw) - - # Apply synonym mapping if available - return _SYNONYM_MAP.get(kw, kw) - - -def _extract_file_key(path: str) -> str: - """ - Extract a file-level key from a chunk's path to determine - whether two chunks belong to the same file. - - Example paths: - "Default_Root/paper.pdf/Section 1/Subsection" → "Default_Root/paper.pdf" - "KB_DATA/reports/annual.docx/Table 1" → "KB_DATA/reports/annual.docx" - - Heuristic: take the path up to and including the first segment - that looks like a filename (has an extension). - """ - if not path: - return "" - - parts = path.replace("\\", "/").split("/") - file_parts = [] - for part in parts: - file_parts.append(part) - # Check if this segment looks like a file (has extension) - if "." in part and not part.startswith("."): - break - - return "/".join(file_parts) - - -# ─── Keyword Inverted Index ────────────────────────────────────────────────── - - -def _build_keyword_index( - chunks: List[Dict[str, Any]], -) -> Dict[str, List[Tuple[str, str]]]: - """ - Build an inverted index: normalized_keyword → [(chunk_id, file_key)]. - - Args: - chunks: List of chunk dicts, each having: - - "chunk_id": str - - "metadata" or "keywords": keyword source - - "path": str - - Returns: - Dict mapping normalized keyword → list of (chunk_id, file_key) tuples. - """ - index: Dict[str, List[Tuple[str, str]]] = defaultdict(list) - - for chunk in chunks: - chunk_id = chunk.get("chunk_id") or chunk.get("know_id", "") - path = chunk.get("path", "") - file_key = _extract_file_key(path) - - # Extract keywords from metadata or top-level - keywords = _get_keywords(chunk) - if not keywords: - continue - - for kw in keywords: - normalized = _normalize_keyword(kw) - if normalized: - index[normalized].append((str(chunk_id), file_key)) - - return dict(index) - - -def _get_keywords(chunk: Dict[str, Any]) -> List[str]: - """ - Extract keywords from a chunk, supporting multiple input formats: - - chunk["metadata"]["keywords"] (list) - - chunk["keywords"] (list or semicolon-separated string) - - chunk["metadata"]["tokens"] or chunk["tokens"] (fallback: jieba word chain) - """ - # Try metadata.keywords first - metadata = chunk.get("metadata", {}) - if isinstance(metadata, dict): - kws = metadata.get("keywords", []) - if isinstance(kws, list) and kws: - return kws - - # Try top-level keywords - kws = chunk.get("keywords", []) - if isinstance(kws, list) and kws: - return kws - if isinstance(kws, str) and kws.strip(): - # Parse semicolon or comma separated - if ";" in kws: - return [k.strip() for k in kws.split(";") if k.strip()] - elif "," in kws: - return [k.strip() for k in kws.split(",") if k.strip()] - return [kws.strip()] - - # ─── Fallback: tokens (jieba word chain) ────────────────────────────── - tokens = _parse_tokens_field( - metadata.get("tokens") if isinstance(metadata, dict) else None - ) - if not tokens: - tokens = _parse_tokens_field(chunk.get("tokens")) - return tokens - - -# Pre-compiled patterns for token noise filtering -_UUID_LIKE_RE = re.compile(r"^[0-9a-f]{4,}$", re.IGNORECASE) -_MARKER_PREFIXES = ("IMAGE_", "TABLE_", "PTXT", "image-", "table-") - - -def _parse_tokens_field(raw) -> List[str]: - """ - Parse the tokens field into a filtered keyword list. - - Accepts: - - List[str]: already parsed (from chunks.json after safe_parse_tokens) - - str with ';': semicolon-separated tokens (new format, matches keywords) - - str with '->': arrow-separated jieba word chain (legacy format) - - str with "['...']": legacy list-repr format - - Filters out noise: single-char tokens, UUIDs, IMAGE_/TABLE_ markers. - """ - if raw is None: - return [] - - # Already a list (from chunks.json) - if isinstance(raw, list): - words = raw - elif isinstance(raw, str): - raw = raw.strip() - if not raw: - return [] - # List-repr format: "['w1;w2;w3']" or "['w1->w2->w3']" - if raw.startswith("[") and raw.endswith("]"): - inner = raw[1:-1].strip() - if (inner.startswith("'") and inner.endswith("'")) or ( - inner.startswith('"') and inner.endswith('"') - ): - inner = inner[1:-1] - raw = inner - # Determine separator: semicolon (new) or arrow (legacy) - if ";" in raw: - words = [t.strip() for t in raw.split(";") if t.strip()] - elif "->" in raw: - words = [t.strip() for t in raw.split("->") if t.strip()] - else: - return [] - else: - return [] - - # Filter noise - filtered = [] - for w in words: - if len(w) <= 1: - continue - if any(w.startswith(p) for p in _MARKER_PREFIXES) or CHUNK_REF_RE.fullmatch(w): - continue - if _UUID_LIKE_RE.match(w): - continue - filtered.append(w) - return filtered - - -# ─── Scoring ───────────────────────────────────────────────────────────────── - - -def _compute_keyword_score( - shared_kws: set, - kws_a: set, - kws_b: set, - weight: float = 1.0, -) -> float: - """ - Compute keyword overlap score using character-length-weighted scoring. - - Longer tokens contribute more: a four-character term contributes twice the - weight of a two-character term. - Formula: score = weight * sum(len(kw) for shared) / min(sum(len) for A, sum(len) for B) - - Args: - shared_kws: Set of shared (normalized) keywords. - kws_a: Full keyword set of chunk A. - kws_b: Full keyword set of chunk B. - weight: Score multiplier. - - Returns: - Float score in [0, weight]. - """ - weighted_a = sum(len(k) for k in kws_a) - weighted_b = sum(len(k) for k in kws_b) - denominator = min(weighted_a, weighted_b) - if denominator == 0: - return 0.0 - weighted_shared = sum(len(k) for k in shared_kws) - return weight * weighted_shared / denominator - - -# ─── Main Entry Point ──────────────────────────────────────────────────────── - - -def build_connections( - chunks: List[Dict[str, Any]], - config: Optional[Dict[str, Any]] = None, -) -> Dict[str, List[Dict[str, Any]]]: - """ - Compute "related" connections between chunks based on keyword overlap. - - Args: - chunks: List of chunk dicts. Each must have: - - chunk_id (or know_id): str - - path: str - - metadata.keywords or keywords: List[str] - config: Optional overrides for DEFAULT_CONFIG. - - Returns: - Dict mapping chunk_id → list of connection dicts: - [{"target": "...", "relation": "related", "score": 0.82, "keywords": ["PPO", "RL"]}] - """ - cfg = {**DEFAULT_CONFIG, **(config or {})} - - min_overlap = cfg["min_keyword_overlap"] - kw_weight = cfg["keyword_score_weight"] - max_conns = cfg["max_connections_per_chunk"] - min_score = cfg["min_score_threshold"] - cross_only = cfg["cross_file_only"] - max_overlap = cfg.get("max_content_overlap", 0.8) - - # Build inverted index - kw_index = _build_keyword_index(chunks) - - # Pre-compute per-chunk data - chunk_data: Dict[ - str, Tuple[str, set] - ] = {} # chunk_id → (file_key, normalized_keywords) - chunk_content: Dict[str, str] = {} # chunk_id → content (for dedup) - for chunk in chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if not cid: - continue - file_key = _extract_file_key(chunk.get("path", "")) - kws = _get_keywords(chunk) - normalized_kws = {_normalize_keyword(k) for k in kws if k} - normalized_kws.discard("") - chunk_data[cid] = (file_key, normalized_kws) - chunk_content[cid] = chunk.get("content") or chunk.get("text", "") - - # For each chunk, find candidates via keyword index - connections: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - - for cid, (file_key, my_kws) in chunk_data.items(): - if not my_kws: - continue - - # Collect candidates and their shared keywords - candidate_shared: Dict[str, set] = defaultdict(set) # target_id → shared_kw set - - for kw in my_kws: - entries = kw_index.get(kw, []) - for target_id, target_file in entries: - if target_id == cid: - continue - if cross_only and target_file == file_key: - continue - candidate_shared[target_id].add(kw) - - # Score and filter candidates - scored: List[Tuple[str, float, set]] = [] - for target_id, shared_kws in candidate_shared.items(): - if len(shared_kws) < min_overlap: - continue - - target_data = chunk_data.get(target_id) - if not target_data: - continue - - _, target_kws = target_data - score = _compute_keyword_score( - shared_kws=shared_kws, - kws_a=my_kws, - kws_b=target_kws, - weight=kw_weight, - ) - if score >= min_score: - # Near-duplicate filter: skip pairs with high character overlap - if max_overlap < 1.0: - src_text = chunk_content.get(cid, "") - tgt_text = chunk_content.get(target_id, "") - if src_text and tgt_text: - char_ratio = SequenceMatcher(None, src_text, tgt_text).ratio() - if char_ratio >= max_overlap: - continue - scored.append((target_id, score, shared_kws)) - - # Sort by score descending, keep top-N - scored.sort(key=lambda x: x[1], reverse=True) - for target_id, score, shared_kws in scored[:max_conns]: - connections[cid].append( - { - "target": target_id, - "relation": "related", - "score": round(score, 4), - "keywords": sorted(shared_kws), - } - ) - - total_edges = sum(len(v) for v in connections.values()) - logger.info( - f"🔗 ConnectTo: {total_edges} connections found " - f"across {len(connections)} chunks " - f"(config: min_overlap={min_overlap}, threshold={min_score})" - ) - - return dict(connections) - - -# ─── Serialization ──────────────────────────────────────────────────────────── - - -def serialize_connections(connections: List[Dict[str, Any]]) -> str: - """ - Serialize a chunk's connection list to a JSON string - for storage in the `connectto` DataFrame column. - - Args: - connections: List of connection dicts. - - Returns: - JSON string, or empty string if no connections. - """ - if not connections: - return "" - return json.dumps(connections, ensure_ascii=False, separators=(",", ":")) - - -def deserialize_connections(raw: Any) -> List[Dict[str, Any]]: - """ - Deserialize the `connectto` column value back to a list of connections. - - Handles: - - JSON array string: '[{"target": "...", ...}]' - - Single target string: 'doc/section-a' - - Empty / NaN / None: returns [] - - Args: - raw: Raw value from DataFrame connectto column. - - Returns: - List of connection dicts. - """ - if raw is None: - return [] - - try: - import pandas as pd - - if pd.isna(raw): - return [] - except (ImportError, TypeError, ValueError): - pass - - raw_str = str(raw).strip() - if not raw_str: - return [] - - # Try JSON parse first - if raw_str.startswith("["): - try: - parsed = json.loads(raw_str) - if isinstance(parsed, list): - return parsed - except json.JSONDecodeError: - pass - - if raw_str: - return [ - {"target": raw_str, "relation": "related", "score": 1.0, "keywords": []} - ] - - return [] - - -# ─── LLM Relation Classification (stub) ────────────────────────────────────── - - -def classify_relation( - summary_a: str, - summary_b: str, - shared_keywords: List[str], - llm_client: Any = None, -) -> Dict[str, Any]: - """ - Classify the specific relation type between two related chunks using LLM. - - This function is a **stub** — the concrete classification prompt and LLM - call logic are TODO. Currently returns "related" for all pairs. - - Args: - summary_a: Summary text of chunk A. - summary_b: Summary text of chunk B. - shared_keywords: Keywords they share. - llm_client: Optional LLM client for making classification calls. - - Returns: - Dict with keys: - - "relation": str (from RELATION_REGISTRY) - - "reason": str (human-readable explanation) - - "confidence": float (0.0 ~ 1.0) - - TODO: Implement classification prompt: - Given two knowledge chunks: - [Chunk A]: {summary_a} - [Chunk B]: {summary_b} - Shared concepts: {shared_keywords} - - Classify their relationship: - - contradicts: A and B describe opposing facts - - causal: A and B have a cause-effect relationship - - extends: B extends or improves upon A - - supports: B provides evidence for A - - same_method: A and B use the same methodology - - same_data: A and B use the same dataset - - related: Related but none of the above - - other: Has a clear relationship not listed above (describe it) - - Return JSON: {"relation": "...", "reason": "...", "confidence": 0.0~1.0} - """ - return { - "relation": "related", - "reason": ( - f"Keyword overlap: {', '.join(shared_keywords)}" - if shared_keywords - else "Keyword overlap" - ), - "confidence": 1.0, - } diff --git a/apps/worker/app/services/connect_builder/graph_builder.py b/apps/worker/app/services/connect_builder/graph_builder.py deleted file mode 100644 index 4803034eb..000000000 --- a/apps/worker/app/services/connect_builder/graph_builder.py +++ /dev/null @@ -1,1538 +0,0 @@ -""" -Knowledge Graph Builder — KB-level knowledge graph assembler (v2.0). - -Assembles a file-level knowledge_graph.json from parsed chunks + connect_builder edges. -Deployed to ~/.knowhere/{kb_id}/ and grows incrementally as more files are parsed. - -Architecture: - - files: per-file summaries (chunks_count, types, top_keywords, importance) - - edges: cross-file relationships (aggregated from chunk-level connections) - - chunk_stats.json: per-chunk usage tracking (hit_count, last_hit, decay) - - Per-file chunks.json: full chunk data lives in subdirectories - -Usage: - # One-stop API (recommended) - graph = build_and_deploy(chunks, kb_id="my_kb", parsed_output_dir=add_dir) - - # Manual: first build - graph = build_knowledge_graph(all_chunks, connections, kb_id="my_kb") - - # Manual: incremental update - graph = update_knowledge_graph(existing_graph, new_chunks, existing_chunks) -""" - -import json -import math -import os -import re -from collections import defaultdict -from datetime import datetime, timezone -from typing import Any, Dict, List, Optional, Tuple - -from app.services.connect_builder.builder import ( - DEFAULT_CONFIG, - _compute_keyword_score, - _extract_file_key, - _get_keywords, - _normalize_keyword, -) -from loguru import logger - -from shared.utils.chunk_refs import CHUNK_REF_PATTERN - -# ─── Tree Construction ─────────────────────────────────────────────────────── - - -def _build_tree_from_paths(paths: List[str]) -> Dict[str, Any]: - """ - Rebuild hierarchical tree from chunk path list. - - Args: - paths: List of chunk paths, e.g. ["Default_Root/report.pdf/Section 1/1.1", ...] - - Returns: - Nested dict tree rooted at Default_Root. - """ - root: Dict[str, Any] = {} - for path in paths: - if not path: - continue - nodes = [n.strip() for n in path.split("/") if n.strip()] - current = root - for node in nodes: - if node not in current: - current[node] = {} - current = current[node] - return root - - -def _merge_tree(base: Dict[str, Any], addition: Dict[str, Any]) -> Dict[str, Any]: - """ - Deep-merge two tree dicts. Addition is merged INTO base (in-place). - - Args: - base: Existing tree. - addition: New tree to merge in. - - Returns: - The merged base dict (same reference). - """ - for key, value in addition.items(): - if key in base and isinstance(base[key], dict) and isinstance(value, dict): - _merge_tree(base[key], value) - else: - base[key] = value - return base - - -# ─── Node Extraction ───────────────────────────────────────────────────────── - - -def _chunks_to_nodes( - chunks: List[Dict[str, Any]], - content_preview_len: int = 200, -) -> List[Dict[str, Any]]: - """ - Extract node metadata from chunks for the knowledge graph. - - Args: - chunks: List of normalized chunk dicts. - content_preview_len: Max characters for content_preview. - - Returns: - List of node dicts with: id, type, path, summary, keywords, content_preview. - """ - nodes = [] - for chunk in chunks: - chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if not chunk_id: - continue - - content = chunk.get("content") or chunk.get("text", "") - metadata = chunk.get("metadata", {}) - if not isinstance(metadata, dict): - metadata = {} - - # Extract keywords from metadata or top-level - keywords = metadata.get("keywords", []) - if not keywords: - keywords = chunk.get("keywords", []) - if isinstance(keywords, str): - keywords = [k.strip() for k in keywords.split(";") if k.strip()] - - node = { - "id": chunk_id, - "type": chunk.get("type", "text"), - "path": chunk.get("path", ""), - "summary": metadata.get("summary") or chunk.get("summary", ""), - "keywords": keywords, - "content_preview": content[:content_preview_len] if content else "", - } - nodes.append(node) - - return nodes - - -# ─── Edge Extraction ───────────────────────────────────────────────────────── - - -def _connections_to_edges( - connections: Dict[str, List[Dict[str, Any]]], -) -> List[Dict[str, Any]]: - """ - Convert connect_builder output to deduplicated edge list. - connect_builder produces bidirectional entries (A→B and B→A); - we deduplicate to keep only one edge per pair. - - Args: - connections: Output from build_connections(), mapping chunk_id → list of connections. - - Returns: - List of edge dicts: {source, target, relation, score, shared_keywords}. - """ - seen_pairs: set = set() - edges = [] - - for source_id, conn_list in connections.items(): - for conn in conn_list: - target_id = conn.get("target", "") - pair_key = tuple(sorted([source_id, target_id])) - if pair_key in seen_pairs: - continue - seen_pairs.add(pair_key) - - edges.append( - { - "source": source_id, - "target": target_id, - "relation": conn.get("relation", "related"), - "score": conn.get("score", 0.0), - "shared_keywords": conn.get("keywords", []), - } - ) - - return edges - - -def _merge_related_connections_into_chunks( - chunks: List[Dict[str, Any]], - connections: Dict[str, List[Dict[str, Any]]], -) -> None: - """Backfill related connections into chunk metadata without touching embeds.""" - if not chunks or not connections: - return - - chunk_map = { - str(chunk.get("chunk_id") or chunk.get("know_id", "")): chunk - for chunk in chunks - if chunk.get("chunk_id") or chunk.get("know_id") - } - - for chunk_id, conn_list in connections.items(): - chunk = chunk_map.get(str(chunk_id)) - if not chunk: - continue - metadata = chunk.setdefault("metadata", {}) - if not isinstance(metadata, dict): - metadata = {} - chunk["metadata"] = metadata - existing = metadata.get("connect_to", []) - if not isinstance(existing, list): - existing = [] - - merged = [] - seen = set() - for item in existing: - if not isinstance(item, dict): - continue - key = ( - str(item.get("target") or ""), - str(item.get("relation") or "related"), - str(item.get("ref") or ""), - ) - if key in seen: - continue - seen.add(key) - merged.append(item) - - for conn in conn_list: - if not isinstance(conn, dict): - continue - if conn.get("relation", "related") != "related": - continue - key = ( - str(conn.get("target") or ""), - "related", - "", - ) - if key in seen: - continue - seen.add(key) - merged.append( - { - "target": conn.get("target", ""), - "relation": "related", - "score": conn.get("score", 0.0), - "keywords": conn.get("keywords", []), - } - ) - - metadata["connect_to"] = merged - - -def _save_chunks_by_source_file(kb_dir: str, chunks: List[Dict[str, Any]]) -> None: - """Persist grouped chunks.json files after metadata backfill.""" - grouped: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - for chunk in chunks: - source_file = chunk.get("_source_file") - if not source_file: - continue - cleaned = dict(chunk) - cleaned.pop("_source_file", None) - grouped[str(source_file)].append(cleaned) - - for source_file, source_chunks in grouped.items(): - output_path = os.path.join(kb_dir, source_file, "chunks.json") - _save_chunks(source_chunks, output_path) - - -# ─── Incremental Matching ──────────────────────────────────────────────────── - - -def _incremental_connections( - new_chunks: List[Dict[str, Any]], - existing_chunks: List[Dict[str, Any]], - config: Optional[Dict[str, Any]] = None, -) -> Dict[str, List[Dict[str, Any]]]: - """ - Match ONLY new_chunks ↔ existing_chunks (skip existing ↔ existing). - Reuses connect_builder scoring functions. - - Complexity: O(new × existing) instead of O(all²). - - Args: - new_chunks: Newly parsed file's chunks. - existing_chunks: All previously known chunks. - config: Optional config overrides (same keys as connect_builder.DEFAULT_CONFIG). - - Returns: - Dict mapping chunk_id → list of connection dicts (same format as build_connections). - """ - from difflib import SequenceMatcher - - cfg = {**DEFAULT_CONFIG, **(config or {})} - min_overlap = cfg["min_keyword_overlap"] - kw_weight = cfg["keyword_score_weight"] - max_conns = cfg["max_connections_per_chunk"] - min_score = cfg["min_score_threshold"] - cross_only = cfg["cross_file_only"] - max_content_overlap = cfg.get("max_content_overlap", 0.8) - - # Pre-compute keyword sets for new chunks - new_data: Dict[str, Tuple[str, set, str]] = {} # id → (file_key, kw_set, content) - for chunk in new_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if not cid: - continue - file_key = _extract_file_key(chunk.get("path", "")) - kws = _get_keywords(chunk) - normalized = {_normalize_keyword(k) for k in kws if k} - normalized.discard("") - content = chunk.get("content") or chunk.get("text", "") - new_data[cid] = (file_key, normalized, content) - - # Pre-compute keyword sets for existing chunks - existing_data: Dict[str, Tuple[str, set, str]] = {} - for chunk in existing_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if not cid: - continue - file_key = _extract_file_key(chunk.get("path", "")) - kws = _get_keywords(chunk) - normalized = {_normalize_keyword(k) for k in kws if k} - normalized.discard("") - content = chunk.get("content") or chunk.get("text", "") - existing_data[cid] = (file_key, normalized, content) - - # Build keyword index for existing chunks only - existing_kw_index: Dict[str, List[str]] = defaultdict(list) # kw → [chunk_id] - for cid, (_, kw_set, _) in existing_data.items(): - for kw in kw_set: - existing_kw_index[kw].append(cid) - - connections: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - - # For each new chunk, find candidates in existing chunks - for new_id, (new_file, new_kws, new_content) in new_data.items(): - if not new_kws: - continue - - candidate_shared: Dict[str, set] = defaultdict(set) - for kw in new_kws: - for existing_id in existing_kw_index.get(kw, []): - if cross_only: - existing_file = existing_data[existing_id][0] - if existing_file == new_file: - continue - candidate_shared[existing_id].add(kw) - - # Score and filter - scored: List[Tuple[str, float, set]] = [] - for existing_id, shared_kws in candidate_shared.items(): - if len(shared_kws) < min_overlap: - continue - - existing_kws = existing_data[existing_id][1] - score = _compute_keyword_score( - shared_kws=shared_kws, - kws_a=new_kws, - kws_b=existing_kws, - weight=kw_weight, - ) - if score >= min_score: - # Near-duplicate filter - if max_content_overlap < 1.0: - existing_content = existing_data[existing_id][2] - if new_content and existing_content: - ratio = SequenceMatcher( - None, new_content, existing_content - ).ratio() - if ratio >= max_content_overlap: - continue - scored.append((existing_id, score, shared_kws)) - - scored.sort(key=lambda x: x[1], reverse=True) - for existing_id, score, shared_kws in scored[:max_conns]: - conn = { - "target": existing_id, - "relation": "related", - "score": round(score, 4), - "keywords": sorted(shared_kws), - } - connections[new_id].append(conn) - # Bidirectional: also add reverse edge - connections[existing_id].append( - { - "target": new_id, - "relation": "related", - "score": round(score, 4), - "keywords": sorted(shared_kws), - } - ) - - total = sum(len(v) for v in connections.values()) - logger.info( - f"🔗 Incremental connections: {total} new edges " - f"between {len(new_data)} new chunks and {len(existing_data)} existing chunks" - ) - - return dict(connections) - - -# ─── File-Level Aggregation (v2.0) ─────────────────────────────────────────── - -# Token filtering — same logic as text_utils._is_meaningful_token -_CN_EN_NUM_RE = re.compile(r"[\u4e00-\u9fff]|[A-Za-z]+|\d+(?:\.\d+)?") -_CHUNK_MARKER_RE = re.compile( - rf"{CHUNK_REF_PATTERN}|image-\d+|table-\d+", - re.IGNORECASE, -) - - -def _is_meaningful_token(token: str) -> bool: - """Check if a token is worth keeping (same logic as text_utils).""" - if not _CN_EN_NUM_RE.search(token): - return False - if len(token) == 1: - return False - if re.fullmatch(r"\d+(?:\.\d+)?", token): - return False - return True - - -def _extract_tokens_from_content(content: str) -> List[str]: - """Extract meaningful tokens from content using jieba (regex fallback).""" - content = _CHUNK_MARKER_RE.sub("", content) - # Strip HTML tags and entities (table chunks contain raw HTML) - content = re.sub(r"<[^>]+>", " ", content) - content = re.sub(r"&\w+;", " ", content) - try: - import jieba - - if hasattr(jieba, "lcut"): - raw = jieba.lcut(content) - else: - raw = list(jieba.cut(content)) - except ImportError: - raw = re.split(r"[\s,;，；。！？、\-/]+", content) - return [t for t in raw if _is_meaningful_token(t)] - - -def _get_chunk_keywords(chunk: Dict[str, Any]) -> List[str]: - """Get keywords for a chunk; falls back to tokens from content if empty.""" - keywords = _get_keywords(chunk) - meaningful = [k for k in keywords if _is_meaningful_token(k)] - if meaningful: - return meaningful - content = chunk.get("content") or chunk.get("text", "") - if not content: - return [] - tokens = _extract_tokens_from_content(content) - seen = set() - unique = [] - for t in tokens: - normalized = _normalize_keyword(t) - if normalized and normalized not in seen: - seen.add(normalized) - unique.append(normalized) - return unique - - -def _compute_tfidf_top_keywords( - file_chunks: Dict[str, List[Dict[str, Any]]], - top_k: int = 6, -) -> Dict[str, List[str]]: - """ - TF-IDF top keywords per file. - TF = chunks in file containing keyword. IDF = log(total_files / files_with_keyword). - """ - total_files = len(file_chunks) - if total_files == 0: - return {} - - file_kw_tf: Dict[str, Dict[str, int]] = {} - doc_freq: Dict[str, int] = defaultdict(int) - - for fk, chunks in file_chunks.items(): - kw_count: Dict[str, int] = defaultdict(int) - file_kw_set: set = set() - for chunk in chunks: - for kw in _get_chunk_keywords(chunk): - normalized = _normalize_keyword(kw) - if normalized: - kw_count[normalized] += 1 - file_kw_set.add(normalized) - file_kw_tf[fk] = dict(kw_count) - for kw in file_kw_set: - doc_freq[kw] += 1 - - result: Dict[str, List[str]] = {} - for fk, kw_count in file_kw_tf.items(): - scored = [] - for kw, tf in kw_count.items(): - if total_files == 1: - score = tf # Single-file KB: pure frequency - else: - idf = ( - math.log(total_files / doc_freq[kw]) - if doc_freq[kw] < total_files - else 0.1 - ) - score = tf * idf - scored.append((score, tf, kw)) - scored.sort(key=lambda x: (x[0], x[1]), reverse=True) - result[fk] = [kw for _, _, kw in scored[:top_k]] - - return result - - -def _compute_file_importance( - chunk_ids: List[str], - chunk_stats: Dict[str, Dict[str, Any]], - half_life_days: float = 30.0, - alpha: float = 0.7, - beta: float = 0.3, -) -> float: - """importance = α × usage_heat + β × freshness""" - if not chunk_ids: - return 0.0 - total_relevance = 0.0 - earliest_created = None - for cid in chunk_ids: - stat = chunk_stats.get(cid, {}) - hc = stat.get("hit_count", 0) - lh = stat.get("last_hit") - ca = stat.get("created_at") - if hc > 0 and lh: - total_relevance += relevance_score(hc, lh, half_life_days) - if ca and (earliest_created is None or ca < earliest_created): - earliest_created = ca - usage_heat = total_relevance / len(chunk_ids) - freshness = ( - relevance_score(1, earliest_created, half_life_days) - if earliest_created - else 1.0 - ) - return round(alpha * usage_heat + beta * freshness, 4) - - -def _aggregate_file_level_edges( - chunk_edges: List[Dict[str, Any]], - chunk_to_file: Dict[str, str], - chunk_paths: Optional[Dict[str, str]] = None, - max_top_connections: int = 10, -) -> List[Dict[str, Any]]: - """ - Aggregate chunk-level edges into file-level edges. - Shows top_connections with readable chunk names instead of raw keywords. - """ - if chunk_paths is None: - chunk_paths = {} - - pair_data: Dict[Tuple[str, str], Dict[str, List[Dict[str, Any]]]] = {} - for edge in chunk_edges: - src_id = str(edge.get("source", "") or "") - tgt_id = str(edge.get("target", "") or "") - sf = chunk_to_file.get(src_id, "") - tf = chunk_to_file.get(tgt_id, "") - if not sf or not tf or sf == tf: - continue - src_path = chunk_paths.get(src_id) or src_id - tgt_path = chunk_paths.get(tgt_id) or tgt_id - src_name = src_path.rsplit("/", 1)[-1] if "/" in src_path else src_path - tgt_name = tgt_path.rsplit("/", 1)[-1] if "/" in tgt_path else tgt_path - - if sf <= tf: - pk: Tuple[str, str] = (sf, tf) - connection = { - "source_chunk": src_name, - "source_id": src_id, - "target_chunk": tgt_name, - "target_id": tgt_id, - "relation": edge.get("relation", "related"), - "score": edge.get("score", 0), - } - else: - pk = (tf, sf) - connection = { - "source_chunk": tgt_name, - "source_id": tgt_id, - "target_chunk": src_name, - "target_id": src_id, - "relation": edge.get("relation", "related"), - "score": edge.get("score", 0), - } - - if pk not in pair_data: - pair_data[pk] = {"connections": []} - pair_data[pk]["connections"].append(connection) - - file_edges = [] - for (f1, f2), data in pair_data.items(): - conns = data["connections"] - # Sort by score desc, take top N - conns.sort(key=lambda x: x["score"], reverse=True) - scores = [c["score"] for c in conns] - file_edges.append( - { - "source": f1, - "target": f2, - "connection_count": len(conns), - "avg_score": round(sum(scores) / len(scores), 4) if scores else 0, - "top_connections": conns[:max_top_connections], - } - ) - file_edges.sort(key=lambda x: x["connection_count"], reverse=True) - return file_edges - - -# ─── Main API ──────────────────────────────────────────────────────────────── - - -def _get_source_file(chunk: Dict[str, Any]) -> str: - """ - Get the source document file for a chunk. - Uses `_source_file` tag (injected by build_and_deploy / _load_all_chunks_from_kb) - for correct grouping of images/tables with their parent document. - Falls back to `_extract_file_key` for backwards compatibility. - """ - sf = chunk.get("_source_file") - if sf: - return sf - return _extract_file_key(chunk.get("path", "")) - - -def build_knowledge_graph( - all_chunks: List[Dict[str, Any]], - connections: Dict[str, List[Dict[str, Any]]], - kb_id: str = "", - chunk_stats: Optional[Dict[str, Dict[str, Any]]] = None, - file_summaries: Optional[Dict[str, str]] = None, -) -> Dict[str, Any]: - """Build a file-level knowledge graph (v2.0).""" - if chunk_stats is None: - chunk_stats = {} - - file_chunks: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - chunk_to_file: Dict[str, str] = {} - chunk_paths: Dict[str, str] = {} - for chunk in all_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - fk = _get_source_file(chunk) - if fk: - file_chunks[fk].append(chunk) - if cid: - chunk_to_file[cid] = fk - chunk_paths[cid] = chunk.get("path", "") - - file_keywords = _compute_tfidf_top_keywords(file_chunks) - chunk_edges = _connections_to_edges(connections) - file_edges = _aggregate_file_level_edges(chunk_edges, chunk_to_file, chunk_paths) - - files_dict = {} - for fk, chunks in file_chunks.items(): - types_count: Dict[str, int] = defaultdict(int) - cids = [] - for c in chunks: - types_count[c.get("type", "text")] += 1 - cid = str(c.get("chunk_id") or c.get("know_id", "")) - if cid: - cids.append(cid) - files_dict[fk] = { - "chunks_count": len(chunks), - "types": dict(types_count), - "top_keywords": file_keywords.get(fk, []), - "top_summary": (file_summaries or {}).get(fk, ""), - "importance": _compute_file_importance(cids, chunk_stats), - "created_at": datetime.now(timezone.utc).isoformat(), - } - - total_chunks = sum(f["chunks_count"] for f in files_dict.values()) - graph = { - "version": "2.0", - "updated_at": datetime.now(timezone.utc).isoformat(), - "kb_id": kb_id, - "stats": { - "total_files": len(files_dict), - "total_chunks": total_chunks, - "total_cross_file_edges": len(file_edges), - }, - "files": files_dict, - "edges": file_edges, - } - logger.info( - f"📊 Knowledge graph built: " - f"{graph['stats']['total_files']} files, " - f"{graph['stats']['total_chunks']} chunks, " - f"{graph['stats']['total_cross_file_edges']} edges" - ) - return graph - - -def update_knowledge_graph( - existing_graph: Dict[str, Any], - new_chunks: List[Dict[str, Any]], - existing_chunks: List[Dict[str, Any]], - kb_id: str = "", - connect_config: Optional[Dict[str, Any]] = None, - chunk_stats: Optional[Dict[str, Dict[str, Any]]] = None, - file_summaries: Optional[Dict[str, str]] = None, - new_connections: Optional[Dict[str, List[Dict[str, Any]]]] = None, -) -> Dict[str, Any]: - """Incrementally update a file-level knowledge graph with new chunks.""" - if chunk_stats is None: - chunk_stats = {} - - all_combined = existing_chunks + new_chunks - file_chunks: Dict[str, List[Dict[str, Any]]] = defaultdict(list) - chunk_to_file: Dict[str, str] = {} - chunk_paths: Dict[str, str] = {} - for chunk in all_combined: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - fk = _get_source_file(chunk) - if fk: - file_chunks[fk].append(chunk) - if cid: - chunk_to_file[cid] = fk - chunk_paths[cid] = chunk.get("path", "") - - file_keywords = _compute_tfidf_top_keywords(file_chunks) - - if new_connections is None: - new_connections = _incremental_connections( - new_chunks=new_chunks, - existing_chunks=existing_chunks, - config=connect_config, - ) - new_chunk_edges = _connections_to_edges(new_connections) - existing_file_edges = existing_graph.get("edges", []) - new_file_edges = _aggregate_file_level_edges( - new_chunk_edges, chunk_to_file, chunk_paths - ) - - # Merge file edges - merged_map: Dict[Tuple[str, str], Dict] = {} - for edge in existing_file_edges + new_file_edges: - pk = tuple(sorted([edge["source"], edge["target"]])) - if pk not in merged_map: - merged_map[pk] = edge - else: - old = merged_map[pk] - # Merge connections, dedup by chunk pair - all_conns = old.get("top_connections", []) + edge.get("top_connections", []) - seen = set() - deduped = [] - for c in all_conns: - pair = (c.get("source_chunk", ""), c.get("target_chunk", "")) - if pair not in seen: - seen.add(pair) - deduped.append(c) - deduped.sort(key=lambda x: x.get("score", 0), reverse=True) - tc = old["connection_count"] + edge["connection_count"] - scores = [c.get("score", 0) for c in deduped] - avg = sum(scores) / len(scores) if scores else 0 - merged_map[pk] = { - "source": pk[0], - "target": pk[1], - "connection_count": tc, - "avg_score": round(avg, 4), - "top_connections": deduped[:10], - } - all_file_edges = sorted( - merged_map.values(), key=lambda x: x["connection_count"], reverse=True - ) - - existing_files = existing_graph.get("files", {}) - files_dict = {} - new_file_count = 0 - for fk, chunks in file_chunks.items(): - types_count: Dict[str, int] = defaultdict(int) - cids = [] - for c in chunks: - types_count[c.get("type", "text")] += 1 - cid = str(c.get("chunk_id") or c.get("know_id", "")) - if cid: - cids.append(cid) - created_at = existing_files.get(fk, {}).get( - "created_at", datetime.now(timezone.utc).isoformat() - ) - if fk not in existing_files: - new_file_count += 1 - files_dict[fk] = { - "chunks_count": len(chunks), - "types": dict(types_count), - "top_keywords": file_keywords.get(fk, []), - "top_summary": (file_summaries or {}).get(fk, "") or existing_files.get(fk, {}).get("top_summary", ""), - "importance": _compute_file_importance(cids, chunk_stats), - "created_at": created_at, - } - - total_chunks = sum(f["chunks_count"] for f in files_dict.values()) - graph = { - "version": "2.0", - "updated_at": datetime.now(timezone.utc).isoformat(), - "kb_id": kb_id or existing_graph.get("kb_id", ""), - "stats": { - "total_files": len(files_dict), - "total_chunks": total_chunks, - "total_cross_file_edges": len(all_file_edges), - }, - "files": files_dict, - "edges": all_file_edges, - } - logger.info( - f"📊 Knowledge graph updated: " - f"+{new_file_count} files → " - f"total {graph['stats']['total_files']} files, " - f"{graph['stats']['total_chunks']} chunks, " - f"{graph['stats']['total_cross_file_edges']} edges" - ) - return graph - - -# ─── Configuration ──────────────────────────────────────────────────────────── - -KNOWHERE_HOME = os.path.expanduser(os.environ.get("KNOWHERE_HOME", "~/.knowhere")) - - -def _get_kb_dir(kb_id: str) -> str: - """Get the knowledge base directory path.""" - return os.path.join(KNOWHERE_HOME, kb_id) - - -def _get_kg_path(kb_id: str) -> str: - """Get the knowledge_graph.json path for a KB.""" - return os.path.join(_get_kb_dir(kb_id), "knowledge_graph.json") - - -def _get_stats_path(kb_id: str) -> str: - """Get the chunk_stats.json path for a KB.""" - return os.path.join(_get_kb_dir(kb_id), "chunk_stats.json") - - -def _empty_knowledge_graph(kb_id: str) -> Dict[str, Any]: - """Build an empty v2 knowledge graph for a KB with no local chunks.""" - return { - "version": "2.0", - "updated_at": datetime.now(timezone.utc).isoformat(), - "kb_id": kb_id, - "stats": { - "total_files": 0, - "total_chunks": 0, - "total_cross_file_edges": 0, - }, - "files": {}, - "edges": [], - } - - -# ─── Chunk Usage Tracking ───────────────────────────────────────────────────── - - -def load_chunk_stats(kb_id: str) -> Dict[str, Dict[str, Any]]: - """Load chunk usage stats from chunk_stats.json.""" - path = _get_stats_path(kb_id) - if not os.path.exists(path): - return {} - try: - with open(path, "r", encoding="utf-8") as f: - return json.load(f) - except (json.JSONDecodeError, IOError): - return {} - - -def record_chunk_hits( - kb_id: str, - chunk_ids: List[str], -) -> None: - """ - Record that chunks were accessed (returned in search results). - Updates hit_count and last_hit for each chunk. - - Args: - kb_id: Knowledge base ID. - chunk_ids: List of chunk IDs that were hit. - """ - stats = load_chunk_stats(kb_id) - now = datetime.now(timezone.utc).isoformat() - - for cid in chunk_ids: - if cid not in stats: - stats[cid] = { - "hit_count": 0, - "first_hit": now, - "last_hit": now, - "created_at": now, - } - stats[cid]["hit_count"] += 1 - stats[cid]["last_hit"] = now - - path = _get_stats_path(kb_id) - os.makedirs(os.path.dirname(path), exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - json.dump(stats, f, ensure_ascii=False, indent=2) - - -def relevance_score( - hit_count: int, - last_hit_iso: str, - half_life_days: float = 30.0, -) -> float: - """ - Compute relevance score with exponential decay. - Higher hit_count + more recent access → higher score. - - Args: - hit_count: Number of times this chunk was accessed. - last_hit_iso: ISO timestamp of last access. - half_life_days: Days until relevance halves. - - Returns: - Decay-weighted score. - """ - try: - last_hit_dt = datetime.fromisoformat(last_hit_iso) - days_since = (datetime.now(timezone.utc) - last_hit_dt).total_seconds() / 86400 - except (ValueError, TypeError): - days_since = 0 - - decay = math.exp(-0.693 * days_since / half_life_days) - return hit_count * decay - - -# ─── File I/O ───────────────────────────────────────────────────────────────── - - -def save_knowledge_graph(graph: Dict[str, Any], output_path: str) -> str: - """Save knowledge graph to a JSON file.""" - os.makedirs(os.path.dirname(output_path), exist_ok=True) - with open(output_path, "w", encoding="utf-8") as f: - json.dump(graph, f, ensure_ascii=False, indent=2) - logger.info(f"💾 Knowledge graph saved: {output_path}") - return output_path - - -def load_knowledge_graph(path: str) -> Optional[Dict[str, Any]]: - """Load an existing knowledge graph from JSON file.""" - if not os.path.exists(path): - return None - try: - with open(path, "r", encoding="utf-8") as f: - return json.load(f) - except (json.JSONDecodeError, IOError) as e: - logger.warning(f"Failed to load knowledge graph from {path}: {e}") - return None - - -def _save_chunks(chunks: List[Dict[str, Any]], output_path: str) -> None: - """Save chunks data to a JSON file in the standard format.""" - os.makedirs(os.path.dirname(output_path), exist_ok=True) - with open(output_path, "w", encoding="utf-8") as f: - json.dump({"chunks": chunks}, f, ensure_ascii=False, indent=2) - - -def _load_chunks(path: str) -> List[Dict[str, Any]]: - """Load chunks from a stored chunks.json file.""" - if not os.path.exists(path): - return [] - try: - with open(path, "r", encoding="utf-8") as f: - data = json.load(f) - if isinstance(data, dict) and "chunks" in data: - return data["chunks"] - if isinstance(data, list): - return data - except (json.JSONDecodeError, IOError): - pass - return [] - - -def extract_chunks_from_graph(graph: Dict[str, Any]) -> List[Dict[str, Any]]: - """ - Reconstruct minimal chunk dicts from graph for incremental matching. - This is a last-resort fallback when subdirectory chunks.json files are unavailable. - v2.0: no chunk-level data in graph; returns empty list. - Legacy: falls back to node_index or nodes array. - """ - chunks = [] - # v2.0: files dict doesn't store chunk IDs, return empty - if graph.get("version", "").startswith("2."): - return chunks - # Legacy v1.x: handle node_index - node_index = graph.get("node_index", {}) - if node_index: - for chunk_id, file_key in node_index.items(): - chunks.append( - { - "chunk_id": chunk_id, - "path": file_key, - "content": "", - "metadata": {"keywords": []}, - } - ) - return chunks - # Legacy: handle old nodes array - for node in graph.get("nodes", []): - chunks.append( - { - "chunk_id": node["id"], - "path": node.get("path", ""), - "content": node.get("content_preview", ""), - "metadata": {"keywords": node.get("keywords", [])}, - } - ) - return chunks - - -# ─── Chunk ID Dedup ────────────────────────────────────────────────────────── - - -def _dedup_chunks_by_content( - new_chunks: List[Dict[str, Any]], - existing_chunks: List[Dict[str, Any]], -) -> List[Dict[str, Any]]: - """ - Filter new_chunks: discard any whose chunk_id already exists in existing_chunks. - - Since all parsers now generate deterministic know_id (content-hash based), - identical content always produces the same chunk_id. Simple set comparison - replaces the old strip+hash pipeline. - - Returns: - List of new chunks that have no chunk_id duplicate in existing_chunks. - """ - existing_ids = { - str(c.get("chunk_id") or c.get("know_id", "")) for c in existing_chunks - } - existing_ids.discard("") - - deduped = [] - skipped = 0 - for chunk in new_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if cid and cid in existing_ids: - skipped += 1 - else: - deduped.append(chunk) - - if skipped > 0: - logger.info( - f"📊 chunk dedup: {skipped} duplicate chunks skipped " - f"(by chunk_id), {len(deduped)} chunks to add" - ) - return deduped - - -def _load_all_chunks_from_kb(kb_dir: str) -> List[Dict[str, Any]]: - """ - Load all chunks from per-file chunks.json files under a KB directory. - Tags each chunk with _source_file = subdirectory name (= source document). - """ - all_chunks = [] - for entry in os.listdir(kb_dir): - entry_path = os.path.join(kb_dir, entry) - if not os.path.isdir(entry_path): - continue - chunks_file = os.path.join(entry_path, "chunks.json") - if os.path.isfile(chunks_file): - loaded = _load_chunks(chunks_file) - for chunk in loaded: - chunk["_source_file"] = entry - all_chunks.extend(loaded) - return all_chunks - - -def _source_files_from_chunks(chunks: List[Dict[str, Any]]) -> set[str]: - """Return the source-file set represented by loaded KB chunks.""" - return { - str(chunk.get("_source_file") or "").strip() - for chunk in chunks - if str(chunk.get("_source_file") or "").strip() - } - - -def _prune_chunk_stats(kb_id: str, chunks: List[Dict[str, Any]]) -> None: - """Remove chunk_stats entries whose chunks no longer exist on disk.""" - stats_path = _get_stats_path(kb_id) - if not os.path.exists(stats_path): - return - - stats = load_chunk_stats(kb_id) - live_chunk_ids = { - str(chunk.get("chunk_id") or chunk.get("know_id", "")) - for chunk in chunks - if chunk.get("chunk_id") or chunk.get("know_id") - } - pruned = {cid: data for cid, data in stats.items() if cid in live_chunk_ids} - if len(pruned) == len(stats): - return - - os.makedirs(os.path.dirname(stats_path), exist_ok=True) - with open(stats_path, "w", encoding="utf-8") as f: - json.dump(pruned, f, ensure_ascii=False, indent=2) - logger.info( - f"📊 Chunk stats pruned: {len(stats) - len(pruned)} stale chunks removed" - ) - - -def sync_knowledge_graph_with_local_files( - kb_id: str, - connect_config: Optional[Dict[str, Any]] = None, - summary_use_llm: bool = False, -) -> Dict[str, Any]: - """Synchronize knowledge_graph.json with current ~/.knowhere/{kb_id} files. - - This is intentionally a no-op when graph files match on-disk document - directories. If a user manually deletes a local parsed document directory, - the graph is rebuilt from remaining chunks and stale chunk_stats entries are - removed. - """ - from app.services.connect_builder.builder import build_connections - from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries - - kb_dir = _get_kb_dir(kb_id) - kg_path = _get_kg_path(kb_id) - os.makedirs(kb_dir, exist_ok=True) - - existing_graph = load_knowledge_graph(kg_path) - chunks_on_disk = _load_all_chunks_from_kb(kb_dir) - disk_files = _source_files_from_chunks(chunks_on_disk) - graph_files = set((existing_graph or {}).get("files", {}).keys()) - - if existing_graph is not None and graph_files == disk_files: - _prune_chunk_stats(kb_id, chunks_on_disk) - return existing_graph - - removed_files = sorted(graph_files - disk_files) - added_files = sorted(disk_files - graph_files) - logger.info( - "📊 Syncing Knowledge Graph with local files: " - f"removed={removed_files}, added={added_files}" - ) - - if not chunks_on_disk: - graph = _empty_knowledge_graph(kb_id) - save_knowledge_graph(graph, kg_path) - _prune_chunk_stats(kb_id, []) - return graph - - try: - file_summaries = enrich_doc_nav_summaries( - kb_dir=kb_dir, - source_file=None, - use_llm=summary_use_llm, - ) - except Exception as e: - logger.warning(f"sync enrich_doc_nav_summaries failed: {e}") - file_summaries = {} - - stats = load_chunk_stats(kb_id) - connections = build_connections(chunks_on_disk, connect_config) - _merge_related_connections_into_chunks(chunks_on_disk, connections) - _save_chunks_by_source_file(kb_dir, chunks_on_disk) - graph = build_knowledge_graph( - all_chunks=chunks_on_disk, - connections=connections, - kb_id=kb_id, - chunk_stats=stats, - file_summaries=file_summaries, - ) - save_knowledge_graph(graph, kg_path) - _prune_chunk_stats(kb_id, chunks_on_disk) - return graph - - -# ─── MCP Auto-Registration ─────────────────────────────────────────────────── - - -def _get_mcp_server_path() -> str: - """Get the absolute path to the MCP server script. - - Points to the consolidated knowhere-mcp/server.py (unified server - with both Cloud API and local search tools). - """ - # Navigate from graph_builder.py → project root → knowhere-mcp/server.py - # graph_builder.py is at: apps/worker/app/services/connect_builder/ - project_root = os.path.normpath( - os.path.join( - os.path.dirname(os.path.abspath(__file__)), - "..", - "..", - "..", - "..", - "..", - ) - ) - return os.path.join(project_root, "knowhere-mcp", "server.py") - - -def _auto_register_mcp() -> None: - """ - Detect installed Agent products and auto-register the knowhere MCP server. - Only runs on first deploy (when ~/.knowhere/ is freshly created). - - Supported products: - - Cursor: ~/.cursor/mcp.json - - Claude Code: ~/.claude.json (project-level) or ~/.claude/claude_code_config.json - """ - mcp_server_path = os.path.normpath(_get_mcp_server_path()) - home = os.path.expanduser("~") - - knowhere_mcp_entry = { - "command": "python3", - "args": [mcp_server_path], - "env": { - "KNOWHERE_API_KEY": os.environ.get("KNOWHERE_API_KEY", ""), - }, - } - - registered = [] - - # ── Cursor ──────────────────────────────────────────────────────────── - cursor_mcp = os.path.join(home, ".cursor", "mcp.json") - if os.path.isdir(os.path.join(home, ".cursor")): - try: - existing = {} - if os.path.exists(cursor_mcp): - with open(cursor_mcp, "r") as f: - existing = json.load(f) - - servers = existing.get("mcpServers", {}) - # Update even if "knowhere" exists (to point to new server) - if "knowhere" not in servers or "mcp/knowhere_mcp_server" in str( - servers.get("knowhere", {}).get("args", []) - ): - servers["knowhere"] = knowhere_mcp_entry - existing["mcpServers"] = servers - with open(cursor_mcp, "w") as f: - json.dump(existing, f, indent=2) - registered.append("Cursor") - except Exception as e: - logger.debug(f"Cursor MCP registration skipped: {e}") - - # ── Claude Code ─────────────────────────────────────────────────────── - claude_config = os.path.join(home, ".claude.json") - if os.path.exists(claude_config) or os.path.isdir(os.path.join(home, ".claude")): - try: - existing = {} - if os.path.exists(claude_config): - with open(claude_config, "r") as f: - existing = json.load(f) - - servers = existing.get("mcpServers", {}) - if "knowhere" not in servers or "mcp/knowhere_mcp_server" in str( - servers.get("knowhere", {}).get("args", []) - ): - servers["knowhere"] = knowhere_mcp_entry - existing["mcpServers"] = servers - with open(claude_config, "w") as f: - json.dump(existing, f, indent=2) - registered.append("Claude Code") - except Exception as e: - logger.debug(f"Claude Code MCP registration skipped: {e}") - - if registered: - logger.info(f"🔌 MCP auto-registered for: {', '.join(registered)}") - else: - logger.debug("No Agent products detected for MCP auto-registration") - - -# ─── doc_nav section extraction for GraphNode persistence ──────────────────── - - - - - -# ─── One-Stop API ───────────────────────────────────────────────────────────── - - -def build_and_deploy( - chunks: List[Dict[str, Any]], - kb_id: str, - parsed_output_dir: Optional[str] = None, - connect_config: Optional[Dict[str, Any]] = None, - rebuild_all: bool = True, - summary_use_llm: bool = False, -) -> Dict[str, Any]: - """ - One-stop knowledge graph build/update + deploy to ~/.knowhere/ + MCP register. - - This is the main entry point for callers (parse services, debug scripts, etc). - Callers just provide chunks + kb_id; everything else is automatic. - - Flow: - 1. If parsed_output_dir provided → copy full parsed output to ~/.knowhere/{kb_id}/data/ - 2. Check if ~/.knowhere/{kb_id}/knowledge_graph.json exists - - No → build_knowledge_graph() (full build) - - rebuild_all=True → scan KB dir for existing files, merge with new chunks - - rebuild_all=False → only use the new chunks (ignore previous files) - - Yes → update_knowledge_graph() (incremental) - 3. Save knowledge_graph.json to ~/.knowhere/{kb_id}/ - 4. On first-ever deploy → _auto_register_mcp() - - Args: - chunks: Parsed chunks from the current file. - kb_id: Knowledge base identifier (e.g. dataset name). - parsed_output_dir: Path to the parsed output directory (add_dir) containing - images, tables, doc_nav.json etc. If provided, its contents are - copied to ~/.knowhere/{kb_id}/data/{dirname}/. - connect_config: Optional config overrides for connect_builder. - rebuild_all: When knowledge_graph.json is missing, whether to scan the - KB directory for existing chunk data and include them in the full - rebuild. Defaults to True. Set to False to only process the new - chunks (legacy behavior). - summary_use_llm: If True, use LLM to generate coherent hierarchical - summaries (slow, costs API tokens). If False (default), use - lightweight title enumeration (e.g. "This section covers: Section 1, - Section 2"). Only affects `top_summary` and `summary` fields. - - Returns: - The knowledge graph dict. - """ - import shutil - - from app.services.connect_builder.builder import build_connections - - kg_path = _get_kg_path(kb_id) - kb_dir = _get_kb_dir(kb_id) - - # Detect if this is a first-ever deploy (for MCP registration) - first_deploy = not os.path.exists(KNOWHERE_HOME) - - # Ensure directory exists - os.makedirs(kb_dir, exist_ok=True) - - # Load existing state BEFORE deploy (to avoid counting new file's chunks twice) - # Determine source_file early so we can exclude it from existing_chunks - source_file = ( - os.path.basename(parsed_output_dir) - if parsed_output_dir and os.path.isdir(parsed_output_dir) - else None - ) - existing_graph = load_knowledge_graph(kg_path) - - # Load chunks once and reuse — avoids the double-load where - # sync_knowledge_graph_with_local_files internally calls - # _load_all_chunks_from_kb and then we call it again. - all_on_disk: List[Dict[str, Any]] = [] - if existing_graph is not None: - all_on_disk = _load_all_chunks_from_kb(kb_dir) - disk_files = _source_files_from_chunks(all_on_disk) - graph_files = set(existing_graph.get("files", {}).keys()) - - if graph_files != disk_files: - # Disk state diverged from graph → full sync rebuild. - # sync will reload chunks internally (it backfills connect_to - # metadata), so we reload afterwards to pick up changes. - existing_graph = sync_knowledge_graph_with_local_files( - kb_id=kb_id, - connect_config=connect_config, - summary_use_llm=summary_use_llm, - ) - all_on_disk = _load_all_chunks_from_kb(kb_dir) - else: - # Files match — fast-path: just prune stale chunk_stats. - _prune_chunk_stats(kb_id, all_on_disk) - - if existing_graph is not None: - if not all_on_disk: - existing_chunks = extract_chunks_from_graph(existing_graph) - else: - # Exclude chunks from the current source file — they may already - # be on disk if parsed_output_dir is inside kb_dir (debug_parse). - # Without this filter, _dedup_chunks_by_content would treat them - # as "existing" and skip the incremental update entirely. - existing_chunks = ( - [c for c in all_on_disk if c.get("_source_file") != source_file] - if source_file - else all_on_disk - ) - else: - existing_chunks = [] - - # ── Deploy parsed output (images, tables, hierarchy, etc.) ── - if parsed_output_dir and os.path.isdir(parsed_output_dir) and source_file is not None: - deploy_target = os.path.join(kb_dir, source_file) - - # Skip copy if parsed output is already in the target location - parsed_abs = os.path.normpath(os.path.abspath(parsed_output_dir)) - target_abs = os.path.normpath(os.path.abspath(deploy_target)) - if parsed_abs == target_abs: - logger.info( - f"📂 Parsed output already at target: {deploy_target} (skip copy)" - ) - else: - if os.path.exists(deploy_target): - shutil.rmtree(deploy_target) - shutil.copytree(parsed_output_dir, deploy_target) - # Delete ZIP files from deployed directory (no longer needed) - import glob - - for zip_file in glob.glob(os.path.join(deploy_target, "*.zip")): - os.remove(zip_file) - logger.info(f"📂 Parsed output deployed: {deploy_target}") - - # Tag all chunks with source document for correct file-level grouping - if source_file: - for chunk in chunks: - chunk["_source_file"] = source_file - - # ── Generate hierarchical summaries ── - from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries - try: - file_summaries = enrich_doc_nav_summaries( - kb_dir=kb_dir, - source_file=source_file, - use_llm=summary_use_llm, - ) - except Exception as e: - logger.warning(f"doc_nav summary enrichment failed: {e}") - file_summaries = {} - - # Load chunk_stats for importance calculation - stats = load_chunk_stats(kb_id) - - if existing_graph is None: - # ── First build: full ── - if rebuild_all: - # Scan KB dir for existing chunk data (deploy already happened, - # so the new file's chunks are on disk if parsed_output_dir was given). - all_on_disk = _load_all_chunks_from_kb(kb_dir) - if source_file and all_on_disk: - # New file already deployed → all_on_disk includes it, no merge needed - all_chunks = all_on_disk - else: - # New file not deployed to disk (no parsed_output_dir), - # or KB dir was empty → merge in-memory chunks with disk data. - # Dedup by chunk_id to prevent double-counting. - seen_ids = { - str(c.get("chunk_id") or c.get("know_id", "")) for c in all_on_disk - } - extra = [ - c - for c in chunks - if str(c.get("chunk_id") or c.get("know_id", "")) not in seen_ids - ] - all_chunks = all_on_disk + extra - # Full rebuild: generate summaries for ALL files, not just source_file - from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries as _enrich_nav - try: - all_nav_summaries = _enrich_nav( - kb_dir=kb_dir, - source_file=None, - use_llm=summary_use_llm, - ) - file_summaries.update(all_nav_summaries) - except Exception as e: - logger.warning(f"Full rebuild enrich_doc_nav_summaries failed: {e}") - logger.info( - f"📊 rebuild Knowledge Graph " - f"(rebuild_all=True, {len(all_chunks)} chunks from KB dir) ..." - ) - else: - all_chunks = chunks - logger.info( - "📊 rebuild Knowledge Graph (rebuild_all=False, new chunks only) ..." - ) - - connections = build_connections(all_chunks, connect_config) - _merge_related_connections_into_chunks(all_chunks, connections) - _save_chunks_by_source_file(kb_dir, all_chunks) - stats_chunks = all_chunks - graph = build_knowledge_graph( - all_chunks=all_chunks, - connections=connections, - kb_id=kb_id, - chunk_stats=stats, - file_summaries=file_summaries, - ) - else: - # ── Incremental update ── - # Content-hash dedup: discard new chunks identical to existing ones - # to preserve established graph edges and relationships. - deduped_new = _dedup_chunks_by_content(chunks, existing_chunks) - if len(deduped_new) == 0: - logger.info( - "📊 All new chunks are duplicates of existing data, " - "skipping incremental update" - ) - stats_chunks = existing_chunks - graph = existing_graph - # Still inject summaries even if chunks unchanged - for fk, fdata in graph.get("files", {}).items(): - if file_summaries and fk in file_summaries and not fdata.get("top_summary"): - fdata["top_summary"] = file_summaries[fk] - else: - logger.info( - f"📊 incremental update Knowledge Graph " - f"({len(deduped_new)} new, {len(chunks) - len(deduped_new)} skipped) ..." - ) - related_connections = _incremental_connections( - new_chunks=deduped_new, - existing_chunks=existing_chunks, - config=connect_config, - ) - _merge_related_connections_into_chunks( - existing_chunks + deduped_new, related_connections - ) - _save_chunks_by_source_file(kb_dir, existing_chunks + deduped_new) - stats_chunks = existing_chunks + deduped_new - graph = update_knowledge_graph( - existing_graph=existing_graph, - new_chunks=deduped_new, - existing_chunks=existing_chunks, - kb_id=kb_id, - connect_config=connect_config, - chunk_stats=stats, - file_summaries=file_summaries, - new_connections=related_connections, - ) - - # Save graph - save_knowledge_graph(graph, kg_path) - - # Initialize chunk_stats.json with created_at for all new chunks - stats_path = _get_stats_path(kb_id) - now = datetime.now(timezone.utc).isoformat() - updated = False - for chunk in stats_chunks: - cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) - if cid and cid not in stats: - stats[cid] = { - "hit_count": 0, - "first_hit": None, - "last_hit": None, - "created_at": now, - } - updated = True - if updated: - os.makedirs(os.path.dirname(stats_path), exist_ok=True) - with open(stats_path, "w", encoding="utf-8") as f: - json.dump(stats, f, ensure_ascii=False, indent=2) - logger.info(f"📊 Chunk stats initialized: {len(stats)} chunks tracked") - - logger.info( - f"✅ Knowledge Graph deployed to {kb_dir}: " - f"{graph['stats']['total_files']} files, " - f"{graph['stats']['total_chunks']} chunks, " - f"{graph['stats']['total_cross_file_edges']} edges" - ) - - # Auto-register MCP on first deploy - if first_deploy: - try: - _auto_register_mcp() - except Exception as e: - logger.debug(f"MCP auto-registration skipped: {e}") - - return graph diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index cac4b38bb..6d1e507c8 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -8,8 +8,6 @@ Usage (standalone): from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries enrich_doc_nav_summaries(kb_dir, source_file="report.pdf") - -Called by graph_builder.build_and_deploy() after file deploy, before KG build. """ import json @@ -131,10 +129,9 @@ def ensure_doc_nav_json( if os.path.exists(nav_path) and not overwrite: return nav_path - # Re-use ZipResultService's builder to keep the format canonical - from shared.services.storage.zip_result_service import ZipResultService - svc = ZipResultService() - doc_nav = svc._build_doc_nav(chunks, source_file_name) + from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder + + doc_nav = ZipResultSchemaBuilder().build_doc_nav(chunks, source_file_name) _save_doc_nav(file_dir, doc_nav) return nav_path diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py new file mode 100644 index 000000000..2af9ce690 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py @@ -0,0 +1,252 @@ +"""Discovery and document selection phase for agentic retrieval.""" +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document +from shared.services.retrieval.agentic import tools +from shared.services.retrieval.agentic.budget import BudgetExceeded +from shared.services.retrieval.agentic.trace import TraceRecorder +from shared.services.retrieval.agentic.types import AgentState, CandidateDoc, ToolResult +from shared.services.retrieval.llm_adapter import LLMFn + + +async def run_initial_discovery( + db: AsyncSession, + *, + state: AgentState, + trace: TraceRecorder, + trace_enabled: bool, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int, + signal_paths: list[str] | None, + filter_mode: str, + channels: list[str] | None, + channel_weights: dict[str, float] | None, + bootstrap_llm_fn: LLMFn | None, +) -> list[dict[str, Any]]: + discovery_kwargs: dict[str, Any] = { + "user_id": user_id, + "namespace": namespace, + "query": query, + "top_k": top_k, + "exclude_document_ids": exclude_document_ids, + "exclude_sections": exclude_sections, + "data_type": data_type, + "signal_paths": signal_paths, + "filter_mode": filter_mode, + "channels": channels, + "channel_weights": channel_weights, + } + + logger.info(" agentic: Phase 1 — discovery + document selection") + discovery_result = await tools.bottom_discovery(db, **discovery_kwargs) + state.step_count += 1 + discovery_rows = ( + discovery_result.payload.get("fused_rows", []) + if discovery_result.status != "error" + else [] + ) + state.discovery_top_doc_ids = ( + discovery_result.payload.get("top_doc_ids", []) + if discovery_result.status != "error" + else [] + ) + + if trace_enabled: + trace.record_step( + "bottom_discovery", + discovery_result, + decision_reason="phase_1_mandatory", + ) + + logger.info( + f" agentic step {state.step_count}: bottom_discovery " + f"status={discovery_result.status} latency={discovery_result.latency_ms}ms" + ) + + if bootstrap_llm_fn is not None: + await _select_documents( + db, + state=state, + trace=trace, + trace_enabled=trace_enabled, + user_id=user_id, + namespace=namespace, + query=query, + exclude_document_ids=exclude_document_ids, + bootstrap_llm_fn=bootstrap_llm_fn, + ) + + return discovery_rows + + +async def register_discovery_documents( + db: AsyncSession, + *, + state: AgentState, + discovery_by_doc: dict[str, list[dict[str, Any]]], +) -> None: + selected_doc_ids = {doc.document_id for doc in state.selected_docs} + for doc_id in discovery_by_doc: + if doc_id in selected_doc_ids or doc_id in state.ever_explored_doc_ids: + continue + doc_stmt = ( + select(Document.document_id, Document.source_file_name, Document.current_job_result_id) + .where(Document.document_id == doc_id) + ) + doc_result = await db.execute(doc_stmt) + row_data = doc_result.first() + if row_data is None: + continue + did, fname, job_result_id = row_data + state.selected_docs.append( + CandidateDoc( + document_id=did, + source_file_name=fname or did, + confidence=0.4, + reason="discovery_auto (not in KG selection)", + source="discovery_auto", + ) + ) + state.doc_id_to_name[did] = fname or did + if job_result_id: + state.doc_job_map[did] = job_result_id + + +async def select_revision_documents( + db: AsyncSession, + *, + state: AgentState, + trace: TraceRecorder, + trace_enabled: bool, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: list[str], + bootstrap_llm_fn: LLMFn, + revision_hint: str, +) -> str | None: + try: + kg_result = await tools.kg_document_select( + db, + user_id=user_id, + namespace=namespace, + query=query, + llm_fn=bootstrap_llm_fn, + exclude_document_ids=list(set(exclude_document_ids)), + revision_hint=revision_hint, + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) + except BudgetExceeded: + logger.info(" agentic: bootstrap budget exhausted during revision doc selection") + if trace_enabled: + trace.record_budget_stop("bootstrap_exhausted") + return "bootstrap_budget" + state.step_count += 1 + _append_selected_docs(state, kg_result) + return None + + +async def _select_documents( + db: AsyncSession, + *, + state: AgentState, + trace: TraceRecorder, + trace_enabled: bool, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: list[str], + bootstrap_llm_fn: LLMFn, +) -> None: + try: + kg_result = await tools.kg_document_select( + db, + user_id=user_id, + namespace=namespace, + query=query, + llm_fn=bootstrap_llm_fn, + exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)), + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) + except BudgetExceeded: + logger.info(" agentic: bootstrap budget exhausted during document selection") + if trace_enabled: + trace.record_budget_stop("bootstrap_exhausted") + kg_result = ToolResult( + status="no_confident_doc", + payload={"reason": "bootstrap budget exhausted"}, + ) + state.step_count += 1 + + if trace_enabled: + trace.record_step( + "kg_document_select", + kg_result, + decision_reason="phase_1_doc_selection", + ) + + _append_selected_docs(state, kg_result) + if not state.selected_docs and state.discovery_top_doc_ids: + await _append_discovery_hints(db, state=state) + + logger.info( + f" agentic step {state.step_count}: kg_document_select " + f"status={kg_result.status} docs={len(state.selected_docs)} " + f"latency={kg_result.latency_ms}ms" + ) + + +async def _append_discovery_hints(db: AsyncSession, *, state: AgentState) -> None: + hint_ids = [ + doc_id + for doc_id in state.discovery_top_doc_ids + if doc_id not in state.ever_explored_doc_ids + ] + if not hint_ids: + return + doc_stmt = ( + select(Document.document_id, Document.source_file_name, Document.current_job_result_id) + .where(Document.document_id.in_(hint_ids)) + ) + doc_result = await db.execute(doc_stmt) + for doc_id, source_file_name, job_result_id in doc_result.all(): + state.selected_docs.append( + CandidateDoc( + document_id=doc_id, + source_file_name=source_file_name or doc_id, + confidence=0.5, + reason="discovery_hint (KG returned 0)", + source="discovery_hint", + ) + ) + state.doc_id_to_name[doc_id] = source_file_name or doc_id + if job_result_id: + state.doc_job_map[doc_id] = job_result_id + + +def _append_selected_docs(state: AgentState, kg_result: ToolResult) -> None: + if kg_result.status != "selected_docs": + return + for doc_data in kg_result.payload.get("candidate_docs", []): + state.selected_docs.append( + CandidateDoc( + document_id=doc_data.get("document_id", ""), + source_file_name=doc_data.get("source_file_name", ""), + confidence=doc_data.get("confidence", 0.0), + reason=doc_data.get("reason", ""), + source=doc_data.get("source", ""), + ) + ) + state.doc_id_to_name.update(kg_result.payload.get("doc_id_to_name", {})) + state.doc_job_map.update(kg_result.payload.get("doc_job_map", {})) diff --git a/packages/shared-python/shared/services/retrieval/agentic/document_navigation.py b/packages/shared-python/shared/services/retrieval/agentic/document_navigation.py new file mode 100644 index 000000000..14b3d5f63 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/document_navigation.py @@ -0,0 +1,414 @@ +"""Per-document navigation for agentic retrieval.""" +from __future__ import annotations + +from typing import Any, cast + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic import tools +from shared.services.retrieval.agentic.budget import BudgetExceeded +from shared.services.retrieval.agentic.evidence import reconcile_deferred_assets +from shared.services.retrieval.agentic.runtime import AgentLlmBudget +from shared.services.retrieval.agentic.trace import TraceRecorder +from shared.services.retrieval.agentic.types import ( + AgentRunConfig, + AgentState, + CandidateDoc, + DocTreeNode, + ToolResult, +) +from shared.services.retrieval.llm_adapter import LLMFn + + +class DocumentNavigationRunner: + def __init__( + self, + *, + db: AsyncSession, + state: AgentState, + trace: TraceRecorder, + trace_enabled: bool, + user_id: str, + namespace: str, + query: str, + config: AgentRunConfig, + discovery_by_doc: dict[str, list[dict[str, Any]]], + llm_fn: LLMFn | None, + llm_budget: AgentLlmBudget, + ) -> None: + self._db = db + self._state = state + self._trace = trace + self._trace_enabled = trace_enabled + self._user_id = user_id + self._namespace = namespace + self._query = query + self._config = config + self._discovery_by_doc = discovery_by_doc + self._llm_fn = llm_fn + self._llm_budget = llm_budget + + async def navigate_selected_documents(self, *, revision_hint: str | None) -> None: + logger.info( + f" agentic: Phase 2 — navigating {len(self._state.selected_docs)} documents" + ) + for doc in self._state.selected_docs: + if self._state.elapsed_ms >= self._config.latency_budget_ms: + logger.info(" agentic: latency budget hit during Phase 2, stopping") + break + await self._navigate_document(doc, revision_hint=revision_hint) + + async def _navigate_document( + self, + doc: CandidateDoc, + *, + revision_hint: str | None, + ) -> None: + job_result_id = self._state.doc_job_map.get(doc.document_id, "") + if not job_result_id: + logger.info(f" agentic: skipping doc {doc.document_id} — no job_result_id") + self._state.ever_explored_doc_ids.add(doc.document_id) + return + + doc_name = doc.source_file_name or self._state.doc_id_to_name.get(doc.document_id, "") + is_discovery_only_doc = doc.source == "discovery_auto" + root = DocTreeNode(scope_path=None) + doc_pending_assets: list[dict[str, Any]] = [] + + if not is_discovery_only_doc: + doc_pending_assets = await self._navigate_bfs( + doc=doc, + root=root, + doc_name=doc_name, + job_result_id=job_result_id, + revision_hint=revision_hint, + ) + + await self._hydrate_discovery_hints( + doc=doc, + root=root, + doc_name=doc_name, + revision_hint=revision_hint, + ) + + if not is_discovery_only_doc and doc_pending_assets: + self._reconcile_pending_assets( + doc=doc, + root=root, + doc_name=doc_name, + doc_pending_assets=doc_pending_assets, + ) + + if doc.document_id in self._state.doc_trees: + self._state.doc_trees[doc.document_id].merge(root) + else: + self._state.doc_trees[doc.document_id] = root + self._state.ever_explored_doc_ids.add(doc.document_id) + if self._state.ledger is not None: + self._state.ledger.mark_explored(docs=1) + + async def _navigate_bfs( + self, + *, + doc: CandidateDoc, + root: DocTreeNode, + doc_name: str, + job_result_id: str, + revision_hint: str | None, + ) -> list[dict[str, Any]]: + doc_exclude: set[str] = { + key.split("::", 1)[1] + for key in self._state.seen_section_keys + if key.startswith(f"{doc.document_id}::") + } if self._state.seen_section_keys else set() + pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)] + doc_pending_assets: list[dict[str, Any]] = [] + + while pending: + if self._state.elapsed_ms >= self._config.latency_budget_ms: + break + + scope, parent_node, depth = pending.pop(0) + if depth >= self._config.max_nav_depth: + continue + if self._llm_fn is None: + break + if self._state.ledger and self._state.ledger.status("planning") in ("CRITICAL", "EXHAUSTED"): + logger.info(" agentic: planning budget critical, ending BFS for current doc") + break + + doc_llm_fn = self._llm_budget.for_document( + cast(LLMFn, self._llm_fn), + doc_id=doc.document_id, + depth=depth, + ) + try: + action, asset_tools, step_node, drill_paths = await tools.navigate_step( + self._db, + document_id=doc.document_id, + job_result_id=job_result_id, + query=self._query, + llm_fn=doc_llm_fn, + user_id=self._user_id, + namespace=self._namespace, + doc_name=doc_name, + scope_path=scope, + exclude_paths=doc_exclude, + revision_hint=revision_hint if depth == 0 else None, + budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None, + ) + except BudgetExceeded: + logger.info(" agentic: planning budget exhausted during navigation") + if self._trace_enabled: + self._trace.record_budget_stop("planning_exhausted") + break + self._state.step_count += 1 + + await self._collect_assets( + doc=doc, + scope=scope, + step_node=step_node, + asset_tools=asset_tools, + pending_assets=doc_pending_assets, + round_scope="nav", + ) + _merge_step_node(parent_node, step_node) + _update_excluded_leaf_paths(doc_exclude, step_node, drill_paths) + _queue_drill_paths(pending, parent_node, drill_paths, depth) + parent_node.reparent_leaf_content() + self._record_navigation_step( + doc=doc, + scope=scope, + depth=depth, + action=action, + asset_tools=asset_tools, + step_node=step_node, + drill_paths=drill_paths, + ) + if self._state.ledger is not None: + self._state.ledger.mark_explored( + chunks=sum(len(chunks) for chunks in step_node.leaf_content.values()), + ) + + return doc_pending_assets + + async def _collect_assets( + self, + *, + doc: CandidateDoc, + scope: str | list[str] | None, + step_node: DocTreeNode, + asset_tools: list[str], + pending_assets: list[dict[str, Any]], + round_scope: str, + ) -> None: + selected_asset_scopes = list(step_node.confidence.keys()) + asset_scope = selected_asset_scopes or scope + for asset_tool in asset_tools: + if asset_tool not in ("FIND_IMAGES", "FIND_TABLES"): + continue + asset_type = "image" if asset_tool == "FIND_IMAGES" else "table" + asset_chunks = await tools.asset_filter_step( + self._db, + document_id=doc.document_id, + job_result_id=self._state.doc_job_map.get(doc.document_id, ""), + scope_path=asset_scope, + asset_type=asset_type, + ) + if asset_chunks: + pending_assets.extend(asset_chunks) + + scope_display = ( + asset_scope + if isinstance(asset_scope, list) + else (asset_scope or "root") + ) + if self._trace_enabled: + self._trace.record_step( + "asset_filter_step", + ToolResult( + status="filtered" if asset_chunks else "empty", + payload={ + "document_id": doc.document_id, + "scope": scope_display, + "navigation_scope": scope if isinstance(scope, str) else (scope or "root"), + "asset_type": asset_type, + "chunks_found": len(asset_chunks) if asset_chunks else 0, + }, + ), + decision_reason=f"asset_{round_scope}_{doc.source_file_name}", + ) + logger.info( + f" agentic step {self._state.step_count}: asset_filter_step " + f'doc="{doc.source_file_name}" scope={scope_display} ' + f"type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}" + ) + + async def _hydrate_discovery_hints( + self, + *, + doc: CandidateDoc, + root: DocTreeNode, + doc_name: str, + revision_hint: str | None, + ) -> None: + doc_hints = self._discovery_by_doc.get(doc.document_id, []) + if not doc_hints or self._llm_fn is None: + return + if self._state.elapsed_ms >= self._config.latency_budget_ms: + return + + discovery_exclude_paths = { + key.split("::", 1)[1] + for key in root.collect_all_paths(doc.document_id) + } + doc_discovery_llm_fn = self._llm_budget.for_discovery( + cast(LLMFn, self._llm_fn), + doc_id=doc.document_id, + low_priority=root.has_content(), + ) + try: + discovery_node = await tools.discovery_select_step( + self._db, + document_id=doc.document_id, + query=self._query, + llm_fn=doc_discovery_llm_fn, + user_id=self._user_id, + namespace=self._namespace, + doc_name=doc_name, + discovery_hints=doc_hints, + exclude_paths=discovery_exclude_paths, + revision_hint=revision_hint, + budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None, + ) + except BudgetExceeded: + logger.info(" agentic: planning budget exhausted during discovery selection") + if self._trace_enabled: + self._trace.record_budget_stop("planning_exhausted") + discovery_node = DocTreeNode(scope_path=None) + self._state.step_count += 1 + + if self._trace_enabled: + self._trace.record_step( + "discovery_select_step", + ToolResult( + status="selected" if discovery_node.has_content() else "empty", + payload={ + "document_id": doc.document_id, + "hints_count": len(doc_hints), + "hydrated_count": len(discovery_node.leaf_content), + }, + ), + decision_reason=f"discovery_{doc.source_file_name}", + ) + root.merge(discovery_node) + if self._state.ledger is not None: + self._state.ledger.mark_explored( + chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()), + ) + + def _reconcile_pending_assets( + self, + *, + doc: CandidateDoc, + root: DocTreeNode, + doc_name: str, + doc_pending_assets: list[dict[str, Any]], + ) -> None: + if doc_name and not root.children and not any( + item.get("path") == doc_name for item in root.outline_items + ): + root.outline_items.insert(0, {"path": doc_name, "level": 0}) + reconcile_deferred_assets(root, doc_pending_assets) + if self._trace_enabled: + self._trace.record_step( + "deferred_asset_reconcile", + ToolResult( + status="reconciled", + payload={ + "document_id": doc.document_id, + "pending_count": len(doc_pending_assets), + "placed_count": sum( + 1 for asset in doc_pending_assets + if str(asset.get("chunk_id") or "") in { + str(row.get("chunk_id") or "") + for row in root.flatten_chunk_rows() + } + ), + }, + ), + decision_reason=f"deferred_reconcile_{doc.source_file_name}", + ) + + def _record_navigation_step( + self, + *, + doc: CandidateDoc, + scope: str | list[str] | None, + depth: int, + action: str, + asset_tools: list[str], + step_node: DocTreeNode, + drill_paths: list[dict[str, Any]], + ) -> None: + if self._trace_enabled: + self._trace.record_step( + "navigate_step", + ToolResult( + status=f"{action.lower()}" + (" (content)" if step_node.has_content() else ""), + payload={ + "document_id": doc.document_id, + "scope": scope if isinstance(scope, str) else (scope or "root"), + "depth": depth, + "action": action, + "asset_tools": asset_tools, + "outline_count": len(step_node.outline_items), + "leaf_count": len(step_node.leaf_content), + "pending_drills": len(drill_paths), + }, + ), + decision_reason=f"nav_d{depth}_{doc.source_file_name}", + ) + scope_log = scope if isinstance(scope, str) else (", ".join(scope) if scope else "root") + logger.info( + f" agentic step {self._state.step_count}: navigate_step " + f'doc="{doc.source_file_name}" scope={scope_log} ' + f"depth={depth} action={action} tools={asset_tools} " + f"outline={len(step_node.outline_items)} " + f"leaves={len(step_node.leaf_content)} " + f"drills={len(drill_paths)}" + ) + + +def _merge_step_node(parent_node: DocTreeNode, step_node: DocTreeNode) -> None: + parent_node.outline_items = step_node.outline_items + for leaf_path, chunks in step_node.leaf_content.items(): + parent_node.add_leaf_chunks(leaf_path, chunks) + parent_node.confidence = step_node.confidence + + +def _update_excluded_leaf_paths( + doc_exclude: set[str], + step_node: DocTreeNode, + drill_paths: list[dict[str, Any]], +) -> None: + drill_path_set = {str(selection["path"]) for selection in drill_paths} + for leaf_path in step_node.leaf_content: + if leaf_path not in drill_path_set: + doc_exclude.add(leaf_path) + + +def _queue_drill_paths( + pending: list[tuple[str | list[str] | None, DocTreeNode, int]], + parent_node: DocTreeNode, + drill_paths: list[dict[str, Any]], + depth: int, +) -> None: + if not drill_paths: + return + for selection in drill_paths: + child = DocTreeNode(scope_path=selection["path"]) + parent_node.children[selection["path"]] = child + batch_scope = [selection["path"] for selection in drill_paths] + pending.append((batch_scope, parent_node, depth + 1)) diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index f5bbd4889..c8fdfa19a 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -18,90 +18,38 @@ from __future__ import annotations import os -import json from typing import Any, cast from loguru import logger -from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document, DocumentChunk - -from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger, BudgetPoolName +from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger +from shared.services.retrieval.agentic.discovery_phase import ( + register_discovery_documents, + run_initial_discovery, + select_revision_documents, +) +from shared.services.retrieval.agentic.document_navigation import DocumentNavigationRunner from shared.services.retrieval.agentic.evidence import ( build_asset_url_map as _build_asset_url_map, collect_media_chunks_all as _collect_media_chunks_all, - reconcile_deferred_assets as _reconcile_deferred_assets, render_evidence as _render_evidence, trim_evidence_to_budget as _trim_evidence_to_budget, with_context_prompt_projection as _with_context_prompt_projection, ) +from shared.services.retrieval.agentic.runtime import ( + AgentLlmBudget, + build_config_from_env as _build_config_from_env, + load_budget_inventory as _load_budget_inventory, +) from shared.services.retrieval.agentic.trace import TraceRecorder from shared.services.retrieval.agentic.types import ( AgentRunConfig, AgentState, AgenticResult, - CandidateDoc, - DocTreeNode, ToolResult, ) from shared.services.retrieval.llm_adapter import LLMFn -from shared.services.retrieval.llm_adapter import current_llm_usage -from shared.utils.token_estimate import estimate_tokens - - - - -def _build_config_from_env() -> AgentRunConfig: - """Read agent config from environment, with sensible defaults.""" - return AgentRunConfig( - max_revisions=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_REVISIONS', '2')), - max_nav_depth=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_NAV_DEPTH', '3')), - latency_budget_ms=int(os.environ.get('RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS', '12000')), - token_budget_total=int(os.environ.get('RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL', '40000')), - planning_ratio=float(os.environ.get('RETRIEVAL_AGENTIC_PLANNING_RATIO', '0.5')), - bootstrap_budget=int(os.environ.get('RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET', '2000')), - per_doc_min_share=int(os.environ.get('RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE', '1500')), - inventory_aware=os.environ.get('RETRIEVAL_AGENTIC_INVENTORY_AWARE', 'true') == 'true', - ) - - -def _stringify_llm_input(prompt: Any) -> str: - if isinstance(prompt, str): - return prompt - try: - return json.dumps(prompt, ensure_ascii=False, default=str) - except Exception: - return str(prompt) - - - -async def _load_budget_inventory( - db: AsyncSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], -) -> tuple[int, int, dict[str, int]]: - stmt = ( - select(Document.document_id, func.count(DocumentChunk.id)) - .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), - ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .group_by(Document.document_id) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - - result = await db.execute(stmt) - doc_chunks = {str(doc_id): int(count or 0) for doc_id, count in result.all()} - return sum(doc_chunks.values()), len(doc_chunks), doc_chunks - class RetrievalAgent: @@ -121,82 +69,6 @@ class RetrievalAgent: If ``llm_fn`` is None, the run returns discovery-only results. """ - async def _call_llm_with_budget( - self, - state: AgentState, - llm_fn: LLMFn, - prompt: Any, - *, - pool: BudgetPoolName, - doc_id: str | None = None, - priority: str = 'normal', - ) -> str: - ledger = state.ledger - if ledger is None: - return await llm_fn(prompt) - - prompt_text = _stringify_llm_input(prompt) - est = estimate_tokens(prompt_text) - reserved = await ledger.try_reserve( - pool, - est, - doc_id=doc_id, - priority='low' if priority == 'low' else 'normal', - ) - if not reserved: - raise BudgetExceeded(f'{pool} budget exhausted') - - try: - response = await llm_fn(prompt) - except Exception: - await ledger.refund(pool, est=est, doc_id=doc_id) - raise - - usage = current_llm_usage.get() or {} - actual = int(usage.get('prompt_tokens') or est) - await ledger.commit(pool, actual=actual, est=est, doc_id=doc_id) - return response - - def _budgeted_doc_llm_fn( - self, - state: AgentState, - llm_fn: LLMFn, - *, - doc_id: str, - depth: int, - ) -> LLMFn: - async def _call(prompt): - return await self._call_llm_with_budget( - state, - llm_fn, - prompt, - pool='planning', - doc_id=doc_id, - priority='low' if depth >= 2 else 'normal', - ) - - return _call - - def _budgeted_discovery_llm_fn( - self, - state: AgentState, - llm_fn: LLMFn, - *, - doc_id: str, - low_priority: bool, - ) -> LLMFn: - async def _call(prompt): - return await self._call_llm_with_budget( - state, - llm_fn, - prompt, - pool='planning', - doc_id=doc_id, - priority='low' if low_priority else 'normal', - ) - - return _call - async def run( self, db: AsyncSession, @@ -225,7 +97,6 @@ async def run( errors are captured in trace and the best available result is returned. """ - from shared.services.retrieval.agentic import tools from shared.services.retrieval.agentic.policy import ( attempt_answer, estimate_attempt_answer_prompt_tokens, @@ -280,135 +151,32 @@ async def run( if llm_fn is None: logger.warning('agentic: no llm_fn provided — running discovery-only mode') - planning_llm_fn: LLMFn | None = None bootstrap_llm_fn: LLMFn | None = None context_llm_fn: LLMFn | None = None + llm_budget = AgentLlmBudget(state) if llm_fn is not None: - base_llm_fn = llm_fn + bootstrap_llm_fn = llm_budget.for_pool(llm_fn, pool='bootstrap') + context_llm_fn = llm_budget.for_pool(llm_fn, pool='context') - async def _planning_llm_call(prompt): - return await self._call_llm_with_budget( - state, base_llm_fn, prompt, pool='planning' - ) - - async def _bootstrap_llm_call(prompt): - return await self._call_llm_with_budget( - state, base_llm_fn, prompt, pool='bootstrap' - ) - - async def _context_llm_call(prompt): - return await self._call_llm_with_budget( - state, base_llm_fn, prompt, pool='context' - ) - - planning_llm_fn = _planning_llm_call - bootstrap_llm_fn = _bootstrap_llm_call - context_llm_fn = _context_llm_call - - # Shared kwargs for bottom_discovery - discovery_kwargs: dict[str, Any] = { - 'user_id': user_id, - 'namespace': namespace, - 'query': query, - 'top_k': top_k, - 'exclude_document_ids': exclude_document_ids, - 'exclude_sections': exclude_sections, - 'data_type': data_type, - 'signal_paths': signal_paths, - 'filter_mode': filter_mode, - 'channels': channels, - 'channel_weights': channel_weights, - } - - # ══════════════════════════════════════════════════════════════════ - # Phase 1: Discovery + Document Selection - # ══════════════════════════════════════════════════════════════════ - logger.info(' agentic: Phase 1 — discovery + document selection') - - # 1a. Bottom discovery (always runs) - discovery_result = await tools.bottom_discovery(db, **discovery_kwargs) - state.step_count += 1 - discovery_rows = discovery_result.payload.get('fused_rows', []) if discovery_result.status != 'error' else [] - state.discovery_top_doc_ids = discovery_result.payload.get('top_doc_ids', []) if discovery_result.status != 'error' else [] - - if trace_enabled: - trace.record_step( - 'bottom_discovery', discovery_result, - decision_reason='phase_1_mandatory', - ) - - logger.info( - f' agentic step {state.step_count}: bottom_discovery ' - f'status={discovery_result.status} latency={discovery_result.latency_ms}ms' + discovery_rows = await run_initial_discovery( + db, + state=state, + trace=trace, + trace_enabled=trace_enabled, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + bootstrap_llm_fn=bootstrap_llm_fn, ) - # 1b. KG document selection (requires LLM) - if bootstrap_llm_fn is not None: - try: - kg_result = await tools.kg_document_select( - db, - user_id=user_id, - namespace=namespace, - query=query, - llm_fn=bootstrap_llm_fn, - exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)), - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: bootstrap budget exhausted during document selection') - if trace_enabled: - trace.record_budget_stop('bootstrap_exhausted') - kg_result = ToolResult( - status='no_confident_doc', - payload={'reason': 'bootstrap budget exhausted'}, - ) - state.step_count += 1 - - if trace_enabled: - trace.record_step( - 'kg_document_select', kg_result, - decision_reason='phase_1_doc_selection', - ) - - if kg_result.status == 'selected_docs': - for doc_data in kg_result.payload.get('candidate_docs', []): - state.selected_docs.append(CandidateDoc( - document_id=doc_data.get('document_id', ''), - source_file_name=doc_data.get('source_file_name', ''), - confidence=doc_data.get('confidence', 0.0), - reason=doc_data.get('reason', ''), - source=doc_data.get('source', ''), - )) - state.doc_id_to_name.update(kg_result.payload.get('doc_id_to_name', {})) - state.doc_job_map.update(kg_result.payload.get('doc_job_map', {})) - - # If KG returned nothing, use discovery hints - if not state.selected_docs and state.discovery_top_doc_ids: - hint_ids = [d for d in state.discovery_top_doc_ids if d not in state.ever_explored_doc_ids] - if hint_ids: - doc_stmt = ( - select(Document.document_id, Document.source_file_name, Document.current_job_result_id) - .where(Document.document_id.in_(hint_ids)) - ) - doc_result = await db.execute(doc_stmt) - for did, fname, jrid in doc_result.all(): - state.selected_docs.append(CandidateDoc( - document_id=did, - source_file_name=fname or did, - confidence=0.5, - reason='discovery_hint (KG returned 0)', - source='discovery_hint', - )) - state.doc_id_to_name[did] = fname or did - if jrid: - state.doc_job_map[did] = jrid - - logger.info( - f' agentic step {state.step_count}: kg_document_select ' - f'status={kg_result.status} docs={len(state.selected_docs)} ' - f'latency={kg_result.latency_ms}ms' - ) - # If no LLM or no docs selected, return discovery rows directly if not state.selected_docs: logger.info('agentic: no documents selected — returning discovery results') @@ -440,38 +208,17 @@ async def _context_llm_call(prompt): router_used='agentic_discovery_only', ) - # ══════════════════════════════════════════════════════════════════ - # Discovery → Navigation integration - # Group discovery_rows by document for post-BFS discovery selection - # ══════════════════════════════════════════════════════════════════ discovery_by_doc: dict[str, list[dict[str, Any]]] = {} for row in discovery_rows: doc_id = row.get('document_id', '') if doc_id: discovery_by_doc.setdefault(doc_id, []).append(row) - # Auto-register B-class docs (discovery-only, not selected by KG) - selected_doc_ids = {d.document_id for d in state.selected_docs} - for doc_id in discovery_by_doc: - if doc_id not in selected_doc_ids and doc_id not in state.ever_explored_doc_ids: - doc_stmt = ( - select(Document.document_id, Document.source_file_name, Document.current_job_result_id) - .where(Document.document_id == doc_id) - ) - doc_result = await db.execute(doc_stmt) - row_data = doc_result.first() - if row_data: - did, fname, jrid = row_data - state.selected_docs.append(CandidateDoc( - document_id=did, - source_file_name=fname or did, - confidence=0.4, - reason='discovery_auto (not in KG selection)', - source='discovery_auto', - )) - state.doc_id_to_name[did] = fname or did - if jrid: - state.doc_job_map[did] = jrid + await register_discovery_documents( + db, + state=state, + discovery_by_doc=discovery_by_doc, + ) if state.ledger is not None: await state.ledger.allocate_doc_caps({ @@ -493,286 +240,20 @@ async def _context_llm_call(prompt): stop_reason = 'latency_budget' break - # ── Phase 2: Per-Document Navigation ──────────────────────── - logger.info( - f' agentic: Phase 2 (round {round_idx}) — ' - f'navigating {len(state.selected_docs)} documents' + navigation_runner = DocumentNavigationRunner( + db=db, + state=state, + trace=trace, + trace_enabled=trace_enabled, + user_id=user_id, + namespace=namespace, + query=query, + config=config, + discovery_by_doc=discovery_by_doc, + llm_fn=llm_fn, + llm_budget=llm_budget, ) - - for doc in state.selected_docs: - if state.elapsed_ms >= config.latency_budget_ms: - logger.info(' agentic: latency budget hit during Phase 2, stopping') - break - - job_result_id = state.doc_job_map.get(doc.document_id, '') - if not job_result_id: - logger.info(f' agentic: skipping doc {doc.document_id} — no job_result_id') - state.ever_explored_doc_ids.add(doc.document_id) - continue - - doc_name = doc.source_file_name or state.doc_id_to_name.get(doc.document_id, '') - - # B-class docs (discovery_auto) skip BFS, go to discovery_select - is_b_class = doc.source == 'discovery_auto' - - if not is_b_class: - # Build exclude_paths for this doc from seen_section_keys - # Starts with revision-carried paths, then accumulates - # leaf paths hydrated during THIS BFS round to prevent - # re-selection in deeper drill-downs. - doc_exclude: set[str] = { - key.split('::', 1)[1] - for key in state.seen_section_keys - if key.startswith(f'{doc.document_id}::') - } if state.seen_section_keys else set() - - # BFS queue: (scope_path(s), parent_node, depth) - # scope can be: None (root), str, or list[str] (multi-scope) - root = DocTreeNode(scope_path=None) - pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)] - doc_pending_assets: list[dict] = [] # deferred asset reconcile - - while pending: - if state.elapsed_ms >= config.latency_budget_ms: - break - - scope, parent_node, depth = pending.pop(0) - if depth >= config.max_nav_depth: - continue - - if planning_llm_fn is None: - break - if state.ledger and state.ledger.status('planning') in ('CRITICAL', 'EXHAUSTED'): - logger.info(' agentic: planning budget critical, ending BFS for current doc') - break - - doc_llm_fn = self._budgeted_doc_llm_fn( - state, - cast(LLMFn, llm_fn), - doc_id=doc.document_id, - depth=depth, - ) - - # ★ Unified navigate step (supports multi-scope batching) - try: - action, asset_tools, step_node, drill_paths = await tools.navigate_step( - db, - document_id=doc.document_id, - job_result_id=job_result_id, - query=query, - llm_fn=doc_llm_fn, - user_id=user_id, - namespace=namespace, - doc_name=doc_name, - scope_path=scope, - exclude_paths=doc_exclude, - revision_hint=revision_hint if depth == 0 else None, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: planning budget exhausted during navigation') - if trace_enabled: - trace.record_budget_stop('planning_exhausted') - break - state.step_count += 1 - - # ★ Asset collection (deferred reconcile) — runs if LLM selected tools. - # If this navigation call selected sections, bind asset tools to - # those selections; otherwise keep the current scope (STOP/root). - selected_asset_scopes = list(step_node.confidence.keys()) - asset_scope = selected_asset_scopes or scope - for asset_tool in asset_tools: - if asset_tool not in ('FIND_IMAGES', 'FIND_TABLES'): - continue - asset_type = 'image' if asset_tool == 'FIND_IMAGES' else 'table' - asset_chunks = await tools.asset_filter_step( - db, - document_id=doc.document_id, - job_result_id=job_result_id, - scope_path=asset_scope, - asset_type=asset_type, - ) - if asset_chunks: - doc_pending_assets.extend(asset_chunks) - - scope_display = ( - asset_scope if isinstance(asset_scope, list) - else (asset_scope or 'root') - ) - if trace_enabled: - trace.record_step( - 'asset_filter_step', ToolResult( - status='filtered' if asset_chunks else 'empty', - payload={ - 'document_id': doc.document_id, - 'scope': scope_display, - 'navigation_scope': scope if isinstance(scope, str) else (scope or 'root'), - 'asset_type': asset_type, - 'chunks_found': len(asset_chunks) if asset_chunks else 0, - }, - ), - decision_reason=f'asset_r{round_idx}_d{depth}_{doc.source_file_name}', - ) - - logger.info( - f' agentic step {state.step_count}: asset_filter_step ' - f'doc="{doc.source_file_name}" scope={scope_display} ' - f'type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}' - ) - - # Merge step result into parent node - parent_node.outline_items = step_node.outline_items - for leaf_path, chunks in step_node.leaf_content.items(): - parent_node.add_leaf_chunks(leaf_path, chunks) - parent_node.confidence = step_node.confidence - - # Accumulate hydrated leaf paths into doc_exclude - drill_path_set = {sel['path'] for sel in drill_paths} - for leaf_path in step_node.leaf_content: - if leaf_path not in drill_path_set: - doc_exclude.add(leaf_path) - - # Queue non-leaf selections as a SINGLE batched item - # (all drill paths expand simultaneously in the next call) - if drill_paths: - for sel in drill_paths: - child = DocTreeNode(scope_path=sel['path']) - parent_node.children[sel['path']] = child - batch_scope = [sel['path'] for sel in drill_paths] - pending.append((batch_scope, parent_node, depth + 1)) - - # Re-parent leaf paths that belong to a child's subtree - parent_node.reparent_leaf_content() - - if trace_enabled: - trace.record_step( - 'navigate_step', ToolResult( - status=f'{action.lower()}' + (' (content)' if step_node.has_content() else ''), - payload={ - 'document_id': doc.document_id, - 'scope': scope if isinstance(scope, str) else (scope or 'root'), - 'depth': depth, - 'action': action, - 'asset_tools': asset_tools, - 'outline_count': len(step_node.outline_items), - 'leaf_count': len(step_node.leaf_content), - 'pending_drills': len(drill_paths), - }, - ), - decision_reason=f'nav_r{round_idx}_d{depth}_{doc.source_file_name}', - ) - - scope_log = scope if isinstance(scope, str) else (', '.join(scope) if scope else 'root') - logger.info( - f' agentic step {state.step_count}: navigate_step ' - f'doc="{doc.source_file_name}" scope={scope_log} ' - f'depth={depth} action={action} tools={asset_tools} ' - f'outline={len(step_node.outline_items)} ' - f'leaves={len(step_node.leaf_content)} ' - f'drills={len(drill_paths)}' - ) - if state.ledger is not None: - state.ledger.mark_explored( - chunks=sum(len(chunks) for chunks in step_node.leaf_content.values()), - ) - else: - # B-class: no BFS, create empty root - root = DocTreeNode(scope_path=None) - - # ── Post-BFS: Discovery selection step ───────────────────── - doc_hints = discovery_by_doc.get(doc.document_id, []) - if doc_hints and planning_llm_fn is not None and state.elapsed_ms < config.latency_budget_ms: - discovery_exclude_paths = { - key.split('::', 1)[1] - for key in root.collect_all_paths(doc.document_id) - } - doc_discovery_llm_fn = self._budgeted_discovery_llm_fn( - state, - cast(LLMFn, llm_fn), - doc_id=doc.document_id, - low_priority=root.has_content(), - ) - try: - discovery_node = await tools.discovery_select_step( - db, - document_id=doc.document_id, - query=query, - llm_fn=doc_discovery_llm_fn, - user_id=user_id, - namespace=namespace, - doc_name=doc_name, - discovery_hints=doc_hints, - exclude_paths=discovery_exclude_paths, - revision_hint=revision_hint, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: planning budget exhausted during discovery selection') - if trace_enabled: - trace.record_budget_stop('planning_exhausted') - discovery_node = DocTreeNode(scope_path=None) - state.step_count += 1 - - if trace_enabled: - trace.record_step( - 'discovery_select_step', ToolResult( - status='selected' if discovery_node.has_content() else 'empty', - payload={ - 'document_id': doc.document_id, - 'hints_count': len(doc_hints), - 'hydrated_count': len(discovery_node.leaf_content), - }, - ), - decision_reason=f'discovery_r{round_idx}_{doc.source_file_name}', - ) - - # Merge discovery results into BFS tree - root.merge(discovery_node) - if state.ledger is not None: - state.ledger.mark_explored( - chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()), - ) - - # ── Deferred asset reconcile ────────────────────────────── - # Assets were collected across all BFS depths but NOT placed - # into the tree yet. Now that the final navigated paths are - # known (BFS + discovery), filter and place only those assets - # whose owner path matches a navigated leaf. - if not is_b_class and doc_pending_assets: - # Inject doc file name as a visible root-level path, but - # ONLY when BFS stopped at root (no children = STOP action). - if doc_name and not root.children and not any( - item.get('path') == doc_name for item in root.outline_items - ): - root.outline_items.insert(0, {'path': doc_name, 'level': 0}) - _reconcile_deferred_assets(root, doc_pending_assets) - if trace_enabled: - trace.record_step( - 'deferred_asset_reconcile', ToolResult( - status='reconciled', - payload={ - 'document_id': doc.document_id, - 'pending_count': len(doc_pending_assets), - 'placed_count': sum( - 1 for a in doc_pending_assets - if str(a.get('chunk_id') or '') in { - str(r.get('chunk_id') or '') - for r in root.flatten_chunk_rows() - } - ), - }, - ), - decision_reason=f'deferred_reconcile_r{round_idx}_{doc.source_file_name}', - ) - - # Merge or store doc tree - if doc.document_id in state.doc_trees: - state.doc_trees[doc.document_id].merge(root) - else: - state.doc_trees[doc.document_id] = root - state.ever_explored_doc_ids.add(doc.document_id) - if state.ledger is not None: - state.ledger.mark_explored(docs=1) + await navigation_runner.navigate_selected_documents(revision_hint=revision_hint) # ── Phase 3: Render evidence + attempt_answer ──────────────── budget_snapshot_before_answer = state.ledger.snapshot() if state.ledger else None @@ -845,9 +326,7 @@ async def _context_llm_call(prompt): ] async def vlm_context_call(prompt, _vlm_fn=vlm_fn): - return await self._call_llm_with_budget( - state, cast(LLMFn, _vlm_fn), prompt, pool='context' - ) + return await llm_budget.call(cast(LLMFn, _vlm_fn), prompt, pool='context') # Auto-trigger attempt_answer (VLM if images present) try: @@ -914,36 +393,21 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn): if bootstrap_llm_fn is None: stop_reason = 'no_llm' break - try: - kg_result = await tools.kg_document_select( - db, - user_id=user_id, - namespace=namespace, - query=query, - llm_fn=bootstrap_llm_fn, - exclude_document_ids=list(set(exclude_document_ids)), - revision_hint=revision_hint, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: bootstrap budget exhausted during revision doc selection') - if trace_enabled: - trace.record_budget_stop('bootstrap_exhausted') - stop_reason = 'bootstrap_budget' + revision_stop_reason = await select_revision_documents( + db, + state=state, + trace=trace, + trace_enabled=trace_enabled, + user_id=user_id, + namespace=namespace, + query=query, + exclude_document_ids=exclude_document_ids, + bootstrap_llm_fn=bootstrap_llm_fn, + revision_hint=revision_hint, + ) + if revision_stop_reason is not None: + stop_reason = revision_stop_reason break - state.step_count += 1 - - if kg_result.status == 'selected_docs': - for doc_data in kg_result.payload.get('candidate_docs', []): - state.selected_docs.append(CandidateDoc( - document_id=doc_data.get('document_id', ''), - source_file_name=doc_data.get('source_file_name', ''), - confidence=doc_data.get('confidence', 0.0), - reason=doc_data.get('reason', ''), - source=doc_data.get('source', ''), - )) - state.doc_id_to_name.update(kg_result.payload.get('doc_id_to_name', {})) - state.doc_job_map.update(kg_result.payload.get('doc_job_map', {})) if not state.selected_docs: logger.info(' agentic: revision found no new docs — stopping') diff --git a/packages/shared-python/shared/services/retrieval/agentic/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/runtime.py new file mode 100644 index 000000000..d658cd562 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/runtime.py @@ -0,0 +1,146 @@ +"""Runtime setup helpers for agentic retrieval.""" +from __future__ import annotations + +import json +import os +from typing import Any + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk +from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetPoolName +from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState +from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage +from shared.utils.token_estimate import estimate_tokens + + +def build_config_from_env() -> AgentRunConfig: + return AgentRunConfig( + max_revisions=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_REVISIONS", "2")), + max_nav_depth=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_NAV_DEPTH", "3")), + latency_budget_ms=int(os.environ.get("RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS", "12000")), + token_budget_total=int(os.environ.get("RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL", "40000")), + planning_ratio=float(os.environ.get("RETRIEVAL_AGENTIC_PLANNING_RATIO", "0.5")), + bootstrap_budget=int(os.environ.get("RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET", "2000")), + per_doc_min_share=int(os.environ.get("RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE", "1500")), + inventory_aware=os.environ.get("RETRIEVAL_AGENTIC_INVENTORY_AWARE", "true") == "true", + ) + + +async def load_budget_inventory( + db: AsyncSession, + *, + user_id: str, + namespace: str, + exclude_document_ids: list[str], +) -> tuple[int, int, dict[str, int]]: + stmt = ( + select(Document.document_id, func.count(DocumentChunk.id)) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .group_by(Document.document_id) + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + + result = await db.execute(stmt) + doc_chunks = {str(doc_id): int(count or 0) for doc_id, count in result.all()} + return sum(doc_chunks.values()), len(doc_chunks), doc_chunks + + +class AgentLlmBudget: + def __init__(self, state: AgentState) -> None: + self._state = state + + async def call( + self, + llm_fn: LLMFn, + prompt: Any, + *, + pool: BudgetPoolName, + doc_id: str | None = None, + priority: str = "normal", + ) -> str: + ledger = self._state.ledger + if ledger is None: + return await llm_fn(prompt) + + prompt_text = _stringify_llm_input(prompt) + est = estimate_tokens(prompt_text) + reserved = await ledger.try_reserve( + pool, + est, + doc_id=doc_id, + priority="low" if priority == "low" else "normal", + ) + if not reserved: + raise BudgetExceeded(f"{pool} budget exhausted") + + try: + response = await llm_fn(prompt) + except Exception: + await ledger.refund(pool, est=est, doc_id=doc_id) + raise + + usage = current_llm_usage.get() or {} + actual = int(usage.get("prompt_tokens") or est) + await ledger.commit(pool, actual=actual, est=est, doc_id=doc_id) + return response + + def for_pool(self, llm_fn: LLMFn, *, pool: BudgetPoolName) -> LLMFn: + async def _call(prompt: Any) -> str: + return await self.call(llm_fn, prompt, pool=pool) + + return _call + + def for_document( + self, + llm_fn: LLMFn, + *, + doc_id: str, + depth: int, + ) -> LLMFn: + async def _call(prompt: Any) -> str: + return await self.call( + llm_fn, + prompt, + pool="planning", + doc_id=doc_id, + priority="low" if depth >= 2 else "normal", + ) + + return _call + + def for_discovery( + self, + llm_fn: LLMFn, + *, + doc_id: str, + low_priority: bool, + ) -> LLMFn: + async def _call(prompt: Any) -> str: + return await self.call( + llm_fn, + prompt, + pool="planning", + doc_id=doc_id, + priority="low" if low_priority else "normal", + ) + + return _call + + +def _stringify_llm_input(prompt: Any) -> str: + if isinstance(prompt, str): + return prompt + try: + return json.dumps(prompt, ensure_ascii=False, default=str) + except Exception: + return str(prompt) diff --git a/packages/shared-python/shared/services/retrieval/graph_service.py b/packages/shared-python/shared/services/retrieval/graph_service.py index 8244a8a72..622f321d2 100644 --- a/packages/shared-python/shared/services/retrieval/graph_service.py +++ b/packages/shared-python/shared/services/retrieval/graph_service.py @@ -18,7 +18,7 @@ _SECTION_EXCLUSION_PAGE_MULTIPLIER = 2 -# ── Keyword overlap config (aligned with connect_builder DEFAULT_CONFIG) ── +# ── Keyword overlap config for document-level publication graph ── _MIN_KEYWORD_OVERLAP = 3 _KEYWORD_SCORE_WEIGHT = 1.0 _MIN_SCORE_THRESHOLD = 0.8 @@ -54,7 +54,7 @@ def is_excluded_section( return False -# ── Keyword extraction & scoring (aligned with connect_builder/builder.py) ── +# ── Keyword extraction & scoring for document-level publication graph ── def _normalize_keyword(keyword: str) -> str: """Normalize a keyword: lowercase, strip, collapse spaces.""" @@ -63,7 +63,7 @@ def _normalize_keyword(keyword: str) -> str: def _extract_keywords_from_chunk_metadata(meta: dict) -> list[str]: - """Extract keywords from chunk metadata, same logic as builder._get_keywords.""" + """Extract keywords from chunk metadata.""" if not isinstance(meta, dict): return [] # Try metadata.keywords @@ -81,7 +81,7 @@ def _compute_tfidf_keywords( chunk_metadata_list: list[dict[str, Any]], top_k: int = 10, ) -> list[str]: - """Compute TF-IDF keywords from chunk metadata, aligned with graph_builder.""" + """Compute TF-IDF keywords from chunk metadata.""" df_count: dict[str, int] = {} tf_count: dict[str, int] = {} total = len(chunk_metadata_list) or 1 @@ -112,7 +112,7 @@ def _compute_keyword_score( kws_b: set[str], weight: float = 1.0, ) -> float: - """Character-length-weighted keyword overlap score (aligned with builder.py). + """Character-length-weighted keyword overlap score. Longer tokens contribute more: '施工现场'(4) has 2x weight of '交底'(2). Formula: score = weight * sum(len(kw) for shared) / min(sum(len) for A, sum(len) for B) @@ -169,7 +169,6 @@ class DocumentGraphService: - Only document-level nodes (no section nodes) - Document nodes carry rich metadata: top_keywords, chunks_count, types, top_summary - Edges are keyword-overlap-based cross-document connections with meaningful scores - - Edge scoring uses connect_builder DEFAULT_CONFIG thresholds """ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, document_id: str, job_result_id: str) -> None: @@ -237,9 +236,8 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d ) db.flush() - # ── Keyword-overlap-based cross-document edges (aligned with KB edges) ── - # Only create edges where keyword overlap score >= threshold, - # matching connect_builder DEFAULT_CONFIG parameters. + # ── Keyword-overlap-based cross-document edges ── + # Only create edges where keyword overlap score >= threshold. other_doc_nodes = list( db.execute( select(GraphNode) diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py new file mode 100644 index 000000000..410bfa111 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_result_schema.py @@ -0,0 +1,345 @@ +"""Schema projection for Knowhere ZIP result packages.""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from shared.services.chunks.chunk_connections import ( + build_resource_target_map, + convert_refs_to_embed_connections, + merge_connections, + normalize_connect_to_targets, + parse_relationship_refs, +) +from shared.utils.text_utils import truncate_content_preview +from shared.utils.utc_now import utc_now_naive + + +class ZipResultSchemaBuilder: + def calculate_statistics(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]: + total_chunks = len(chunks) + text_chunks = 0 + image_chunks = 0 + table_chunks = 0 + + for chunk in chunks: + chunk_type = chunk.get("type", "") + raw_type = str(chunk_type).strip() + normalized_type = raw_type.split("\n", 1)[0].lower() + if normalized_type == "image": + image_chunks += 1 + elif normalized_type == "table": + table_chunks += 1 + else: + text_chunks += 1 + + return { + "total_chunks": total_chunks, + "text_chunks": text_chunks, + "image_chunks": image_chunks, + "table_chunks": table_chunks, + "total_pages": None, + } + + def format_chunks( + self, + chunks: List[Dict[str, Any]], + image_files_map: Dict[str, Dict[str, Any]], + table_files_map: Dict[str, Dict[str, Any]], + ) -> List[Dict[str, Any]]: + resource_target_map = build_resource_target_map( + chunks, + image_files_map=image_files_map, + table_files_map=table_files_map, + ) + + formatted = [] + for chunk in chunks: + chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id")) + chunk_type_str = chunk.get("type", "") + raw_type = str(chunk_type_str).strip() + normalized_type = raw_type.split("\n", 1)[0].lower() + img_info = image_files_map.get(chunk_id) + + if normalized_type == "image": + chunk_type = "image" + elif normalized_type == "table": + chunk_type = "table" + else: + chunk_type = "text" + + content = chunk.get("text") or chunk.get("content", "") + path = chunk.get("path", "") + existing_metadata = chunk.get("metadata", {}) + metadata = { + "length": existing_metadata.get("length") or len(content), + "summary": existing_metadata.get("summary") or chunk.get("summary", ""), + "page_nums": existing_metadata.get("page_nums", []), + } + document_top_summary = str( + existing_metadata.get("document_top_summary") or "" + ).strip() + if document_top_summary: + metadata["document_top_summary"] = document_top_summary + + if chunk_type == "text": + metadata["tokens"] = existing_metadata.get("tokens") or chunk.get( + "tokens", 0 + ) + metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( + "keywords", [] + ) + relationship_refs = parse_relationship_refs( + chunk.get("type_raw") or chunk_type_str, + str(content), + ) + embed_connections = convert_refs_to_embed_connections( + relationship_refs, resource_target_map + ) + related_connections = normalize_connect_to_targets( + existing_metadata.get("connect_to") + or chunk.get("connect_to") + or chunk.get("connectto"), + resource_target_map, + ) + metadata["connect_to"] = merge_connections( + embed_connections, related_connections + ) + + elif chunk_type == "image": + if img_info: + metadata["file_path"] = img_info["file_path"] + metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( + "keywords", [] + ) + metadata["tokens"] = [] + + elif chunk_type == "table": + file_path = existing_metadata.get("file_path") + if not file_path: + table_info = table_files_map.get(chunk_id) + if table_info: + file_path = table_info["file_path"] + else: + table_name = ( + path.split("/")[-1] + if "/" in path + else f"table_{chunk_id}.html" + ) + file_path = f"tables/{table_name}" + + metadata["file_path"] = file_path + metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( + "keywords", [] + ) + metadata["tokens"] = [] + + formatted.append( + { + "chunk_id": chunk_id, + "type": chunk_type, + "content": content, + "path": path, + "metadata": metadata, + } + ) + + return formatted + + def generate_manifest( + self, + *, + job_id: str, + data_id: Optional[str], + source_file_name: str, + statistics: Dict[str, Any], + job_metadata: Dict[str, Any], + hierarchy: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + return { + "version": "2.0", + "job_id": job_id, + "data_id": data_id, + "source_file_name": source_file_name, + "processing_date": utc_now_naive().isoformat() + "Z", + "processing": { + "page_count": job_metadata.get("page_count"), + "billing_status": job_metadata.get("billing_status"), + "cost": { + "micro_dollars": job_metadata.get("billing_amount_micro_dollars"), + "credits": job_metadata.get("billing_credits"), + }, + "timing": { + "started_at": job_metadata.get("processing_started_at"), + "completed_at": job_metadata.get("processing_completed_at"), + "duration_ms": job_metadata.get("processing_duration_ms"), + }, + }, + "statistics": statistics, + "HIERARCHY": hierarchy or {}, + } + + def build_hierarchy_dict( + self, + sections: List[Dict[str, Any]], + ) -> Dict[str, Any]: + hierarchy: Dict[str, Any] = {} + title_counts: Dict[str, int] = {} + + for section in sections: + raw_title = str(section.get("title") or "").strip() + if not raw_title: + continue + + title_counts[raw_title] = title_counts.get(raw_title, 0) + 1 + title = ( + raw_title + if title_counts[raw_title] == 1 + else f"{raw_title} ({title_counts[raw_title]})" + ) + hierarchy[title] = self.build_hierarchy_dict( + section.get("children") or [] + ) + + return hierarchy + + def build_doc_nav( + self, + formatted_chunks: List[Dict[str, Any]], + source_file_name: str, + ) -> Dict[str, Any]: + text_chunks: List[Dict[str, Any]] = [] + image_resources: List[Dict[str, Any]] = [] + table_resources: List[Dict[str, Any]] = [] + + stats = { + "total_chunks": 0, + "text_chunks": 0, + "image_chunks": 0, + "table_chunks": 0, + "max_depth": 0, + } + + for formatted_chunk in formatted_chunks: + chunk_type = formatted_chunk.get("type", "text") + path = formatted_chunk.get("path", "") + metadata = formatted_chunk.get("metadata") or {} + summary_raw = (metadata.get("summary") or "").strip() + content_raw = (formatted_chunk.get("content") or "").strip() + summary = " ".join(summary_raw.split()) if summary_raw else "" + content_preview = truncate_content_preview(content_raw) if content_raw else "" + + stats["total_chunks"] += 1 + if chunk_type == "image": + stats["image_chunks"] += 1 + image_resources.append( + { + "path": path, + "summary": summary or content_preview, + } + ) + elif chunk_type == "table": + stats["table_chunks"] += 1 + table_resources.append( + { + "path": path, + "summary": summary or content_preview, + } + ) + else: + stats["text_chunks"] += 1 + text_chunks.append( + { + "path": path, + "summary": summary or content_preview, + } + ) + + sections = self._build_section_tree(text_chunks) + stats["max_depth"] = _max_depth(sections) + return { + "version": "1.0", + "file_name": source_file_name or "", + "stats": stats, + "sections": sections, + "resources": { + "images": image_resources, + "tables": table_resources, + }, + } + + def _build_section_tree( + self, + text_chunks: List[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + root_children: Dict[str, dict] = {} + + for chunk in text_chunks: + path = chunk.get("path", "") + parts = [part.strip() for part in path.split("/") if part.strip()] + section_parts = parts[2:] if len(parts) > 2 else [] + + if not section_parts: + key = "__root__" + if key not in root_children: + root_children[key] = { + "title": "Root", + "path": "/".join(parts[:2]) if len(parts) >= 2 else path, + "summary": chunk.get("summary", ""), + "chunk_count": 0, + "_children_map": {}, + } + root_children[key]["chunk_count"] += 1 + if not root_children[key]["summary"]: + root_children[key]["summary"] = chunk.get("summary", "") + continue + + current_level = root_children + full_section_path_parts = parts[:2] + for index, part in enumerate(section_parts): + full_section_path_parts.append(part) + if part not in current_level: + current_level[part] = { + "title": part, + "path": "/".join(full_section_path_parts), + "summary": "", + "chunk_count": 0, + "_children_map": {}, + } + node = current_level[part] + if index == len(section_parts) - 1: + node["chunk_count"] += 1 + if not node["summary"]: + node["summary"] = chunk.get("summary", "") + current_level = node["_children_map"] + + return _section_tree_to_output(root_children) + + +def _max_depth(nodes: list, depth: int = 1) -> int: + max_depth = depth if nodes else 0 + for node in nodes: + max_depth = max(max_depth, _max_depth(node.get("children", []), depth + 1)) + return max_depth + + +def _section_tree_to_output( + children_map: Dict[str, dict], + level: int = 1, +) -> List[Dict[str, Any]]: + result = [] + for node in children_map.values(): + children = _section_tree_to_output(node["_children_map"], level + 1) + total_chunks = node["chunk_count"] + sum( + child.get("chunk_count", 0) for child in children + ) + result.append( + { + "title": node["title"], + "path": node["path"], + "level": level, + "summary": node["summary"], + "chunk_count": total_chunks, + "children": children, + } + ) + return result diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index 68eb1b0fc..033c2db25 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -13,29 +13,20 @@ from loguru import logger from PIL import Image -from shared.services.chunks.chunk_connections import ( - build_resource_target_map, - convert_refs_to_embed_connections, - merge_connections, - normalize_connect_to_targets, - parse_relationship_refs, -) -from shared.utils.text_utils import truncate_content_preview - import pandas as pd from shared.core.exceptions.domain_exceptions import ( KnowhereException, StorageServiceException, ) -from shared.utils.utc_now import utc_now_naive +from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder class ZipResultService: """ZIP Result Package Generation Service""" def __init__(self): - pass + self._schema = ZipResultSchemaBuilder() def generate_zip_package( self, @@ -85,10 +76,10 @@ def generate_zip_package( table_files_map = {tb["id"]: tb for tb in table_files_info} # Convert chunks data format (using file info) - formatted_chunks = self._format_chunks( + formatted_chunks = self._schema.format_chunks( chunks, image_files_map, table_files_map ) - statistics = self._calculate_statistics(formatted_chunks) + statistics = self._schema.calculate_statistics(formatted_chunks) doc_nav: Dict[str, Any] = {} hierarchy: Dict[str, Any] = {} @@ -136,8 +127,8 @@ def generate_zip_package( # 5. Generate doc_nav.json — unified navigation file try: - doc_nav = self._build_doc_nav(formatted_chunks, source_file_name) - hierarchy = self._build_hierarchy_dict(doc_nav.get("sections", [])) + doc_nav = self._schema.build_doc_nav(formatted_chunks, source_file_name) + hierarchy = self._schema.build_hierarchy_dict(doc_nav.get("sections", [])) doc_nav_json = json.dumps(doc_nav, ensure_ascii=False, indent=2) zip_file.writestr("doc_nav.json", doc_nav_json.encode("utf-8")) logger.info("Added doc_nav.json") @@ -145,7 +136,7 @@ def generate_zip_package( logger.warning(f"generate doc_nav.json fail {e}") # 6. Generate manifest.json (checksum not included, stored in database) - manifest = self._generate_manifest( + manifest = self._schema.generate_manifest( job_id=job_id, data_id=data_id, source_file_name=source_file_name, @@ -179,184 +170,6 @@ def generate_zip_package( original_exception=e, ) - def _calculate_statistics(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]: - """Calculate statistics""" - total_chunks = len(chunks) - text_chunks = 0 - image_chunks = 0 - table_chunks = 0 - - for chunk in chunks: - chunk_type = chunk.get("type", "") - raw_type = str(chunk_type).strip() - normalized_type = raw_type.split("\n", 1)[0].lower() - if normalized_type == "image": - image_chunks += 1 - elif normalized_type == "table": - table_chunks += 1 - else: - text_chunks += 1 - - return { - "total_chunks": total_chunks, - "text_chunks": text_chunks, - "image_chunks": image_chunks, - "table_chunks": table_chunks, - "total_pages": None, # Cannot determine page count at this point - } - - def _format_chunks( - self, - chunks: List[Dict[str, Any]], - image_files_map: Dict[str, Dict[str, Any]], - table_files_map: Dict[str, Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Convert chunks data to ZIP specification format""" - resource_target_map = build_resource_target_map( - chunks, - image_files_map=image_files_map, - table_files_map=table_files_map, - ) - - formatted = [] - for chunk in chunks: - chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id")) - chunk_type_str = chunk.get("type", "") - raw_type = str(chunk_type_str).strip() - normalized_type = raw_type.split("\n", 1)[0].lower() - img_info = image_files_map.get(chunk_id) - - # Determine chunk type - if normalized_type == "image": - chunk_type = "image" - elif normalized_type == "table": - chunk_type = "table" - else: - chunk_type = "text" - - # Get content - content = chunk.get("text") or chunk.get("content", "") - - # Use original path directly to match kb.csv - path = chunk.get("path", "") - - # Get or build base metadata - existing_metadata = chunk.get("metadata", {}) - metadata = { - "length": existing_metadata.get("length") or len(content), - "summary": existing_metadata.get("summary") or chunk.get("summary", ""), - "page_nums": existing_metadata.get("page_nums", []), - } - document_top_summary = str( - existing_metadata.get("document_top_summary") or "" - ).strip() - if document_top_summary: - metadata["document_top_summary"] = document_top_summary - - # Add type-specific fields - if chunk_type == "text": - metadata["tokens"] = existing_metadata.get("tokens") or chunk.get( - "tokens", 0 - ) - metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( - "keywords", [] - ) - - # Convert in-text resource refs into embeds edges. - relationship_refs = parse_relationship_refs( - chunk.get("type_raw") or chunk_type_str, - str(content), - ) - embed_connections = convert_refs_to_embed_connections( - relationship_refs, resource_target_map - ) - related_connections = normalize_connect_to_targets( - existing_metadata.get("connect_to") - or chunk.get("connect_to") - or chunk.get("connectto"), - resource_target_map, - ) - metadata["connect_to"] = merge_connections( - embed_connections, related_connections - ) - - elif chunk_type == "image": - if img_info: - metadata["file_path"] = img_info["file_path"] - # Unified schema: include keywords and tokens for all chunk types - metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( - "keywords", [] - ) - metadata["tokens"] = [] - - elif chunk_type == "table": - # Get table info from existing_metadata or table_files_map - file_path = existing_metadata.get("file_path") - - if not file_path: - # Get table info from table_files_map - tb_info = table_files_map.get(chunk_id) - if tb_info: - file_path = tb_info["file_path"] - else: - # Extract from path or use default - tbl_name = ( - path.split("/")[-1] - if "/" in path - else f"table_{chunk_id}.html" - ) - file_path = f"tables/{tbl_name}" - - metadata["file_path"] = file_path - # Unified schema: include keywords and tokens for all chunk types - metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( - "keywords", [] - ) - metadata["tokens"] = [] - - formatted_chunk = { - "chunk_id": chunk_id, - "type": chunk_type, - "content": content, - "path": path, - "metadata": metadata, - } - formatted.append(formatted_chunk) - - return formatted - - def _clean_path(self, path: str) -> str: - """Clean path, keep only logical path""" - if not path: - return "/" - - # Remove filesystem path prefix - # Example: .-->users-->KB_DATA_xxx-->dir-->file.pdf-->chapter-->section - # Should extract: chapter-->section - - # Find the last .pdf, .docx, etc. file extension - import re - - # Match filename pattern (with extension) - file_pattern = r"[^/]+\.(pdf|docx|doc|txt|md|xlsx|xls|pptx|ppt)" - match = re.search(file_pattern, path, re.IGNORECASE) - - if match: - # Extract the part after filename - path_after_file = path[match.end() :] - # Clean path separators - path_after_file = path_after_file.replace("-->", "/").strip("/") - if path_after_file: - return path_after_file - - # If no file pattern found, try to clean common prefixes - path = path.replace("-->", "/") - # Remove leading path separators and empty segments - path = "/".join( - [p for p in path.split("/") if p and p not in ["", ".", "users"]] - ) - return path if path else "/" - def _collect_image_files( self, chunks: List[Dict[str, Any]], images_dir: str ) -> List[Dict[str, Any]]: @@ -642,41 +455,6 @@ def resolve_source_path( return table_files - def _generate_manifest( - self, - job_id: str, - data_id: Optional[str], - source_file_name: str, - statistics: Dict[str, Any], - job_metadata: Dict[str, Any], - hierarchy: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Generate manifest.json""" - manifest = { - "version": "2.0", - "job_id": job_id, - "data_id": data_id, - "source_file_name": source_file_name, - "processing_date": utc_now_naive().isoformat() + "Z", - "processing": { - "page_count": job_metadata.get("page_count"), - "billing_status": job_metadata.get("billing_status"), - "cost": { - "micro_dollars": job_metadata.get("billing_amount_micro_dollars"), - "credits": job_metadata.get("billing_credits"), - }, - "timing": { - "started_at": job_metadata.get("processing_started_at"), - "completed_at": job_metadata.get("processing_completed_at"), - "duration_ms": job_metadata.get("processing_duration_ms"), - }, - }, - "statistics": statistics, - "HIERARCHY": hierarchy or {}, - } - - return manifest - def _calculate_zip_checksum(self, zip_file_path: str) -> str: """Calculate SHA-256 checksum of ZIP file""" sha256_hash = hashlib.sha256() @@ -684,193 +462,3 @@ def _calculate_zip_checksum(self, zip_file_path: str) -> str: for byte_block in iter(lambda: f.read(4096), b""): sha256_hash.update(byte_block) return sha256_hash.hexdigest().lower() - - def _build_hierarchy_dict( - self, - sections: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """Build a title-only nested hierarchy from doc_nav sections.""" - hierarchy: Dict[str, Any] = {} - title_counts: Dict[str, int] = {} - - for section in sections: - raw_title = str(section.get("title") or "").strip() - if not raw_title: - continue - - title_counts[raw_title] = title_counts.get(raw_title, 0) + 1 - title = ( - raw_title - if title_counts[raw_title] == 1 - else f"{raw_title} ({title_counts[raw_title]})" - ) - hierarchy[title] = self._build_hierarchy_dict( - section.get("children") or [] - ) - - return hierarchy - - def _build_doc_nav( - self, - formatted_chunks: List[Dict[str, Any]], - source_file_name: str, - ) -> Dict[str, Any]: - """Build doc_nav.json — unified navigation file. - - Structured file serving both human demo and LLM navigation. - - The output contains: - - ``sections``: tree of text sections with summaries and chunk counts. - - ``resources``: flat lists of image/table chunks with summaries. - - ``stats``: chunk counts by type. - - Each leaf section carries a ``summary`` derived from: - 1. chunk.metadata.summary (LLM-generated, highest quality) - 2. chunk.content[:300] (fallback truncation) - - Non-leaf section summaries are left empty at this stage and are - populated later by ``summary_builder.enrich_doc_nav_summaries``. - """ - # ── Separate text chunks from resource chunks ── - text_chunks: List[Dict[str, Any]] = [] - image_resources: List[Dict[str, Any]] = [] - table_resources: List[Dict[str, Any]] = [] - - stats = {"total_chunks": 0, "text_chunks": 0, "image_chunks": 0, "table_chunks": 0, "max_depth": 0} - - for fc in formatted_chunks: - ctype = fc.get("type", "text") - path = fc.get("path", "") - meta = fc.get("metadata") or {} - summary_raw = (meta.get("summary") or "").strip() - content_raw = (fc.get("content") or "").strip() - # Normalize whitespace - summary = " ".join(summary_raw.split()) if summary_raw else "" - content_preview = truncate_content_preview(content_raw) if content_raw else "" - - stats["total_chunks"] += 1 - - if ctype == "image": - stats["image_chunks"] += 1 - image_resources.append({ - "path": path, - "summary": summary or content_preview, - }) - elif ctype == "table": - stats["table_chunks"] += 1 - table_resources.append({ - "path": path, - "summary": summary or content_preview, - }) - else: - stats["text_chunks"] += 1 - text_chunks.append({ - "path": path, - "summary": summary or content_preview, - }) - - # ── Build section tree from text chunk paths ── - # Each text chunk path looks like: "kb_root/filename.pdf/Section/Subsection" - # We strip the kb_root and filename prefix to get relative section paths. - sections = self._build_section_tree(text_chunks) - - # Compute max depth - def _max_depth(nodes: list, d: int = 1) -> int: - m = d if nodes else 0 - for n in nodes: - m = max(m, _max_depth(n.get("children", []), d + 1)) - return m - - stats["max_depth"] = _max_depth(sections) - - return { - "version": "1.0", - "file_name": source_file_name or "", - "stats": stats, - "sections": sections, - "resources": { - "images": image_resources, - "tables": table_resources, - }, - } - - def _build_section_tree( - self, - text_chunks: List[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Build a tree of sections from flat text chunk paths. - - Each text chunk has a ``path`` like ``"kb/file.pdf/Sec1/Sub1"``. - We extract section parts (after kb_root + filename) and build a - tree using ``children`` arrays. - - Returns a list of top-level section nodes. - """ - # Internal tree node: {title, summary, chunk_count, children: {title: node}} - root_children: Dict[str, dict] = {} # ordered dict of top-level titles - - for chunk in text_chunks: - path = chunk.get("path", "") - parts = [p.strip() for p in path.split("/") if p.strip()] - # Skip kb_root + filename → section parts start at index 2 - section_parts = parts[2:] if len(parts) > 2 else [] - - if not section_parts: - # Root-level chunk (no section hierarchy) - key = "__root__" - if key not in root_children: - root_children[key] = { - "title": "Root", - "path": "/".join(parts[:2]) if len(parts) >= 2 else path, - "summary": chunk.get("summary", ""), - "chunk_count": 0, - "_children_map": {}, - } - root_children[key]["chunk_count"] += 1 - # Use the first chunk's summary for root - if not root_children[key]["summary"]: - root_children[key]["summary"] = chunk.get("summary", "") - continue - - # Walk the tree, creating nodes as needed - current_level = root_children - full_section_path_parts = parts[:2] # start with kb_root/filename - for i, part in enumerate(section_parts): - full_section_path_parts.append(part) - if part not in current_level: - current_level[part] = { - "title": part, - "path": "/".join(full_section_path_parts), - "summary": "", - "chunk_count": 0, - "_children_map": {}, - } - node = current_level[part] - if i == len(section_parts) - 1: - # Leaf — this is the chunk's actual section - node["chunk_count"] += 1 - if not node["summary"]: - node["summary"] = chunk.get("summary", "") - current_level = node["_children_map"] - - # Convert internal tree to output format - def _to_output(children_map: Dict[str, dict], level: int = 1) -> List[Dict[str, Any]]: - result = [] - for node in children_map.values(): - children = _to_output(node["_children_map"], level + 1) - # Compute total chunk_count including descendants - total_chunks = node["chunk_count"] + sum( - c.get("chunk_count", 0) for c in children - ) - out = { - "title": node["title"], - "path": node["path"], - "level": level, - "summary": node["summary"], - "chunk_count": total_chunks, - "children": children, - } - result.append(out) - return result - - return _to_output(root_children) From c8c2eefe77f150d330be5cff821b7ca405952238 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:31:34 +0800 Subject: [PATCH 24/40] refactor deepen webhook delivery modules --- apps/api/app/api/v1/routes/webhook.py | 17 +- .../tests/contract/test_webhook_contract.py | 11 +- .../services/webhook/delivery_client.py | 151 +++++++++ .../shared/services/webhook/dispatcher.py | 314 ++---------------- .../shared/services/webhook/event_delivery.py | 127 +++++++ .../services/webhook/payload_enrichment.py | 50 +++ .../services/webhook/qstash_publisher.py | 20 +- .../services/webhook/secret_resolver.py | 60 ++++ .../shared/services/webhook/signing.py | 34 ++ 9 files changed, 458 insertions(+), 326 deletions(-) create mode 100644 packages/shared-python/shared/services/webhook/delivery_client.py create mode 100644 packages/shared-python/shared/services/webhook/event_delivery.py create mode 100644 packages/shared-python/shared/services/webhook/payload_enrichment.py create mode 100644 packages/shared-python/shared/services/webhook/secret_resolver.py create mode 100644 packages/shared-python/shared/services/webhook/signing.py diff --git a/apps/api/app/api/v1/routes/webhook.py b/apps/api/app/api/v1/routes/webhook.py index 18ba4abc8..4e155cd11 100644 --- a/apps/api/app/api/v1/routes/webhook.py +++ b/apps/api/app/api/v1/routes/webhook.py @@ -160,24 +160,17 @@ async def trigger_webhook( ), ) - # Use dispatcher to send synchronously dispatcher = get_webhook_dispatcher() - # Pass db session and is_manual=True to handle logging internally - ( - success, - status_code, - duration_ms, - error_message, - ) = await dispatcher._send_webhook(db=db, event=event, is_manual=True) + delivery_result = await dispatcher.send_manual_webhook(db=db, event=event) # 5. Return response return WebhookTriggerResponse( - success=success, - status_code=status_code, + success=delivery_result.success, + status_code=delivery_result.status_code, response_body=None, # Dispatcher doesn't return response body - duration_ms=duration_ms, + duration_ms=delivery_result.duration_ms, delivery_id=None, # Manual trigger doesn't create delivery log - error_message=error_message, + error_message=delivery_result.error_message, ) except KnowhereException: diff --git a/apps/api/tests/contract/test_webhook_contract.py b/apps/api/tests/contract/test_webhook_contract.py index 92f064d3c..890379af1 100644 --- a/apps/api/tests/contract/test_webhook_contract.py +++ b/apps/api/tests/contract/test_webhook_contract.py @@ -1,6 +1,7 @@ import importlib from collections.abc import Callable from contextlib import AbstractAsyncContextManager +from types import SimpleNamespace from typing import cast from uuid import uuid4 @@ -163,11 +164,15 @@ async def test_should_trigger_a_webhook_for_an_owned_terminal_job_with_a_matchin event_id: str = "" class FakeDispatcher: - async def _send_webhook(self, db, event, is_manual: bool = False): - assert is_manual is True + async def send_manual_webhook(self, db, event): assert event.id == event_id assert event.job_id == job_id - return True, 202, 118, None + return SimpleNamespace( + success=True, + status_code=202, + duration_ms=118, + error_message=None, + ) async with developer_api_client_factory() as api_client: job_id = await _insert_webhook_job(user_id="local-dev-user", status="done") diff --git a/packages/shared-python/shared/services/webhook/delivery_client.py b/packages/shared-python/shared/services/webhook/delivery_client.py new file mode 100644 index 000000000..232f9a875 --- /dev/null +++ b/packages/shared-python/shared/services/webhook/delivery_client.py @@ -0,0 +1,151 @@ +"""Pinned HTTP delivery for outbound webhooks.""" + +import asyncio +import time +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from loguru import logger + +from shared.utils.pinned_outbound_http import send_pinned_outbound_request +from shared.utils.url_security import validate_http_url_and_resolve_ip_async + +HTTP_TIMEOUT_SECONDS = 10 + + +@dataclass(frozen=True) +class WebhookDeliveryTarget: + target_url: str + pinned_ip: str + + +@dataclass(frozen=True) +class WebhookDeliveryResult: + success: bool + status_code: int | None + duration_ms: int + error_message: str | None + + +@dataclass(frozen=True) +class WebhookTargetValidation: + target: WebhookDeliveryTarget | None + failure: WebhookDeliveryResult | None + + +class WebhookDeliveryClient: + """Validate and send direct webhook HTTP requests with DNS pinning.""" + + async def validate_target( + self, *, event_id: str, target_url: str + ) -> WebhookTargetValidation: + validation = await validate_http_url_and_resolve_ip_async(target_url) + + if not validation.is_valid: + logger.warning( + f"SSRF validation failed: event_id={event_id}, " + f"error={validation.error_message}" + ) + return WebhookTargetValidation( + target=None, + failure=WebhookDeliveryResult( + success=False, + status_code=400, + duration_ms=0, + error_message=f"SSRF: {validation.error_message}", + ), + ) + + if not validation.validated_ip: + return WebhookTargetValidation( + target=None, + failure=WebhookDeliveryResult( + success=False, + status_code=400, + duration_ms=0, + error_message="SSRF validation did not return a pinned IP", + ), + ) + + return WebhookTargetValidation( + target=WebhookDeliveryTarget( + target_url=target_url, + pinned_ip=validation.validated_ip, + ), + failure=None, + ) + + async def post_json( + self, + *, + event_id: str, + target: WebhookDeliveryTarget, + payload: Mapping[str, Any], + headers: Mapping[str, str], + ) -> WebhookDeliveryResult: + start_time = time.time() + + try: + response = await send_pinned_outbound_request( + method="POST", + url=target.target_url, + pinned_ip=target.pinned_ip, + timeout_seconds=HTTP_TIMEOUT_SECONDS, + headers=headers, + json_body=payload, + ) + duration_ms = int((time.time() - start_time) * 1000) + + if 200 <= response.status < 300: + logger.info( + f"Webhook delivered: event_id={event_id}, status={response.status}" + ) + return WebhookDeliveryResult( + success=True, + status_code=response.status, + duration_ms=duration_ms, + error_message=None, + ) + + if 300 <= response.status < 400: + logger.warning( + f"Webhook redirect blocked (SSRF protection): " + f"event_id={event_id}, status={response.status}" + ) + return WebhookDeliveryResult( + success=False, + status_code=response.status, + duration_ms=duration_ms, + error_message=f"Redirect blocked: HTTP {response.status}", + ) + + logger.warning( + f"Webhook failed: event_id={event_id}, status={response.status}" + ) + return WebhookDeliveryResult( + success=False, + status_code=response.status, + duration_ms=duration_ms, + error_message=f"HTTP {response.status}", + ) + + except asyncio.TimeoutError: + duration_ms = int((time.time() - start_time) * 1000) + logger.error(f"Webhook timeout: event_id={event_id}") + return WebhookDeliveryResult( + success=False, + status_code=None, + duration_ms=duration_ms, + error_message="Connection timeout", + ) + + except Exception as error: + duration_ms = int((time.time() - start_time) * 1000) + logger.error(f"Webhook error: event_id={event_id}, error={error}") + return WebhookDeliveryResult( + success=False, + status_code=None, + duration_ms=duration_ms, + error_message=str(error), + ) diff --git a/packages/shared-python/shared/services/webhook/dispatcher.py b/packages/shared-python/shared/services/webhook/dispatcher.py index 1c5890f4f..bf6ada8e9 100644 --- a/packages/shared-python/shared/services/webhook/dispatcher.py +++ b/packages/shared-python/shared/services/webhook/dispatcher.py @@ -1,19 +1,13 @@ """ Webhook Dispatcher Service -Dispatches webhook events via HTTP requests with HMAC signing and delivery logging. -Called by Celery task for async processing. +Dispatches webhook events with retry policy. Direct HTTP delivery details live +behind WebhookEventDelivery. """ -import asyncio -import hashlib -import hmac -import json import threading -import time -import uuid from datetime import datetime, timezone -from typing import Any, Dict, Optional, Tuple +from typing import Optional from loguru import logger from sqlalchemy import select @@ -21,25 +15,12 @@ # Use standard db context - run_async_task handles the loop reuse from shared.core.database import get_db_context -from shared.core.exceptions.domain_exceptions import ( - SystemSettingInvalidException, - SystemSettingMissingException, -) from shared.core.exceptions.webhook_exceptions import WebhookDeliveryException -from shared.models.database.job import Job from shared.models.database.webhook import WebhookEvent, WebhookEventStatus -from shared.models.database.webhook_log import WebhookLog -from shared.services.jobs.result_delivery import JobResultDeliveryResolver -from shared.utils.pinned_outbound_http import ( - send_pinned_outbound_request, -) -from shared.utils.url_security import ( - HTTPURLValidationResult, - validate_http_url_and_resolve_ip_async, -) +from shared.services.webhook.delivery_client import WebhookDeliveryResult +from shared.services.webhook.event_delivery import WebhookEventDelivery # Constants -HTTP_TIMEOUT_SECONDS = 10 MAX_ATTEMPTS = 6 @@ -55,6 +36,9 @@ class WebhookDispatcher: 5. On failure, signals the caller to schedule retry """ + def __init__(self, event_delivery: WebhookEventDelivery | None = None) -> None: + self._event_delivery = event_delivery or WebhookEventDelivery() + async def dispatch(self, event_id: str) -> bool: """ Dispatch a webhook event. @@ -88,27 +72,24 @@ async def dispatch(self, event_id: str) -> bool: await self._mark_failed(db, event) return True # ACK - # 4. Dispatch the webhook - # Logging is now handled inside _send_webhook - success, status_code, duration_ms, error_message = await self._send_webhook( + delivery_result = await self._event_delivery.send( db=db, event=event, is_manual=False ) - # 6. Handle result (Logging already done) - if success: + if delivery_result.success: await self._mark_delivered(db, event) return True # Success else: # Determine if error is retryable # Retryable: 5xx, timeout (None), 429 (rate limit) # NOT retryable: 4xx (except 429) - client errors won't be fixed by retrying - is_retryable = self._is_retryable_error(status_code) + is_retryable = self._is_retryable_error(delivery_result.status_code) if not is_retryable: # Permanent failure - don't retry logger.warning( f"WebhookEvent permanent failure (non-retryable): " - f"event_id={event_id}, status={status_code}" + f"event_id={event_id}, status={delivery_result.status_code}" ) await self._mark_failed(db, event) return True # ACK - no point retrying @@ -122,9 +103,11 @@ async def dispatch(self, event_id: str) -> bool: # Raise exception so Celery task will retry raise WebhookDeliveryException( - internal_message=f"Webhook delivery failed: {error_message}", + internal_message=( + f"Webhook delivery failed: {delivery_result.error_message}" + ), retryable=True, - status_code=status_code, + status_code=delivery_result.status_code, ) async def mark_event_failed(self, event_id: str) -> None: @@ -148,210 +131,11 @@ async def _fetch_event( ) return result.scalar_one_or_none() - async def _send_webhook( - self, db: AsyncSession, event: WebhookEvent, is_manual: bool = False - ) -> Tuple[bool, Optional[int], int, Optional[str]]: - """ - Send HTTP POST request to webhook target and log the attempt. - - Args: - db: Database session - event: WebhookEvent object - is_manual: True if manually triggered (adds 'trigger': 'manual' to payload) - - Returns: - Tuple of (success, status_code, duration_ms, error_message) - """ - - # Generate attempt ID - attempt_id = str(uuid.uuid4()) - - # SSRF Protection - validation: HTTPURLValidationResult = await validate_http_url_and_resolve_ip_async( - event.target_url, - ) - if not validation.is_valid: - logger.warning( - f"SSRF validation failed: event_id={event.id}, error={validation.error_message}" - ) - return False, 400, 0, f"SSRF: {validation.error_message}" - - # Enrich payload with job result data at delivery time - enriched_payload = await self._enrich_payload(event) - - # Add manual mark if requested - if is_manual: - enriched_payload["trigger"] = "manual" - - # Helper to get user_id from job - async def _get_job_owner(job_id: str) -> Optional[str]: - result = await db.execute(select(Job.user_id).where(Job.job_id == job_id)) - return result.scalar_one_or_none() - - # Resolve secret (Lazy creation) - secret = None - try: - user_id = await _get_job_owner(event.job_id) - if user_id: - secret = await self._resolve_secret(db, user_id, event.target_url) - else: - logger.warning( - f"Could not resolve secret: Job {event.job_id} has no user_id" - ) - except (SystemSettingMissingException, SystemSettingInvalidException) as e: - logger.error(f"Configuration error during secret resolution: {e}") - # Return 424 (Failed Dependency) to ensure it's treated as a non-retryable error - return False, 424, 0, f"Configuration Error: {e}" - except Exception as e: - logger.error(f"Secret resolution failed: {e}") - - if not secret: - logger.error(f"No secret found or created/resolved for event {event.id}") - # Default to non-retryable error for any secret resolution failure - return False, 424, 0, "Secret resolution failed" - - # Sign payload - signature = self._sign_payload(enriched_payload, secret) - - # Build headers - headers = { - "Content-Type": "application/json", - "X-Knowhere-Signature": signature, - "X-Knowhere-Attempt-ID": attempt_id, - "User-Agent": "Knowhere-Webhook/1.0", - } - - start_time = time.time() - status_code = None - error_message = None - success = False - - try: - pinned_ip = validation.validated_ip - if not pinned_ip: - return False, 400, 0, "SSRF validation did not return a pinned IP" - - response = await send_pinned_outbound_request( - method="POST", - url=event.target_url, - pinned_ip=pinned_ip, - timeout_seconds=HTTP_TIMEOUT_SECONDS, - headers=headers, - json_body=enriched_payload, - ) - duration_ms = int((time.time() - start_time) * 1000) - status_code = response.status - - if 200 <= response.status < 300: - logger.info( - f"Webhook delivered: event_id={event.id}, status={response.status}" - ) - success = True - elif 300 <= response.status < 400: - logger.warning( - f"Webhook redirect blocked (SSRF protection): " - f"event_id={event.id}, status={response.status}" - ) - error_message = f"Redirect blocked: HTTP {response.status}" - success = False - else: - logger.warning( - f"Webhook failed: event_id={event.id}, status={response.status}" - ) - error_message = f"HTTP {response.status}" - success = False - - except asyncio.TimeoutError: - duration_ms = int((time.time() - start_time) * 1000) - logger.error(f"Webhook timeout: event_id={event.id}") - error_message = "Connection timeout" - success = False - - except Exception as e: - duration_ms = int((time.time() - start_time) * 1000) - logger.error(f"Webhook error: event_id={event.id}, error={e}") - error_message = str(e) - success = False - - # Log delivery attempt - # If manual, event_id is None to avoid FK violation - log_event_id = None if is_manual else event.id - - try: - # Combine headers and payload - combined_payload = {"header": headers, "payload": enriched_payload} - - log = WebhookLog( - job_id=event.job_id, - event_id=log_event_id, - webhook_url=event.target_url, - attempt_number=event.attempts + 1, - request_payload=combined_payload, - signature=signature, - idempotency_key=str(uuid.uuid4()), - response_status_code=status_code, - error_message=error_message, - duration_ms=duration_ms, - ) - db.add(log) - # If auto-commit is needed? - # Dispatcher.dispatch uses passed 'db' session which is managed by 'async with get_db_context()'. - # It commits inside _mark_delivered etc. - # We should probably commit/flush here to persist log even if update fails? - await db.commit() - - except Exception as e: - logger.error(f"Failed to log webhook delivery: {e}") - - return success, status_code, duration_ms, error_message - - async def _enrich_payload(self, event: WebhookEvent) -> Dict[str, Any]: - """ - Enrich webhook payload with job result data at delivery time. - - For job.completed events: - - Adds result_url (fresh download URL for result zip) - - Adds result (inline payload with checksum/statistics) - - This ensures download URLs are generated fresh (they expire) - and data is current at delivery time. - """ - payload = dict(event.payload) # Copy to avoid mutating stored payload - - # Only enrich completion events - if payload.get("event") != "job.completed": - return payload - - try: - # Fetch job with result - from sqlalchemy.orm import selectinload - - from shared.models.database.job import Job - - async with get_db_context() as db: - result = await db.execute( - select(Job) - .options(selectinload(Job.job_result)) - .where(Job.job_id == event.job_id) - ) - job = result.scalar_one_or_none() - - if not job or not job.job_result: - logger.warning( - f"Job or result not found for enrichment: job_id={event.job_id}" - ) - return payload - - payload = JobResultDeliveryResolver().enrich_payload( - payload, - job_result=job.job_result, - ) - - except Exception as e: - logger.error(f"Failed to enrich payload for event {event.id}: {e}") - # Continue with original payload if enrichment fails - - return payload + async def send_manual_webhook( + self, db: AsyncSession, event: WebhookEvent + ) -> WebhookDeliveryResult: + """Send a webhook immediately for an operator-triggered retry.""" + return await self._event_delivery.send(db=db, event=event, is_manual=True) def _is_retryable_error(self, status_code: Optional[int]) -> bool: """ @@ -387,62 +171,6 @@ def _is_retryable_error(self, status_code: Optional[int]) -> bool: # Examples: 400 Bad Request, 401 Unauthorized, 404 Not Found return False - async def _resolve_secret( - self, db: AsyncSession, user_id: str, endpoint: str - ) -> Optional[str]: - """ - Resolve webhook secret using repository (Lazy creation). - - 1. Try to get existing active secret for user/endpoint. - 2. If not found, create a new one. - 3. Decrypt and return the raw secret string. - """ - try: - # Import here to avoid circular dependency with WebhookDispatcher - from shared.repositories.webhook_secret_repository import ( - WebhookSecretRepository, - ) - - repo = WebhookSecretRepository() - secret_obj = await repo.get_or_create_secret(db, user_id, endpoint=endpoint) - - # Update usage timestamp - if secret_obj: - secret_obj.last_used_at = datetime.now(timezone.utc).replace( - tzinfo=None - ) - db.add(secret_obj) - # We don't commit here to avoid side effects if the caller aborts, - # but the session will eventually be committed by the caller. - - # Decrypt - return repo.decrypt_secret(secret_obj) - except (SystemSettingMissingException, SystemSettingInvalidException): - # Re-raise configuration errors so they can be handled as non-retryable - raise - except Exception as e: - logger.error(f"Failed to resolve/create secret for user {user_id}: {e}") - return None - - def _sign_payload(self, payload: Dict[str, Any], secret: str) -> str: - """ - Generate timestamped HMAC-SHA256 signature. - - Format: t=,v1= - Signed content: "{timestamp}.{json_payload}" - - This prevents replay attacks by binding the signature to the current time. - """ - timestamp = int(time.time()) - payload_str = json.dumps(payload, separators=(",", ":")) - signed_content = f"{timestamp}.{payload_str}" - - signature = hmac.new( - secret.encode("utf-8"), signed_content.encode("utf-8"), hashlib.sha256 - ).hexdigest() - - return f"t={timestamp},v1={signature}" - async def _mark_delivered(self, db: AsyncSession, event: WebhookEvent) -> None: """Mark event as delivered.""" event.status = WebhookEventStatus.DELIVERED diff --git a/packages/shared-python/shared/services/webhook/event_delivery.py b/packages/shared-python/shared/services/webhook/event_delivery.py new file mode 100644 index 000000000..21b5509fa --- /dev/null +++ b/packages/shared-python/shared/services/webhook/event_delivery.py @@ -0,0 +1,127 @@ +"""Direct WebhookEvent delivery attempt orchestration.""" + +import uuid +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + SystemSettingInvalidException, + SystemSettingMissingException, +) +from shared.models.database.webhook import WebhookEvent +from shared.models.database.webhook_log import WebhookLog +from shared.services.webhook.delivery_client import ( + WebhookDeliveryClient, + WebhookDeliveryResult, +) +from shared.services.webhook.payload_enrichment import WebhookPayloadEnricher +from shared.services.webhook.secret_resolver import WebhookSecretResolver +from shared.services.webhook.signing import build_webhook_headers + + +class WebhookEventDelivery: + """Send one direct webhook attempt and persist its delivery log.""" + + def __init__( + self, + *, + client: WebhookDeliveryClient | None = None, + enricher: WebhookPayloadEnricher | None = None, + secret_resolver: WebhookSecretResolver | None = None, + ) -> None: + self._client = client or WebhookDeliveryClient() + self._enricher = enricher or WebhookPayloadEnricher() + self._secret_resolver = secret_resolver or WebhookSecretResolver() + + async def send( + self, *, db: AsyncSession, event: WebhookEvent, is_manual: bool = False + ) -> WebhookDeliveryResult: + attempt_id = str(uuid.uuid4()) + target_validation = await self._client.validate_target( + event_id=event.id, + target_url=event.target_url, + ) + if target_validation.failure: + return target_validation.failure + if not target_validation.target: + return WebhookDeliveryResult( + success=False, + status_code=400, + duration_ms=0, + error_message="Webhook target validation failed", + ) + + payload = await self._enricher.enrich(event) + if is_manual: + payload["trigger"] = "manual" + + secret, secret_error = await self._resolve_secret(db, event) + if not secret: + logger.error(f"No secret found or created/resolved for event {event.id}") + return WebhookDeliveryResult( + success=False, + status_code=424, + duration_ms=0, + error_message=secret_error or "Secret resolution failed", + ) + + headers = build_webhook_headers( + payload=payload, + secret=secret, + attempt_id=attempt_id, + ) + result = await self._client.post_json( + event_id=event.id, + target=target_validation.target, + payload=payload, + headers=headers, + ) + await self._log_attempt( + db=db, + event=event, + is_manual=is_manual, + headers=headers, + payload=payload, + result=result, + ) + return result + + async def _resolve_secret( + self, db: AsyncSession, event: WebhookEvent + ) -> tuple[str | None, str | None]: + try: + return await self._secret_resolver.resolve_for_event(db, event), None + except (SystemSettingMissingException, SystemSettingInvalidException) as error: + logger.error(f"Configuration error during secret resolution: {error}") + return None, f"Configuration Error: {error}" + + async def _log_attempt( + self, + *, + db: AsyncSession, + event: WebhookEvent, + is_manual: bool, + headers: dict[str, str], + payload: dict[str, Any], + result: WebhookDeliveryResult, + ) -> None: + try: + log = WebhookLog( + job_id=event.job_id, + event_id=None if is_manual else event.id, + webhook_url=event.target_url, + attempt_number=event.attempts + 1, + request_payload={"header": headers, "payload": payload}, + signature=headers["X-Knowhere-Signature"], + idempotency_key=str(uuid.uuid4()), + response_status_code=result.status_code, + error_message=result.error_message, + duration_ms=result.duration_ms, + ) + db.add(log) + await db.commit() + + except Exception as error: + logger.error(f"Failed to log webhook delivery: {error}") diff --git a/packages/shared-python/shared/services/webhook/payload_enrichment.py b/packages/shared-python/shared/services/webhook/payload_enrichment.py new file mode 100644 index 000000000..2e8d1b73d --- /dev/null +++ b/packages/shared-python/shared/services/webhook/payload_enrichment.py @@ -0,0 +1,50 @@ +"""Delivery-time webhook payload enrichment.""" + +from collections.abc import Mapping +from typing import Any, cast + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.orm import selectinload + +from shared.core.database import get_db_context +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent +from shared.services.jobs.result_delivery import JobResultDeliveryResolver + + +class WebhookPayloadEnricher: + """Add fresh Job Result delivery metadata to webhook payloads.""" + + def __init__(self, resolver: JobResultDeliveryResolver | None = None) -> None: + self._resolver = resolver or JobResultDeliveryResolver() + + async def enrich(self, event: WebhookEvent) -> dict[str, Any]: + payload = dict(cast(Mapping[str, Any], event.payload)) + + if payload.get("event") != "job.completed": + return payload + + try: + async with get_db_context() as db: + result = await db.execute( + select(Job) + .options(selectinload(Job.job_result)) + .where(Job.job_id == event.job_id) + ) + job = result.scalar_one_or_none() + + if not job or not job.job_result: + logger.warning( + f"Job or result not found for enrichment: job_id={event.job_id}" + ) + return payload + + return self._resolver.enrich_payload( + payload, + job_result=job.job_result, + ) + + except Exception as error: + logger.error(f"Failed to enrich payload for event {event.id}: {error}") + return payload diff --git a/packages/shared-python/shared/services/webhook/qstash_publisher.py b/packages/shared-python/shared/services/webhook/qstash_publisher.py index 03279d89e..2c7dbf33e 100644 --- a/packages/shared-python/shared/services/webhook/qstash_publisher.py +++ b/packages/shared-python/shared/services/webhook/qstash_publisher.py @@ -11,10 +11,7 @@ from __future__ import annotations -import hashlib -import hmac import json -import time from dataclasses import dataclass from typing import Any, Dict, Optional @@ -24,6 +21,7 @@ from shared.core.exceptions.domain_exceptions import QStashServiceException from shared.models.database.webhook import WebhookEventStatus from shared.services.jobs.result_delivery import JobResultDeliveryResolver +from shared.services.webhook.signing import sign_webhook_payload from shared.utils.url_security import ( validate_http_url_and_resolve_ip, ) @@ -134,7 +132,7 @@ def publish_event(self, event_id: str) -> Optional[str]: return None # Sign payload with our HMAC - signature = self._sign_payload(payload, secret) + signature = sign_webhook_payload(payload, secret) # Publish to QStash try: @@ -346,20 +344,6 @@ def _resolve_secret(self, db: Any, user_id: str, endpoint: str) -> Optional[str] db.add(secret_obj) return fernet.decrypt(secret_obj.secret_encrypted) - @staticmethod - def _sign_payload(payload: Dict[str, Any], secret: str) -> str: - """Generate HMAC-SHA256 signature matching the existing Knowhere format.""" - timestamp = int(time.time()) - payload_str = json.dumps(payload, separators=(",", ":")) - signed_content = f"{timestamp}.{payload_str}" - sig = hmac.new( - secret.encode("utf-8"), - signed_content.encode("utf-8"), - hashlib.sha256, - ).hexdigest() - return f"t={timestamp},v1={sig}" - - # Module-level singleton _publisher: Optional[QStashWebhookPublisher] = None diff --git a/packages/shared-python/shared/services/webhook/secret_resolver.py b/packages/shared-python/shared/services/webhook/secret_resolver.py new file mode 100644 index 000000000..d0a648ef3 --- /dev/null +++ b/packages/shared-python/shared/services/webhook/secret_resolver.py @@ -0,0 +1,60 @@ +"""Webhook secret resolution for direct deliveries.""" + +from datetime import datetime, timezone + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ( + SystemSettingInvalidException, + SystemSettingMissingException, +) +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent +from shared.repositories.webhook_secret_repository import WebhookSecretRepository + + +class WebhookSecretResolver: + """Resolve the active endpoint secret for a WebhookEvent delivery.""" + + def __init__(self, repository: WebhookSecretRepository | None = None) -> None: + self._repository = repository or WebhookSecretRepository() + + async def resolve_for_event( + self, db: AsyncSession, event: WebhookEvent + ) -> str | None: + user_id = await self._get_job_owner(db, event.job_id) + if not user_id: + logger.warning(f"Could not resolve secret: Job {event.job_id} has no user_id") + return None + + return await self.resolve_for_endpoint( + db, + user_id=user_id, + endpoint=event.target_url, + ) + + async def resolve_for_endpoint( + self, db: AsyncSession, *, user_id: str, endpoint: str + ) -> str | None: + try: + secret = await self._repository.get_or_create_secret( + db, user_id, endpoint=endpoint + ) + + if secret: + secret.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None) + db.add(secret) + + return self._repository.decrypt_secret(secret) + + except (SystemSettingMissingException, SystemSettingInvalidException): + raise + except Exception as error: + logger.error(f"Failed to resolve/create secret for user {user_id}: {error}") + return None + + async def _get_job_owner(self, db: AsyncSession, job_id: str) -> str | None: + result = await db.execute(select(Job.user_id).where(Job.job_id == job_id)) + return result.scalar_one_or_none() diff --git a/packages/shared-python/shared/services/webhook/signing.py b/packages/shared-python/shared/services/webhook/signing.py new file mode 100644 index 000000000..15c521946 --- /dev/null +++ b/packages/shared-python/shared/services/webhook/signing.py @@ -0,0 +1,34 @@ +"""Webhook request signing.""" + +import hashlib +import hmac +import json +import time +from collections.abc import Mapping +from typing import Any + + +def sign_webhook_payload(payload: Mapping[str, Any], secret: str) -> str: + """Generate the timestamped Knowhere webhook HMAC signature.""" + timestamp = int(time.time()) + payload_text = json.dumps(payload, separators=(",", ":")) + signed_content = f"{timestamp}.{payload_text}" + signature = hmac.new( + secret.encode("utf-8"), + signed_content.encode("utf-8"), + hashlib.sha256, + ).hexdigest() + + return f"t={timestamp},v1={signature}" + + +def build_webhook_headers( + *, payload: Mapping[str, Any], secret: str, attempt_id: str +) -> dict[str, str]: + """Build signed HTTP headers for a direct webhook delivery attempt.""" + return { + "Content-Type": "application/json", + "X-Knowhere-Signature": sign_webhook_payload(payload, secret), + "X-Knowhere-Attempt-ID": attempt_id, + "User-Agent": "Knowhere-Webhook/1.0", + } From 35493104e35fd67fab5d64d717febe506e9ac992 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:42:02 +0800 Subject: [PATCH 25/40] refactor split retrieval hydration modules --- .../retrieval/agentic/navigation_tools.py | 6 +- .../shared/services/retrieval/assets.py | 2 +- .../services/retrieval/connected_hydration.py | 97 ++++ .../services/retrieval/execution_routes.py | 6 +- .../shared/services/retrieval/hydration.py | 549 ------------------ .../services/retrieval/path_hydration.py | 271 +++++++++ .../services/retrieval/reference_hydration.py | 128 ++++ .../services/retrieval/response_projection.py | 2 +- .../services/retrieval/result_assembly.py | 80 +++ .../shared/services/retrieval/row_utils.py | 82 +++ 10 files changed, 664 insertions(+), 559 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/connected_hydration.py delete mode 100644 packages/shared-python/shared/services/retrieval/hydration.py create mode 100644 packages/shared-python/shared/services/retrieval/path_hydration.py create mode 100644 packages/shared-python/shared/services/retrieval/reference_hydration.py create mode 100644 packages/shared-python/shared/services/retrieval/result_assembly.py create mode 100644 packages/shared-python/shared/services/retrieval/row_utils.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py index 8d489f8be..251f1aae5 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py @@ -28,10 +28,8 @@ load_child_sections, ) from shared.services.retrieval.agentic.types import DocTreeNode -from shared.services.retrieval.hydration import ( - hydrate_connected_target_rows, - hydrate_paths_to_rows, -) +from shared.services.retrieval.connected_hydration import hydrate_connected_target_rows +from shared.services.retrieval.path_hydration import hydrate_paths_to_rows from shared.services.retrieval.lexical_text import normalize_section_path from shared.services.retrieval.llm_adapter import LLMFn diff --git a/packages/shared-python/shared/services/retrieval/assets.py b/packages/shared-python/shared/services/retrieval/assets.py index 54868e980..035825430 100644 --- a/packages/shared-python/shared/services/retrieval/assets.py +++ b/packages/shared-python/shared/services/retrieval/assets.py @@ -4,7 +4,7 @@ from loguru import logger -from shared.services.retrieval.hydration import MEDIA_CHUNK_TYPES, normalize_chunk_type +from shared.services.retrieval.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type from shared.services.storage.result_storage import get_result_storage diff --git a/packages/shared-python/shared/services/retrieval/connected_hydration.py b/packages/shared-python/shared/services/retrieval/connected_hydration.py new file mode 100644 index 000000000..c5c2a181b --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/connected_hydration.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy import and_, or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.row_utils import ( + filter_excluded_rows, + iter_connected_target_ids, + normalize_chunk_type, +) + + +async def hydrate_connected_target_rows( + *, + db: AsyncSession | None, + rows: list[dict[str, Any]], + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + if db is None: + return [] + + existing_chunk_ids = { + str(row.get('chunk_id') or '').strip() + for row in rows + if row.get('chunk_id') + } + target_ids_by_revision: dict[tuple[str, str], set[str]] = {} + for row in rows: + if normalize_chunk_type(row.get('chunk_type')) != 'text': + continue + document_id = str(row.get('document_id') or '').strip() + job_result_id = str(row.get('job_result_id') or '').strip() + if not document_id or not job_result_id: + continue + for target_id in iter_connected_target_ids(row): + if target_id in existing_chunk_ids: + continue + target_ids_by_revision.setdefault((document_id, job_result_id), set()).add( + target_id + ) + + if not target_ids_by_revision: + return [] + + revision_filters = [ + and_( + DocumentChunk.document_id == document_id, + DocumentChunk.job_result_id == job_result_id, + DocumentChunk.chunk_id.in_(sorted(target_ids)), + ) + for (document_id, job_result_id), target_ids in target_ids_by_revision.items() + if target_ids + ] + if not revision_filters: + return [] + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(or_(*revision_filters)) + .order_by(DocumentChunk.sort_order) + ) + result = await db.execute(stmt) + + hydrated_rows: list[dict[str, Any]] = [] + for document, chunk, section, job_result in result.all(): + section_path = section.section_path if section else None + hydrated_rows.append( + { + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 0.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'sort_order': chunk.sort_order, + } + ) + + return filter_excluded_rows( + hydrated_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) diff --git a/packages/shared-python/shared/services/retrieval/execution_routes.py b/packages/shared-python/shared/services/retrieval/execution_routes.py index 14dcc619f..efa032d23 100644 --- a/packages/shared-python/shared/services/retrieval/execution_routes.py +++ b/packages/shared-python/shared/services/retrieval/execution_routes.py @@ -10,10 +10,8 @@ from shared.services.retrieval.channels import content_channel, path_channel, term_channel from shared.services.retrieval.graph_service import GraphQueryService -from shared.services.retrieval.hydration import ( - assemble_retrieval_results, - hydrate_referenced_chunk_rows, -) +from shared.services.retrieval.reference_hydration import hydrate_referenced_chunk_rows +from shared.services.retrieval.result_assembly import assemble_retrieval_results from shared.services.retrieval.ranking import rank_retrieval_candidates from shared.services.retrieval.response_projection import ( attach_citation, diff --git a/packages/shared-python/shared/services/retrieval/hydration.py b/packages/shared-python/shared/services/retrieval/hydration.py deleted file mode 100644 index 2b6821f96..000000000 --- a/packages/shared-python/shared/services/retrieval/hydration.py +++ /dev/null @@ -1,549 +0,0 @@ -from __future__ import annotations - -import re -from typing import Any - -from loguru import logger -from sqlalchemy import and_, or_, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import Document, DocumentChunk, DocumentSection -from shared.models.database.job_result import JobResult -from shared.services.retrieval.graph_service import is_excluded_section -from shared.services.retrieval.lexical_text import normalize_section_path -from shared.services.retrieval.scoring import get_row_path - -MEDIA_CHUNK_TYPES = {'image', 'table'} -PUBLIC_RESULT_FIELDS = { - 'chunk_type', 'content', 'score', 'asset_url', -} -PUBLIC_SOURCE_FIELDS = { - 'document_id', 'source_file_name', 'section_path', -} - -ReferenceLookupKey = tuple[str, str, str, str] - -_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]') - - -def clean_content(content: str) -> str: - return _PATH_REF_RE.sub('', content).strip() - - -def normalize_chunk_type(raw: object) -> str: - return str(raw or '').strip().split('\n', 1)[0].lower() - - -def is_media_chunk(row: dict[str, Any]) -> bool: - return normalize_chunk_type(row.get('chunk_type')) in MEDIA_CHUNK_TYPES - - -def build_reference_lookup_key( - *, - document_id: object, - chunk_id: object, - section_path: object = '', - file_path: object = '', -) -> ReferenceLookupKey: - return ( - str(document_id or '').strip(), - str(chunk_id or '').strip(), - str(section_path or '').strip(), - str(file_path or '').strip(), - ) - - -def filter_excluded_rows( - rows: list[dict[str, Any]], - *, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - filtered: list[dict[str, Any]] = [] - excluded_documents = set(exclude_document_ids) - for row in rows: - document_id = row.get('document_id') - if document_id in excluded_documents: - continue - if is_excluded_section( - document_id=document_id, - section_path=row.get('section_path'), - exclude_sections=exclude_sections, - ): - continue - filtered.append(row) - return filtered - - -def iter_connected_target_ids(row: dict[str, Any]) -> list[str]: - metadata = row.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - return [] - - target_ids: list[str] = [] - for item in metadata.get('connect_to') or []: - if not isinstance(item, dict): - continue - target_id = str(item.get('target') or '').strip() - if target_id: - target_ids.append(target_id) - return target_ids - - -async def hydrate_connected_target_rows( - *, - db: AsyncSession | None, - rows: list[dict[str, Any]], - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - if db is None: - return [] - - existing_chunk_ids = { - str(row.get('chunk_id') or '').strip() - for row in rows - if row.get('chunk_id') - } - target_ids_by_revision: dict[tuple[str, str], set[str]] = {} - for row in rows: - if normalize_chunk_type(row.get('chunk_type')) != 'text': - continue - document_id = str(row.get('document_id') or '').strip() - job_result_id = str(row.get('job_result_id') or '').strip() - if not document_id or not job_result_id: - continue - for target_id in iter_connected_target_ids(row): - if target_id in existing_chunk_ids: - continue - target_ids_by_revision.setdefault((document_id, job_result_id), set()).add(target_id) - - if not target_ids_by_revision: - return [] - - revision_filters = [ - and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == job_result_id, - DocumentChunk.chunk_id.in_(sorted(target_ids)), - ) - for (document_id, job_result_id), target_ids in target_ids_by_revision.items() - if target_ids - ] - if not revision_filters: - return [] - - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(or_(*revision_filters)) - .order_by(DocumentChunk.sort_order) - ) - result = await db.execute(stmt) - - hydrated_rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): - section_path = section.section_path if section else None - hydrated_rows.append( - { - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': 0.0, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'sort_order': chunk.sort_order, - } - ) - - return filter_excluded_rows( - hydrated_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - - -async def hydrate_referenced_chunk_rows( - *, - db: AsyncSession | None, - user_id: str, - namespace: str, - refs: list[dict[str, Any]], -) -> list[dict[str, Any]]: - if db is None or not refs: - return [] - - ref_keys = [ - build_reference_lookup_key( - document_id=ref.get('document_id'), - chunk_id=ref.get('chunk_id'), - section_path=ref.get('section_path'), - file_path=ref.get('file_path'), - ) - for ref in refs - ] - ref_keys = [key for key in ref_keys if key[0] and key[1]] - if not ref_keys: - return [] - - document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) - chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), - ) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(Document.document_id.in_(document_ids)) - .where(DocumentChunk.chunk_id.in_(chunk_ids)) - .order_by(DocumentChunk.sort_order) - ) - result = await db.execute(stmt) - - rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} - rows_by_base_key: dict[tuple[str, str], list[dict[str, Any]]] = {} - for document, chunk, section, job_result in result.all(): - row = { - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section.section_path if section else None, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': 1.0, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'source_chunk_path': chunk.source_chunk_path, - 'sort_order': chunk.sort_order, - } - key = build_reference_lookup_key( - document_id=row['document_id'], - chunk_id=row['chunk_id'], - section_path=row['section_path'], - file_path=row['file_path'], - ) - rows_by_key[key] = row - rows_by_base_key.setdefault((key[0], key[1]), []).append(row) - - rows: list[dict[str, Any]] = [] - seen_keys: set[ReferenceLookupKey] = set() - for key in ref_keys: - row = rows_by_key.get(key) - if row is None: - candidates = rows_by_base_key.get((key[0], key[1]), []) - row = next( - ( - candidate for candidate in candidates - if key[2] and str(candidate.get('section_path') or '').strip() == key[2] - ), - None, - ) - if row is None: - row = next( - ( - candidate for candidate in candidates - if build_reference_lookup_key( - document_id=candidate.get('document_id'), - chunk_id=candidate.get('chunk_id'), - section_path=candidate.get('section_path'), - file_path=candidate.get('file_path'), - ) - not in seen_keys - ), - None, - ) - if row is not None: - row_key = build_reference_lookup_key( - document_id=row.get('document_id'), - chunk_id=row.get('chunk_id'), - section_path=row.get('section_path'), - file_path=row.get('file_path'), - ) - if row_key in seen_keys: - continue - seen_keys.add(row_key) - rows.append(row) - return rows - - -async def assemble_retrieval_results( - *, - db: AsyncSession | None = None, - rows: list[dict[str, Any]], - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None = None, -) -> list[dict[str, Any]]: - filtered_rows = filter_excluded_rows( - rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if allowed_chunk_types is not None: - filtered_rows = [ - row for row in filtered_rows - if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types - ] - hydrated_rows = await hydrate_connected_target_rows( - db=db, - rows=filtered_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - rows_by_chunk_id = { - str(row.get('chunk_id') or ''): row - for row in [*filtered_rows, *hydrated_rows] - if row.get('chunk_id') - } - - embedded_targets: set[str] = set() - for row in filtered_rows: - for target_id in iter_connected_target_ids(row): - if target_id in rows_by_chunk_id: - embedded_targets.add(target_id) - - assembled: list[dict[str, Any]] = [] - for row in filtered_rows: - if row.get('chunk_id') in embedded_targets: - continue - metadata = row.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - metadata = {} - assembled_row = dict(row) - base_content = str(row.get('content') or '') - if normalize_chunk_type(row.get('chunk_type')) == 'text': - connected_targets: list[tuple[int, str]] = [] - for target_id in iter_connected_target_ids(row): - target_row = rows_by_chunk_id.get(target_id) - if not target_row: - continue - if normalize_chunk_type(target_row.get('chunk_type')) != 'table': - continue - target_content = str(target_row.get('content') or '').strip() - if target_content: - sort_key = int(target_row.get('sort_order', 0) or 0) - connected_targets.append((sort_key, target_content)) - connected_targets.sort(key=lambda x: x[0]) - related_parts = [content for _, content in connected_targets] - if base_content and related_parts: - assembled_row['content'] = '\n\n'.join([base_content, *related_parts]) - else: - assembled_row['content'] = base_content - else: - assembled_row['content'] = base_content - assembled_row['content'] = clean_content(assembled_row['content']) - assembled.append(assembled_row) - return assembled - - -async def hydrate_paths_to_rows( - db: AsyncSession, - *, - path_selections: list[dict[str, Any]], - user_id: str, - namespace: str, - document_id: str | None = None, -) -> list[dict[str, Any]]: - """Load full chunk rows by section_path or source_chunk_path.""" - if not path_selections: - return [] - - confidence_by_path: dict[str, float] = {} - mode_by_path: dict[str, str] = {} - ordered_paths: list[str] = [] - for item in path_selections: - raw_path = str(item.get('path') or '').strip() - path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path - if not path: - continue - confidence = float(item.get('confidence', 0.0) or 0.0) - hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower() - if path not in confidence_by_path: - ordered_paths.append(path) - confidence_by_path[path] = confidence - mode_by_path[path] = hydrate_mode - else: - confidence_by_path[path] = max(confidence_by_path[path], confidence) - if not ordered_paths: - return [] - - outline_paths = [p for p in ordered_paths if mode_by_path.get(p) == 'outline'] - chunk_paths = [p for p in ordered_paths if mode_by_path.get(p) != 'outline'] - - rows: list[dict[str, Any]] = [] - - if outline_paths: - outline_section_filters = [ - DocumentSection.section_path == path - for path in outline_paths - ] - outline_stmt = ( - select(Document, DocumentSection) - .join( - DocumentSection, - (DocumentSection.document_id == Document.document_id) - & (DocumentSection.job_result_id == Document.current_job_result_id), - ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(or_(*outline_section_filters)) - ) - if document_id: - outline_stmt = outline_stmt.where(Document.document_id == document_id) - outline_result = await db.execute(outline_stmt) - for document, section in outline_result.all(): - agent_score = confidence_by_path.get(section.section_path, 0.0) - summary_text = (section.summary or '').strip() - title_text = (section.section_title or '').strip() - content = f'[Outline] {title_text}' - if summary_text: - content += f'\n{summary_text}' - rows.append({ - 'document_id': document.document_id, - 'chunk_id': f'outline_{section.section_id}', - 'section_id': section.section_id, - 'section_path': section.section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': 'outline', - 'content': content, - 'score': agent_score, - 'agent_score': agent_score, - 'file_path': None, - 'chunk_metadata': {}, - 'job_result_id': section.job_result_id, - 'job_id': None, - 'source_chunk_path': None, - 'sort_order': section.sort_order, - 'hydrate_mode': 'outline', - }) - - if chunk_paths: - section_path_filters = [] - self_only_paths = {p for p in chunk_paths if mode_by_path.get(p) == 'self_only'} - for path in chunk_paths: - section_path_filters.append(DocumentSection.section_path == path) - if path not in self_only_paths: - section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) - - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), - ) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where( - or_( - *section_path_filters, - DocumentChunk.source_chunk_path.in_(chunk_paths), - ) - ) - ) - if document_id: - stmt = stmt.where(Document.document_id == document_id) - result = await db.execute(stmt) - - mode_allowed_types: dict[str, set[str] | None] = { - 'chunks': None, - 'self_only': None, - 'assets_only': {'image', 'table'}, - 'image_only': {'image'}, - 'table_only': {'table'}, - } - - seen_paths: set[str] = set() - for document, chunk, section, job_result in result.all(): - row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' - if row_path in seen_paths: - continue - - matched_path = row_path - if section and section.section_path not in confidence_by_path: - matched_path = next( - ( - path for path in chunk_paths - if section.section_path == path or section.section_path.startswith(f'{path} / ') - ), - row_path, - ) - - path_mode = mode_by_path.get(matched_path, 'chunks') - allowed_types = mode_allowed_types.get(path_mode) - if allowed_types is not None: - chunk_type_lower = (chunk.chunk_type or '').strip().lower() - if chunk_type_lower not in allowed_types: - continue - - seen_paths.add(row_path) - agent_score = confidence_by_path.get(matched_path, 0.0) - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section.section_path if section else None, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': agent_score, - 'agent_score': agent_score, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'source_chunk_path': chunk.source_chunk_path, - 'sort_order': chunk.sort_order, - 'hydrate_mode': path_mode, - }) - - path_order = {path: index for index, path in enumerate(ordered_paths)} - - def _row_sort_key(row: dict[str, Any]) -> int: - row_path = get_row_path(row) - if row_path in path_order: - return path_order[row_path] - for path, index in path_order.items(): - if row_path.startswith(f'{path} / '): - return index - return 10**9 - - rows.sort(key=_row_sort_key) - hydrated_paths = {get_row_path(row) for row in rows} - resolved_inputs = { - path for path in ordered_paths - if path in hydrated_paths - or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths) - } - resolved_inputs |= set(outline_paths) - missed = len(ordered_paths) - len(resolved_inputs) - if missed > 0: - missing_paths = [path for path in ordered_paths if path not in resolved_inputs] - logger.warning( - f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); ' - f'missing[:5]={missing_paths[:5]}' - ) - else: - logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved') - return rows diff --git a/packages/shared-python/shared/services/retrieval/path_hydration.py b/packages/shared-python/shared/services/retrieval/path_hydration.py new file mode 100644 index 000000000..d973f8a8b --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/path_hydration.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.lexical_text import normalize_section_path +from shared.services.retrieval.scoring import get_row_path + + +async def hydrate_paths_to_rows( + db: AsyncSession, + *, + path_selections: list[dict[str, Any]], + user_id: str, + namespace: str, + document_id: str | None = None, +) -> list[dict[str, Any]]: + """Load full chunk rows by section_path or source_chunk_path.""" + if not path_selections: + return [] + + confidence_by_path: dict[str, float] = {} + mode_by_path: dict[str, str] = {} + ordered_paths: list[str] = [] + for item in path_selections: + raw_path = str(item.get('path') or '').strip() + path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path + if not path: + continue + confidence = float(item.get('confidence', 0.0) or 0.0) + hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower() + if path not in confidence_by_path: + ordered_paths.append(path) + confidence_by_path[path] = confidence + mode_by_path[path] = hydrate_mode + else: + confidence_by_path[path] = max(confidence_by_path[path], confidence) + if not ordered_paths: + return [] + + outline_paths = [path for path in ordered_paths if mode_by_path.get(path) == 'outline'] + chunk_paths = [path for path in ordered_paths if mode_by_path.get(path) != 'outline'] + + rows: list[dict[str, Any]] = [] + + if outline_paths: + rows.extend( + await _hydrate_outline_paths( + db, + outline_paths=outline_paths, + confidence_by_path=confidence_by_path, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + ) + + if chunk_paths: + rows.extend( + await _hydrate_chunk_paths( + db, + chunk_paths=chunk_paths, + confidence_by_path=confidence_by_path, + mode_by_path=mode_by_path, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + ) + + _sort_rows_by_selection_order(rows, ordered_paths) + _log_hydration_resolution(rows=rows, ordered_paths=ordered_paths, outline_paths=outline_paths) + return rows + + +async def _hydrate_outline_paths( + db: AsyncSession, + *, + outline_paths: list[str], + confidence_by_path: dict[str, float], + user_id: str, + namespace: str, + document_id: str | None, +) -> list[dict[str, Any]]: + outline_section_filters = [ + DocumentSection.section_path == path + for path in outline_paths + ] + outline_stmt = ( + select(Document, DocumentSection) + .join( + DocumentSection, + (DocumentSection.document_id == Document.document_id) + & (DocumentSection.job_result_id == Document.current_job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(or_(*outline_section_filters)) + ) + if document_id: + outline_stmt = outline_stmt.where(Document.document_id == document_id) + + rows: list[dict[str, Any]] = [] + outline_result = await db.execute(outline_stmt) + for document, section in outline_result.all(): + agent_score = confidence_by_path.get(section.section_path, 0.0) + summary_text = (section.summary or '').strip() + title_text = (section.section_title or '').strip() + content = f'[Outline] {title_text}' + if summary_text: + content += f'\n{summary_text}' + rows.append({ + 'document_id': document.document_id, + 'chunk_id': f'outline_{section.section_id}', + 'section_id': section.section_id, + 'section_path': section.section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': 'outline', + 'content': content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': None, + 'chunk_metadata': {}, + 'job_result_id': section.job_result_id, + 'job_id': None, + 'source_chunk_path': None, + 'sort_order': section.sort_order, + 'hydrate_mode': 'outline', + }) + return rows + + +async def _hydrate_chunk_paths( + db: AsyncSession, + *, + chunk_paths: list[str], + confidence_by_path: dict[str, float], + mode_by_path: dict[str, str], + user_id: str, + namespace: str, + document_id: str | None, +) -> list[dict[str, Any]]: + section_path_filters = [] + self_only_paths = {path for path in chunk_paths if mode_by_path.get(path) == 'self_only'} + for path in chunk_paths: + section_path_filters.append(DocumentSection.section_path == path) + if path not in self_only_paths: + section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where( + or_( + *section_path_filters, + DocumentChunk.source_chunk_path.in_(chunk_paths), + ) + ) + ) + if document_id: + stmt = stmt.where(Document.document_id == document_id) + result = await db.execute(stmt) + + rows: list[dict[str, Any]] = [] + seen_paths: set[str] = set() + for document, chunk, section, job_result in result.all(): + row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' + if row_path in seen_paths: + continue + + matched_path = row_path + if section and section.section_path not in confidence_by_path: + matched_path = next( + ( + path for path in chunk_paths + if section.section_path == path + or section.section_path.startswith(f'{path} / ') + ), + row_path, + ) + + path_mode = mode_by_path.get(matched_path, 'chunks') + allowed_types = _get_allowed_types_for_mode(path_mode) + if allowed_types is not None: + chunk_type_lower = (chunk.chunk_type or '').strip().lower() + if chunk_type_lower not in allowed_types: + continue + + seen_paths.add(row_path) + agent_score = confidence_by_path.get(matched_path, 0.0) + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, + 'hydrate_mode': path_mode, + }) + return rows + + +def _get_allowed_types_for_mode(path_mode: str) -> set[str] | None: + mode_allowed_types: dict[str, set[str] | None] = { + 'chunks': None, + 'self_only': None, + 'assets_only': {'image', 'table'}, + 'image_only': {'image'}, + 'table_only': {'table'}, + } + return mode_allowed_types.get(path_mode) + + +def _sort_rows_by_selection_order(rows: list[dict[str, Any]], ordered_paths: list[str]) -> None: + path_order = {path: index for index, path in enumerate(ordered_paths)} + + def row_sort_key(row: dict[str, Any]) -> int: + row_path = get_row_path(row) + if row_path in path_order: + return path_order[row_path] + for path, index in path_order.items(): + if row_path.startswith(f'{path} / '): + return index + return 10**9 + + rows.sort(key=row_sort_key) + + +def _log_hydration_resolution( + *, rows: list[dict[str, Any]], ordered_paths: list[str], outline_paths: list[str] +) -> None: + hydrated_paths = {get_row_path(row) for row in rows} + resolved_inputs = { + path for path in ordered_paths + if path in hydrated_paths + or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths) + } + resolved_inputs |= set(outline_paths) + missed = len(ordered_paths) - len(resolved_inputs) + if missed > 0: + missing_paths = [path for path in ordered_paths if path not in resolved_inputs] + logger.warning( + f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); ' + f'missing[:5]={missing_paths[:5]}' + ) + else: + logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved') diff --git a/packages/shared-python/shared/services/retrieval/reference_hydration.py b/packages/shared-python/shared/services/retrieval/reference_hydration.py new file mode 100644 index 000000000..299024e2e --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/reference_hydration.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.row_utils import ( + ReferenceLookupKey, + build_reference_lookup_key, +) + + +async def hydrate_referenced_chunk_rows( + *, + db: AsyncSession | None, + user_id: str, + namespace: str, + refs: list[dict[str, Any]], +) -> list[dict[str, Any]]: + if db is None or not refs: + return [] + + ref_keys = [ + build_reference_lookup_key( + document_id=ref.get('document_id'), + chunk_id=ref.get('chunk_id'), + section_path=ref.get('section_path'), + file_path=ref.get('file_path'), + ) + for ref in refs + ] + ref_keys = [key for key in ref_keys if key[0] and key[1]] + if not ref_keys: + return [] + + document_ids = sorted({document_id for document_id, _, _, _ in ref_keys}) + chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys}) + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(Document.document_id.in_(document_ids)) + .where(DocumentChunk.chunk_id.in_(chunk_ids)) + .order_by(DocumentChunk.sort_order) + ) + result = await db.execute(stmt) + + rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} + rows_by_base_key: dict[tuple[str, str], list[dict[str, Any]]] = {} + for document, chunk, section, job_result in result.all(): + row = { + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 1.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, + } + key = build_reference_lookup_key( + document_id=row['document_id'], + chunk_id=row['chunk_id'], + section_path=row['section_path'], + file_path=row['file_path'], + ) + rows_by_key[key] = row + rows_by_base_key.setdefault((key[0], key[1]), []).append(row) + + rows: list[dict[str, Any]] = [] + seen_keys: set[ReferenceLookupKey] = set() + for key in ref_keys: + row = rows_by_key.get(key) + if row is None: + candidates = rows_by_base_key.get((key[0], key[1]), []) + row = next( + ( + candidate + for candidate in candidates + if key[2] + and str(candidate.get('section_path') or '').strip() == key[2] + ), + None, + ) + if row is None: + row = next( + ( + candidate + for candidate in candidates + if build_reference_lookup_key( + document_id=candidate.get('document_id'), + chunk_id=candidate.get('chunk_id'), + section_path=candidate.get('section_path'), + file_path=candidate.get('file_path'), + ) + not in seen_keys + ), + None, + ) + if row is not None: + row_key = build_reference_lookup_key( + document_id=row.get('document_id'), + chunk_id=row.get('chunk_id'), + section_path=row.get('section_path'), + file_path=row.get('file_path'), + ) + if row_key in seen_keys: + continue + seen_keys.add(row_key) + rows.append(row) + return rows diff --git a/packages/shared-python/shared/services/retrieval/response_projection.py b/packages/shared-python/shared/services/retrieval/response_projection.py index 71c092267..f12638a5c 100644 --- a/packages/shared-python/shared/services/retrieval/response_projection.py +++ b/packages/shared-python/shared/services/retrieval/response_projection.py @@ -3,7 +3,7 @@ from typing import Any from shared.services.retrieval.assets import enrich_rows_with_retrieval_asset_urls -from shared.services.retrieval.hydration import ( +from shared.services.retrieval.row_utils import ( PUBLIC_RESULT_FIELDS, PUBLIC_SOURCE_FIELDS, ) diff --git a/packages/shared-python/shared/services/retrieval/result_assembly.py b/packages/shared-python/shared/services/retrieval/result_assembly.py new file mode 100644 index 000000000..1ef948078 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/result_assembly.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.connected_hydration import hydrate_connected_target_rows +from shared.services.retrieval.row_utils import ( + clean_content, + filter_excluded_rows, + iter_connected_target_ids, + normalize_chunk_type, +) + + +async def assemble_retrieval_results( + *, + db: AsyncSession | None = None, + rows: list[dict[str, Any]], + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + allowed_chunk_types: set[str] | None = None, +) -> list[dict[str, Any]]: + filtered_rows = filter_excluded_rows( + rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + if allowed_chunk_types is not None: + filtered_rows = [ + row for row in filtered_rows + if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types + ] + hydrated_rows = await hydrate_connected_target_rows( + db=db, + rows=filtered_rows, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + rows_by_chunk_id = { + str(row.get('chunk_id') or ''): row + for row in [*filtered_rows, *hydrated_rows] + if row.get('chunk_id') + } + + embedded_targets: set[str] = set() + for row in filtered_rows: + for target_id in iter_connected_target_ids(row): + if target_id in rows_by_chunk_id: + embedded_targets.add(target_id) + + assembled: list[dict[str, Any]] = [] + for row in filtered_rows: + if row.get('chunk_id') in embedded_targets: + continue + assembled_row = dict(row) + base_content = str(row.get('content') or '') + if normalize_chunk_type(row.get('chunk_type')) == 'text': + connected_targets: list[tuple[int, str]] = [] + for target_id in iter_connected_target_ids(row): + target_row = rows_by_chunk_id.get(target_id) + if not target_row: + continue + if normalize_chunk_type(target_row.get('chunk_type')) != 'table': + continue + target_content = str(target_row.get('content') or '').strip() + if target_content: + sort_key = int(target_row.get('sort_order', 0) or 0) + connected_targets.append((sort_key, target_content)) + connected_targets.sort(key=lambda item: item[0]) + related_parts = [content for _, content in connected_targets] + if base_content and related_parts: + assembled_row['content'] = '\n\n'.join([base_content, *related_parts]) + else: + assembled_row['content'] = base_content + else: + assembled_row['content'] = base_content + assembled_row['content'] = clean_content(assembled_row['content']) + assembled.append(assembled_row) + return assembled diff --git a/packages/shared-python/shared/services/retrieval/row_utils.py b/packages/shared-python/shared/services/retrieval/row_utils.py new file mode 100644 index 000000000..86afa461d --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/row_utils.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import re +from typing import Any + +from shared.services.retrieval.graph_service import is_excluded_section + +MEDIA_CHUNK_TYPES = {'image', 'table'} +PUBLIC_RESULT_FIELDS = { + 'chunk_type', 'content', 'score', 'asset_url', +} +PUBLIC_SOURCE_FIELDS = { + 'document_id', 'source_file_name', 'section_path', +} + +ReferenceLookupKey = tuple[str, str, str, str] + +_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]') + + +def clean_content(content: str) -> str: + return _PATH_REF_RE.sub('', content).strip() + + +def normalize_chunk_type(raw: object) -> str: + return str(raw or '').strip().split('\n', 1)[0].lower() + + +def is_media_chunk(row: dict[str, Any]) -> bool: + return normalize_chunk_type(row.get('chunk_type')) in MEDIA_CHUNK_TYPES + + +def build_reference_lookup_key( + *, + document_id: object, + chunk_id: object, + section_path: object = '', + file_path: object = '', +) -> ReferenceLookupKey: + return ( + str(document_id or '').strip(), + str(chunk_id or '').strip(), + str(section_path or '').strip(), + str(file_path or '').strip(), + ) + + +def filter_excluded_rows( + rows: list[dict[str, Any]], + *, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + filtered: list[dict[str, Any]] = [] + excluded_documents = set(exclude_document_ids) + for row in rows: + document_id = row.get('document_id') + if document_id in excluded_documents: + continue + if is_excluded_section( + document_id=document_id, + section_path=row.get('section_path'), + exclude_sections=exclude_sections, + ): + continue + filtered.append(row) + return filtered + + +def iter_connected_target_ids(row: dict[str, Any]) -> list[str]: + metadata = row.get('chunk_metadata') or {} + if not isinstance(metadata, dict): + return [] + + target_ids: list[str] = [] + for item in metadata.get('connect_to') or []: + if not isinstance(item, dict): + continue + target_id = str(item.get('target') or '').strip() + if target_id: + target_ids.append(target_id) + return target_ids From 7519c3ac41da3d1b0e8571cf2eb203e3a540f7e3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:54:13 +0800 Subject: [PATCH 26/40] refactor split retrieval graph modules --- .../shared/services/retrieval/__init__.py | 3 +- .../shared/services/retrieval/channels.py | 2 +- .../services/retrieval/execution_routes.py | 2 +- .../services/retrieval/graph_keywords.py | 100 +++++ .../services/retrieval/graph_query_service.py | 208 ++++++++++ .../services/retrieval/graph_service.py | 363 +++--------------- .../shared/services/retrieval/row_utils.py | 2 +- .../services/retrieval/scoped_corpus.py | 2 +- .../services/retrieval/section_filters.py | 25 ++ 9 files changed, 383 insertions(+), 324 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/graph_keywords.py create mode 100644 packages/shared-python/shared/services/retrieval/graph_query_service.py create mode 100644 packages/shared-python/shared/services/retrieval/section_filters.py diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py index 832e81eb6..4c1a56a95 100644 --- a/packages/shared-python/shared/services/retrieval/__init__.py +++ b/packages/shared-python/shared/services/retrieval/__init__.py @@ -6,7 +6,8 @@ invalidate_retrieval_cache_namespaces, set_cached_retrieval_query_result, ) -from .graph_service import DocumentGraphService, GraphQueryService, GraphScope +from .graph_query_service import GraphQueryService +from .graph_service import DocumentGraphService, GraphScope from .hit_stats_service import record_retrieval_hits from .llm_adapter import create_retrieval_llm_fn, create_retrieval_planner_fn diff --git a/packages/shared-python/shared/services/retrieval/channels.py b/packages/shared-python/shared/services/retrieval/channels.py index a597307ca..ec1764c94 100644 --- a/packages/shared-python/shared/services/retrieval/channels.py +++ b/packages/shared-python/shared/services/retrieval/channels.py @@ -12,7 +12,7 @@ from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.graph_service import is_excluded_section +from shared.services.retrieval.section_filters import is_excluded_section from shared.utils.text_utils import tokenize_for_retrieval diff --git a/packages/shared-python/shared/services/retrieval/execution_routes.py b/packages/shared-python/shared/services/retrieval/execution_routes.py index efa032d23..23d16261f 100644 --- a/packages/shared-python/shared/services/retrieval/execution_routes.py +++ b/packages/shared-python/shared/services/retrieval/execution_routes.py @@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.channels import content_channel, path_channel, term_channel -from shared.services.retrieval.graph_service import GraphQueryService +from shared.services.retrieval.graph_query_service import GraphQueryService from shared.services.retrieval.reference_hydration import hydrate_referenced_chunk_rows from shared.services.retrieval.result_assembly import assemble_retrieval_results from shared.services.retrieval.ranking import rank_retrieval_candidates diff --git a/packages/shared-python/shared/services/retrieval/graph_keywords.py b/packages/shared-python/shared/services/retrieval/graph_keywords.py new file mode 100644 index 000000000..5769bd355 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/graph_keywords.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import math +import re +from typing import Any + +MIN_KEYWORD_OVERLAP = 3 +KEYWORD_SCORE_WEIGHT = 1.0 +MIN_SCORE_THRESHOLD = 0.8 + + +def normalize_keyword(keyword: str) -> str: + """Normalize a keyword: lowercase, strip, collapse spaces.""" + keyword = keyword.lower().strip() + return re.sub(r'\s+', ' ', keyword) + + +def extract_keywords_from_chunk_metadata(meta: dict) -> list[str]: + """Extract keywords from chunk metadata.""" + if not isinstance(meta, dict): + return [] + + keywords = meta.get('keywords', []) + if isinstance(keywords, list) and keywords: + return [str(keyword) for keyword in keywords if keyword] + + tokens = meta.get('tokens', []) + if isinstance(tokens, list) and tokens: + return [str(token) for token in tokens if token and len(str(token)) > 1] + + return [] + + +def compute_tfidf_keywords( + chunk_metadata_list: list[dict[str, Any]], + top_k: int = 10, +) -> list[str]: + """Compute TF-IDF keywords from chunk metadata.""" + df_count: dict[str, int] = {} + tf_count: dict[str, int] = {} + total = len(chunk_metadata_list) or 1 + for meta in chunk_metadata_list: + keywords = extract_keywords_from_chunk_metadata(meta) + seen: set[str] = set() + for keyword in keywords: + if len(str(keyword)) <= 1 or re.match(r'^\d+[.,%]*$', str(keyword)): + continue + normalized = normalize_keyword(str(keyword)) + if not normalized: + continue + tf_count[normalized] = tf_count.get(normalized, 0) + 1 + if normalized not in seen: + df_count[normalized] = df_count.get(normalized, 0) + 1 + seen.add(normalized) + scored = [ + (term, freq * (math.log(total / (df_count.get(term, 1))) + 1)) + for term, freq in tf_count.items() + ] + scored.sort(key=lambda item: item[1], reverse=True) + return [term for term, _ in scored[:top_k]] + + +def compute_keyword_score( + shared_keywords: set[str], + keywords_a: set[str], + keywords_b: set[str], + weight: float = 1.0, +) -> float: + """Character-length-weighted keyword overlap score.""" + weighted_a = sum(len(keyword) for keyword in keywords_a) + weighted_b = sum(len(keyword) for keyword in keywords_b) + denominator = min(weighted_a, weighted_b) + if denominator == 0: + return 0.0 + weighted_shared = sum(len(keyword) for keyword in shared_keywords) + return weight * weighted_shared / denominator + + +def get_normalized_keyword_set(chunk_metadata_list: list[dict[str, Any]]) -> set[str]: + """Collect all normalized keywords from chunk metadata for a document.""" + result: set[str] = set() + for meta in chunk_metadata_list: + for keyword in extract_keywords_from_chunk_metadata(meta): + normalized = normalize_keyword(str(keyword)) + if normalized and len(normalized) > 1 and not re.match( + r'^\d+[.,%]*$', normalized + ): + result.add(normalized) + return result + + +def extract_document_top_summary(chunk_metadata_list: list[dict[str, Any]]) -> str: + """Read the parser-injected top summary from chunk metadata.""" + for meta in chunk_metadata_list: + if not isinstance(meta, dict): + continue + summary = str(meta.get('document_top_summary') or '').strip() + if summary: + return summary + return '' diff --git a/packages/shared-python/shared/services/retrieval/graph_query_service.py b/packages/shared-python/shared/services/retrieval/graph_query_service.py new file mode 100644 index 000000000..de912a0ca --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/graph_query_service.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import Any + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.section_filters import is_excluded_section + +_SECTION_EXCLUSION_PAGE_MULTIPLIER = 2 + + +def _build_lexical_match_predicate(query: str): + like = f'%{query}%' + return ( + DocumentChunk.content_lexical_text.ilike(like) + | DocumentChunk.path_lexical_text.ilike(like) + ) + + +class GraphQueryService: + """Read-side graph routing before canonical chunk hydration.""" + + async def find_entry_documents( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: Iterable[str] = (), + exclude_sections: Iterable[dict[str, str]] = (), + ) -> list[str]: + query_lc = query.lower().strip() + excluded_document_ids = set(exclude_document_ids) + + if query_lc: + section_matches = await self._find_documents_by_section( + db, + user_id=user_id, + namespace=namespace, + query=query_lc, + exclude_document_ids=excluded_document_ids, + exclude_sections=exclude_sections, + ) + if section_matches: + return section_matches + + return await self._find_documents_by_content( + db, + user_id=user_id, + namespace=namespace, + query=query_lc, + exclude_document_ids=excluded_document_ids, + ) + + async def _find_documents_by_section( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: set[str], + exclude_sections: Iterable[dict[str, str]], + ) -> list[str]: + like = f'%{query}%' + stmt = ( + select(DocumentSection.document_id) + .join( + Document, + (Document.document_id == DocumentSection.document_id) + & (Document.current_job_result_id == DocumentSection.job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where( + DocumentSection.section_title.ilike(like) + | DocumentSection.section_path.ilike(like) + ) + .distinct() + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + for item in exclude_sections or (): + if not isinstance(item, dict): + continue + excluded_document_id = str(item.get('document_id') or '').strip() + excluded_path = str(item.get('section_path') or '').strip() + if excluded_document_id and excluded_path: + stmt = stmt.where( + ~( + (DocumentSection.document_id == excluded_document_id) + & ( + (DocumentSection.section_path == excluded_path) + | DocumentSection.section_path.like(f'{excluded_path} / %') + ) + ) + ) + + result = await db.execute(stmt) + return [document_id for (document_id,) in result.all()] + + async def _find_documents_by_content( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + exclude_document_ids: set[str], + ) -> list[str]: + like = f'%{query}%' + stmt = ( + select(Document.document_id) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(DocumentChunk.content_lexical_text.ilike(like)) + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + result = await db.execute(stmt) + seen: list[str] = [] + for (document_id,) in result.all(): + if document_id and document_id not in seen: + seen.append(document_id) + return seen + + async def collect_candidate_chunks( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + entry_document_ids: Sequence[str], + query: str, + top_k: int, + exclude_sections: Iterable[dict[str, str]] = (), + ) -> list[dict[str, Any]]: + if not entry_document_ids: + return [] + page_size = top_k + if exclude_sections: + page_size = max(top_k, top_k * _SECTION_EXCLUSION_PAGE_MULTIPLIER) + base_stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin( + DocumentSection, + DocumentSection.section_id == DocumentChunk.section_id, + ) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(Document.document_id.in_(list(entry_document_ids))) + .where(_build_lexical_match_predicate(query)) + .order_by(DocumentChunk.sort_order) + ) + rows: list[dict[str, Any]] = [] + offset = 0 + while len(rows) < top_k: + result = await db.execute(base_stmt.limit(page_size).offset(offset)) + result_rows = result.all() + if not result_rows: + break + for document, chunk, section, job_result in result_rows: + section_path = section.section_path if section else None + if is_excluded_section( + document_id=document.document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': 2.0, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + }) + if len(rows) >= top_k: + break + if len(result_rows) < page_size: + break + offset += page_size + return rows diff --git a/packages/shared-python/shared/services/retrieval/graph_service.py b/packages/shared-python/shared/services/retrieval/graph_service.py index 622f321d2..0d96a76e3 100644 --- a/packages/shared-python/shared/services/retrieval/graph_service.py +++ b/packages/shared-python/shared/services/retrieval/graph_service.py @@ -1,160 +1,31 @@ from __future__ import annotations import logging -import math -import re from collections import defaultdict from dataclasses import dataclass -from typing import Any, Iterable, Sequence from sqlalchemy import delete, or_, select -from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session -from shared.models.database.document import Document, DocumentChunk, DocumentSection, GraphEdge, GraphNode -from shared.models.database.job_result import JobResult +from shared.models.database.document import ( + Document, + DocumentChunk, + GraphEdge, + GraphNode, +) +from shared.services.retrieval.graph_keywords import ( + KEYWORD_SCORE_WEIGHT, + MIN_KEYWORD_OVERLAP, + MIN_SCORE_THRESHOLD, + compute_keyword_score, + compute_tfidf_keywords, + extract_document_top_summary, + get_normalized_keyword_set, + normalize_keyword, +) logger = logging.getLogger(__name__) -_SECTION_EXCLUSION_PAGE_MULTIPLIER = 2 - -# ── Keyword overlap config for document-level publication graph ── -_MIN_KEYWORD_OVERLAP = 3 -_KEYWORD_SCORE_WEIGHT = 1.0 -_MIN_SCORE_THRESHOLD = 0.8 -_CROSS_FILE_ONLY = True -_MAX_CONTENT_OVERLAP = 0.8 - - -def _build_lexical_match_predicate(query: str): - like = f'%{query}%' - return ( - DocumentChunk.content_lexical_text.ilike(like) - | DocumentChunk.path_lexical_text.ilike(like) - ) - - -def is_excluded_section( - *, - document_id: str | None, - section_path: str | None, - exclude_sections: Iterable[dict[str, str]], -) -> bool: - document_id = str(document_id or '').strip() - section_path = str(section_path or '').strip() - if not document_id or not section_path: - return False - for item in exclude_sections: - if not isinstance(item, dict): - continue - exc_doc = str(item.get('document_id') or '').strip() - exc_path = str(item.get('section_path') or '').strip() - if document_id == exc_doc and (section_path == exc_path or section_path.startswith(exc_path + ' / ')): - return True - return False - - -# ── Keyword extraction & scoring for document-level publication graph ── - -def _normalize_keyword(keyword: str) -> str: - """Normalize a keyword: lowercase, strip, collapse spaces.""" - kw = keyword.lower().strip() - return re.sub(r'\s+', ' ', kw) - - -def _extract_keywords_from_chunk_metadata(meta: dict) -> list[str]: - """Extract keywords from chunk metadata.""" - if not isinstance(meta, dict): - return [] - # Try metadata.keywords - kws = meta.get('keywords', []) - if isinstance(kws, list) and kws: - return [str(k) for k in kws if k] - # Fallback: tokens - tokens = meta.get('tokens', []) - if isinstance(tokens, list) and tokens: - return [str(t) for t in tokens if t and len(str(t)) > 1] - return [] - - -def _compute_tfidf_keywords( - chunk_metadata_list: list[dict[str, Any]], - top_k: int = 10, -) -> list[str]: - """Compute TF-IDF keywords from chunk metadata.""" - df_count: dict[str, int] = {} - tf_count: dict[str, int] = {} - total = len(chunk_metadata_list) or 1 - for meta in chunk_metadata_list: - kws = _extract_keywords_from_chunk_metadata(meta) - seen: set[str] = set() - for k in kws: - if len(str(k)) <= 1 or re.match(r'^\d+[.,%]*$', str(k)): - continue - lower = _normalize_keyword(str(k)) - if not lower: - continue - tf_count[lower] = tf_count.get(lower, 0) + 1 - if lower not in seen: - df_count[lower] = df_count.get(lower, 0) + 1 - seen.add(lower) - scored = [ - (term, freq * (math.log(total / (df_count.get(term, 1))) + 1)) - for term, freq in tf_count.items() - ] - scored.sort(key=lambda x: x[1], reverse=True) - return [s[0] for s in scored[:top_k]] - - -def _compute_keyword_score( - shared_kws: set[str], - kws_a: set[str], - kws_b: set[str], - weight: float = 1.0, -) -> float: - """Character-length-weighted keyword overlap score. - - Longer tokens contribute more: '施工现场'(4) has 2x weight of '交底'(2). - Formula: score = weight * sum(len(kw) for shared) / min(sum(len) for A, sum(len) for B) - """ - weighted_a = sum(len(k) for k in kws_a) - weighted_b = sum(len(k) for k in kws_b) - denominator = min(weighted_a, weighted_b) - if denominator == 0: - return 0.0 - weighted_shared = sum(len(k) for k in shared_kws) - return weight * weighted_shared / denominator - - -def _get_normalized_keyword_set(chunk_metadata_list: list[dict[str, Any]]) -> set[str]: - """Collect all normalized keywords from chunk metadata for a document.""" - result: set[str] = set() - for meta in chunk_metadata_list: - for k in _extract_keywords_from_chunk_metadata(meta): - normalized = _normalize_keyword(str(k)) - if normalized and len(normalized) > 1 and not re.match(r'^\d+[.,%]*$', normalized): - result.add(normalized) - return result - - -def _extract_document_top_summary( - chunk_metadata_list: list[dict[str, Any]], - section_titles: Sequence[str], -) -> str: - """Extract document_top_summary from chunk metadata. - - The summary is injected by kb_tasks.py via load_nav_top_summary() - at parse time, so it should always be present. If missing, return empty - string rather than fabricating a low-quality fallback. - """ - for meta in chunk_metadata_list: - if not isinstance(meta, dict): - continue - summary = str(meta.get('document_top_summary') or '').strip() - if summary: - return summary - return '' - @dataclass class GraphScope: @@ -171,7 +42,15 @@ class DocumentGraphService: - Edges are keyword-overlap-based cross-document connections with meaningful scores """ - def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, document_id: str, job_result_id: str) -> None: + def publish_document_graph( + self, + db: Session, + *, + user_id: str, + namespace: str, + document_id: str, + job_result_id: str, + ) -> None: document = db.execute( select(Document).where(Document.document_id == document_id) ).scalar_one_or_none() @@ -189,29 +68,22 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d chunk_metadata_list = [row[1] or {} for row in chunk_meta_rows] # Compute document-level metadata (aligned with KB knowledge_graph.json files dict) - top_keywords = _compute_tfidf_keywords(chunk_metadata_list) - new_doc_kws = _get_normalized_keyword_set(chunk_metadata_list) + top_keywords = compute_tfidf_keywords(chunk_metadata_list) + new_doc_kws = get_normalized_keyword_set(chunk_metadata_list) types_breakdown: dict[str, int] = defaultdict(int) for chunk_type, _ in chunk_meta_rows: types_breakdown[chunk_type or 'text'] += 1 chunks_count = len(chunk_meta_rows) - sections = [ - section_title - for section_title in db.execute( - select(DocumentSection.section_title) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.section_level <= 2) - .order_by(DocumentSection.sort_order) - ).scalars() - if section_title is not None - ] - top_summary = _extract_document_top_summary(chunk_metadata_list, sections) + top_summary = extract_document_top_summary(chunk_metadata_list) # ── Clean up old graph data for this document ── - self.remove_document_graph(db, scope=GraphScope(user_id=user_id, namespace=namespace), document_id=document_id) + self.remove_document_graph( + db, + scope=GraphScope(user_id=user_id, namespace=namespace), + document_id=document_id, + ) # ── Create document-level node (no section nodes — aligned with KB KG) ── document_node_id = f"doc:{document_id}" @@ -255,7 +127,7 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d # Build normalized keyword sets for comparison peer_kws: set[str] = set() for k in peer_keywords: - normalized = _normalize_keyword(str(k)) + normalized = normalize_keyword(str(k)) if normalized: peer_kws.add(normalized) @@ -264,17 +136,17 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d # Find shared keywords shared_kws = new_doc_kws & peer_kws - if len(shared_kws) < _MIN_KEYWORD_OVERLAP: + if len(shared_kws) < MIN_KEYWORD_OVERLAP: continue # Compute character-length-weighted score - score = _compute_keyword_score( - shared_kws=shared_kws, - kws_a=new_doc_kws, - kws_b=peer_kws, - weight=_KEYWORD_SCORE_WEIGHT, + score = compute_keyword_score( + shared_keywords=shared_kws, + keywords_a=new_doc_kws, + keywords_b=peer_kws, + weight=KEYWORD_SCORE_WEIGHT, ) - if score < _MIN_SCORE_THRESHOLD: + if score < MIN_SCORE_THRESHOLD: continue # Create edge with meaningful weight and metadata @@ -305,7 +177,9 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d f"keywords={len(top_keywords)} chunks={chunks_count}" ) - def remove_document_graph(self, db: Session, *, scope: GraphScope | None, document_id: str) -> None: + def remove_document_graph( + self, db: Session, *, scope: GraphScope | None, document_id: str + ) -> None: document_node_id = f"doc:{document_id}" edge_delete = delete(GraphEdge).where( or_( @@ -327,152 +201,3 @@ def remove_document_graph(self, db: Session, *, scope: GraphScope | None, docume db.execute(edge_delete) db.execute(node_delete) db.flush() - - -class GraphQueryService: - """Read-side graph service for document routing before canonical chunk hydration.""" - - async def find_entry_documents( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - exclude_document_ids: Iterable[str] = (), - exclude_sections: Iterable[dict[str, str]] = (), - ) -> list[str]: - query_lc = query.lower().strip() - exclude_document_ids = set(exclude_document_ids) - - if query_lc: - like = f'%{query_lc}%' - stmt = ( - select(DocumentSection.document_id) - .join(Document, (Document.document_id == DocumentSection.document_id) & (Document.current_job_result_id == DocumentSection.job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where( - DocumentSection.section_title.ilike(like) - | DocumentSection.section_path.ilike(like) - ) - .distinct() - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - for exc in (exclude_sections or ()): - if not isinstance(exc, dict): - continue - exc_doc = str(exc.get('document_id') or '').strip() - exc_path = str(exc.get('section_path') or '').strip() - if exc_doc and exc_path: - stmt = stmt.where( - ~((DocumentSection.document_id == exc_doc) & ( - (DocumentSection.section_path == exc_path) | - DocumentSection.section_path.like(f'{exc_path} / %') - )) - ) - result = await db.execute(stmt) - seen = [row[0] for row in result.all()] - if seen: - return seen - - return await self._find_documents_by_content( - db, - user_id=user_id, - namespace=namespace, - query=query_lc, - exclude_document_ids=exclude_document_ids, - ) - - async def _find_documents_by_content( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - exclude_document_ids: set[str], - ) -> list[str]: - like = f'%{query}%' - stmt = ( - select(Document.document_id) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(DocumentChunk.content_lexical_text.ilike(like)) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - result = await db.execute(stmt) - seen: list[str] = [] - for (doc_id,) in result.all(): - if doc_id and doc_id not in seen: - seen.append(doc_id) - return seen - - async def collect_candidate_chunks( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - entry_document_ids: Sequence[str], - query: str, - top_k: int, - exclude_sections: Iterable[dict[str, str]] = (), - ) -> list[dict[str, Any]]: - if not entry_document_ids: - return [] - page_size = top_k - if exclude_sections: - page_size = max(top_k, top_k * _SECTION_EXCLUSION_PAGE_MULTIPLIER) - base_stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where(Document.document_id.in_(list(entry_document_ids))) - .where(_build_lexical_match_predicate(query)) - .order_by(DocumentChunk.sort_order) - ) - rows = [] - offset = 0 - while len(rows) < top_k: - result = await db.execute(base_stmt.limit(page_size).offset(offset)) - result_rows = result.all() - if not result_rows: - break - for document, chunk, section, job_result in result_rows: - section_path = section.section_path if section else None - if is_excluded_section( - document_id=document.document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': 2.0, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - }) - if len(rows) >= top_k: - break - if len(result_rows) < page_size: - break - offset += page_size - return rows diff --git a/packages/shared-python/shared/services/retrieval/row_utils.py b/packages/shared-python/shared/services/retrieval/row_utils.py index 86afa461d..0784830b0 100644 --- a/packages/shared-python/shared/services/retrieval/row_utils.py +++ b/packages/shared-python/shared/services/retrieval/row_utils.py @@ -3,7 +3,7 @@ import re from typing import Any -from shared.services.retrieval.graph_service import is_excluded_section +from shared.services.retrieval.section_filters import is_excluded_section MEDIA_CHUNK_TYPES = {'image', 'table'} PUBLIC_RESULT_FIELDS = { diff --git a/packages/shared-python/shared/services/retrieval/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/scoped_corpus.py index 604400253..012ba7129 100644 --- a/packages/shared-python/shared/services/retrieval/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/scoped_corpus.py @@ -7,7 +7,7 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult -from shared.services.retrieval.graph_service import is_excluded_section +from shared.services.retrieval.section_filters import is_excluded_section async def count_scoped_chunks( diff --git a/packages/shared-python/shared/services/retrieval/section_filters.py b/packages/shared-python/shared/services/retrieval/section_filters.py new file mode 100644 index 000000000..f74b7f3c3 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/section_filters.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from collections.abc import Iterable + + +def is_excluded_section( + *, + document_id: str | None, + section_path: str | None, + exclude_sections: Iterable[dict[str, str]], +) -> bool: + document_id = str(document_id or '').strip() + section_path = str(section_path or '').strip() + if not document_id or not section_path: + return False + for item in exclude_sections: + if not isinstance(item, dict): + continue + exc_doc = str(item.get('document_id') or '').strip() + exc_path = str(item.get('section_path') or '').strip() + if document_id == exc_doc and ( + section_path == exc_path or section_path.startswith(exc_path + ' / ') + ): + return True + return False From f5d8d3b2dd2f1cc9dbfb85af6583576c4070a83e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 17:11:29 +0800 Subject: [PATCH 27/40] refactor extract agentic selection hydration --- .../retrieval/agentic/navigation_tools.py | 128 +----------------- .../retrieval/agentic/selection_hydration.py | 104 ++++++++++++++ 2 files changed, 109 insertions(+), 123 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py index 251f1aae5..23edec823 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py @@ -15,7 +15,6 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import DocumentChunk, DocumentSection -from shared.services.retrieval.agentic import asset_tools from shared.services.retrieval.agentic.budget import BudgetExceeded from shared.services.retrieval.agentic.prompts import ( ACTION_PROMPT, @@ -27,9 +26,10 @@ format_items_for_llm, load_child_sections, ) +from shared.services.retrieval.agentic.selection_hydration import ( + hydrate_path_selections_into_node, +) from shared.services.retrieval.agentic.types import DocTreeNode -from shared.services.retrieval.connected_hydration import hydrate_connected_target_rows -from shared.services.retrieval.path_hydration import hydrate_paths_to_rows from shared.services.retrieval.lexical_text import normalize_section_path from shared.services.retrieval.llm_adapter import LLMFn @@ -141,7 +141,7 @@ async def navigate_step( "hydrate_mode": "self_only", }) - await _hydrate_selections_into_node( + await hydrate_path_selections_into_node( db, node=node, path_selections=path_selections, @@ -216,7 +216,7 @@ async def discovery_select_step( root_path_selections=root_path_selections, node=node, ) - await _hydrate_discovery_selections_into_node( + await hydrate_path_selections_into_node( db, node=node, path_selections=path_selections, @@ -341,56 +341,6 @@ def _build_navigation_prompt( return prompt -async def _hydrate_selections_into_node( - db: AsyncSession, - *, - node: DocTreeNode, - path_selections: list[dict[str, Any]], - user_id: str, - namespace: str, - document_id: str, - job_result_id: str, -) -> None: - chunks = await hydrate_paths_to_rows( - db, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - if not chunks: - return - - connected = await hydrate_connected_target_rows( - db=db, - rows=chunks, - exclude_document_ids=[], - exclude_sections=[], - ) - if connected: - owner_map = asset_tools.build_connected_owner_map(chunks) - for chunk in connected: - if not chunk.get("owner_section_path"): - chunk["owner_section_path"] = owner_map.get(str(chunk.get("chunk_id") or "")) - chunks = chunks + connected - - root_map = await asset_tools.resolve_root_asset_owners( - db, - document_id=document_id, - job_result_id=job_result_id, - chunks=chunks, - ) - if root_map: - for chunk in chunks: - if chunk.get("owner_section_path"): - continue - chunk_id = str(chunk.get("chunk_id") or "") - if chunk_id in root_map: - chunk["owner_section_path"] = root_map[chunk_id] - - _add_chunks_to_node(node, chunks) - - def _project_discovery_hints( hints: list[dict[str, Any]], *, @@ -492,71 +442,3 @@ def _build_discovery_path_selections( }) return path_selections - - -async def _hydrate_discovery_selections_into_node( - db: AsyncSession, - *, - node: DocTreeNode, - path_selections: list[dict[str, Any]], - user_id: str, - namespace: str, - document_id: str, -) -> None: - chunks = await hydrate_paths_to_rows( - db, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - if not chunks: - return - - connected = await hydrate_connected_target_rows( - db=db, - rows=chunks, - exclude_document_ids=[], - exclude_sections=[], - ) - if connected: - owner_map = asset_tools.build_connected_owner_map(chunks) - for chunk in connected: - if not chunk.get("owner_section_path"): - chunk["owner_section_path"] = owner_map.get(str(chunk.get("chunk_id") or "")) - chunks = chunks + connected - - job_result_id = next( - (str(chunk["job_result_id"]) for chunk in chunks if chunk.get("job_result_id")), - None, - ) - root_map = ( - await asset_tools.resolve_root_asset_owners( - db, - document_id=document_id, - job_result_id=job_result_id, - chunks=chunks, - ) - if job_result_id - else {} - ) - if root_map: - for chunk in chunks: - if chunk.get("owner_section_path"): - continue - chunk_id = str(chunk.get("chunk_id") or "") - if chunk_id in root_map: - chunk["owner_section_path"] = root_map[chunk_id] - - _add_chunks_to_node(node, chunks) - - -def _add_chunks_to_node(node: DocTreeNode, chunks: list[dict[str, Any]]) -> None: - for chunk in chunks: - real_path = ( - chunk.get("owner_section_path") - or chunk.get("section_path") - or chunk.get("source_chunk_path") - ) - if real_path: - node.add_leaf_chunks(str(real_path), [chunk]) diff --git a/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py b/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py new file mode 100644 index 000000000..2fa86c690 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic import asset_tools +from shared.services.retrieval.agentic.types import DocTreeNode +from shared.services.retrieval.connected_hydration import hydrate_connected_target_rows +from shared.services.retrieval.path_hydration import hydrate_paths_to_rows + + +async def hydrate_path_selections_into_node( + db: AsyncSession, + *, + node: DocTreeNode, + path_selections: list[dict[str, Any]], + user_id: str, + namespace: str, + document_id: str, + job_result_id: str | None = None, +) -> None: + chunks = await hydrate_paths_to_rows( + db, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + if not chunks: + return + + chunks = await _append_connected_asset_targets(db, chunks) + resolved_job_result_id = job_result_id or _find_job_result_id(chunks) + if resolved_job_result_id: + await _attach_root_asset_owners( + db, + document_id=document_id, + job_result_id=resolved_job_result_id, + chunks=chunks, + ) + + add_chunks_to_node(node, chunks) + + +async def _append_connected_asset_targets( + db: AsyncSession, chunks: list[dict[str, Any]] +) -> list[dict[str, Any]]: + connected = await hydrate_connected_target_rows( + db=db, + rows=chunks, + exclude_document_ids=[], + exclude_sections=[], + ) + if not connected: + return chunks + + owner_map = asset_tools.build_connected_owner_map(chunks) + for chunk in connected: + if not chunk.get("owner_section_path"): + chunk["owner_section_path"] = owner_map.get(str(chunk.get("chunk_id") or "")) + return [*chunks, *connected] + + +async def _attach_root_asset_owners( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + chunks: list[dict[str, Any]], +) -> None: + root_map = await asset_tools.resolve_root_asset_owners( + db, + document_id=document_id, + job_result_id=job_result_id, + chunks=chunks, + ) + if not root_map: + return + + for chunk in chunks: + if chunk.get("owner_section_path"): + continue + chunk_id = str(chunk.get("chunk_id") or "") + if chunk_id in root_map: + chunk["owner_section_path"] = root_map[chunk_id] + + +def _find_job_result_id(chunks: list[dict[str, Any]]) -> str | None: + return next( + (str(chunk["job_result_id"]) for chunk in chunks if chunk.get("job_result_id")), + None, + ) + + +def add_chunks_to_node(node: DocTreeNode, chunks: list[dict[str, Any]]) -> None: + for chunk in chunks: + real_path = ( + chunk.get("owner_section_path") + or chunk.get("section_path") + or chunk.get("source_chunk_path") + ) + if real_path: + node.add_leaf_chunks(str(real_path), [chunk]) From 0db18742f5cefac6d23c43ee1b4258fc2008a8b8 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 17:28:45 +0800 Subject: [PATCH 28/40] refactor extract legacy retrieval route --- .../tests/contract/test_retrieval_contract.py | 6 +- .../services/retrieval/execution_plan.py | 6 +- .../services/retrieval/execution_routes.py | 365 +----------------- .../shared/services/retrieval/legacy_route.py | 347 +++++++++++++++++ .../shared/services/retrieval/route_types.py | 35 ++ 5 files changed, 393 insertions(+), 366 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/legacy_route.py create mode 100644 packages/shared-python/shared/services/retrieval/route_types.py diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 3ccf8b5d0..cb99fd689 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -405,15 +405,15 @@ async def fake_graph_routing(*_args: object, **_kwargs: object) -> list[dict[str return [] monkeypatch.setattr( - "shared.services.retrieval.execution_routes.path_channel", + "shared.services.retrieval.legacy_route.path_channel", fake_path_channel, ) monkeypatch.setattr( - "shared.services.retrieval.execution_routes.content_channel", + "shared.services.retrieval.legacy_route.content_channel", fake_content_channel, ) monkeypatch.setattr( - "shared.services.retrieval.execution_routes.list_graph_routed_chunks", + "shared.services.retrieval.legacy_route.list_graph_routed_chunks", fake_graph_routing, ) diff --git a/packages/shared-python/shared/services/retrieval/execution_plan.py b/packages/shared-python/shared/services/retrieval/execution_plan.py index 96e0e15c7..71669f3ed 100644 --- a/packages/shared-python/shared/services/retrieval/execution_plan.py +++ b/packages/shared-python/shared/services/retrieval/execution_plan.py @@ -10,10 +10,8 @@ get_cached_retrieval_query_result, set_cached_retrieval_query_result, ) -from shared.services.retrieval.execution_routes import ( - RetrievalRouteContext, - run_retrieval_route, -) +from shared.services.retrieval.execution_routes import run_retrieval_route +from shared.services.retrieval.route_types import RetrievalRouteContext from shared.services.retrieval.hit_stats_recorder import ( schedule_retrieval_hit_stats_update, ) diff --git a/packages/shared-python/shared/services/retrieval/execution_routes.py b/packages/shared-python/shared/services/retrieval/execution_routes.py index 23d16261f..ee5516159 100644 --- a/packages/shared-python/shared/services/retrieval/execution_routes.py +++ b/packages/shared-python/shared/services/retrieval/execution_routes.py @@ -1,97 +1,24 @@ from __future__ import annotations import os -import time -from dataclasses import dataclass -from typing import Any from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.channels import content_channel, path_channel, term_channel -from shared.services.retrieval.graph_query_service import GraphQueryService +from shared.services.retrieval.legacy_route import run_legacy_retrieval_route from shared.services.retrieval.reference_hydration import hydrate_referenced_chunk_rows from shared.services.retrieval.result_assembly import assemble_retrieval_results -from shared.services.retrieval.ranking import rank_retrieval_candidates from shared.services.retrieval.response_projection import ( attach_citation, enrich_referenced_chunks_with_asset_urls, ) +from shared.services.retrieval.route_types import ( + RetrievalRouteContext, + RetrievalRouteOutcome, +) from shared.services.retrieval.scoped_corpus import ( count_scoped_chunks, load_all_scoped_chunks, ) -from shared.services.retrieval.scoring import ( - get_row_path, - merge_channels_rrf, - merge_same_section_rows, - normalize_row_scores, -) -from shared.services.retrieval.settings import ( - CHANNEL_WEIGHT_CONTENT, - CHANNEL_WEIGHT_PATH, - CHANNEL_WEIGHT_TERM, - INTERNAL_RECALL_K_MULTIPLIER, -) - - -@dataclass(frozen=True) -class RetrievalRouteContext: - db: AsyncSession - user_id: str - namespace: str - query: str - top_k: int - exclude_document_ids: list[str] - exclude_sections: list[dict[str, str]] - allowed_chunk_types: set[str] | None - data_type: int - signal_paths: list[str] | None - filter_mode: str - channels: list[str] | None - channel_weights: dict[str, float] | None - threshold: float - effective_recall_k: int - use_agentic: bool | None - - -@dataclass(frozen=True) -class RetrievalRouteOutcome: - response: dict[str, Any] - hit_stats_results: list[dict[str, Any]] - completion_label: str - completion_count: int - completion_detail: str - - -async def list_graph_routed_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - service = GraphQueryService() - entry_document_ids = await service.find_entry_documents( - db, - user_id=user_id, - namespace=namespace, - query=query, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - return await service.collect_candidate_chunks( - db, - user_id=user_id, - namespace=namespace, - entry_document_ids=entry_document_ids, - query=query, - top_k=top_k * INTERNAL_RECALL_K_MULTIPLIER, - exclude_sections=exclude_sections, - ) async def run_retrieval_route( @@ -104,7 +31,7 @@ async def run_retrieval_route( if _should_use_agentic_route(context.use_agentic): return await _run_agentic_route(context) - return await _run_legacy_route(context) + return await run_legacy_retrieval_route(context) async def _try_run_small_kb_route( @@ -243,283 +170,3 @@ async def _run_agentic_route( completion_count=len(enriched_refs), completion_detail=completion_detail, ) - - -async def _run_legacy_route( - context: RetrievalRouteContext, -) -> RetrievalRouteOutcome: - active_channels = set(context.channels) if context.channels else { - "path", - "content", - "term", - } - logger.info( - f"\n PHASE 1: Bottom-Layer Discovery " - f"(channels={sorted(active_channels)})" - ) - logger.info(f" effective_recall_k={context.effective_recall_k}") - - path_rows = await _load_path_rows(context, active_channels) - content_rows = await _load_content_rows(context, active_channels) - term_rows = await _load_term_rows(context, active_channels) - - fused_rows = _fuse_legacy_rows( - context=context, - path_rows=path_rows, - content_rows=content_rows, - term_rows=term_rows, - ) - router_used = "discovery_only" - agent_rows: list[dict[str, Any]] = [] - - logger.info("\n PHASE 2: Legacy Graph Routing") - try: - agent_rows = await list_graph_routed_chunks( - context.db, - user_id=context.user_id, - namespace=context.namespace, - query=context.query, - top_k=context.top_k, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - ) - if agent_rows: - router_used = "discovery+graph" - logger.info(f" Graph routing: {len(agent_rows)} rows") - except Exception as exc: - logger.error(f" Graph routing failed (ignored): {exc}") - agent_rows = [] - - if agent_rows: - normalize_row_scores( - agent_rows, - source_field="score", - target_field="agent_score", - default=0.5, - ) - - ranked_rows = await rank_retrieval_candidates( - context.db, - user_id=context.user_id, - namespace=context.namespace, - discovery_rows=fused_rows, - routed_rows=agent_rows, - top_k=context.top_k, - ) - if ranked_rows: - logger.info(f"\n Unified candidate ranking: {len(ranked_rows)} rows") - for index, row in enumerate(ranked_rows[:10]): - logger.info( - " " - f"[{index}] evidence={row.get('evidence_score', 0.0):.4f} " - f"discovery={row.get('discovery_score', 0.0):.4f} " - f"agent={row.get('agent_score', 0.0):.4f} " - f"path={get_row_path(row)}" - ) - - assembled_rows = await assemble_retrieval_results( - db=context.db, - rows=ranked_rows, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - allowed_chunk_types=context.allowed_chunk_types, - ) - results = [attach_citation(row) for row in assembled_rows] - response = { - "namespace": context.namespace, - "query": context.query, - "router_used": router_used, - "results": results, - } - return RetrievalRouteOutcome( - response=response, - hit_stats_results=results, - completion_label="RETRIEVAL", - completion_count=len(results), - completion_detail=f"results | router={router_used}", - ) - - -async def _load_path_rows( - context: RetrievalRouteContext, - active_channels: set[str], -) -> list[dict[str, Any]]: - if "path" not in active_channels: - return [] - - start_time = time.monotonic() - rows = await path_channel( - context.db, - user_id=context.user_id, - namespace=context.namespace, - query=context.query, - top_k=context.effective_recall_k, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - allowed_chunk_types=context.allowed_chunk_types, - signal_paths=context.signal_paths, - filter_mode=context.filter_mode, - ) - elapsed_ms = round((time.monotonic() - start_time) * 1000) - logger.info(f"\n path_channel: {len(rows)} rows in {elapsed_ms}ms") - for index, row in enumerate(rows[:5]): - logger.info( - f" [{index}] score={row.get('score', 0):.4f} " - f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " - f"type={row.get('chunk_type', '?')}" - ) - if len(rows) > 5: - logger.info(f" ... and {len(rows) - 5} more") - return rows - - -async def _load_content_rows( - context: RetrievalRouteContext, - active_channels: set[str], -) -> list[dict[str, Any]]: - if "content" not in active_channels: - return [] - - start_time = time.monotonic() - rows = await content_channel( - context.db, - user_id=context.user_id, - namespace=context.namespace, - query=context.query, - top_k=context.effective_recall_k, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - allowed_chunk_types=context.allowed_chunk_types, - signal_paths=context.signal_paths, - filter_mode=context.filter_mode, - ) - elapsed_ms = round((time.monotonic() - start_time) * 1000) - logger.info(f"\n content_channel: {len(rows)} rows in {elapsed_ms}ms") - for index, row in enumerate(rows[:5]): - logger.info( - f" [{index}] score={row.get('score', 0):.4f} " - f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " - f"content={str(row.get('content', ''))[:80]}" - ) - if len(rows) > 5: - logger.info(f" ... and {len(rows) - 5} more") - return rows - - -async def _load_term_rows( - context: RetrievalRouteContext, - active_channels: set[str], -) -> list[dict[str, Any]]: - if "term" not in active_channels: - return [] - - start_time = time.monotonic() - rows = await term_channel( - context.db, - user_id=context.user_id, - namespace=context.namespace, - query=context.query, - top_k=context.effective_recall_k, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - allowed_chunk_types=context.allowed_chunk_types, - signal_paths=context.signal_paths, - filter_mode=context.filter_mode, - ) - elapsed_ms = round((time.monotonic() - start_time) * 1000) - logger.info(f"\n term_channel: {len(rows)} rows in {elapsed_ms}ms") - for index, row in enumerate(rows[:5]): - logger.info( - f" [{index}] score={row.get('score', 0):.4f} " - f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " - f"type={row.get('chunk_type', '?')}" - ) - if len(rows) > 5: - logger.info(f" ... and {len(rows) - 5} more") - return rows - - -def _fuse_legacy_rows( - *, - context: RetrievalRouteContext, - path_rows: list[dict[str, Any]], - content_rows: list[dict[str, Any]], - term_rows: list[dict[str, Any]], -) -> list[dict[str, Any]]: - default_weights = { - "path": CHANNEL_WEIGHT_PATH, - "content": CHANNEL_WEIGHT_CONTENT, - "term": CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(context.channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) - - if channel_lists: - fused_rows = merge_channels_rrf( - channel_lists, - weight_list, - context.effective_recall_k, - ) - else: - fused_rows = [] - logger.info( - f"\n RRF Fusion: {len(fused_rows)} rows from " - f"{len(channel_lists)} channels " - f"(weights={dict(zip(['path', 'content', 'term'][:len(weight_list)], weight_list))})" - ) - for index, row in enumerate(fused_rows[:5]): - logger.info( - f" [{index}] rrf_score={row.get('score', 0):.4f} " - f"path={row.get('section_path', '') or row.get('source_chunk_path', '')}" - ) - if len(fused_rows) > 5: - logger.info(f" ... and {len(fused_rows) - 5} more") - - pre_merge = len(fused_rows) - fused_rows = merge_same_section_rows(fused_rows) - if len(fused_rows) != pre_merge: - logger.info(f"retrieval: section_merge={pre_merge}->{len(fused_rows)}") - - if context.channel_weights is not None: - logger.debug(f"retrieval: channel_weights={context.channel_weights}") - - fused_rows = _filter_rows_by_threshold(fused_rows, context) - if fused_rows: - normalize_row_scores( - fused_rows, - source_field="score", - target_field="discovery_score", - default=0.5, - ) - - return fused_rows - - -def _filter_rows_by_threshold( - rows: list[dict[str, Any]], - context: RetrievalRouteContext, -) -> list[dict[str, Any]]: - if context.threshold <= 0.0 or not rows: - return rows - - pre_count = len(rows) - filtered_rows = [ - row for row in rows if row.get("score", 0.0) >= context.threshold - ] - logger.info( - f"retrieval: threshold_filter={pre_count}->{len(filtered_rows)} " - f"(threshold={context.threshold})" - ) - return filtered_rows diff --git a/packages/shared-python/shared/services/retrieval/legacy_route.py b/packages/shared-python/shared/services/retrieval/legacy_route.py new file mode 100644 index 000000000..e8642840b --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/legacy_route.py @@ -0,0 +1,347 @@ +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.channels import content_channel, path_channel, term_channel +from shared.services.retrieval.graph_query_service import GraphQueryService +from shared.services.retrieval.ranking import rank_retrieval_candidates +from shared.services.retrieval.response_projection import attach_citation +from shared.services.retrieval.result_assembly import assemble_retrieval_results +from shared.services.retrieval.route_types import ( + RetrievalRouteContext, + RetrievalRouteOutcome, +) +from shared.services.retrieval.scoring import ( + get_row_path, + merge_channels_rrf, + merge_same_section_rows, + normalize_row_scores, +) +from shared.services.retrieval.settings import ( + CHANNEL_WEIGHT_CONTENT, + CHANNEL_WEIGHT_PATH, + CHANNEL_WEIGHT_TERM, + INTERNAL_RECALL_K_MULTIPLIER, +) + + +async def run_legacy_retrieval_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome: + active_channels = set(context.channels) if context.channels else { + "path", + "content", + "term", + } + logger.info( + f"\n PHASE 1: Bottom-Layer Discovery " + f"(channels={sorted(active_channels)})" + ) + logger.info(f" effective_recall_k={context.effective_recall_k}") + + path_rows = await _load_path_rows(context, active_channels) + content_rows = await _load_content_rows(context, active_channels) + term_rows = await _load_term_rows(context, active_channels) + + fused_rows = _fuse_legacy_rows( + context=context, + path_rows=path_rows, + content_rows=content_rows, + term_rows=term_rows, + ) + router_used, graph_rows = await _run_legacy_graph_routing(context) + + ranked_rows = await rank_retrieval_candidates( + context.db, + user_id=context.user_id, + namespace=context.namespace, + discovery_rows=fused_rows, + routed_rows=graph_rows, + top_k=context.top_k, + ) + _log_ranked_rows(ranked_rows) + + assembled_rows = await assemble_retrieval_results( + db=context.db, + rows=ranked_rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) + results = [attach_citation(row) for row in assembled_rows] + response = { + "namespace": context.namespace, + "query": context.query, + "router_used": router_used, + "results": results, + } + return RetrievalRouteOutcome( + response=response, + hit_stats_results=results, + completion_label="RETRIEVAL", + completion_count=len(results), + completion_detail=f"results | router={router_used}", + ) + + +async def list_graph_routed_chunks( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], +) -> list[dict[str, Any]]: + service = GraphQueryService() + entry_document_ids = await service.find_entry_documents( + db, + user_id=user_id, + namespace=namespace, + query=query, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + return await service.collect_candidate_chunks( + db, + user_id=user_id, + namespace=namespace, + entry_document_ids=entry_document_ids, + query=query, + top_k=top_k * INTERNAL_RECALL_K_MULTIPLIER, + exclude_sections=exclude_sections, + ) + + +async def _run_legacy_graph_routing( + context: RetrievalRouteContext, +) -> tuple[str, list[dict[str, Any]]]: + logger.info("\n PHASE 2: Legacy Graph Routing") + try: + graph_rows = await list_graph_routed_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.top_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + ) + if graph_rows: + logger.info(f" Graph routing: {len(graph_rows)} rows") + normalize_row_scores( + graph_rows, + source_field="score", + target_field="agent_score", + default=0.5, + ) + return "discovery+graph", graph_rows + except Exception as exc: + logger.error(f" Graph routing failed (ignored): {exc}") + + return "discovery_only", [] + + +async def _load_path_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "path" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await path_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n path_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"type={row.get('chunk_type', '?')}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +async def _load_content_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "content" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await content_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n content_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"content={str(row.get('content', ''))[:80]}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +async def _load_term_rows( + context: RetrievalRouteContext, + active_channels: set[str], +) -> list[dict[str, Any]]: + if "term" not in active_channels: + return [] + + start_time = time.monotonic() + rows = await term_channel( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + ) + elapsed_ms = round((time.monotonic() - start_time) * 1000) + logger.info(f"\n term_channel: {len(rows)} rows in {elapsed_ms}ms") + for index, row in enumerate(rows[:5]): + logger.info( + f" [{index}] score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} " + f"type={row.get('chunk_type', '?')}" + ) + if len(rows) > 5: + logger.info(f" ... and {len(rows) - 5} more") + return rows + + +def _fuse_legacy_rows( + *, + context: RetrievalRouteContext, + path_rows: list[dict[str, Any]], + content_rows: list[dict[str, Any]], + term_rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + default_weights = { + "path": CHANNEL_WEIGHT_PATH, + "content": CHANNEL_WEIGHT_CONTENT, + "term": CHANNEL_WEIGHT_TERM, + } + effective_weights = {**default_weights, **(context.channel_weights or {})} + + channel_lists: list[list[dict[str, Any]]] = [] + weight_list: list[float] = [] + + if path_rows: + channel_lists.append(path_rows) + weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) + if content_rows: + channel_lists.append(content_rows) + weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) + if term_rows: + channel_lists.append(term_rows) + weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) + + if channel_lists: + fused_rows = merge_channels_rrf( + channel_lists, + weight_list, + context.effective_recall_k, + ) + else: + fused_rows = [] + logger.info( + f"\n RRF Fusion: {len(fused_rows)} rows from " + f"{len(channel_lists)} channels " + f"(weights={dict(zip(['path', 'content', 'term'][:len(weight_list)], weight_list))})" + ) + for index, row in enumerate(fused_rows[:5]): + logger.info( + f" [{index}] rrf_score={row.get('score', 0):.4f} " + f"path={row.get('section_path', '') or row.get('source_chunk_path', '')}" + ) + if len(fused_rows) > 5: + logger.info(f" ... and {len(fused_rows) - 5} more") + + pre_merge = len(fused_rows) + fused_rows = merge_same_section_rows(fused_rows) + if len(fused_rows) != pre_merge: + logger.info(f"retrieval: section_merge={pre_merge}->{len(fused_rows)}") + + if context.channel_weights is not None: + logger.debug(f"retrieval: channel_weights={context.channel_weights}") + + fused_rows = _filter_rows_by_threshold(fused_rows, context) + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + + return fused_rows + + +def _filter_rows_by_threshold( + rows: list[dict[str, Any]], + context: RetrievalRouteContext, +) -> list[dict[str, Any]]: + if context.threshold <= 0.0 or not rows: + return rows + + pre_count = len(rows) + filtered_rows = [ + row for row in rows if row.get("score", 0.0) >= context.threshold + ] + logger.info( + f"retrieval: threshold_filter={pre_count}->{len(filtered_rows)} " + f"(threshold={context.threshold})" + ) + return filtered_rows + + +def _log_ranked_rows(ranked_rows: list[dict[str, Any]]) -> None: + if not ranked_rows: + return + + logger.info(f"\n Unified candidate ranking: {len(ranked_rows)} rows") + for index, row in enumerate(ranked_rows[:10]): + logger.info( + " " + f"[{index}] evidence={row.get('evidence_score', 0.0):.4f} " + f"discovery={row.get('discovery_score', 0.0):.4f} " + f"agent={row.get('agent_score', 0.0):.4f} " + f"path={get_row_path(row)}" + ) diff --git a/packages/shared-python/shared/services/retrieval/route_types.py b/packages/shared-python/shared/services/retrieval/route_types.py new file mode 100644 index 000000000..f1aa94160 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/route_types.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + + +@dataclass(frozen=True) +class RetrievalRouteContext: + db: AsyncSession + user_id: str + namespace: str + query: str + top_k: int + exclude_document_ids: list[str] + exclude_sections: list[dict[str, str]] + allowed_chunk_types: set[str] | None + data_type: int + signal_paths: list[str] | None + filter_mode: str + channels: list[str] | None + channel_weights: dict[str, float] | None + threshold: float + effective_recall_k: int + use_agentic: bool | None + + +@dataclass(frozen=True) +class RetrievalRouteOutcome: + response: dict[str, Any] + hit_stats_results: list[dict[str, Any]] + completion_label: str + completion_count: int + completion_detail: str From 4517c4358b7418ba800200fe23aed49db1215810 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:11:08 +0800 Subject: [PATCH 29/40] refactor deepen qstash webhook publisher --- .../test_webhook_recovery_contract.py | 24 +- .../shared/services/webhook/qstash_client.py | 145 ++++++++++ .../shared/services/webhook/qstash_payload.py | 40 +++ .../services/webhook/qstash_publisher.py | 267 ++---------------- .../webhook/qstash_secret_resolver.py | 77 +++++ 5 files changed, 305 insertions(+), 248 deletions(-) create mode 100644 packages/shared-python/shared/services/webhook/qstash_client.py create mode 100644 packages/shared-python/shared/services/webhook/qstash_payload.py create mode 100644 packages/shared-python/shared/services/webhook/qstash_secret_resolver.py diff --git a/apps/worker/tests/contract/test_webhook_recovery_contract.py b/apps/worker/tests/contract/test_webhook_recovery_contract.py index 569946a14..606e25818 100644 --- a/apps/worker/tests/contract/test_webhook_recovery_contract.py +++ b/apps/worker/tests/contract/test_webhook_recovery_contract.py @@ -179,8 +179,8 @@ def publish(self, **kwargs: Any) -> SimpleNamespace: ), ) monkeypatch.setattr( - publisher, - "_get_client", + publisher._client_adapter, + "get_client", lambda: SimpleNamespace(message=FakeMessageClient()), ) @@ -323,6 +323,8 @@ def test_should_publish_completed_webhook_with_result_delivery_payload( monkeypatch: MonkeyPatch, ) -> None: _, qstash_publisher, engine = _load_worker_modules() + from shared.core.config import app_config + from shared.services.jobs.result_delivery import JobResultDeliveryResolver from shared.services.storage.job_file_storage import JobFileStorage user_id = f"worker-user-{uuid4().hex[:12]}" @@ -358,7 +360,6 @@ def generate_presigned_url( ) return f"signed://{bucket}/{key}?expires={expiration}" - publisher = qstash_publisher.QStashWebhookPublisher() monkeypatch.setattr( qstash_publisher, "validate_http_url_and_resolve_ip", @@ -370,12 +371,7 @@ def generate_presigned_url( ), ) monkeypatch.setattr( - publisher, - "_get_client", - lambda: SimpleNamespace(message=FakeMessageClient()), - ) - monkeypatch.setattr( - qstash_publisher.JobResultDeliveryResolver, + JobResultDeliveryResolver, "__init__", lambda self: setattr( self, @@ -383,6 +379,12 @@ def generate_presigned_url( JobFileStorage(storage_adapter=FakeStorageAdapter()), ), ) + publisher = qstash_publisher.QStashWebhookPublisher() + monkeypatch.setattr( + publisher._client_adapter, + "get_client", + lambda: SimpleNamespace(message=FakeMessageClient()), + ) now = _utc_now() with engine.begin() as connection: @@ -424,7 +426,7 @@ def generate_presigned_url( { "key": result_s3_key, "expiration": 3600, - "bucket": qstash_publisher.app_config.S3_RESULTS_BUCKET, + "bucket": app_config.S3_RESULTS_BUCKET, "method": "GET", "headers": None, } @@ -435,7 +437,7 @@ def generate_presigned_url( assert published_payload["job_id"] == job_id assert published_payload["result"] == {"checksum": "contract-checksum"} assert published_payload["result_url"] == ( - f"signed://{qstash_publisher.app_config.S3_RESULTS_BUCKET}/{result_s3_key}" + f"signed://{app_config.S3_RESULTS_BUCKET}/{result_s3_key}" "?expires=3600" ) diff --git a/packages/shared-python/shared/services/webhook/qstash_client.py b/packages/shared-python/shared/services/webhook/qstash_client.py new file mode 100644 index 000000000..78547d04a --- /dev/null +++ b/packages/shared-python/shared/services/webhook/qstash_client.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Optional + +from loguru import logger + +from shared.core.config import app_config +from shared.core.exceptions.domain_exceptions import QStashServiceException +from shared.models.database.webhook import WebhookEventStatus + + +@dataclass(frozen=True) +class QStashDeliveryStatus: + """Terminal delivery status observed from QStash logs.""" + + status: str + response_status_code: Optional[int] + response_body: Optional[str] + error_message: Optional[str] + + +class QStashClientAdapter: + """Upstash QStash client adapter for webhook publication and log lookup.""" + + def __init__(self) -> None: + self._client: Any = None + + def get_client(self) -> Any: + """Lazily initialize the QStash client.""" + if self._client is None: + try: + from qstash import QStash + except ImportError as exc: + raise QStashServiceException( + internal_message=( + "qstash package is required for QStash webhook delivery. " + "Install it with: pip install qstash" + ), + operation="initialize_client", + original_exception=exc, + ) from exc + + token = app_config.QSTASH_TOKEN + if not token: + raise QStashServiceException( + internal_message="QSTASH_TOKEN is not configured", + operation="initialize_client", + ) + + self._client = QStash(token, base_url=app_config.QSTASH_BASE_URL) + return self._client + + def publish_webhook( + self, + *, + target_url: str, + payload: dict[str, Any], + signature: str, + event_id: str, + ) -> Optional[str]: + """Call the QStash publish API.""" + headers = { + "Content-Type": "application/json", + "X-Knowhere-Signature": signature, + "X-Knowhere-Event-ID": event_id, + "User-Agent": "Knowhere-Webhook/1.0", + } + + callback_url = app_config.qstash_callback_url + failure_callback_url = app_config.qstash_failure_callback_url + if not callback_url or not failure_callback_url: + raise QStashServiceException( + internal_message=( + "QSTASH_CALLBACK_BASE_URL must be configured for QStash " + "webhook delivery" + ), + operation="publish_webhook", + ) + + publish_kwargs: dict[str, Any] = { + "url": target_url, + "body": json.dumps(payload, separators=(",", ":")), + "headers": headers, + "retries": app_config.QSTASH_MAX_RETRIES, + "content_type": "application/json", + "retry_delay": _get_retry_delay_expression(), + "callback": callback_url, + "failure_callback": failure_callback_url, + "deduplication_id": event_id, + "label": "knowhere-webhook", + } + + response = self.get_client().message.publish(**publish_kwargs) + + message_id = getattr(response, "message_id", None) + if message_id is None and isinstance(response, dict): + message_id = response.get("messageId") or response.get("message_id") + + return message_id + + def get_terminal_delivery_status( + self, + qstash_message_id: str, + ) -> Optional[QStashDeliveryStatus]: + """Read QStash logs for a terminal destination delivery state.""" + try: + from qstash.log import LogState + + response = self.get_client().log.list( + filter={"message_id": qstash_message_id}, + count=20, + ) + except Exception as exc: + logger.warning( + f"QStash delivery status lookup failed: " + f"message_id={qstash_message_id}, error={exc}" + ) + return None + + terminal_logs = sorted(response.logs, key=lambda log: log.time, reverse=True) + for log in terminal_logs: + if log.state == LogState.DELIVERED: + return QStashDeliveryStatus( + status=WebhookEventStatus.DELIVERED, + response_status_code=log.response_status, + response_body=log.response_body, + error_message=log.error, + ) + + if log.state == LogState.FAILED: + return QStashDeliveryStatus( + status=WebhookEventStatus.FAILED, + response_status_code=log.response_status, + response_body=log.response_body, + error_message=log.error, + ) + + return None + + +def _get_retry_delay_expression() -> str: + # Approximate exponential backoff: 1m, 10m, ~100m, ~100m, ~100m. + return "pow(10, min(retried, 2)) * 60000" diff --git a/packages/shared-python/shared/services/webhook/qstash_payload.py b/packages/shared-python/shared/services/webhook/qstash_payload.py new file mode 100644 index 000000000..e454c076b --- /dev/null +++ b/packages/shared-python/shared/services/webhook/qstash_payload.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger +from sqlalchemy import select +from sqlalchemy.orm import selectinload + +from shared.models.database.job import Job +from shared.services.jobs.result_delivery import JobResultDeliveryResolver + + +class QStashPayloadEnricher: + """Sync payload enricher for QStash webhook publication.""" + + def __init__(self, resolver: JobResultDeliveryResolver | None = None) -> None: + self._resolver = resolver or JobResultDeliveryResolver() + + def enrich(self, db: Any, event: Any) -> dict[str, Any]: + payload = dict(event.payload) + if payload.get("event") != "job.completed": + return payload + + try: + result = db.execute( + select(Job) + .options(selectinload(Job.job_result)) + .where(Job.job_id == event.job_id) + ) + job = result.scalar_one_or_none() + if not job or not job.job_result: + return payload + + return self._resolver.enrich_payload( + payload, + job_result=job.job_result, + ) + except Exception as exc: + logger.error(f"Failed to enrich payload for event {event.id}: {exc}") + return payload diff --git a/packages/shared-python/shared/services/webhook/qstash_publisher.py b/packages/shared-python/shared/services/webhook/qstash_publisher.py index 2c7dbf33e..7a0580cb6 100644 --- a/packages/shared-python/shared/services/webhook/qstash_publisher.py +++ b/packages/shared-python/shared/services/webhook/qstash_publisher.py @@ -11,62 +11,39 @@ from __future__ import annotations -import json -from dataclasses import dataclass -from typing import Any, Dict, Optional +from typing import Optional from loguru import logger - -from shared.core.config import app_config -from shared.core.exceptions.domain_exceptions import QStashServiceException -from shared.models.database.webhook import WebhookEventStatus -from shared.services.jobs.result_delivery import JobResultDeliveryResolver +from sqlalchemy import select + +from shared.core.database_sync import get_sync_db_context +from shared.models.database.job import Job +from shared.models.database.webhook import WebhookEvent, WebhookEventStatus +from shared.services.webhook.qstash_client import ( + QStashClientAdapter, + QStashDeliveryStatus, +) +from shared.services.webhook.qstash_payload import QStashPayloadEnricher +from shared.services.webhook.qstash_secret_resolver import QStashSecretResolver from shared.services.webhook.signing import sign_webhook_payload from shared.utils.url_security import ( validate_http_url_and_resolve_ip, ) -@dataclass(frozen=True) -class QStashDeliveryStatus: - """Terminal delivery status observed from QStash logs.""" - - status: str - response_status_code: Optional[int] - response_body: Optional[str] - error_message: Optional[str] - - class QStashWebhookPublisher: """Publishes webhook events to customer endpoints via QStash.""" - def __init__(self) -> None: - self._client: Any = None - - def _get_client(self) -> Any: - """Lazily initialize the QStash client.""" - if self._client is None: - try: - from qstash import QStash - except ImportError as exc: - raise QStashServiceException( - internal_message=( - "qstash package is required for QStash webhook delivery. " - "Install it with: pip install qstash" - ), - operation="initialize_client", - original_exception=exc, - ) from exc - - token = app_config.QSTASH_TOKEN - if not token: - raise QStashServiceException( - internal_message="QSTASH_TOKEN is not configured", - operation="initialize_client", - ) - - self._client = QStash(token, base_url=app_config.QSTASH_BASE_URL) - return self._client + def __init__( + self, + *, + client_adapter: QStashClientAdapter | None = None, + payload_enricher: QStashPayloadEnricher | None = None, + secret_resolver: QStashSecretResolver | None = None, + ) -> None: + self._client_adapter = client_adapter or QStashClientAdapter() + self._payload_enricher = payload_enricher or QStashPayloadEnricher() + self._secret_resolver = secret_resolver or QStashSecretResolver() def publish_event(self, event_id: str) -> Optional[str]: """Publish a webhook event via QStash. @@ -76,12 +53,6 @@ def publish_event(self, event_id: str) -> Optional[str]: Returns the QStash message_id on success, or None on failure. """ - from sqlalchemy import select - - from shared.core.database_sync import get_sync_db_context - from shared.models.database.job import Job - from shared.models.database.webhook import WebhookEvent - with get_sync_db_context() as db: event = db.execute( select(WebhookEvent).where(WebhookEvent.id == event_id) @@ -108,10 +79,8 @@ def publish_event(self, event_id: str) -> Optional[str]: db.commit() return None - # Enrich payload (presigned S3 URL for completed jobs) - payload = self._enrich_payload(db, event) + payload = self._payload_enricher.enrich(db, event) - # Resolve signing secret user_id = db.execute( select(Job.user_id).where(Job.job_id == event.job_id) ).scalar_one_or_none() @@ -122,7 +91,11 @@ def publish_event(self, event_id: str) -> Optional[str]: db.commit() return None - secret = self._resolve_secret(db, str(user_id), event.target_url) + secret = self._secret_resolver.resolve( + db, + user_id=str(user_id), + endpoint=event.target_url, + ) if not secret: logger.error( f"QStash publish: secret resolution failed for event {event_id}" @@ -131,12 +104,10 @@ def publish_event(self, event_id: str) -> Optional[str]: db.commit() return None - # Sign payload with our HMAC signature = sign_webhook_payload(payload, secret) - # Publish to QStash try: - message_id = self._publish_to_qstash( + message_id = self._client_adapter.publish_webhook( target_url=event.target_url, payload=payload, signature=signature, @@ -158,191 +129,13 @@ def publish_event(self, event_id: str) -> Optional[str]: ) return message_id - def _publish_to_qstash( - self, - target_url: str, - payload: Dict[str, Any], - signature: str, - event_id: str, - ) -> Optional[str]: - """Call the QStash publish API.""" - headers = { - "Content-Type": "application/json", - "X-Knowhere-Signature": signature, - "X-Knowhere-Event-ID": event_id, - "User-Agent": "Knowhere-Webhook/1.0", - } - - # Approximate exponential backoff: 1m, 10m, ~100m, ~100m, ~100m - # pow(10, min(retried, 2)) * 60000 → 60s, 600s, 6000s capped - retry_delay_expression = "pow(10, min(retried, 2)) * 60000" - - callback_url = app_config.qstash_callback_url - failure_callback_url = app_config.qstash_failure_callback_url - if not callback_url or not failure_callback_url: - raise QStashServiceException( - internal_message=( - "QSTASH_CALLBACK_BASE_URL must be configured for QStash " - "webhook delivery" - ), - operation="publish_webhook", - ) - - client = self._get_client() - - publish_kwargs: Dict[str, Any] = { - "url": target_url, - "body": json.dumps(payload, separators=(",", ":")), - "headers": headers, - "retries": app_config.QSTASH_MAX_RETRIES, - "content_type": "application/json", - "retry_delay": retry_delay_expression, - "callback": callback_url, - "failure_callback": failure_callback_url, - "deduplication_id": event_id, - "label": "knowhere-webhook", - } - - response = client.message.publish(**publish_kwargs) - - message_id = getattr(response, "message_id", None) - if message_id is None and isinstance(response, dict): - message_id = response.get("messageId") or response.get("message_id") - - return message_id - def get_terminal_delivery_status( self, qstash_message_id: str, ) -> Optional[QStashDeliveryStatus]: """Read QStash logs for a terminal destination delivery state.""" - try: - from qstash.log import LogState - - response = self._get_client().log.list( - filter={"message_id": qstash_message_id}, - count=20, - ) - except Exception as exc: - logger.warning( - f"QStash delivery status lookup failed: " - f"message_id={qstash_message_id}, error={exc}" - ) - return None - - terminal_logs = sorted(response.logs, key=lambda log: log.time, reverse=True) - for log in terminal_logs: - if log.state == LogState.DELIVERED: - return QStashDeliveryStatus( - status=WebhookEventStatus.DELIVERED, - response_status_code=log.response_status, - response_body=log.response_body, - error_message=log.error, - ) - - if log.state == LogState.FAILED: - return QStashDeliveryStatus( - status=WebhookEventStatus.FAILED, - response_status_code=log.response_status, - response_body=log.response_body, - error_message=log.error, - ) - - return None - - def _enrich_payload(self, db: Any, event: Any) -> Dict[str, Any]: - """Enrich the webhook payload (e.g., generate fresh presigned S3 URL).""" - from sqlalchemy import select - from sqlalchemy.orm import selectinload - - from shared.models.database.job import Job - - payload = dict(event.payload) - if payload.get("event") != "job.completed": - return payload - - try: - result = db.execute( - select(Job) - .options(selectinload(Job.job_result)) - .where(Job.job_id == event.job_id) - ) - job = result.scalar_one_or_none() - if not job or not job.job_result: - return payload - - payload = JobResultDeliveryResolver().enrich_payload( - payload, - job_result=job.job_result, - ) - except Exception as exc: - logger.error(f"Failed to enrich payload for event {event.id}: {exc}") - - return payload - - def _resolve_secret(self, db: Any, user_id: str, endpoint: str) -> Optional[str]: - """Resolve the webhook signing secret for a user/endpoint.""" - from datetime import datetime, timezone - - from sqlalchemy import and_, select - - from shared.core.exceptions.domain_exceptions import ( - SystemSettingInvalidException, - SystemSettingMissingException, - ) - from shared.models.database.webhook_secret import ( - WebhookSecret, - WebhookSecretStatus, - ) - from shared.services.encryption import get_fernet_service - - try: - fernet = get_fernet_service() - except (SystemSettingMissingException, SystemSettingInvalidException) as exc: - logger.error(f"Configuration error during secret resolution: {exc}") - return None - - # Try endpoint-specific secret first, then global - secret_obj = None - if endpoint: - result = db.execute( - select(WebhookSecret).where( - and_( - WebhookSecret.user_id == user_id, - WebhookSecret.endpoint == endpoint, - WebhookSecret.status == WebhookSecretStatus.ACTIVE, - ) - ) - ) - secret_obj = result.scalar_one_or_none() - - if secret_obj is None: - result = db.execute( - select(WebhookSecret).where( - and_( - WebhookSecret.user_id == user_id, - WebhookSecret.endpoint.is_(None), - WebhookSecret.status == WebhookSecretStatus.ACTIVE, - ) - ) - ) - secret_obj = result.scalar_one_or_none() - - if secret_obj is None: - raw_secret = fernet.generate_webhook_secret() - secret_obj = WebhookSecret( - user_id=user_id, - endpoint=endpoint, - secret_encrypted=fernet.encrypt(raw_secret), - status=WebhookSecretStatus.ACTIVE, - ) - db.add(secret_obj) - db.commit() - db.refresh(secret_obj) + return self._client_adapter.get_terminal_delivery_status(qstash_message_id) - secret_obj.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None) - db.add(secret_obj) - return fernet.decrypt(secret_obj.secret_encrypted) # Module-level singleton _publisher: Optional[QStashWebhookPublisher] = None diff --git a/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py b/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py new file mode 100644 index 000000000..fb8174883 --- /dev/null +++ b/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any, Optional + +from loguru import logger +from sqlalchemy import and_, select + +from shared.core.exceptions.domain_exceptions import ( + SystemSettingInvalidException, + SystemSettingMissingException, +) +from shared.models.database.webhook_secret import ( + WebhookSecret, + WebhookSecretStatus, +) +from shared.services.encryption import get_fernet_service + + +class QStashSecretResolver: + """Sync webhook secret resolver for QStash publication.""" + + def resolve(self, db: Any, *, user_id: str, endpoint: str) -> Optional[str]: + try: + fernet = get_fernet_service() + except (SystemSettingMissingException, SystemSettingInvalidException) as exc: + logger.error(f"Configuration error during secret resolution: {exc}") + return None + + secret = self._find_active_secret(db, user_id=user_id, endpoint=endpoint) + if secret is None: + raw_secret = fernet.generate_webhook_secret() + secret = WebhookSecret( + user_id=user_id, + endpoint=endpoint, + secret_encrypted=fernet.encrypt(raw_secret), + status=WebhookSecretStatus.ACTIVE, + ) + db.add(secret) + db.commit() + db.refresh(secret) + + secret.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None) + db.add(secret) + return fernet.decrypt(secret.secret_encrypted) + + def _find_active_secret( + self, + db: Any, + *, + user_id: str, + endpoint: str, + ) -> WebhookSecret | None: + if endpoint: + result = db.execute( + select(WebhookSecret).where( + and_( + WebhookSecret.user_id == user_id, + WebhookSecret.endpoint == endpoint, + WebhookSecret.status == WebhookSecretStatus.ACTIVE, + ) + ) + ) + secret = result.scalar_one_or_none() + if secret is not None: + return secret + + result = db.execute( + select(WebhookSecret).where( + and_( + WebhookSecret.user_id == user_id, + WebhookSecret.endpoint.is_(None), + WebhookSecret.status == WebhookSecretStatus.ACTIVE, + ) + ) + ) + return result.scalar_one_or_none() From 5fe4961cc4f98870e878c31ecaefa96a484a1523 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:16:16 +0800 Subject: [PATCH 30/40] refactor extract retrieval lexical ranker --- .../shared/services/retrieval/channels.py | 62 ++-------------- .../services/retrieval/lexical_ranker.py | 74 +++++++++++++++++++ 2 files changed, 81 insertions(+), 55 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/lexical_ranker.py diff --git a/packages/shared-python/shared/services/retrieval/channels.py b/packages/shared-python/shared/services/retrieval/channels.py index ec1764c94..35966ed23 100644 --- a/packages/shared-python/shared/services/retrieval/channels.py +++ b/packages/shared-python/shared/services/retrieval/channels.py @@ -8,12 +8,14 @@ from typing import Any -from loguru import logger from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession +from shared.services.retrieval.lexical_ranker import ( + rank_rows_by_bm25, + tokenize_query_for_ranker, +) from shared.services.retrieval.section_filters import is_excluded_section -from shared.utils.text_utils import tokenize_for_retrieval _SCOPED_CORPUS_CTE = """ @@ -193,56 +195,6 @@ async def content_channel( ) -def _tokenize_query(query: str) -> list[str]: - return tokenize_for_retrieval(query, dedupe=True) - - -def _bm25_rerank( - rows: list[dict[str, Any]], - query_tokens: list[str], - *, - search_field: str, -) -> list[dict[str, Any]]: - """Rank matching rows with BM25 over pre-tokenized search text.""" - try: - from rank_bm25 import BM25Okapi - except ImportError: - logger.warning("rank_bm25 not installed, skipping BM25 re-rank") - ranked_rows: list[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = [token for token in str(row.get(search_field) or "").split() if token] - overlap = len(query_token_set.intersection(tokens)) - if overlap <= 0: - continue - row["score"] = float(overlap) - ranked_rows.append(row) - ranked_rows.sort(key=lambda r: r["score"], reverse=True) - return ranked_rows - - corpus: list[list[str]] = [] - ranked_rows: list[dict[str, Any]] = [] - query_token_set = set(query_tokens) - for row in rows: - tokens = [token for token in str(row.get(search_field) or "").split() if token] - if not tokens or not query_token_set.intersection(tokens): - continue - corpus.append(tokens) - ranked_rows.append(row) - - if not corpus or not query_tokens: - return [] - - bm25 = BM25Okapi(corpus) - scores = bm25.get_scores(query_tokens) - - for i, row in enumerate(ranked_rows): - row["score"] = float(scores[i]) - - ranked_rows.sort(key=lambda r: r["score"], reverse=True) - return ranked_rows - - async def _bm25_channel( db: AsyncSession, *, @@ -260,7 +212,7 @@ async def _bm25_channel( if search_field not in {"content_search_text", "path_search_text"}: raise ValueError(f"Unsupported search_field: {search_field}") - query_tokens = _tokenize_query(query) + query_tokens = tokenize_query_for_ranker(query) if not query_tokens: return [] @@ -287,7 +239,7 @@ async def _bm25_channel( rows = [_row_to_dict(r) for r in result.all()] rows = _filter_excluded_sections(rows, exclude_sections) - ranked_rows = _bm25_rerank(rows, query_tokens, search_field=search_field) + ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field) return ranked_rows[:top_k] @@ -312,7 +264,7 @@ async def term_channel( Note: top_k is already effective_recall_k from app_service. """ query_lower = query.lower().strip() - query_tokens = _tokenize_query(query) + query_tokens = tokenize_query_for_ranker(query) if not query_lower or not query_tokens: return [] diff --git a/packages/shared-python/shared/services/retrieval/lexical_ranker.py b/packages/shared-python/shared/services/retrieval/lexical_ranker.py new file mode 100644 index 000000000..8b50cd404 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/lexical_ranker.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from shared.utils.text_utils import tokenize_for_retrieval + + +def tokenize_query_for_ranker(query: str) -> list[str]: + return tokenize_for_retrieval(query, dedupe=True) + + +def rank_rows_by_bm25( + rows: list[dict[str, Any]], + query_tokens: list[str], + *, + search_field: str, +) -> list[dict[str, Any]]: + """Rank matching rows with BM25 over pre-tokenized search text.""" + try: + from rank_bm25 import BM25Okapi + except ImportError: + return _rank_rows_by_token_overlap( + rows, + query_tokens, + search_field=search_field, + ) + + corpus: list[list[str]] = [] + ranked_rows: list[dict[str, Any]] = [] + query_token_set = set(query_tokens) + for row in rows: + tokens = _get_search_tokens(row, search_field=search_field) + if not tokens or not query_token_set.intersection(tokens): + continue + corpus.append(tokens) + ranked_rows.append(row) + + if not corpus or not query_tokens: + return [] + + bm25 = BM25Okapi(corpus) + scores = bm25.get_scores(query_tokens) + + for index, row in enumerate(ranked_rows): + row["score"] = float(scores[index]) + + ranked_rows.sort(key=lambda row: row["score"], reverse=True) + return ranked_rows + + +def _rank_rows_by_token_overlap( + rows: list[dict[str, Any]], + query_tokens: list[str], + *, + search_field: str, +) -> list[dict[str, Any]]: + logger.warning("rank_bm25 not installed, skipping BM25 re-rank") + ranked_rows: list[dict[str, Any]] = [] + query_token_set = set(query_tokens) + for row in rows: + tokens = _get_search_tokens(row, search_field=search_field) + overlap = len(query_token_set.intersection(tokens)) + if overlap <= 0: + continue + row["score"] = float(overlap) + ranked_rows.append(row) + ranked_rows.sort(key=lambda row: row["score"], reverse=True) + return ranked_rows + + +def _get_search_tokens(row: dict[str, Any], *, search_field: str) -> list[str]: + return [token for token in str(row.get(search_field) or "").split() if token] From 42f4aec28d6328fb8042f93d6a6feb0741af17cb Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:21:31 +0800 Subject: [PATCH 31/40] refactor extract agentic section counts --- .../retrieval/agentic/section_counts.py | 174 +++++++++++++++++ .../retrieval/agentic/section_tree.py | 175 +----------------- 2 files changed, 178 insertions(+), 171 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/section_counts.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_counts.py b/packages/shared-python/shared/services/retrieval/agentic/section_counts.py new file mode 100644 index 000000000..74ba465d3 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/section_counts.py @@ -0,0 +1,174 @@ +"""Section count aggregation for agentic navigation.""" +from __future__ import annotations + +from sqlalchemy import case, func, literal_column, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import DocumentChunk + + +async def attach_section_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + all_sections: dict[str, dict], + items_by_path: dict[str, dict], +) -> None: + """Attach direct chunk and connected asset counts to visible section items.""" + scope_item_sids = { + item["section_id"] + for item in items_by_path.values() + if item["show_summary"] + } + all_section_ids = [meta["section_id"] for meta in all_sections.values()] + if not all_section_ids or not scope_item_sids: + return + + section_id_counts = await _load_direct_chunk_counts( + db, + document_id=document_id, + job_result_id=job_result_id, + all_section_ids=all_section_ids, + ) + + sid_to_path = {meta["section_id"]: path for path, meta in all_sections.items()} + for section_id, (text_count, image_count, table_count) in section_id_counts.items(): + chunk_path = sid_to_path.get(section_id, "") + if not chunk_path: + continue + + for item_path, item in items_by_path.items(): + if not item["show_summary"]: + continue + if chunk_path == item_path or chunk_path.startswith(item_path + " / "): + item["chunk_count"] += text_count + item["image_count"] += image_count + item["table_count"] += table_count + + await _attach_connected_asset_counts( + db, + document_id=document_id, + job_result_id=job_result_id, + items_by_path=items_by_path, + sid_to_path=sid_to_path, + ) + + +async def _load_direct_chunk_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + all_section_ids: list[str], +) -> dict[str, tuple[int, int, int]]: + chunk_stmt = ( + select( + DocumentChunk.section_id, + func.count( + case( + (DocumentChunk.chunk_type.notin_(["image", "table"]), literal_column("1")), + ) + ).label("text_count"), + func.count( + case( + (DocumentChunk.chunk_type == "image", literal_column("1")), + ) + ).label("image_count"), + func.count( + case( + (DocumentChunk.chunk_type == "table", literal_column("1")), + ) + ).label("table_count"), + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(all_section_ids)) + .group_by(DocumentChunk.section_id) + ) + chunk_rows = (await db.execute(chunk_stmt)).all() + return { + section_id: (int(text_count), int(image_count), int(table_count)) + for section_id, text_count, image_count, table_count in chunk_rows + } + + +async def _attach_connected_asset_counts( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + items_by_path: dict[str, dict], + sid_to_path: dict[str, str], +) -> None: + scope_items_with_zero_assets = [ + item + for item in items_by_path.values() + if item["show_summary"] and item["image_count"] == 0 and item["table_count"] == 0 + ] + if not scope_items_with_zero_assets: + return + + scope_section_ids = { + item["section_id"] + for item in items_by_path.values() + if item.get("section_id") + } + if not scope_section_ids: + return + + connect_stmt = ( + select( + DocumentChunk.section_id, + DocumentChunk.chunk_metadata, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(scope_section_ids))) + .where(DocumentChunk.chunk_type == "text") + ) + connect_result = (await db.execute(connect_stmt)).all() + + section_target_ids: dict[str, set[str]] = {} + for section_id, metadata in connect_result: + if not isinstance(metadata, dict): + continue + for connection in metadata.get("connect_to") or []: + target_id = connection.get("target", "") + if target_id: + section_target_ids.setdefault(section_id, set()).add(target_id) + + if not section_target_ids: + return + + all_target_ids: set[str] = set() + for target_ids in section_target_ids.values(): + all_target_ids.update(target_ids) + + target_type_stmt = ( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_id.in_(list(all_target_ids))) + .where(DocumentChunk.chunk_type.in_(["image", "table"])) + ) + target_type_result = (await db.execute(target_type_stmt)).all() + target_types = {chunk_id: chunk_type for chunk_id, chunk_type in target_type_result} + + for section_id, target_ids in section_target_ids.items(): + ref_path = sid_to_path.get(section_id, "") + if not ref_path: + continue + referenced_images = sum(1 for target_id in target_ids if target_types.get(target_id) == "image") + referenced_tables = sum(1 for target_id in target_ids if target_types.get(target_id) == "table") + if referenced_images == 0 and referenced_tables == 0: + continue + for item_path, item in items_by_path.items(): + if not item["show_summary"]: + continue + if ref_path == item_path or ref_path.startswith(item_path + " / "): + item["image_count"] += referenced_images + item["table_count"] += referenced_tables diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py index 8c395c07c..2f8b53526 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py +++ b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py @@ -2,10 +2,11 @@ from __future__ import annotations from loguru import logger -from sqlalchemy import func, select +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import DocumentChunk, DocumentSection +from shared.models.database.document import DocumentSection +from shared.services.retrieval.agentic.section_counts import attach_section_counts from shared.services.retrieval.lexical_text import normalize_section_path, split_section_path from shared.utils.text_utils import truncate_content_preview @@ -93,7 +94,7 @@ async def load_child_sections( if not items_by_path: return [] - await _attach_chunk_counts( + await attach_section_counts( db, document_id=document_id, job_result_id=job_result_id, @@ -219,174 +220,6 @@ def _resolve_allowed_depths(items_by_path: dict[str, dict], scope_list: list[str return allowed_set -async def _attach_chunk_counts( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - all_sections: dict[str, dict], - items_by_path: dict[str, dict], -) -> None: - scope_item_sids = { - item["section_id"] - for item in items_by_path.values() - if item["show_summary"] - } - all_section_ids = [meta["section_id"] for meta in all_sections.values()] - if not all_section_ids or not scope_item_sids: - return - - section_id_counts = await _load_direct_chunk_counts( - db, - document_id=document_id, - job_result_id=job_result_id, - all_section_ids=all_section_ids, - ) - - sid_to_path = {meta["section_id"]: path for path, meta in all_sections.items()} - for section_id, (text_count, image_count, table_count) in section_id_counts.items(): - chunk_path = sid_to_path.get(section_id, "") - if not chunk_path: - continue - - for item_path, item in items_by_path.items(): - if not item["show_summary"]: - continue - if chunk_path == item_path or chunk_path.startswith(item_path + " / "): - item["chunk_count"] += text_count - item["image_count"] += image_count - item["table_count"] += table_count - - await _attach_connected_asset_counts( - db, - document_id=document_id, - job_result_id=job_result_id, - items_by_path=items_by_path, - sid_to_path=sid_to_path, - ) - - -async def _load_direct_chunk_counts( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - all_section_ids: list[str], -) -> dict[str, tuple[int, int, int]]: - from sqlalchemy import case, literal_column - - chunk_stmt = ( - select( - DocumentChunk.section_id, - func.count( - case( - (DocumentChunk.chunk_type.notin_(["image", "table"]), literal_column("1")), - ) - ).label("text_count"), - func.count( - case( - (DocumentChunk.chunk_type == "image", literal_column("1")), - ) - ).label("image_count"), - func.count( - case( - (DocumentChunk.chunk_type == "table", literal_column("1")), - ) - ).label("table_count"), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_section_ids)) - .group_by(DocumentChunk.section_id) - ) - chunk_rows = (await db.execute(chunk_stmt)).all() - return { - section_id: (int(text_count), int(image_count), int(table_count)) - for section_id, text_count, image_count, table_count in chunk_rows - } - - -async def _attach_connected_asset_counts( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - items_by_path: dict[str, dict], - sid_to_path: dict[str, str], -) -> None: - scope_items_with_zero_assets = [ - item - for item in items_by_path.values() - if item["show_summary"] and item["image_count"] == 0 and item["table_count"] == 0 - ] - if not scope_items_with_zero_assets: - return - - scope_section_ids = { - item["section_id"] - for item in items_by_path.values() - if item.get("section_id") - } - if not scope_section_ids: - return - - connect_stmt = ( - select( - DocumentChunk.section_id, - DocumentChunk.chunk_metadata, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(scope_section_ids))) - .where(DocumentChunk.chunk_type == "text") - ) - connect_result = (await db.execute(connect_stmt)).all() - - section_target_ids: dict[str, set[str]] = {} - for section_id, metadata in connect_result: - if not isinstance(metadata, dict): - continue - for connection in metadata.get("connect_to") or []: - target_id = connection.get("target", "") - if target_id: - section_target_ids.setdefault(section_id, set()).add(target_id) - - if not section_target_ids: - return - - all_target_ids: set[str] = set() - for target_ids in section_target_ids.values(): - all_target_ids.update(target_ids) - - target_type_stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_id.in_(list(all_target_ids))) - .where(DocumentChunk.chunk_type.in_(["image", "table"])) - ) - target_type_result = (await db.execute(target_type_stmt)).all() - target_types = {chunk_id: chunk_type for chunk_id, chunk_type in target_type_result} - - for section_id, target_ids in section_target_ids.items(): - ref_path = sid_to_path.get(section_id, "") - if not ref_path: - continue - referenced_images = sum(1 for target_id in target_ids if target_types.get(target_id) == "image") - referenced_tables = sum(1 for target_id in target_ids if target_types.get(target_id) == "table") - if referenced_images == 0 and referenced_tables == 0: - continue - for item_path, item in items_by_path.items(): - if not item["show_summary"]: - continue - if ref_path == item_path or ref_path.startswith(item_path + " / "): - item["image_count"] += referenced_images - item["table_count"] += referenced_tables - - def _mark_leaf_and_selectable( sorted_items: list[dict], *, From e079d6a127e8ab8d5cd7dded94edb8221e6e5246 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:25:04 +0800 Subject: [PATCH 32/40] refactor extract agentic section prompt projection --- .../retrieval/agentic/navigation_tools.py | 6 +- .../agentic/section_prompt_projection.py | 59 +++++++++++++++++++ .../retrieval/agentic/section_tree.py | 56 ------------------ 3 files changed, 61 insertions(+), 60 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py index 23edec823..9844ed7e4 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py @@ -22,10 +22,8 @@ format_budget_block, parse_action_response, ) -from shared.services.retrieval.agentic.section_tree import ( - format_items_for_llm, - load_child_sections, -) +from shared.services.retrieval.agentic.section_prompt_projection import format_items_for_llm +from shared.services.retrieval.agentic.section_tree import load_child_sections from shared.services.retrieval.agentic.selection_hydration import ( hydrate_path_selections_into_node, ) diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py b/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py new file mode 100644 index 000000000..b7a251805 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py @@ -0,0 +1,59 @@ +"""Prompt projection for agentic section navigation.""" +from __future__ import annotations + +from shared.utils.text_utils import truncate_content_preview + + +def format_items_for_llm( + items: list[dict], + max_chars: int = 20000, +) -> tuple[str, bool]: + """Format section items with hierarchy, selectability, counts, and summaries.""" + if not items: + return "(no items available)", False + + full_text = "\n".join(_render_item(item, include_summary=True) for item in items) + if len(full_text) <= max_chars: + return full_text, False + + slim_text = "\n".join(_render_item(item, include_summary=False) for item in items) + return slim_text[:max_chars], True + + +def _render_item(item: dict, include_summary: bool) -> str: + level = item.get("level", 1) + show_summary = item.get("show_summary", True) + is_leaf = item.get("is_leaf", False) + leaf_tag = " [Leaf]" if is_leaf else "" + path = item.get("path", "") + summary = item.get("summary") or "" + + counts_str = "" + if show_summary: + count_parts: list[str] = [] + chunk_count = item.get("chunk_count", 0) + if chunk_count > 0: + count_parts.append(f"text={chunk_count}") + image_count = item.get("image_count", 0) + if image_count > 0: + count_parts.append(f"image={image_count}") + table_count = item.get("table_count", 0) + if table_count > 0: + count_parts.append(f"table={table_count}") + counts_str = f' [{" ".join(count_parts)}]' if count_parts else "" + + indent = " " * (level - 1) + prefix = "▸" if level == 1 else "└" + level_tag = f"[L{level}]" + select_tag = "[SELECT] " if item.get("selectable", False) else "" + + lines = [ + f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}' + ] + + if include_summary and show_summary and summary: + sub_indent = " " * level + clipped = truncate_content_preview(summary, head=80, tail=0) + lines.append(f"{sub_indent}{clipped}") + + return "\n".join(lines) diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py index 2f8b53526..53507f35d 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py +++ b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py @@ -8,7 +8,6 @@ from shared.models.database.document import DocumentSection from shared.services.retrieval.agentic.section_counts import attach_section_counts from shared.services.retrieval.lexical_text import normalize_section_path, split_section_path -from shared.utils.text_utils import truncate_content_preview async def load_child_sections( @@ -111,22 +110,6 @@ async def load_child_sections( return sorted_items -def format_items_for_llm( - items: list[dict], - max_chars: int = 20000, -) -> tuple[str, bool]: - """Format section items with hierarchy, selectability, counts, and summaries.""" - if not items: - return "(no items available)", False - - full_text = "\n".join(_render_item(item, include_summary=True) for item in items) - if len(full_text) <= max_chars: - return full_text, False - - slim_text = "\n".join(_render_item(item, include_summary=False) for item in items) - return slim_text[:max_chars], True - - def _select_scope_items( all_sections: dict[str, dict], *, @@ -246,42 +229,3 @@ def _mark_leaf_and_selectable( else: for item in sorted_items: item["selectable"] = item.get("show_summary", True) - - -def _render_item(item: dict, include_summary: bool) -> str: - level = item.get("level", 1) - show_summary = item.get("show_summary", True) - is_leaf = item.get("is_leaf", False) - leaf_tag = " [Leaf]" if is_leaf else "" - path = item.get("path", "") - summary = item.get("summary") or "" - - counts_str = "" - if show_summary: - count_parts: list[str] = [] - chunk_count = item.get("chunk_count", 0) - if chunk_count > 0: - count_parts.append(f"text={chunk_count}") - image_count = item.get("image_count", 0) - if image_count > 0: - count_parts.append(f"image={image_count}") - table_count = item.get("table_count", 0) - if table_count > 0: - count_parts.append(f"table={table_count}") - counts_str = f' [{" ".join(count_parts)}]' if count_parts else "" - - indent = " " * (level - 1) - prefix = "▸" if level == 1 else "└" - level_tag = f"[L{level}]" - select_tag = "[SELECT] " if item.get("selectable", False) else "" - - lines = [ - f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}' - ] - - if include_summary and show_summary and summary: - sub_indent = " " * level - clipped = truncate_content_preview(summary, head=80, tail=0) - lines.append(f"{sub_indent}{clipped}") - - return "\n".join(lines) From 7f4c504ab55abadfa0505b95426a939d1c7ff91a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:29:18 +0800 Subject: [PATCH 33/40] refactor move agentic asset availability --- .../services/retrieval/agentic/asset_tools.py | 70 ++++++++++++++++- .../retrieval/agentic/navigation_tools.py | 78 ++----------------- 2 files changed, 75 insertions(+), 73 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py b/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py index 395120e9d..25f97866c 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py @@ -4,7 +4,8 @@ from typing import Any from loguru import logger -from sqlalchemy import select +from sqlalchemy import func as sa_func +from sqlalchemy import or_, select from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection @@ -31,6 +32,73 @@ def build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, st return owner_map +async def count_assets_under_scope( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + scope_paths: list[str], +) -> tuple[int, int]: + scope_section_stmt = ( + select(DocumentSection.section_id) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + ) + if scope_paths: + scope_filters = [] + for scope in scope_paths: + scope_filters.append(DocumentSection.section_path == scope) + scope_filters.append(DocumentSection.section_path.like(f"{scope} / %")) + scope_section_stmt = scope_section_stmt.where(or_(*scope_filters)) + scope_section_ids = await db.execute(scope_section_stmt) + all_section_ids = [row[0] for row in scope_section_ids.all()] + + if not all_section_ids: + return 0, 0 + + count_stmt = ( + select( + DocumentChunk.chunk_type, + sa_func.count(DocumentChunk.id), + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(all_section_ids)) + .where(DocumentChunk.chunk_type.in_(["image", "table"])) + .group_by(DocumentChunk.chunk_type) + ) + count_result = await db.execute(count_stmt) + + total_images = 0 + total_tables = 0 + for chunk_type, count in count_result.all(): + if chunk_type == "image": + total_images = count + elif chunk_type == "table": + total_tables = count + return total_images, total_tables + + +def build_asset_tools_block(total_images: int, total_tables: int) -> str: + if total_images <= 0 and total_tables <= 0: + return "" + + tools_lines = ["\nOptional asset tools (usable with NAVIGATE or STOP):\n"] + if total_images > 0: + tools_lines.append( + f" FIND_IMAGES — Extract image/chart assets under the current scope ({total_images} available).\n" + ) + if total_tables > 0: + tools_lines.append( + f" FIND_TABLES — Extract table/data assets under the current scope ({total_tables} available).\n" + ) + tools_lines.append( + " Note: with NAVIGATE selections, asset tools are limited to the selected sections; " + "with STOP or no selections, they use the current scope.\n" + ) + return "".join(tools_lines) + + async def resolve_root_asset_owners( db: AsyncSession, *, diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py index 9844ed7e4..1f6bbfe7a 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py @@ -10,11 +10,12 @@ from typing import Any from loguru import logger -from sqlalchemy import func as sa_func -from sqlalchemy import or_, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import DocumentChunk, DocumentSection +from shared.services.retrieval.agentic.asset_tools import ( + build_asset_tools_block, + count_assets_under_scope, +) from shared.services.retrieval.agentic.budget import BudgetExceeded from shared.services.retrieval.agentic.prompts import ( ACTION_PROMPT, @@ -74,13 +75,13 @@ async def navigate_step( selectable = { item["path"]: item for item in items if item.get("selectable", False) } - total_images, total_tables = await _count_assets_under_scope( + total_images, total_tables = await count_assets_under_scope( db, document_id=document_id, job_result_id=job_result_id, scope_paths=scope_paths, ) - tools_block = _build_tools_block(total_images, total_tables) + tools_block = build_asset_tools_block(total_images, total_tables) items_text, overflowed = format_items_for_llm(items) prompt = _build_navigation_prompt( @@ -237,73 +238,6 @@ async def discovery_select_step( return node -async def _count_assets_under_scope( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - scope_paths: list[str], -) -> tuple[int, int]: - scope_section_stmt = ( - select(DocumentSection.section_id) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - ) - if scope_paths: - scope_filters = [] - for scope in scope_paths: - scope_filters.append(DocumentSection.section_path == scope) - scope_filters.append(DocumentSection.section_path.like(f"{scope} / %")) - scope_section_stmt = scope_section_stmt.where(or_(*scope_filters)) - scope_section_ids = await db.execute(scope_section_stmt) - all_section_ids = [row[0] for row in scope_section_ids.all()] - - if not all_section_ids: - return 0, 0 - - count_stmt = ( - select( - DocumentChunk.chunk_type, - sa_func.count(DocumentChunk.id), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_section_ids)) - .where(DocumentChunk.chunk_type.in_(["image", "table"])) - .group_by(DocumentChunk.chunk_type) - ) - count_result = await db.execute(count_stmt) - - total_images = 0 - total_tables = 0 - for chunk_type, count in count_result.all(): - if chunk_type == "image": - total_images = count - elif chunk_type == "table": - total_tables = count - return total_images, total_tables - - -def _build_tools_block(total_images: int, total_tables: int) -> str: - if total_images <= 0 and total_tables <= 0: - return "" - - tools_lines = ["\nOptional asset tools (usable with NAVIGATE or STOP):\n"] - if total_images > 0: - tools_lines.append( - f" FIND_IMAGES — Extract image/chart assets under the current scope ({total_images} available).\n" - ) - if total_tables > 0: - tools_lines.append( - f" FIND_TABLES — Extract table/data assets under the current scope ({total_tables} available).\n" - ) - tools_lines.append( - " Note: with NAVIGATE selections, asset tools are limited to the selected sections; " - "with STOP or no selections, they use the current scope.\n" - ) - return "".join(tools_lines) - - def _build_navigation_prompt( *, document_id: str, From 9a11351c64acaa4ce11629bc5004642f7b492970 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:33:24 +0800 Subject: [PATCH 34/40] refactor split agentic discovery selection --- .../retrieval/agentic/discovery_selection.py | 206 ++++++++++++++++++ .../retrieval/agentic/navigation_tools.py | 190 ---------------- .../services/retrieval/agentic/tools.py | 9 +- 3 files changed, 213 insertions(+), 192 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py new file mode 100644 index 000000000..d54df4e5f --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py @@ -0,0 +1,206 @@ +"""Post-navigation discovery selection for agentic retrieval.""" +from __future__ import annotations + +import time +from typing import Any + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agentic.budget import BudgetExceeded +from shared.services.retrieval.agentic.prompts import ( + DISCOVERY_SELECT_PROMPT, + format_budget_block, + parse_action_response, +) +from shared.services.retrieval.agentic.selection_hydration import ( + hydrate_path_selections_into_node, +) +from shared.services.retrieval.agentic.types import DocTreeNode +from shared.services.retrieval.lexical_text import normalize_section_path +from shared.services.retrieval.llm_adapter import LLMFn + + +_MAX_DISCOVERY_PER_DOC = 3 + + +async def discovery_select_step( + db: AsyncSession, + *, + document_id: str, + query: str, + llm_fn: LLMFn, + user_id: str, + namespace: str, + doc_name: str = "", + discovery_hints: list[dict[str, Any]], + exclude_paths: set[str] | None = None, + revision_hint: str | None = None, + budget_snapshot: dict | None = None, +) -> DocTreeNode: + """Select and hydrate discovery-found sections after BFS navigation.""" + node = DocTreeNode(scope_path=None) + if not discovery_hints: + return node + + hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC] + + t0 = time.monotonic() + try: + hint_lines, hint_by_path, root_path_selections = _project_discovery_hints( + hints, + exclude_paths=exclude_paths, + ) + if not hint_lines and not root_path_selections: + return node + + selections: list[dict[str, Any]] = [] + if hint_lines: + prompt = _build_discovery_selection_prompt( + document_id=document_id, + doc_name=doc_name, + query=query, + hint_lines=hint_lines, + revision_hint=revision_hint, + budget_snapshot=budget_snapshot, + ) + response = await llm_fn(prompt) + parsed = parse_action_response(response) + selections = parsed.get("selections", []) + + logger.info( + f' discovery_select_step doc="{doc_name}": ' + f"hints={len(hints)} selections={len(selections)} " + f"root_selections={len(root_path_selections)}" + ) + + path_selections = _build_discovery_path_selections( + selections=selections, + hint_by_path=hint_by_path, + root_path_selections=root_path_selections, + node=node, + ) + await hydrate_path_selections_into_node( + db, + node=node, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f" discovery_select_step done: hydrated={len(node.leaf_content)} " + f"latency={latency}ms" + ) + return node + + except BudgetExceeded: + raise + except Exception as exc: + logger.error(f" discovery_select_step failed for doc={document_id}: {exc}") + return node + + +def _project_discovery_hints( + hints: list[dict[str, Any]], + *, + exclude_paths: set[str] | None, +) -> tuple[list[str], dict[str, dict], list[dict[str, Any]]]: + exclude_set = { + normalize_section_path(path) + for path in (exclude_paths or set()) + if path + } + hint_lines: list[str] = [] + hint_by_path: dict[str, dict] = {} + root_path_selections: list[dict[str, Any]] = [] + for hint in hints: + section_path = normalize_section_path(hint.get("section_path", "")) + if not section_path: + continue + if section_path in exclude_set: + continue + if section_path in hint_by_path: + continue + + hint_by_path[section_path] = hint + if section_path == "Root": + root_path_selections.append({ + "path": section_path, + "confidence": float( + hint.get("discovery_score") or hint.get("score") or 0.7 + ), + "hydrate_mode": "self_only", + }) + continue + + summary = hint.get("summary", "") or "" + hint_lines.append(f'▸ path="{section_path}"') + if summary: + hint_lines.append(f" {summary[:300]}") + + return hint_lines, hint_by_path, root_path_selections + + +def _build_discovery_selection_prompt( + *, + document_id: str, + doc_name: str, + query: str, + hint_lines: list[str], + revision_hint: str | None, + budget_snapshot: dict | None, +) -> str: + revision_context = "" + if revision_hint: + revision_context = ( + "\nIMPORTANT: This is a REVISION round. " + "The previous search attempt failed because:\n" + f'"{revision_hint}"\n' + "Adjust your selection accordingly. " + "If no candidate is relevant, return an EMPTY list [].\n" + ) + + return DISCOVERY_SELECT_PROMPT.format( + doc_name=doc_name or document_id, + budget_block=format_budget_block(budget_snapshot), + items="\n".join(hint_lines), + query=query, + revision_context=revision_context, + ) + + +def _build_discovery_path_selections( + *, + selections: list[dict[str, Any]], + hint_by_path: dict[str, dict], + root_path_selections: list[dict[str, Any]], + node: DocTreeNode, +) -> list[dict[str, Any]]: + valid_selections = [ + selection for selection in selections if selection["path"] in hint_by_path + ] + path_selections = list(root_path_selections) + for selection in valid_selections: + path = selection["path"] + confidence = selection.get("confidence", 0.7) + node.confidence[path] = confidence + path_selections.append({"path": path, "confidence": confidence}) + + if not path_selections and hint_by_path: + fallback_path, fallback_hint = next(iter(hint_by_path.items())) + fallback_confidence = float( + fallback_hint.get("discovery_score") + or fallback_hint.get("score") + or 0.5 + ) + node.confidence[fallback_path] = fallback_confidence + path_selections.append({ + "path": fallback_path, + "confidence": fallback_confidence, + "hydrate_mode": "self_only", + }) + + return path_selections diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py index 1f6bbfe7a..98f8fd424 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py @@ -6,7 +6,6 @@ """ from __future__ import annotations -import time from typing import Any from loguru import logger @@ -19,7 +18,6 @@ from shared.services.retrieval.agentic.budget import BudgetExceeded from shared.services.retrieval.agentic.prompts import ( ACTION_PROMPT, - DISCOVERY_SELECT_PROMPT, format_budget_block, parse_action_response, ) @@ -29,13 +27,9 @@ hydrate_path_selections_into_node, ) from shared.services.retrieval.agentic.types import DocTreeNode -from shared.services.retrieval.lexical_text import normalize_section_path from shared.services.retrieval.llm_adapter import LLMFn -_MAX_DISCOVERY_PER_DOC = 3 - - async def navigate_step( db: AsyncSession, *, @@ -157,87 +151,6 @@ async def navigate_step( except Exception as exc: logger.error(f" navigate_step failed for doc={document_id}: {exc}") return "STOP", [], empty, [] - - -async def discovery_select_step( - db: AsyncSession, - *, - document_id: str, - query: str, - llm_fn: LLMFn, - user_id: str, - namespace: str, - doc_name: str = "", - discovery_hints: list[dict[str, Any]], - exclude_paths: set[str] | None = None, - revision_hint: str | None = None, - budget_snapshot: dict | None = None, -) -> DocTreeNode: - """Select and hydrate discovery-found sections after BFS navigation.""" - node = DocTreeNode(scope_path=None) - if not discovery_hints: - return node - - hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC] - - t0 = time.monotonic() - try: - hint_lines, hint_by_path, root_path_selections = _project_discovery_hints( - hints, - exclude_paths=exclude_paths, - ) - if not hint_lines and not root_path_selections: - return node - - selections: list[dict[str, Any]] = [] - if hint_lines: - prompt = _build_discovery_selection_prompt( - document_id=document_id, - doc_name=doc_name, - query=query, - hint_lines=hint_lines, - revision_hint=revision_hint, - budget_snapshot=budget_snapshot, - ) - response = await llm_fn(prompt) - parsed = parse_action_response(response) - selections = parsed.get("selections", []) - - logger.info( - f' discovery_select_step doc="{doc_name}": ' - f"hints={len(hints)} selections={len(selections)} " - f"root_selections={len(root_path_selections)}" - ) - - path_selections = _build_discovery_path_selections( - selections=selections, - hint_by_path=hint_by_path, - root_path_selections=root_path_selections, - node=node, - ) - await hydrate_path_selections_into_node( - db, - node=node, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f" discovery_select_step done: hydrated={len(node.leaf_content)} " - f"latency={latency}ms" - ) - return node - - except BudgetExceeded: - raise - except Exception as exc: - logger.error(f" discovery_select_step failed for doc={document_id}: {exc}") - return node - - def _build_navigation_prompt( *, document_id: str, @@ -271,106 +184,3 @@ def _build_navigation_prompt( f'"{revision_hint}". Adjust your selections accordingly.' ) return prompt - - -def _project_discovery_hints( - hints: list[dict[str, Any]], - *, - exclude_paths: set[str] | None, -) -> tuple[list[str], dict[str, dict], list[dict[str, Any]]]: - exclude_set = { - normalize_section_path(path) - for path in (exclude_paths or set()) - if path - } - hint_lines: list[str] = [] - hint_by_path: dict[str, dict] = {} - root_path_selections: list[dict[str, Any]] = [] - for hint in hints: - section_path = normalize_section_path(hint.get("section_path", "")) - if not section_path: - continue - if section_path in exclude_set: - continue - if section_path in hint_by_path: - continue - - hint_by_path[section_path] = hint - if section_path == "Root": - root_path_selections.append({ - "path": section_path, - "confidence": float( - hint.get("discovery_score") or hint.get("score") or 0.7 - ), - "hydrate_mode": "self_only", - }) - continue - - summary = hint.get("summary", "") or "" - hint_lines.append(f'▸ path="{section_path}"') - if summary: - hint_lines.append(f" {summary[:300]}") - - return hint_lines, hint_by_path, root_path_selections - - -def _build_discovery_selection_prompt( - *, - document_id: str, - doc_name: str, - query: str, - hint_lines: list[str], - revision_hint: str | None, - budget_snapshot: dict | None, -) -> str: - revision_context = "" - if revision_hint: - revision_context = ( - "\nIMPORTANT: This is a REVISION round. " - "The previous search attempt failed because:\n" - f'"{revision_hint}"\n' - "Adjust your selection accordingly. " - "If no candidate is relevant, return an EMPTY list [].\n" - ) - - return DISCOVERY_SELECT_PROMPT.format( - doc_name=doc_name or document_id, - budget_block=format_budget_block(budget_snapshot), - items="\n".join(hint_lines), - query=query, - revision_context=revision_context, - ) - - -def _build_discovery_path_selections( - *, - selections: list[dict[str, Any]], - hint_by_path: dict[str, dict], - root_path_selections: list[dict[str, Any]], - node: DocTreeNode, -) -> list[dict[str, Any]]: - valid_selections = [ - selection for selection in selections if selection["path"] in hint_by_path - ] - path_selections = list(root_path_selections) - for selection in valid_selections: - path = selection["path"] - confidence = selection.get("confidence", 0.7) - node.confidence[path] = confidence - path_selections.append({"path": path, "confidence": confidence}) - - if not path_selections and hint_by_path: - fallback_path, fallback_hint = next(iter(hint_by_path.items())) - fallback_confidence = float( - fallback_hint.get("discovery_score") - or fallback_hint.get("score") - or 0.5 - ) - node.confidence[fallback_path] = fallback_confidence - path_selections.append({ - "path": fallback_path, - "confidence": fallback_confidence, - "hydrate_mode": "self_only", - }) - - return path_selections diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 29c0c4e73..4fde17d1d 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -9,7 +9,12 @@ from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic import asset_tools, discovery_tools, navigation_tools +from shared.services.retrieval.agentic import ( + asset_tools, + discovery_selection, + discovery_tools, + navigation_tools, +) from shared.services.retrieval.agentic.types import DocTreeNode, ToolResult from shared.services.retrieval.llm_adapter import LLMFn @@ -134,7 +139,7 @@ async def discovery_select_step( revision_hint: str | None = None, budget_snapshot: dict | None = None, ) -> DocTreeNode: - return await navigation_tools.discovery_select_step( + return await discovery_selection.discovery_select_step( db, document_id=document_id, query=query, From b4d2b764ea41da246c26064c4b9fca33ab9b8373 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 19:02:47 +0800 Subject: [PATCH 35/40] refactor share agentic asset scope loading --- .../services/retrieval/agentic/asset_tools.py | 48 +++++++++++-------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py b/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py index 25f97866c..48403e3ca 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py @@ -32,15 +32,15 @@ def build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, st return owner_map -async def count_assets_under_scope( +async def _load_scope_sections( db: AsyncSession, *, document_id: str, job_result_id: str, scope_paths: list[str], -) -> tuple[int, int]: - scope_section_stmt = ( - select(DocumentSection.section_id) +) -> list[tuple[str, str]]: + section_stmt = ( + select(DocumentSection.section_id, DocumentSection.section_path) .where(DocumentSection.document_id == document_id) .where(DocumentSection.job_result_id == job_result_id) ) @@ -49,9 +49,25 @@ async def count_assets_under_scope( for scope in scope_paths: scope_filters.append(DocumentSection.section_path == scope) scope_filters.append(DocumentSection.section_path.like(f"{scope} / %")) - scope_section_stmt = scope_section_stmt.where(or_(*scope_filters)) - scope_section_ids = await db.execute(scope_section_stmt) - all_section_ids = [row[0] for row in scope_section_ids.all()] + section_stmt = section_stmt.where(or_(*scope_filters)) + rows = (await db.execute(section_stmt)).all() + return [(section_id, section_path or "") for section_id, section_path in rows] + + +async def count_assets_under_scope( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + scope_paths: list[str], +) -> tuple[int, int]: + section_rows = await _load_scope_sections( + db, + document_id=document_id, + job_result_id=job_result_id, + scope_paths=scope_paths, + ) + all_section_ids = [section_id for section_id, _section_path in section_rows] if not all_section_ids: return 0, 0 @@ -170,20 +186,12 @@ async def asset_filter_step( else [] ) - section_stmt = ( - select(DocumentSection.section_id, DocumentSection.section_path) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) + section_rows = await _load_scope_sections( + db, + document_id=document_id, + job_result_id=job_result_id, + scope_paths=scope_list, ) - if scope_list: - from sqlalchemy import or_ - - scope_filters = [] - for scope in scope_list: - scope_filters.append(DocumentSection.section_path == scope) - scope_filters.append(DocumentSection.section_path.like(f"{scope} / %")) - section_stmt = section_stmt.where(or_(*scope_filters)) - section_rows = (await db.execute(section_stmt)).all() section_ids = {row[0] for row in section_rows} if not section_ids: From f430026bec4c73d73507eb507bd2430502be0dd9 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 22:55:21 +0800 Subject: [PATCH 36/40] refactor(worker): deepen parser architecture --- CONTEXT.md | 57 + .../services/document_parser/atlas_parser.py | 36 +- .../dataframe_html_renderer.py | 294 ++++ .../services/document_parser/doc_parser.py | 564 ++----- .../document_parser/doc_profile_model.py | 75 + .../document_parser/doc_profile_pdf.py | 676 ++++++++ .../services/document_parser/doc_profiler.py | 769 +-------- .../document_parser/docx_asset_accumulator.py | 83 + .../document_parser/docx_asset_store.py | 63 + .../document_parser/docx_block_stream.py | 315 ++++ .../document_parser/docx_table_html.py | 138 ++ .../document_parser/excel_structure_parser.py | 745 +++++++++ .../document_parser/heading_candidates.py | 488 ++++++ .../document_parser/heading_hierarchy.py | 37 + .../document_parser/heading_llm_executor.py | 610 +++++++ .../services/document_parser/heading_tree.py | 172 ++ .../services/document_parser/html_parser.py | 565 +------ .../services/document_parser/image_parser.py | 28 +- .../services/document_parser/inline_asset.py | 50 + .../services/document_parser/layout_parser.py | 1425 +---------------- .../markdown_deferred_summary.py | 292 ++++ .../document_parser/markdown_parse_state.py | 179 +++ .../app/services/document_parser/md_parser.py | 517 ++---- .../services/document_parser/mineru_client.py | 94 ++ .../document_parser/mineru_pdf_service.py | 370 +---- .../document_parser/mineru_task_polling.py | 240 +++ .../orchestration/format_adapters.py | 221 +++ .../orchestration/format_router.py | 94 ++ .../orchestration/parse_input.py | 29 + .../orchestration/parse_session.py | 132 +- .../orchestration/route_parse.py | 199 +-- .../services/document_parser/parse_service.py | 26 +- .../services/document_parser/parser_rows.py | 61 + .../services/document_parser/pptx_parser.py | 115 +- .../document_parser/rendered_pdf_transform.py | 90 ++ .../document_parser/table_asset_writer.py | 48 + .../services/document_parser/table_parser.py | 884 +--------- .../app/services/document_parser/toc_docx.py | 291 ++++ .../services/document_parser/toc_hierarchy.py | 149 ++ .../services/document_parser/toc_parser.py | 554 +------ ...t_document_parser_architecture_contract.py | 732 +++++++++ 41 files changed, 6821 insertions(+), 5686 deletions(-) create mode 100644 apps/worker/app/services/document_parser/dataframe_html_renderer.py create mode 100644 apps/worker/app/services/document_parser/doc_profile_model.py create mode 100644 apps/worker/app/services/document_parser/doc_profile_pdf.py create mode 100644 apps/worker/app/services/document_parser/docx_asset_accumulator.py create mode 100644 apps/worker/app/services/document_parser/docx_asset_store.py create mode 100644 apps/worker/app/services/document_parser/docx_block_stream.py create mode 100644 apps/worker/app/services/document_parser/docx_table_html.py create mode 100644 apps/worker/app/services/document_parser/excel_structure_parser.py create mode 100644 apps/worker/app/services/document_parser/heading_candidates.py create mode 100644 apps/worker/app/services/document_parser/heading_hierarchy.py create mode 100644 apps/worker/app/services/document_parser/heading_llm_executor.py create mode 100644 apps/worker/app/services/document_parser/heading_tree.py create mode 100644 apps/worker/app/services/document_parser/inline_asset.py create mode 100644 apps/worker/app/services/document_parser/markdown_deferred_summary.py create mode 100644 apps/worker/app/services/document_parser/markdown_parse_state.py create mode 100644 apps/worker/app/services/document_parser/mineru_client.py create mode 100644 apps/worker/app/services/document_parser/mineru_task_polling.py create mode 100644 apps/worker/app/services/document_parser/orchestration/format_adapters.py create mode 100644 apps/worker/app/services/document_parser/orchestration/format_router.py create mode 100644 apps/worker/app/services/document_parser/orchestration/parse_input.py create mode 100644 apps/worker/app/services/document_parser/parser_rows.py create mode 100644 apps/worker/app/services/document_parser/rendered_pdf_transform.py create mode 100644 apps/worker/app/services/document_parser/table_asset_writer.py create mode 100644 apps/worker/app/services/document_parser/toc_docx.py create mode 100644 apps/worker/app/services/document_parser/toc_hierarchy.py create mode 100644 apps/worker/tests/contract/test_document_parser_architecture_contract.py diff --git a/CONTEXT.md b/CONTEXT.md index bafd49c76..f87d42697 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -52,6 +52,35 @@ The retrieval-visible text, image, or table row attached to a Document Section. The workflow that creates a Job, accepts a file or URL source, confirms upload state, and starts parsing work. +### Worker Document Parsing + +The worker-side workflow that turns a source file into parsed DataFrame rows, +parsed assets, and parser debug artifacts before chunk conversion and result +packaging. + +### Parser Input + +The typed worker-side parse request assembled from Job metadata, parser options, +source-file identity, output naming, and storage transform keys. + +### Document Format Routing + +The Worker Document Parsing module that selects one concrete parser adapter for +the source document format while keeping format-specific conversion details out +of the stable parser entrypoint. + +### Rendered PDF Transform + +The Worker Document Parsing module that reuses or creates rendered PDF artifacts +for PDF-backed parsing paths, including PPTX-to-PDF fallback handling, image-only +PDF rendering, temporary PDF materialization, MinerU handoff, and cleanup. + +### Heading Hierarchy + +The Worker Document Parsing module that predicts section levels from Markdown +lines, DOCX blocks, TOC context, layout metadata, heuristics, and optional LLM +inference. + ### Job Admission The policy checks that must pass before a new Job is created: authentication, @@ -230,11 +259,39 @@ exceptions. - `app/api/v1/routes/qstash_callbacks.py` - `app/services/qstash_callback_service.py` +## apps/worker Workflow Ownership + +### Worker Document Parsing + +- `app/services/document_parser/parse_service.py` +- `app/services/document_parser/orchestration/parse_input.py` +- `app/services/document_parser/orchestration/parse_session.py` +- `app/services/document_parser/orchestration/route_parse.py` +- `app/services/document_parser/orchestration/format_router.py` +- `app/services/document_parser/orchestration/format_adapters.py` + +### Rendered PDF Transform + +- `app/services/document_parser/rendered_pdf_transform.py` +- `app/services/document_parser/pptx_pdf_rendering.py` +- `app/services/document_parser/pdf_parser.py` +- `app/services/document_parser/pptx_parser.py` + +### Heading Hierarchy + +- `app/services/document_parser/heading_hierarchy.py` +- `app/services/document_parser/layout_parser.py` +- `app/services/document_parser/md_parser.py` +- `app/services/document_parser/doc_parser.py` + ## Invariants - `apps/api` coordinates workflows. Parsing, publication, retrieval internals, storage mechanics, and state-machine implementation mostly live outside the route modules. +- Worker Document Parsing exposes `checkerboard_inject_parse` as the stable + parser entrypoint; parser option shaping, format routing, rendered PDF + transforms, and heading inference stay behind that entrypoint. - A Job and a Document are not the same thing. Jobs track intake and processing; Documents track retrieval-visible knowledge state. - `current_job_result_id` selects the active revision of a Document. diff --git a/apps/worker/app/services/document_parser/atlas_parser.py b/apps/worker/app/services/document_parser/atlas_parser.py index 2d95b1b5a..fef806add 100644 --- a/apps/worker/app/services/document_parser/atlas_parser.py +++ b/apps/worker/app/services/document_parser/atlas_parser.py @@ -21,6 +21,7 @@ import pandas as pd from app.services.document_parser.dataframe_helpers import process_dup_paths_df from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker from app.services.document_parser.toc_parser import detect_tocs_in_texts from loguru import logger @@ -373,7 +374,7 @@ def _vlm_task(page_num, img_name): ) # ── Phase 3: Build chunks ── - df_list = [] + rows_builder = ParsedRowsBuilder() custom_md_lines = [] skipped_null = 0 used_image_names: set[str] = set() @@ -431,26 +432,20 @@ def _vlm_task(page_num, img_name): else: chunk_path = safe_title - # Build df row (11 columns) # Atlas chunks are image-primary: use IMAGE marker directly. # find_matches_parsing() prepends "PTXT\n" which causes downstream # chunk type classifier to misclassify as "text" instead of "image". - match_type = "image" tokens = tokenize2stw_remove([content], stopwords) - df_list.append( - [ - content, # content - chunk_path, # path - match_type, # type - len(content), # length - "", # keywords - "", # summary - know_id, # know_id - tokens, # tokens - "", # connectto - time_stamp, # addtime - str(page_num), # page_nums - ] + rows_builder.append( + ParsedRow( + content=content, + path=chunk_path, + type="image", + know_id=know_id, + addtime=time_stamp, + tokens=tokens, + page_nums=str(page_num), + ) ) # Build custom md line @@ -468,8 +463,9 @@ def _vlm_task(page_num, img_name): with open(custom_md_path, "w", encoding="utf-8") as f: f.write("\n".join(custom_md_lines)) + df = rows_builder.to_dataframe() logger.info( - f"📐 Atlas pipeline complete: {len(df_list)} chunks created " + f"📐 Atlas pipeline complete: {len(df)} chunks created " f"(skipped {len(toc_page_set)} TOC + {skipped_null} null pages, total {total_pages})" ) @@ -477,9 +473,5 @@ def _vlm_task(page_num, img_name): # Currently atlas chunks are flat with no parent-child relationships. # Future: integrate with hierarchy builder for unified schema. - # ── Build DataFrame ── - all_cols = settings.ALL_DF_COLS.split(",") - - df = pd.DataFrame(df_list, columns=all_cols) df = process_dup_paths_df(df) return df diff --git a/apps/worker/app/services/document_parser/dataframe_html_renderer.py b/apps/worker/app/services/document_parser/dataframe_html_renderer.py new file mode 100644 index 000000000..97da217ce --- /dev/null +++ b/apps/worker/app/services/document_parser/dataframe_html_renderer.py @@ -0,0 +1,294 @@ +from __future__ import annotations + +from typing import List, Union + +import pandas as pd + + +def render_multiindex_thead(columns: pd.MultiIndex, escape: bool = False) -> str: + """Convert MultiIndex columns to an HTML thead with colspan and rowspan.""" + import html as html_lib + + level_count = columns.nlevels + column_count = len(columns) + + grid = [] + for level in range(level_count): + row = [columns.get_level_values(level)[col] for col in range(column_count)] + grid.append(row) + + colspan = [[1] * column_count for _ in range(level_count)] + + for level in range(level_count): + col = 0 + while col < column_count: + span = 1 + while col + span < column_count and grid[level][col] == grid[level][col + span]: + parent_match = True + for parent_level in range(level): + if grid[parent_level][col] != grid[parent_level][col + span]: + parent_match = False + break + if parent_match: + span += 1 + else: + break + colspan[level][col] = span + col += span + + rowspan = [[1] * column_count for _ in range(level_count)] + + for col in range(column_count): + level = 0 + while level < level_count: + span = 1 + while level + span < level_count: + if ( + grid[level][col] == grid[level + span][col] + and colspan[level][col] == colspan[level + span][col] + ): + span += 1 + else: + break + rowspan[level][col] = span + level += span + + covered = [[False] * column_count for _ in range(level_count)] + html_parts = [""] + + for level in range(level_count): + html_parts.append('') + col = 0 + while col < column_count: + if covered[level][col]: + col += 1 + continue + + val = grid[level][col] + val_str = str(val) if val is not None else "" + if escape: + val_str = html_lib.escape(val_str) + + column_span = colspan[level][col] + row_span = rowspan[level][col] + + for row_offset in range(row_span): + for column_offset in range(column_span): + if row_offset > 0 or column_offset > 0: + if ( + level + row_offset < level_count + and col + column_offset < column_count + ): + covered[level + row_offset][col + column_offset] = True + + attrs = [] + if column_span > 1: + attrs.append(f'colspan="{column_span}"') + if row_span > 1: + attrs.append(f'rowspan="{row_span}"') + + attr_str = " " + " ".join(attrs) if attrs else "" + html_parts.append(f"{val_str}") + + col += column_span + + html_parts.append("") + + html_parts.append("") + return "".join(html_parts) + + +def render_tbody_with_row_headers( + tb_df: pd.DataFrame, + row_header_cols: int = 0, + na_rep: str = "—", + escape: bool = False, +) -> str: + """Render a DataFrame body with optional row-header cells and merged spans.""" + import html as html_lib + + if row_header_cols <= 0: + html_parts = [""] + for _, row in tb_df.iterrows(): + html_parts.append("") + for val in row: + val_str = na_rep if pd.isna(val) else str(val) + if escape: + val_str = html_lib.escape(val_str) + html_parts.append(f"{val_str}") + html_parts.append("") + html_parts.append("") + return "".join(html_parts) + + row_count = len(tb_df) + column_count = len(tb_df.columns) + + if row_count == 0: + return "" + + grid = [] + for row_idx in range(row_count): + row_values = [] + for col_idx in range(row_header_cols): + val = tb_df.iloc[row_idx, col_idx] + val = na_rep if pd.isna(val) else str(val) + row_values.append(val) + grid.append(row_values) + + colspan = [[1] * row_header_cols for _ in range(row_count)] + + for row_idx in range(row_count): + col_idx = 0 + while col_idx < row_header_cols: + span = 1 + while ( + col_idx + span < row_header_cols + and grid[row_idx][col_idx] == grid[row_idx][col_idx + span] + ): + span += 1 + colspan[row_idx][col_idx] = span + col_idx += span + + rowspan = [[1] * row_header_cols for _ in range(row_count)] + + col_idx = 0 + while col_idx < row_header_cols: + row_idx = 0 + while row_idx < row_count: + if col_idx > 0 and grid[row_idx][col_idx] == grid[row_idx][col_idx - 1]: + row_idx += 1 + continue + + current_colspan = colspan[row_idx][col_idx] + span = 1 + + while row_idx + span < row_count: + if grid[row_idx][col_idx] != grid[row_idx + span][col_idx]: + break + if colspan[row_idx + span][col_idx] != current_colspan: + break + + parent_match = True + for parent_col in range(col_idx): + if grid[row_idx][parent_col] != grid[row_idx + span][parent_col]: + parent_match = False + break + if parent_match: + span += 1 + else: + break + + rowspan[row_idx][col_idx] = span + row_idx += span + col_idx += 1 + + covered = [[False] * row_header_cols for _ in range(row_count)] + + for row_idx in range(row_count): + col_idx = 0 + while col_idx < row_header_cols: + column_span = colspan[row_idx][col_idx] + for offset in range(1, column_span): + if col_idx + offset < row_header_cols: + covered[row_idx][col_idx + offset] = True + col_idx += column_span + + for row_idx in range(row_count): + for col_idx in range(row_header_cols): + if covered[row_idx][col_idx]: + continue + row_span = rowspan[row_idx][col_idx] + for offset in range(1, row_span): + if row_idx + offset < row_count: + column_span = colspan[row_idx][col_idx] + for column_offset in range(column_span): + if col_idx + column_offset < row_header_cols: + covered[row_idx + offset][col_idx + column_offset] = True + + html_parts = [""] + + for row_idx in range(row_count): + html_parts.append("") + + for col_idx in range(row_header_cols): + if covered[row_idx][col_idx]: + continue + + val_str = grid[row_idx][col_idx] + if escape: + val_str = html_lib.escape(val_str) + + row_span = rowspan[row_idx][col_idx] + column_span = colspan[row_idx][col_idx] + + attrs = [] + if row_span > 1: + attrs.append(f'rowspan="{row_span}"') + if column_span > 1: + attrs.append(f'colspan="{column_span}"') + + attr_str = " " + " ".join(attrs) if attrs else "" + html_parts.append(f'{val_str}') + + for col_idx in range(row_header_cols, column_count): + val = tb_df.iloc[row_idx, col_idx] + val_str = na_rep if pd.isna(val) else str(val) + if escape: + val_str = html_lib.escape(val_str) + html_parts.append(f"{val_str}") + + html_parts.append("") + + html_parts.append("") + return "".join(html_parts) + + +def df2html( + tb_df: pd.DataFrame, + *, + index: bool = False, + classes: Union[str, List[str], None] = "table table-striped", + na_rep: str = "—", + escape: bool = False, + row_header_cols: int = 0, +) -> str: + """Convert a DataFrame to an HTML table.""" + class_str = ( + classes if isinstance(classes, str) else " ".join(classes) if classes else "" + ) + + if isinstance(tb_df.columns, pd.MultiIndex): + thead_html = render_multiindex_thead(tb_df.columns, escape=escape) + tbody_html = render_tbody_with_row_headers( + tb_df, row_header_cols, na_rep, escape + ) + return f'{thead_html}{tbody_html}
' + + if row_header_cols <= 0: + table_html = tb_df.to_html( + index=index, + na_rep=na_rep, + classes=classes, + escape=escape, + border=0, + justify="center", + ) + return table_html.replace("\n", "") + + html_parts = [f''] + html_parts.append("") + html_parts.append('') + for col in tb_df.columns: + col_str = str(col) if col is not None else "" + if escape: + import html as html_lib + + col_str = html_lib.escape(col_str) + html_parts.append(f"") + html_parts.append("") + html_parts.append("") + + tbody_html = render_tbody_with_row_headers(tb_df, row_header_cols, na_rep, escape) + html_parts.append(tbody_html) + html_parts.append("
{col_str}
") + return "".join(html_parts) diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index 5146e19df..956cc3b57 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -1,39 +1,38 @@ # pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false -import io import json import os -import shutil -import zipfile import pandas as pd from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.docx_asset_accumulator import DocxAssetAccumulator +from app.services.document_parser.docx_asset_store import DocxAssetStore +from app.services.document_parser.docx_block_stream import iter_block_items +from app.services.document_parser.docx_table_html import table2html from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.inline_asset import ( + build_image_asset_row, + build_table_asset_row, +) +from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder from app.services.document_parser.path_helpers import ( find_matches_parsing, process_path_texts, remove_spaces, ) -from app.services.document_parser.html_parser import table2html +from app.services.document_parser.heading_hierarchy import ( + HeadingHierarchyInput, + predict_heading_hierarchy, +) from app.services.document_parser.image_parser import ( _get_vision_client, ask_image, perceptual_hash, ) -from app.services.document_parser.layout_parser import pred_titles from app.services.document_parser.table_parser import sanitize_table_name_from_header -from app.services.document_parser.toc_parser import ( - build_docx_toc_hierarchies, - detect_doc_tocs, - detect_sdt_toc, -) +from app.services.document_parser.toc_docx import build_docx_toc_hierarchies from app.services.document_parser.txt_parser import postprocess_leaf_dics -from docx import Document -from docx.oxml.table import CT_Tbl -from docx.oxml.text.paragraph import CT_P -from docx.table import Table from docx.text.paragraph import Paragraph from loguru import logger -from lxml import etree from shared.core.config import settings from shared.core.exceptions.domain_exceptions import DocxParsingException @@ -99,7 +98,7 @@ def _find_img_context(headings_stack, max_chars=100): def handle_image( df_list, img_file, - img_dir, + asset_store, headings_stack, current_heading, img_count, @@ -115,19 +114,13 @@ def handle_image( cached = seen_images[img_hash] headings_stack[-1]["content"].append(cached["image_ref"]) df_list.append( - [ - cached["image_ref"], - cached["img_path"], - "image", - len(cached["image_ref"]), - "", - cached["img_summary_field"], - cached["temp_uid"], - "", - "", - time_stamp, - "", - ] + build_image_asset_row( + content=cached["image_ref"], + relative_path=cached["img_path"], + summary=cached["img_summary_field"], + know_id=cached["temp_uid"], + addtime=time_stamp, + ).to_list() ) logger.debug(f"Skipped duplicate image (hash={img_hash[:12]}...)") return headings_stack, df_list, False # False = cache hit, don't increment @@ -142,10 +135,7 @@ def handle_image( raw_img_name = process_path_texts( f"image-{str(img_count + 1)} {current_heading} {last_context}", last=30 ) - img_raw_path = os.path.join(img_dir, f"{raw_img_name}{img_ext}") - - with open(img_raw_path, "wb") as image_file: - image_file.write(img_file["data"]) + raw_image_asset = asset_store.write_image(raw_img_name, img_ext, img_file["data"]) # LLM title + summary (optional, with fallback to last_context) llm_title = None @@ -156,7 +146,10 @@ def handle_image( # TODO: Risk of missing text content if the image is a screenshot of pure text. # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. llm_resp = ask_image( - client, img_dir, [f"{raw_img_name}{img_ext}"], title_text=last_context + client, + asset_store.image_dir, + [f"{raw_img_name}{img_ext}"], + title_text=last_context, ) if llm_resp: llm_title, llm_summary = split_title_summary(llm_resp) @@ -175,8 +168,7 @@ def handle_image( img_name = process_path_texts( f"image-{str(img_count + 1)} {current_heading} {img_title or ''}", last=30 ) - img_path = os.path.join(img_dir, f"{img_name}{img_ext}") - os.rename(img_raw_path, img_path) # if summary fails, renaming is not applied + image_asset = asset_store.rename_image(raw_image_asset, img_name) temp_uid = gen_str_codes(img_hash) @@ -186,8 +178,7 @@ def handle_image( else: img_summary_field = image_index - img_path = f"images/{img_name}{img_ext}" - img_ref = build_chunk_ref(img_path) + img_ref = build_chunk_ref(image_asset.relative_path) # Build image_ref for heading_stack: optional summary + image path ref if img_summary: @@ -197,25 +188,19 @@ def handle_image( headings_stack[-1]["content"].append(image_ref) df_list.append( - [ - image_ref, - img_path, - "image", - len(image_ref), - "", - img_summary_field, - temp_uid, - "", - "", - time_stamp, - "", - ] + build_image_asset_row( + content=image_ref, + relative_path=image_asset.relative_path, + summary=img_summary_field, + know_id=temp_uid, + addtime=time_stamp, + ).to_list() ) # Cache result for document-level dedup if seen_images is not None: seen_images[img_hash] = { - "img_path": img_path, + "img_path": image_asset.relative_path, "image_ref": image_ref, "img_summary_field": img_summary_field, "temp_uid": temp_uid, @@ -274,14 +259,13 @@ def _first_cols_rows(table_block, max_items=10, max_chars=20): def handle_table( df_list, block, - tb_dir, + asset_store, headings_stack, current_heading, table_count, summary_table=False, summary_image=False, cell_images=None, - img_dir=None, img_count=0, seen_images=None, ): @@ -328,9 +312,9 @@ def handle_table( img_name = process_path_texts( f"table-{table_count + 1}-{image_index} {current_heading}", last=30 ) - img_save_path = os.path.join(img_dir, f"{img_name}{img_ext}") - with open(img_save_path, "wb") as f: - f.write(img_data["data"]) + image_asset = asset_store.write_image( + img_name, img_ext, img_data["data"] + ) # LLM summary (optional) img_summary = None @@ -339,7 +323,7 @@ def handle_table( client = _get_vision_client() img_summary = ask_image( client, - img_dir, + asset_store.image_dir, [f"{img_name}{img_ext}"], title_text=current_heading, ) @@ -354,32 +338,25 @@ def handle_table( img_summary_field = ( f"{image_index}\n{img_summary}" if img_summary else image_index ) - relative_img_path = f"images/{img_name}{img_ext}" - img_ref = build_chunk_ref(relative_img_path) + img_ref = build_chunk_ref(image_asset.relative_path) if img_summary: image_ref = f"\n{img_summary}\n{img_ref}\n" else: image_ref = f"\n{img_ref}\n" table_img_entries.append( - [ - image_ref, - relative_img_path, - "image", - len(image_ref), - "", - img_summary_field, - temp_uid, - "", - "", - time_stamp, - "", - ] + build_image_asset_row( + content=image_ref, + relative_path=image_asset.relative_path, + summary=img_summary_field, + know_id=temp_uid, + addtime=time_stamp, + ).to_list() ) # Cache result for document-level dedup if seen_images is not None: seen_images[cell_img_hash] = { - "img_path": relative_img_path, + "img_path": image_asset.relative_path, "image_ref": image_ref, "img_summary_field": img_summary_field, "temp_uid": temp_uid, @@ -391,7 +368,6 @@ def handle_table( f"Extracted {sum(len(v) for v in cell_images.values())} images from table-{table_count + 1} cells" ) - # Generate HTML with image descriptions embedded tb_html_str = table2html( block, cell_image_map=cell_image_map if cell_image_map else None @@ -437,14 +413,8 @@ def handle_table( tb_name = path_handle( f"table-{str(table_count + 1)} {effective_name}", mode="clean_single" ) - tb_path = os.path.join(tb_dir, f"{tb_name}.html") - - with open(tb_path, "w", encoding="utf-8") as f: - f.write(tb_html_str) - - # Use relative path for tables (avoid absolute path in path column) - tb_path = f"tables/{tb_name}.html" - tb_ref = build_chunk_ref(tb_path) + table_asset = asset_store.write_table(tb_name, tb_html_str) + tb_ref = build_chunk_ref(table_asset.relative_path) # Build table_ref for heading_stack: optional LLM summary + table path ref if llm_summary: table_ref = f"\n{llm_summary}\n{tb_ref}\n" @@ -452,323 +422,18 @@ def handle_table( table_ref = f"\n{tb_ref}\n" headings_stack[-1]["content"].append(table_ref) df_list.append( - [ - tb_html_str, - tb_path, - "table", - len(tb_html_str), - tb_keywords, - tb_summary, - temp_uid, - "", - "", - time_stamp, - "", - ] + build_table_asset_row( + content=tb_html_str, + relative_path=table_asset.relative_path, + summary=tb_summary, + keywords=tb_keywords, + know_id=temp_uid, + addtime=time_stamp, + ).to_list() ) return headings_stack, df_list, img_count -def iter_block_items(doc_data): - doc_stream = io.BytesIO(doc_data) - doc = Document(doc_stream) - - # python-docx mapping - p_tbl_map = [] - for child in doc.element.body: - if isinstance(child, CT_P): - p_tbl_map.append(("p", child)) - elif isinstance(child, CT_Tbl): - p_tbl_map.append(("tbl", child)) - - with zipfile.ZipFile(io.BytesIO(doc_data), "r") as docx: - xml = docx.read("word/document.xml") - rels = etree.fromstring(docx.read("word/_rels/document.xml.rels")) - rel_map = { - r.get("Id"): r.get("Target") for r in rels.findall(".//{*}Relationship") - } - ns = { - "a": "http://schemas.openxmlformats.org/drawingml/2006/main", - "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", - "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main", - "v": "urn:schemas-microsoft-com:vml", - "o": "urn:schemas-microsoft-com:office:office", - } - r_ns = "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}" - - root = etree.fromstring(xml) - body = root.find(".//w:body", namespaces=ns) - - ele_num = 1 - map_index = 0 # point to p_tbl_map - toc_field_active = False - - for elem in body.iterchildren(): - if not isinstance(elem.tag, str): - continue - - tag = etree.QName(elem.tag).localname - - # --- SDT (Structured Document Tag) container --- - # TOC generated by MS Word is usually in sdt - if tag == "sdt": - sdt_toc_info = detect_sdt_toc(elem, ns) - is_toc_sdt = sdt_toc_info["is_toc_sdt"] - - sdt_content = elem.find(".//w:sdtContent", namespaces=ns) - if sdt_content is not None: - for p_elem in sdt_content.findall(".//w:p", namespaces=ns): - texts = p_elem.xpath(".//w:t/text()", namespaces=ns) - text = "".join(texts).strip() - - if is_toc_sdt: - label = "TOC-AREA" - toc_info = detect_doc_tocs(p_elem, ns) - else: - toc_info = detect_doc_tocs(p_elem, ns) - if toc_info["is_style"] or toc_info["is_field_start"]: - label = "TOC-AREA" - else: - label = "PTXT" - - if text: - meta = None - if "TOC" in label: - meta = { - "toc_level": toc_info.get("toc_level"), - "toc_outline_level": toc_info.get("outline_level"), - "toc_left_indent": toc_info.get("left_indent"), - "toc_style_name": toc_info.get("style_name"), - "toc_source": "sdt", - } - yield ele_num, text, label, meta - ele_num += 1 - continue - - # --- text paras --- - if tag == "p": - texts = elem.xpath(".//w:t/text()", namespaces=ns) - text = "".join(texts).strip() - - if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "p": - p_obj = Paragraph(p_tbl_map[map_index][1], doc) - else: - p_obj = None - - toc_info = detect_doc_tocs(elem, ns) - if toc_info["is_field_start"]: - toc_field_active = True - - if toc_info["is_style"] or toc_field_active: - label = "TOC-AREA" - else: - label = "PTXT" - - if text or p_obj is not None: - meta = None - if "TOC" in label: - meta = { - "toc_level": toc_info.get("toc_level"), - "toc_outline_level": toc_info.get("outline_level"), - "toc_left_indent": toc_info.get("left_indent"), - "toc_style_name": toc_info.get("style_name"), - "toc_source": "paragraph", - } - yield ele_num, p_obj or text, label, meta - ele_num += 1 - - # images (DrawingML: ) - seen_rids = set() - blips = elem.xpath(".//a:blip", namespaces=ns) - for b in blips: - rid = b.get(f"{r_ns}embed") - if not rid or rid in seen_rids: - continue - target = rel_map.get(rid) - if not target or not target.startswith("media/"): - continue - seen_rids.add(rid) - data = docx.read("word/" + target) - yield ( - ele_num, - None, - "IMAGE", - { - "image_name": target.split("/")[-1], - "from": "paragraph", - "size": len(data), - "data": data, - }, - ) - ele_num += 1 - - # TODO: Re-evaluate VML group extraction strategy. - # Complex VML composite images () are currently skipped because extracting - # piece-by-piece loses textual overlay and positioning. - # Future plan: Use LibreOffice headless conversion to render the entire document - # and map the perfectly rendered images back to the layout via text anchors. - # - # Temporary: detect VML-only paragraphs and inject a placeholder so the - # paragraph isn't silently swallowed, leaving its parent section empty. - if not text and not seen_rids: - # No text and no DrawingML images — check for VML content - vml_groups = elem.xpath(".//v:group", namespaces=ns) - vml_images_check = elem.xpath(".//v:imagedata", namespaces=ns) - if vml_groups or vml_images_check: - vml_placeholder = "[VML graphic \u2014 extraction not yet supported]" - yield ele_num, vml_placeholder, "PTXT", None - ele_num += 1 - logger.debug( - f"Injected VML placeholder for paragraph with " - f"{len(vml_groups)} v:group, {len(vml_images_check)} v:imagedata" - ) - """ - # images (VML: ) — convert to PNG - from PIL import Image as PILImage - - vml_images = elem.xpath(".//v:imagedata", namespaces=ns) - for v in vml_images: - rid = v.get(f"{r_ns}id") - if not rid or rid in seen_rids: - continue - target = rel_map.get(rid) - if not target or not target.startswith("media/"): - continue - seen_rids.add(rid) - raw_data = docx.read("word/" + target) - # Convert to PNG for uniform downstream handling - try: - pil_img = PILImage.open(io.BytesIO(raw_data)) - png_buf = io.BytesIO() - pil_img.save(png_buf, format="PNG") - png_data = png_buf.getvalue() - except Exception as e: - logger.warning(f"Failed to convert VML image to PNG: {e}") - continue - orig_name = target.split("/")[-1] - png_name = os.path.splitext(orig_name)[0] + ".png" - yield ( - ele_num, - None, - "IMAGE", - { - "image_name": png_name, - "from": "paragraph_vml", - "size": len(png_data), - "data": png_data, - }, - ) - ele_num += 1 - """ - map_index += 1 - - if toc_info["is_field_end"]: - toc_field_active = False - - # --- tables --- - elif tag == "tbl": - if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "tbl": - tbl = Table(p_tbl_map[map_index][1], doc) - else: - tbl = Table(elem, doc) - - # Extract images from each cell, keyed by (row_idx, col_idx) - cell_images = {} # {(row_idx, col_idx): [{'image_name', 'data', 'size'}]} - for row_idx, tr in enumerate(elem.findall(".//w:tr", namespaces=ns)): - for col_idx, tc in enumerate(tr.findall(".//w:tc", namespaces=ns)): - cell_seen_rids = set() - imgs_in_cell = [] - # DrawingML images in cell - blips = tc.xpath(".//a:blip", namespaces=ns) - for b in blips: - rid = b.get(f"{r_ns}embed") - if not rid or rid in cell_seen_rids: - continue - target = rel_map.get(rid) - if not target or not target.startswith("media/"): - continue - cell_seen_rids.add(rid) - data = docx.read("word/" + target) - if ( - len(data) < 10 * 1024 - ): # Skip small images (<10KB, likely icons) - continue - imgs_in_cell.append( - { - "image_name": target.split("/")[-1], - "data": data, - "size": len(data), - } - ) - # TODO: VML in tables is temporarily skipped to avoid extracting - # fragmented textless background images. (Same as paragraph VML logic) - """ - # VML images in cell — convert to PNG - from PIL import Image as PILImage - - vml_in_cell = tc.xpath(".//v:imagedata", namespaces=ns) - for v in vml_in_cell: - rid = v.get(f"{r_ns}id") - if not rid or rid in cell_seen_rids: - continue - target = rel_map.get(rid) - if not target or not target.startswith("media/"): - continue - cell_seen_rids.add(rid) - raw_data = docx.read("word/" + target) - try: - pil_img = PILImage.open(io.BytesIO(raw_data)) - png_buf = io.BytesIO() - pil_img.save(png_buf, format="PNG") - png_data = png_buf.getvalue() - except Exception as e: - logger.warning( - f"Failed to convert VML cell image to PNG: {e}" - ) - continue - if len(png_data) < 10 * 1024: - continue - orig_name = target.split("/")[-1] - png_name = os.path.splitext(orig_name)[0] + ".png" - imgs_in_cell.append( - { - "image_name": png_name, - "data": png_data, - "size": len(png_data), - } - ) - """ - if imgs_in_cell: - cell_images[(row_idx, col_idx)] = imgs_in_cell - - yield ele_num, tbl, "TABLE", cell_images if cell_images else None - ele_num += 1 - map_index += 1 - else: - continue - - # --- handle p_tbl_map at the end --- - while map_index < len(p_tbl_map): - tag, node = p_tbl_map[map_index] - if tag == "p": - toc_info = detect_doc_tocs(node, ns) - label = "TOC-AREA" if toc_info["is_style"] else "PTXT" - meta = None - if "TOC" in label: - meta = { - "toc_level": toc_info.get("toc_level"), - "toc_outline_level": toc_info.get("outline_level"), - "toc_left_indent": toc_info.get("left_indent"), - "toc_style_name": toc_info.get("style_name"), - "toc_source": "tail-map", - } - yield ele_num, Paragraph(node, doc), label, meta - elif tag == "tbl": - yield ele_num, Table(node, doc), "TABLE", None - ele_num += 1 - map_index += 1 - - def parse_docx( docx_path, llm_paras, @@ -786,16 +451,8 @@ def parse_docx( headings_stack = [{"level": -1, "content": doc_structure}] current_heading = "" - # Clean old artifacts to prevent accumulation across debug runs. - # In production each job uses a fresh workspace so rmtree never triggers. - tb_dir = os.path.join(output_dir, "tables") - if os.path.isdir(tb_dir): - shutil.rmtree(tb_dir) - os.makedirs(tb_dir, exist_ok=True) - img_dir = os.path.join(output_dir, "images") - if os.path.isdir(img_dir): - shutil.rmtree(img_dir) - os.makedirs(img_dir, exist_ok=True) + asset_store = DocxAssetStore(output_dir) + asset_store.reset() block_tuples = list(iter_block_items(doc_data)) # Record first TOC block position before filtering, for pre-TOC exclusion in pred_titles @@ -829,15 +486,17 @@ def parse_docx( if llm_paras else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) ) - heading_candidates = pred_titles( - heading_infos, - doc_type="docx", - toc_hierarchies=toc_hierarchies or None, - enable_regx=True, - smart_parse=smart_title_parse, - model_name=model_name, - output_dir=output_dir, - first_toc_ele_num=first_toc_ele_num, + heading_candidates = predict_heading_hierarchy( + HeadingHierarchyInput( + infos=heading_infos, + doc_type="docx", + toc_hierarchies=toc_hierarchies or None, + enable_regex=True, + smart_parse=smart_title_parse, + model_name=model_name, + output_dir=output_dir, + first_toc_ele_num=first_toc_ele_num, + ) ) if len(heading_candidates) > 0 and not (heading_candidates["level"] == -1).all(): @@ -853,10 +512,13 @@ def parse_docx( headings_stack.append(new_content) logger.debug("⚠️no headings detected, using file name or mine a heading=>", text) - df_list = [] - table_count = 0 - image_count = 0 - _seen_images: dict[str, dict] = {} # sha256_hex -> cached image info for dedup + asset_accumulator = DocxAssetAccumulator( + asset_store=asset_store, + should_summary_image=llm_paras["summary_image"], + should_summary_table=llm_paras["summary_table"], + image_handler=handle_image, + table_handler=handle_table, + ) logger.debug("Parsing docx file... total_blocks={}", len(block_tuples)) for block_tuple in block_tuples: @@ -894,43 +556,27 @@ def parse_docx( if meta and meta.get("size", 0) < 10 * 1024: continue - headings_stack, df_list, is_new = handle_image( - df_list, + headings_stack = asset_accumulator.append_image( meta, - img_dir, headings_stack, current_heading, - image_count, - llm_paras["summary_image"], - seen_images=_seen_images, ) - if is_new: - image_count += 1 current_heading = last_heading_before_block elif label == "TABLE": # TODO: handle cross-page tables - headings_stack, df_list, image_count = handle_table( - df_list, + headings_stack = asset_accumulator.append_table( block, - tb_dir, headings_stack, current_heading, - table_count, - summary_table=llm_paras["summary_table"], - summary_image=llm_paras["summary_image"], cell_images=meta, - img_dir=img_dir, - img_count=image_count, - seen_images=_seen_images, ) - table_count += 1 current_heading = last_heading_before_block else: # TODO: handle latex, etc. pass - return {"content": doc_structure}, df_list + return {"content": doc_structure}, asset_accumulator.rows() def convert_doc2dics( @@ -998,19 +644,16 @@ def convert_doc2dics( else (relative_root or path_suffix) ) df_list.append( - [ - bottom_content, - know_path, - match_type, - len(bottom_content), - keywords, - summary, - know_id, - bottom_tokens, - "", - time_stamp, - "", - ] + ParsedRow( + content=bottom_content, + path=know_path, + type=match_type, + keywords=keywords, + summary=summary, + know_id=know_id, + tokens=bottom_tokens, + addtime=time_stamp, + ).to_list() ) except KnowhereException: raise @@ -1023,6 +666,23 @@ def convert_doc2dics( original_exception=e, ) - doc_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(",")) + rows_builder = ParsedRowsBuilder() + for row_values in df_list: + rows_builder.append( + ParsedRow( + content=str(row_values[0]), + path=str(row_values[1]), + type=str(row_values[2]), + length=int(row_values[3]), + keywords=str(row_values[4]), + summary=str(row_values[5]), + know_id=str(row_values[6]), + tokens=str(row_values[7]), + connectto=str(row_values[8]), + addtime=str(row_values[9]), + page_nums=str(row_values[10]), + ) + ) + doc_df = rows_builder.to_dataframe() doc_df = process_dup_paths_df(doc_df) return doc_df diff --git a/apps/worker/app/services/document_parser/doc_profile_model.py b/apps/worker/app/services/document_parser/doc_profile_model.py new file mode 100644 index 000000000..ad1a12e65 --- /dev/null +++ b/apps/worker/app/services/document_parser/doc_profile_model.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import gc +import json +import os +from dataclasses import asdict, dataclass, field +from typing import List, Literal, Optional + +from loguru import logger + + +@dataclass +class DocProfile: + """Document profile data contract used by parser routing.""" + + file_type: str = "" + route: Literal["fast", "standard"] = "standard" + decision_band: Literal["safe_fast", "gray_zone", "safe_standard"] = "safe_standard" + scan_type: Optional[Literal["electronic", "scanned", "mixed"]] = None + doc_category: Literal["generic", "atlas", "ppt_converted"] = "generic" + page_count: int = 0 + avg_text_density: float = 0.0 + avg_image_coverage: float = 0.0 + has_tables: bool = False + has_embedded_fonts: bool = False + is_multi_column: bool = False + is_degraded_electronic: bool = False + sample_text: str = "" + has_significant_images: bool = False + significant_image_count: int = 0 + max_image_coverage_on_page: float = 0.0 + pages_with_significant_images: int = 0 + large_image_page_ratio: float = 0.0 + table_signal_pages: int = 0 + table_signal_strength: float = 0.0 + complex_pages: int = 0 + complex_page_ratio: float = 0.0 + max_drawing_count: int = 0 + min_text_density_page: float = 0.0 + text_density_std: float = 0.0 + estimated_fast_benefit: float = 0.0 + estimated_risk_score: float = 0.0 + atlas_candidate: bool = False + page_details: List[dict] = field(default_factory=list) + reasoning: str = "" + + def to_dict(self) -> dict: + data = asdict(self) + data.pop("page_details", None) + data.pop("sample_text", None) + return data + + def summary(self) -> str: + parts = ( + f"[{self.file_type.upper()}] route={self.route}, band={self.decision_band}, " + f"scan={self.scan_type}, category={self.doc_category}, " + f"pages={self.page_count}, text_density={self.avg_text_density:.0f}, " + f"img_coverage={self.avg_image_coverage:.1%}, " + f"risk={self.estimated_risk_score:.2f}, gain={self.estimated_fast_benefit:.2f}" + ) + if self.is_degraded_electronic: + parts += ", degraded=True" + return parts + + +def publish_profile_result(queue, profile: DocProfile) -> None: + gc.collect() + queue.put({"ok": True, "profile": asdict(profile)}) + + +def save_profile_metadata(profile: DocProfile, output_dir: str) -> None: + profile_path = os.path.join(output_dir, "profile.json") + with open(profile_path, "w", encoding="utf-8") as file_obj: + json.dump(profile.to_dict(), file_obj, ensure_ascii=False, indent=2) + logger.debug(f"Profile metadata saved to {profile_path}") diff --git a/apps/worker/app/services/document_parser/doc_profile_pdf.py b/apps/worker/app/services/document_parser/doc_profile_pdf.py new file mode 100644 index 000000000..216613372 --- /dev/null +++ b/apps/worker/app/services/document_parser/doc_profile_pdf.py @@ -0,0 +1,676 @@ +# pyright: reportAttributeAccessIssue=false, reportOperatorIssue=false +from __future__ import annotations + +import math +from typing import Any + +from app.services.document_parser.doc_profile_model import ( + DocProfile, + publish_profile_result, +) +from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from loguru import logger + +# Thresholds +SCAN_TEXT_THRESHOLD = 50 +SCAN_IMAGE_COVERAGE_MIN = 0.6 +SCAN_PAGE_RATIO = 0.7 + +ATLAS_TEXT_THRESHOLD = 200 +ATLAS_CANDIDATE_IMAGE_COVERAGE_MIN = 0.30 +ATLAS_MIN_LANDSCAPE_RATIO = 0.5 # ≥50% of sampled pages must be landscape +ATLAS_MIN_PAGES = 2 # single-page scans (resumes, posters) are not atlases + +FAST_TEXT_THRESHOLD = 500 +MIN_FAST_TEXT_DENSITY_FLOOR = 120 +SAFE_FAST_MAX_PAGE_COUNT = 80 +HARD_STANDARD_PAGE_COUNT = 150 + +MULTI_COL_GAP_RATIO = 0.15 +MULTI_COL_MIN_BLOCKS = 4 + +DEGRADED_SKINNY_ASPECT = 50 +DEGRADED_SKINNY_MAX_H = 30 +DEGRADED_SKINNY_MIN_PER_PAGE = 50 +DEGRADED_PAGE_RATIO = 0.5 + +SIGNIFICANT_IMAGE_AREA_RATIO = 0.12 +MEDIUM_IMAGE_AREA_RATIO = 0.03 +LARGE_IMAGE_PAGE_RATIO = 0.25 +SIGNIFICANT_IMAGE_MIN_DIM = 400 +SIGNIFICANT_IMAGE_MIN_PIXELS = 250_000 + +PROFILE_MAX_NEW_XREFS_PER_PAGE = 30 + +TABLE_DRAWING_LINE_THRESHOLD = 12 +TABLE_DRAWING_STRONG_THRESHOLD = 18 +TABLE_DRAWING_RECT_THRESHOLD = 2 + +SAFE_FAST_MAX_COMPLEX_PAGE_RATIO = 0.05 +SAFE_FAST_MAX_IMAGE_COVERAGE_ON_PAGE = 0.08 +SAFE_FAST_MAX_AVG_IMAGE_COVERAGE = 0.03 +SAFE_FAST_MAX_TEXT_STD = 600.0 +HARD_COMPLEX_PAGE_RATIO = 0.2 +HARD_SIGNIFICANT_IMAGE_PAGES = 3 +HARD_LARGE_IMAGE_PAGE_RATIO = 0.15 + + +def _clamp(value: float, min_value: float = 0.0, max_value: float = 1.0) -> float: + return max(min_value, min(max_value, value)) + + +def _stddev(values: list[float]) -> float: + if not values: + return 0.0 + mean = sum(values) / len(values) + variance = sum((value - mean) ** 2 for value in values) / len(values) + return math.sqrt(variance) + + +def _count_detected_tables(page: Any) -> int: + try: + finder = page.find_tables() + except Exception: + return 0 + + if not finder: + return 0 + + tables = getattr(finder, "tables", finder) + try: + return len(tables) + except TypeError: + return 1 if tables else 0 + + +def _is_stroked_drawing(drawing: dict[str, Any]) -> bool: + stroke_width = drawing.get("width") + return drawing.get("color") is not None or ( + stroke_width is not None and stroke_width > 0 + ) + + +def _estimate_fast_benefit(profile: DocProfile) -> float: + if profile.page_count <= 2: + page_factor = 0.35 + elif profile.page_count <= 10: + page_factor = 0.7 + elif profile.page_count <= SAFE_FAST_MAX_PAGE_COUNT: + page_factor = 1.0 + elif profile.page_count <= HARD_STANDARD_PAGE_COUNT: + page_factor = 0.8 + else: + page_factor = 0.45 + + density_factor = _clamp(profile.avg_text_density / 1200.0) + stability_factor = _clamp( + 1.0 + - (profile.complex_page_ratio * 1.5) + - (profile.large_image_page_ratio * 1.2) + - (profile.table_signal_strength * 0.8) + ) + return _clamp( + (0.35 * page_factor) + (0.40 * density_factor) + (0.25 * stability_factor) + ) + + +def _estimate_risk_score(profile: DocProfile) -> float: + risk = 0.0 + if profile.scan_type != "electronic": + risk += 0.35 + if profile.doc_category != "generic": + risk += 0.20 + if profile.is_multi_column: + risk += 0.20 + if profile.is_degraded_electronic: + risk += 0.20 + if profile.has_tables: + risk += 0.30 + + risk += min(0.20, profile.large_image_page_ratio * 1.2) + risk += min(0.20, profile.complex_page_ratio * 0.8) + risk += min(0.15, profile.table_signal_strength * 0.2) + risk += min(0.12, profile.pages_with_significant_images * 0.04) + + if profile.page_count > HARD_STANDARD_PAGE_COUNT: + risk += 0.10 + + return _clamp(risk) + + +def _classify_route(profile: DocProfile) -> tuple[str, str, float, float, list[str]]: + hard_gate_reasons: list[str] = [] + + if profile.scan_type != "electronic": + hard_gate_reasons.append(f"scan_type={profile.scan_type}") + if profile.doc_category != "generic": + hard_gate_reasons.append(f"doc_category={profile.doc_category}") + if profile.is_multi_column: + hard_gate_reasons.append("multi_column") + if profile.is_degraded_electronic: + hard_gate_reasons.append("degraded_electronic") + if profile.has_tables: + hard_gate_reasons.append( + f"table_signals={profile.table_signal_pages}p/{profile.table_signal_strength:.2f}" + ) + if ( + profile.max_image_coverage_on_page >= LARGE_IMAGE_PAGE_RATIO + or profile.pages_with_significant_images >= HARD_SIGNIFICANT_IMAGE_PAGES + or profile.large_image_page_ratio >= HARD_LARGE_IMAGE_PAGE_RATIO + ): + hard_gate_reasons.append( + "significant_images=" + f"{profile.pages_with_significant_images}p,max={profile.max_image_coverage_on_page:.1%}" + ) + if profile.complex_page_ratio >= HARD_COMPLEX_PAGE_RATIO: + hard_gate_reasons.append(f"complex_pages={profile.complex_page_ratio:.0%}") + if profile.page_count > HARD_STANDARD_PAGE_COUNT: + hard_gate_reasons.append( + f"page_count={profile.page_count}>{HARD_STANDARD_PAGE_COUNT}" + ) + + benefit = _estimate_fast_benefit(profile) + risk = _estimate_risk_score(profile) + + if hard_gate_reasons: + return ( + "standard", + "safe_standard", + benefit, + risk, + [ + "decision=safe_standard: hard gate matched", + "hard_gates=" + ",".join(hard_gate_reasons), + ], + ) + + safe_fast_checks = [ + ( + profile.page_count <= SAFE_FAST_MAX_PAGE_COUNT, + f"page_count={profile.page_count}<={SAFE_FAST_MAX_PAGE_COUNT}", + ), + ( + profile.avg_text_density >= MIN_FAST_TEXT_DENSITY_FLOOR, + "text_density_floor=" + f"{profile.avg_text_density:.0f}>={MIN_FAST_TEXT_DENSITY_FLOOR}", + ), + ( + not profile.has_significant_images, + f"has_significant_images={profile.has_significant_images}", + ), + ( + profile.max_image_coverage_on_page <= SAFE_FAST_MAX_IMAGE_COVERAGE_ON_PAGE, + "max_image_coverage_on_page=" + f"{profile.max_image_coverage_on_page:.1%}<={SAFE_FAST_MAX_IMAGE_COVERAGE_ON_PAGE:.0%}", + ), + ( + profile.avg_image_coverage <= SAFE_FAST_MAX_AVG_IMAGE_COVERAGE, + f"avg_image_coverage={profile.avg_image_coverage:.1%}<={SAFE_FAST_MAX_AVG_IMAGE_COVERAGE:.0%}", + ), + ( + profile.complex_page_ratio <= SAFE_FAST_MAX_COMPLEX_PAGE_RATIO, + f"complex_page_ratio={profile.complex_page_ratio:.0%}<={SAFE_FAST_MAX_COMPLEX_PAGE_RATIO:.0%}", + ), + ( + profile.text_density_std <= SAFE_FAST_MAX_TEXT_STD, + f"text_density_std={profile.text_density_std:.0f}<={SAFE_FAST_MAX_TEXT_STD:.0f}", + ), + ( + risk <= 0.35, + f"estimated_risk_score={risk:.2f}<=0.35", + ), + ] + + failed_checks = [reason for passed, reason in safe_fast_checks if not passed] + if not failed_checks: + return ( + "fast", + "safe_fast", + benefit, + risk, + [ + "decision=safe_fast: low-complexity high-yield pdf", + "safe_fast_checks_passed", + ], + ) + + return ( + "standard", + "gray_zone", + benefit, + risk, + [ + "decision=gray_zone: conservative fallback to standard in phase1", + "borderline=" + ",".join(failed_checks[:4]), + ], + ) + + + +@worker +def _profile_pdf_worker(queue, file_path: str) -> None: + """Child process: analyze PDF features, return profile as dict.""" + import pymupdf + + profile = DocProfile(file_type="pdf") + reasons: list[str] = [] + + try: + doc = pymupdf.open(file_path) + except Exception as exc: + profile.reasoning = f"Cannot open file: {exc}" + publish_profile_result(queue, profile) + return + + profile.page_count = doc.page_count + + if doc.page_count == 0: + profile.reasoning = "Empty file (0 pages)" + doc.close() + del doc + publish_profile_result(queue, profile) + return + + if doc.page_count <= 50: + sample_indices = list(range(doc.page_count)) + else: + step = max(1, doc.page_count // 20) + sample_indices = list(range(0, doc.page_count, step))[:20] + sample_indices = sorted( + set( + sample_indices + + [0, 1, 2] + + [doc.page_count - 3, doc.page_count - 2, doc.page_count - 1] + ) + ) + sample_indices = [idx for idx in sample_indices if 0 <= idx < doc.page_count] + + page_details = [] + text_lengths: list[float] = [] + total_text_len = 0 + total_image_coverage = 0.0 + scanned_pages = 0 + all_text_parts: list[str] = [] + has_any_fonts = False + has_any_tables = False + table_signal_pages = 0 + total_table_signal_strength = 0.0 + multi_col_pages = 0 + landscape_pages = 0 + degraded_pages = 0 + doc_page_sizes = [] + + significant_image_count = 0 + pages_with_significant_images = 0 + large_image_pages = 0 + max_image_coverage_on_page = 0.0 + + complex_pages = 0 + max_drawing_count = 0 + + # Track xrefs already processed across pages to avoid redundant + # get_image_rects() calls on shared/inherited image resources. + # PDFs with shared xrefs (e.g. scanned docs) report ALL document + # images on every page; without dedup this causes O(pages × images) + # content-stream scans. + seen_xrefs: set = set() + + for idx in sample_indices: + page = doc[idx] + page_width = page.rect.width + page_height = page.rect.height + page_area = page_width * page_height + + if page_width > page_height: + landscape_pages += 1 + doc_page_sizes.append((page_width, page_height)) + + text = page.get_text().strip() + text_len = len(text) + text_lengths.append(float(text_len)) + total_text_len += text_len + + if len("".join(all_text_parts)) < 500: + all_text_parts.append(text[:200]) + + images = page.get_images(full=True) + img_total_area = 0.0 + page_significant_image_count = 0 + page_max_rect_ratio = 0.0 + page_medium_image_coverage = 0.0 + skinny_count = 0 + + new_xref_count = 0 + for img in images: + xref = img[0] + img_w, img_h = img[2], img[3] + if ( + img_h > 0 + and img_w / img_h > DEGRADED_SKINNY_ASPECT + and img_h < DEGRADED_SKINNY_MAX_H + ): + skinny_count += 1 + + # ── xref dedup: skip images already analyzed on earlier pages ── + if xref in seen_xrefs: + continue + seen_xrefs.add(xref) + new_xref_count += 1 + # Cap expensive get_image_rects calls per page + if new_xref_count > PROFILE_MAX_NEW_XREFS_PER_PAGE: + continue + + try: + rects = page.get_image_rects(xref) + except Exception: + rects = [] + + for rect in rects: + rect_area = rect.width * rect.height + img_total_area += rect_area + + area_ratio = rect_area / page_area if page_area > 0 else 0.0 + page_max_rect_ratio = max(page_max_rect_ratio, area_ratio) + + is_significant = ( + area_ratio >= SIGNIFICANT_IMAGE_AREA_RATIO + or ( + area_ratio >= 0.05 + and ( + max(img_w, img_h) >= SIGNIFICANT_IMAGE_MIN_DIM + or (img_w * img_h) >= SIGNIFICANT_IMAGE_MIN_PIXELS + ) + ) + or ( + area_ratio >= 0.02 + and (img_w * img_h) >= (SIGNIFICANT_IMAGE_MIN_PIXELS * 2) + ) + ) + + if is_significant: + page_significant_image_count += 1 + elif area_ratio >= MEDIUM_IMAGE_AREA_RATIO: + page_medium_image_coverage += area_ratio + + if skinny_count >= DEGRADED_SKINNY_MIN_PER_PAGE: + degraded_pages += 1 + + img_coverage = img_total_area / page_area if page_area > 0 else 0.0 + img_coverage = min(img_coverage, 1.0) + total_image_coverage += img_coverage + + fonts = page.get_fonts() + if fonts: + has_any_fonts = True + + drawings = page.get_drawings() + drawing_count = len(drawings) + max_drawing_count = max(max_drawing_count, drawing_count) + line_like_items = 0 + horizontal_line_items = 0 + vertical_line_items = 0 + rect_items = 0 + fill_only_rect_items = 0 + for drawing in drawings: + is_stroked = _is_stroked_drawing(drawing) + for item in drawing.get("items", []): + if item[0] == "l": + if is_stroked: + line_like_items += 1 + point_a = item[1] + point_b = item[2] + if abs(point_a.y - point_b.y) <= 2: + horizontal_line_items += 1 + if abs(point_a.x - point_b.x) <= 2: + vertical_line_items += 1 + elif item[0] == "re": + if is_stroked: + rect_items += 1 + line_like_items += 4 + horizontal_line_items += 2 + vertical_line_items += 2 + else: + fill_only_rect_items += 1 + + detected_table_count = _count_detected_tables(page) + drawing_table_signal = line_like_items >= TABLE_DRAWING_LINE_THRESHOLD and ( + (horizontal_line_items >= 2 and vertical_line_items >= 2) + or rect_items >= TABLE_DRAWING_RECT_THRESHOLD + or ( + line_like_items >= TABLE_DRAWING_STRONG_THRESHOLD + and horizontal_line_items >= 3 + and vertical_line_items >= 3 + ) + ) + # NOTE: + # `page.find_tables()` produces too many false positives on Word / Writer + # exported pure-text PDFs, where paragraph background boxes are inferred as + # full-page tables. For Phase 1 fast-path routing, keep `find_tables()` + # only as debug evidence and rely on explicit drawing-grid signals for + # table hard gates. + table_hit = drawing_table_signal + page_table_strength = 0.0 + if drawing_table_signal: + page_table_strength = min( + 1.0, + line_like_items / float(TABLE_DRAWING_STRONG_THRESHOLD), + ) + + if table_hit: + has_any_tables = True + table_signal_pages += 1 + total_table_signal_strength += page_table_strength + + blocks = page.get_text("blocks") + text_blocks = [ + block + for block in blocks + if block[6] == 0 + and (block[2] - block[0]) > 20 + and (block[3] - block[1]) > 10 + ] + + is_multi_col_page = False + if len(text_blocks) >= MULTI_COL_MIN_BLOCKS: + min_x_gap = page.rect.width * MULTI_COL_GAP_RATIO + side_by_side_count = 0 + + for i in range(len(text_blocks)): + for j in range(i + 1, len(text_blocks)): + block_i = text_blocks[i] + block_j = text_blocks[j] + y_overlap = min(block_i[3], block_j[3]) - max( + block_i[1], block_j[1] + ) + if y_overlap <= 0: + continue + x_gap = max(block_j[0] - block_i[2], block_i[0] - block_j[2]) + if x_gap > min_x_gap: + side_by_side_count += 1 + if side_by_side_count >= 3: + is_multi_col_page = True + break + if is_multi_col_page: + break + + if is_multi_col_page: + multi_col_pages += 1 + + is_scan_page = ( + text_len < SCAN_TEXT_THRESHOLD and img_coverage > SCAN_IMAGE_COVERAGE_MIN + ) + if is_scan_page: + scanned_pages += 1 + + page_has_significant_images = ( + page_significant_image_count > 0 or page_medium_image_coverage >= 0.18 + ) + if page_has_significant_images: + pages_with_significant_images += 1 + significant_image_count += page_significant_image_count or 1 + + page_has_large_image = ( + page_max_rect_ratio >= LARGE_IMAGE_PAGE_RATIO or img_coverage >= 0.35 + ) + if page_has_large_image: + large_image_pages += 1 + + max_image_coverage_on_page = max( + max_image_coverage_on_page, page_max_rect_ratio + ) + + is_complex_page = ( + table_hit + or page_has_large_image + or is_multi_col_page + or (page_has_significant_images and text_len < FAST_TEXT_THRESHOLD) + or (drawing_count >= 25 and text_len < FAST_TEXT_THRESHOLD) + ) + if is_complex_page: + complex_pages += 1 + + page_details.append( + { + "page": idx + 1, + "text_len": text_len, + "image_count": len(images), + "img_coverage": round(img_coverage, 3), + "font_count": len(fonts), + "drawing_count": drawing_count, + "line_like_items": line_like_items, + "horizontal_line_items": horizontal_line_items, + "vertical_line_items": vertical_line_items, + "table_hit": table_hit, + "detected_table_count": detected_table_count, + "stroked_rect_count": rect_items, + "fill_only_rect_count": fill_only_rect_items, + "significant_image_count": page_significant_image_count, + "max_image_coverage": round(page_max_rect_ratio, 3), + "is_multi_col_page": is_multi_col_page, + "is_scan_page": is_scan_page, + "is_complex_page": is_complex_page, + "text_block_count": len(text_blocks), + } + ) + + del text_blocks + del blocks + del drawings + del fonts + del images + del page + + doc.close() + del doc + + n_sampled = len(sample_indices) + profile.avg_text_density = total_text_len / n_sampled if n_sampled > 0 else 0.0 + profile.avg_image_coverage = ( + total_image_coverage / n_sampled if n_sampled > 0 else 0.0 + ) + profile.has_embedded_fonts = has_any_fonts + profile.has_tables = has_any_tables + profile.is_multi_column = multi_col_pages > (n_sampled * 0.3) + profile.is_degraded_electronic = degraded_pages > (n_sampled * DEGRADED_PAGE_RATIO) + profile.sample_text = " ".join(all_text_parts)[:500] + profile.page_details = page_details + + profile.has_significant_images = pages_with_significant_images > 0 + profile.significant_image_count = significant_image_count + profile.max_image_coverage_on_page = max_image_coverage_on_page + profile.pages_with_significant_images = pages_with_significant_images + profile.large_image_page_ratio = ( + large_image_pages / n_sampled if n_sampled > 0 else 0.0 + ) + + profile.table_signal_pages = table_signal_pages + profile.table_signal_strength = ( + total_table_signal_strength / n_sampled if n_sampled > 0 else 0.0 + ) + + profile.complex_pages = complex_pages + profile.complex_page_ratio = complex_pages / n_sampled if n_sampled > 0 else 0.0 + profile.max_drawing_count = max_drawing_count + profile.min_text_density_page = min(text_lengths) if text_lengths else 0.0 + profile.text_density_std = _stddev(text_lengths) + + scan_ratio = scanned_pages / n_sampled if n_sampled > 0 else 0.0 + if scan_ratio >= SCAN_PAGE_RATIO: + profile.scan_type = "scanned" + reasons.append( + f"scanned: {scanned_pages}/{n_sampled} sampled pages are scanned ({scan_ratio:.0%})" + ) + elif scanned_pages > 0: + profile.scan_type = "mixed" + reasons.append(f"mixed: {scanned_pages}/{n_sampled} sampled pages are scanned") + else: + profile.scan_type = "electronic" + reasons.append( + f"electronic: sampled pages contain extractable text (avg={profile.avg_text_density:.0f})" + ) + + landscape_ratio = landscape_pages / n_sampled if n_sampled > 0 else 0.0 + + # ── Linear atlas gate: VLM always makes the final call ── + # Any document meeting all 4 conditions is sent for VLM visual confirmation. + # We do NOT heuristically commit here — VLM decides in parse_service. + is_atlas_candidate = ( + profile.avg_text_density + < ATLAS_TEXT_THRESHOLD # text-sparse (< 200 chars/page) + and profile.avg_image_coverage + > ATLAS_CANDIDATE_IMAGE_COVERAGE_MIN # image-heavy (> 30%) + and landscape_ratio >= ATLAS_MIN_LANDSCAPE_RATIO # mostly landscape (>= 50%) + and profile.page_count >= ATLAS_MIN_PAGES # multi-page (>= 2) + ) + if is_atlas_candidate: + profile.doc_category = ( + "generic" # provisional — VLM will promote to "atlas" if confirmed + ) + profile.atlas_candidate = True + reasons.append( + f"atlas_candidate: text={profile.avg_text_density:.0f}<{ATLAS_TEXT_THRESHOLD}, " + f"img={profile.avg_image_coverage:.1%}>{ATLAS_CANDIDATE_IMAGE_COVERAGE_MIN:.0%}, " + f"landscape={landscape_ratio:.0%}>={ATLAS_MIN_LANDSCAPE_RATIO:.0%}, " + f"pages={profile.page_count}>={ATLAS_MIN_PAGES} → VLM confirmation required" + ) + else: + profile.doc_category = "generic" + + if landscape_ratio >= 0.8 and profile.doc_category == "generic": + slide_ratios = [1.333, 1.778, 1.600] + tolerance = 0.05 + ref_page = doc_page_sizes[0] if doc_page_sizes else None + if ref_page: + page_ratio = ref_page[0] / ref_page[1] if ref_page[1] > 0 else 0.0 + is_slide_ratio = any( + abs(page_ratio - ratio) < tolerance for ratio in slide_ratios + ) + if is_slide_ratio: + profile.doc_category = "ppt_converted" + reasons.append( + f"ppt_converted: {landscape_pages}/{n_sampled} landscape, ratio={page_ratio:.2f}" + ) + + route, decision_band, benefit, risk, route_reasons = _classify_route(profile) + profile.route = route + profile.decision_band = decision_band + profile.estimated_fast_benefit = benefit + profile.estimated_risk_score = risk + reasons.extend(route_reasons) + + profile.reasoning = " | ".join(reasons) + publish_profile_result(queue, profile) + + +def profile_pdf(file_path: str) -> DocProfile: + """Profile a PDF by running PyMuPDF analysis in a spawned child process.""" + result = run_in_child_process(_profile_pdf_worker, file_path, timeout=300) + profile = DocProfile(**result["profile"]) + logger.info( + f"[doc-profiler] route={profile.route} band={profile.decision_band} " + f"category={profile.doc_category} scan={profile.scan_type} " + f"pages={profile.page_count} text_density={profile.avg_text_density:.0f} " + f"img_coverage={profile.avg_image_coverage:.1%} risk={profile.estimated_risk_score:.2f} " + f"gain={profile.estimated_fast_benefit:.2f}" + ) + return profile diff --git a/apps/worker/app/services/document_parser/doc_profiler.py b/apps/worker/app/services/document_parser/doc_profiler.py index 8fbefea76..bf728ccf3 100644 --- a/apps/worker/app/services/document_parser/doc_profiler.py +++ b/apps/worker/app/services/document_parser/doc_profiler.py @@ -1,4 +1,3 @@ -# pyright: reportAttributeAccessIssue=false, reportOperatorIssue=false """ Agentic Document Profiler @@ -10,764 +9,10 @@ profile = profile_document("/path/to/file.pdf") """ -import gc -import json -import math import os -from dataclasses import asdict, dataclass, field -from typing import Any, List, Literal, Optional -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker -from loguru import logger - - -@dataclass -class DocProfile: - """Profile data structure.""" - - # Basic information - file_type: str = "" - - # Routing decision - route: Literal["fast", "standard"] = "standard" - decision_band: Literal["safe_fast", "gray_zone", "safe_standard"] = "safe_standard" - - # Document type - scan_type: Optional[Literal["electronic", "scanned", "mixed"]] = None - doc_category: Literal["generic", "atlas", "ppt_converted"] = "generic" - - # Raw features - page_count: int = 0 - avg_text_density: float = 0.0 - avg_image_coverage: float = 0.0 - has_tables: bool = False - has_embedded_fonts: bool = False - is_multi_column: bool = False - is_degraded_electronic: bool = False - sample_text: str = "" - - # Image complexity - has_significant_images: bool = False - significant_image_count: int = 0 - max_image_coverage_on_page: float = 0.0 - pages_with_significant_images: int = 0 - large_image_page_ratio: float = 0.0 - - # Table complexity - table_signal_pages: int = 0 - table_signal_strength: float = 0.0 - - # Page complexity - complex_pages: int = 0 - complex_page_ratio: float = 0.0 - max_drawing_count: int = 0 - min_text_density_page: float = 0.0 - text_density_std: float = 0.0 - - # Aggregated decision scores - estimated_fast_benefit: float = 0.0 - estimated_risk_score: float = 0.0 - - # Atlas VLM second-pass flag - # True when heuristics suggest atlas-like layout but confidence is not high enough - # to commit without visual confirmation from a VLM. - atlas_candidate: bool = False - - # Page details for debug - page_details: List[dict] = field(default_factory=list) - - # Reasoning - reasoning: str = "" - - def to_dict(self) -> dict: - """Convert to dict (excluding page_details/sample_text to reduce size).""" - data = asdict(self) - data.pop("page_details", None) - data.pop("sample_text", None) - return data - - def summary(self) -> str: - """One-line summary.""" - parts = ( - f"[{self.file_type.upper()}] route={self.route}, band={self.decision_band}, " - f"scan={self.scan_type}, category={self.doc_category}, " - f"pages={self.page_count}, text_density={self.avg_text_density:.0f}, " - f"img_coverage={self.avg_image_coverage:.1%}, " - f"risk={self.estimated_risk_score:.2f}, gain={self.estimated_fast_benefit:.2f}" - ) - if self.is_degraded_electronic: - parts += ", degraded=True" - return parts - - -# Thresholds -SCAN_TEXT_THRESHOLD = 50 -SCAN_IMAGE_COVERAGE_MIN = 0.6 -SCAN_PAGE_RATIO = 0.7 - -ATLAS_TEXT_THRESHOLD = 200 -ATLAS_CANDIDATE_IMAGE_COVERAGE_MIN = 0.30 -ATLAS_MIN_LANDSCAPE_RATIO = 0.5 # ≥50% of sampled pages must be landscape -ATLAS_MIN_PAGES = 2 # single-page scans (resumes, posters) are not atlases - -FAST_TEXT_THRESHOLD = 500 -MIN_FAST_TEXT_DENSITY_FLOOR = 120 -SAFE_FAST_MAX_PAGE_COUNT = 80 -HARD_STANDARD_PAGE_COUNT = 150 - -MULTI_COL_GAP_RATIO = 0.15 -MULTI_COL_MIN_BLOCKS = 4 - -DEGRADED_SKINNY_ASPECT = 50 -DEGRADED_SKINNY_MAX_H = 30 -DEGRADED_SKINNY_MIN_PER_PAGE = 50 -DEGRADED_PAGE_RATIO = 0.5 - -SIGNIFICANT_IMAGE_AREA_RATIO = 0.12 -MEDIUM_IMAGE_AREA_RATIO = 0.03 -LARGE_IMAGE_PAGE_RATIO = 0.25 -SIGNIFICANT_IMAGE_MIN_DIM = 400 -SIGNIFICANT_IMAGE_MIN_PIXELS = 250_000 - -PROFILE_MAX_NEW_XREFS_PER_PAGE = 30 - -TABLE_DRAWING_LINE_THRESHOLD = 12 -TABLE_DRAWING_STRONG_THRESHOLD = 18 -TABLE_DRAWING_RECT_THRESHOLD = 2 - -SAFE_FAST_MAX_COMPLEX_PAGE_RATIO = 0.05 -SAFE_FAST_MAX_IMAGE_COVERAGE_ON_PAGE = 0.08 -SAFE_FAST_MAX_AVG_IMAGE_COVERAGE = 0.03 -SAFE_FAST_MAX_TEXT_STD = 600.0 -HARD_COMPLEX_PAGE_RATIO = 0.2 -HARD_SIGNIFICANT_IMAGE_PAGES = 3 -HARD_LARGE_IMAGE_PAGE_RATIO = 0.15 - - -def _clamp(value: float, min_value: float = 0.0, max_value: float = 1.0) -> float: - return max(min_value, min(max_value, value)) - - -def _stddev(values: list[float]) -> float: - if not values: - return 0.0 - mean = sum(values) / len(values) - variance = sum((value - mean) ** 2 for value in values) / len(values) - return math.sqrt(variance) - - -def _count_detected_tables(page: Any) -> int: - try: - finder = page.find_tables() - except Exception: - return 0 - - if not finder: - return 0 - - tables = getattr(finder, "tables", finder) - try: - return len(tables) - except TypeError: - return 1 if tables else 0 - - -def _is_stroked_drawing(drawing: dict[str, Any]) -> bool: - stroke_width = drawing.get("width") - return drawing.get("color") is not None or ( - stroke_width is not None and stroke_width > 0 - ) - - -def _estimate_fast_benefit(profile: DocProfile) -> float: - if profile.page_count <= 2: - page_factor = 0.35 - elif profile.page_count <= 10: - page_factor = 0.7 - elif profile.page_count <= SAFE_FAST_MAX_PAGE_COUNT: - page_factor = 1.0 - elif profile.page_count <= HARD_STANDARD_PAGE_COUNT: - page_factor = 0.8 - else: - page_factor = 0.45 - - density_factor = _clamp(profile.avg_text_density / 1200.0) - stability_factor = _clamp( - 1.0 - - (profile.complex_page_ratio * 1.5) - - (profile.large_image_page_ratio * 1.2) - - (profile.table_signal_strength * 0.8) - ) - return _clamp( - (0.35 * page_factor) + (0.40 * density_factor) + (0.25 * stability_factor) - ) - - -def _estimate_risk_score(profile: DocProfile) -> float: - risk = 0.0 - if profile.scan_type != "electronic": - risk += 0.35 - if profile.doc_category != "generic": - risk += 0.20 - if profile.is_multi_column: - risk += 0.20 - if profile.is_degraded_electronic: - risk += 0.20 - if profile.has_tables: - risk += 0.30 - - risk += min(0.20, profile.large_image_page_ratio * 1.2) - risk += min(0.20, profile.complex_page_ratio * 0.8) - risk += min(0.15, profile.table_signal_strength * 0.2) - risk += min(0.12, profile.pages_with_significant_images * 0.04) - - if profile.page_count > HARD_STANDARD_PAGE_COUNT: - risk += 0.10 - - return _clamp(risk) - - -def _classify_route(profile: DocProfile) -> tuple[str, str, float, float, list[str]]: - hard_gate_reasons: list[str] = [] - - if profile.scan_type != "electronic": - hard_gate_reasons.append(f"scan_type={profile.scan_type}") - if profile.doc_category != "generic": - hard_gate_reasons.append(f"doc_category={profile.doc_category}") - if profile.is_multi_column: - hard_gate_reasons.append("multi_column") - if profile.is_degraded_electronic: - hard_gate_reasons.append("degraded_electronic") - if profile.has_tables: - hard_gate_reasons.append( - f"table_signals={profile.table_signal_pages}p/{profile.table_signal_strength:.2f}" - ) - if ( - profile.max_image_coverage_on_page >= LARGE_IMAGE_PAGE_RATIO - or profile.pages_with_significant_images >= HARD_SIGNIFICANT_IMAGE_PAGES - or profile.large_image_page_ratio >= HARD_LARGE_IMAGE_PAGE_RATIO - ): - hard_gate_reasons.append( - "significant_images=" - f"{profile.pages_with_significant_images}p,max={profile.max_image_coverage_on_page:.1%}" - ) - if profile.complex_page_ratio >= HARD_COMPLEX_PAGE_RATIO: - hard_gate_reasons.append(f"complex_pages={profile.complex_page_ratio:.0%}") - if profile.page_count > HARD_STANDARD_PAGE_COUNT: - hard_gate_reasons.append( - f"page_count={profile.page_count}>{HARD_STANDARD_PAGE_COUNT}" - ) - - benefit = _estimate_fast_benefit(profile) - risk = _estimate_risk_score(profile) - - if hard_gate_reasons: - return ( - "standard", - "safe_standard", - benefit, - risk, - [ - "decision=safe_standard: hard gate matched", - "hard_gates=" + ",".join(hard_gate_reasons), - ], - ) - - safe_fast_checks = [ - ( - profile.page_count <= SAFE_FAST_MAX_PAGE_COUNT, - f"page_count={profile.page_count}<={SAFE_FAST_MAX_PAGE_COUNT}", - ), - ( - profile.avg_text_density >= MIN_FAST_TEXT_DENSITY_FLOOR, - "text_density_floor=" - f"{profile.avg_text_density:.0f}>={MIN_FAST_TEXT_DENSITY_FLOOR}", - ), - ( - not profile.has_significant_images, - f"has_significant_images={profile.has_significant_images}", - ), - ( - profile.max_image_coverage_on_page <= SAFE_FAST_MAX_IMAGE_COVERAGE_ON_PAGE, - "max_image_coverage_on_page=" - f"{profile.max_image_coverage_on_page:.1%}<={SAFE_FAST_MAX_IMAGE_COVERAGE_ON_PAGE:.0%}", - ), - ( - profile.avg_image_coverage <= SAFE_FAST_MAX_AVG_IMAGE_COVERAGE, - f"avg_image_coverage={profile.avg_image_coverage:.1%}<={SAFE_FAST_MAX_AVG_IMAGE_COVERAGE:.0%}", - ), - ( - profile.complex_page_ratio <= SAFE_FAST_MAX_COMPLEX_PAGE_RATIO, - f"complex_page_ratio={profile.complex_page_ratio:.0%}<={SAFE_FAST_MAX_COMPLEX_PAGE_RATIO:.0%}", - ), - ( - profile.text_density_std <= SAFE_FAST_MAX_TEXT_STD, - f"text_density_std={profile.text_density_std:.0f}<={SAFE_FAST_MAX_TEXT_STD:.0f}", - ), - ( - risk <= 0.35, - f"estimated_risk_score={risk:.2f}<=0.35", - ), - ] - - failed_checks = [reason for passed, reason in safe_fast_checks if not passed] - if not failed_checks: - return ( - "fast", - "safe_fast", - benefit, - risk, - [ - "decision=safe_fast: low-complexity high-yield pdf", - "safe_fast_checks_passed", - ], - ) - - return ( - "standard", - "gray_zone", - benefit, - risk, - [ - "decision=gray_zone: conservative fallback to standard in phase1", - "borderline=" + ",".join(failed_checks[:4]), - ], - ) - - -def _publish_profile_result(queue, profile: DocProfile) -> None: - """Release Python-side wrappers before publishing the profile result.""" - gc.collect() - queue.put({"ok": True, "profile": asdict(profile)}) - - -@worker -def _profile_pdf_worker(queue, file_path: str) -> None: - """Child process: analyze PDF features, return profile as dict.""" - import pymupdf - - profile = DocProfile(file_type="pdf") - reasons: list[str] = [] - - try: - doc = pymupdf.open(file_path) - except Exception as exc: - profile.reasoning = f"Cannot open file: {exc}" - _publish_profile_result(queue, profile) - return - - profile.page_count = doc.page_count - - if doc.page_count == 0: - profile.reasoning = "Empty file (0 pages)" - doc.close() - del doc - _publish_profile_result(queue, profile) - return - - if doc.page_count <= 50: - sample_indices = list(range(doc.page_count)) - else: - step = max(1, doc.page_count // 20) - sample_indices = list(range(0, doc.page_count, step))[:20] - sample_indices = sorted( - set( - sample_indices - + [0, 1, 2] - + [doc.page_count - 3, doc.page_count - 2, doc.page_count - 1] - ) - ) - sample_indices = [idx for idx in sample_indices if 0 <= idx < doc.page_count] - - page_details = [] - text_lengths: list[float] = [] - total_text_len = 0 - total_image_coverage = 0.0 - scanned_pages = 0 - all_text_parts: list[str] = [] - has_any_fonts = False - has_any_tables = False - table_signal_pages = 0 - total_table_signal_strength = 0.0 - multi_col_pages = 0 - landscape_pages = 0 - degraded_pages = 0 - doc_page_sizes = [] - - significant_image_count = 0 - pages_with_significant_images = 0 - large_image_pages = 0 - max_image_coverage_on_page = 0.0 - - complex_pages = 0 - max_drawing_count = 0 - - # Track xrefs already processed across pages to avoid redundant - # get_image_rects() calls on shared/inherited image resources. - # PDFs with shared xrefs (e.g. scanned docs) report ALL document - # images on every page; without dedup this causes O(pages × images) - # content-stream scans. - seen_xrefs: set = set() - - for idx in sample_indices: - page = doc[idx] - page_width = page.rect.width - page_height = page.rect.height - page_area = page_width * page_height - - if page_width > page_height: - landscape_pages += 1 - doc_page_sizes.append((page_width, page_height)) - - text = page.get_text().strip() - text_len = len(text) - text_lengths.append(float(text_len)) - total_text_len += text_len - - if len("".join(all_text_parts)) < 500: - all_text_parts.append(text[:200]) - - images = page.get_images(full=True) - img_total_area = 0.0 - page_significant_image_count = 0 - page_max_rect_ratio = 0.0 - page_medium_image_coverage = 0.0 - skinny_count = 0 - - new_xref_count = 0 - for img in images: - xref = img[0] - img_w, img_h = img[2], img[3] - if ( - img_h > 0 - and img_w / img_h > DEGRADED_SKINNY_ASPECT - and img_h < DEGRADED_SKINNY_MAX_H - ): - skinny_count += 1 - - # ── xref dedup: skip images already analyzed on earlier pages ── - if xref in seen_xrefs: - continue - seen_xrefs.add(xref) - new_xref_count += 1 - # Cap expensive get_image_rects calls per page - if new_xref_count > PROFILE_MAX_NEW_XREFS_PER_PAGE: - continue - - try: - rects = page.get_image_rects(xref) - except Exception: - rects = [] - - for rect in rects: - rect_area = rect.width * rect.height - img_total_area += rect_area - - area_ratio = rect_area / page_area if page_area > 0 else 0.0 - page_max_rect_ratio = max(page_max_rect_ratio, area_ratio) - - is_significant = ( - area_ratio >= SIGNIFICANT_IMAGE_AREA_RATIO - or ( - area_ratio >= 0.05 - and ( - max(img_w, img_h) >= SIGNIFICANT_IMAGE_MIN_DIM - or (img_w * img_h) >= SIGNIFICANT_IMAGE_MIN_PIXELS - ) - ) - or ( - area_ratio >= 0.02 - and (img_w * img_h) >= (SIGNIFICANT_IMAGE_MIN_PIXELS * 2) - ) - ) - - if is_significant: - page_significant_image_count += 1 - elif area_ratio >= MEDIUM_IMAGE_AREA_RATIO: - page_medium_image_coverage += area_ratio - - if skinny_count >= DEGRADED_SKINNY_MIN_PER_PAGE: - degraded_pages += 1 - - img_coverage = img_total_area / page_area if page_area > 0 else 0.0 - img_coverage = min(img_coverage, 1.0) - total_image_coverage += img_coverage - - fonts = page.get_fonts() - if fonts: - has_any_fonts = True - - drawings = page.get_drawings() - drawing_count = len(drawings) - max_drawing_count = max(max_drawing_count, drawing_count) - line_like_items = 0 - horizontal_line_items = 0 - vertical_line_items = 0 - rect_items = 0 - fill_only_rect_items = 0 - for drawing in drawings: - is_stroked = _is_stroked_drawing(drawing) - for item in drawing.get("items", []): - if item[0] == "l": - if is_stroked: - line_like_items += 1 - point_a = item[1] - point_b = item[2] - if abs(point_a.y - point_b.y) <= 2: - horizontal_line_items += 1 - if abs(point_a.x - point_b.x) <= 2: - vertical_line_items += 1 - elif item[0] == "re": - if is_stroked: - rect_items += 1 - line_like_items += 4 - horizontal_line_items += 2 - vertical_line_items += 2 - else: - fill_only_rect_items += 1 - - detected_table_count = _count_detected_tables(page) - drawing_table_signal = line_like_items >= TABLE_DRAWING_LINE_THRESHOLD and ( - (horizontal_line_items >= 2 and vertical_line_items >= 2) - or rect_items >= TABLE_DRAWING_RECT_THRESHOLD - or ( - line_like_items >= TABLE_DRAWING_STRONG_THRESHOLD - and horizontal_line_items >= 3 - and vertical_line_items >= 3 - ) - ) - # NOTE: - # `page.find_tables()` produces too many false positives on Word / Writer - # exported pure-text PDFs, where paragraph background boxes are inferred as - # full-page tables. For Phase 1 fast-path routing, keep `find_tables()` - # only as debug evidence and rely on explicit drawing-grid signals for - # table hard gates. - table_hit = drawing_table_signal - page_table_strength = 0.0 - if drawing_table_signal: - page_table_strength = min( - 1.0, - line_like_items / float(TABLE_DRAWING_STRONG_THRESHOLD), - ) - - if table_hit: - has_any_tables = True - table_signal_pages += 1 - total_table_signal_strength += page_table_strength - - blocks = page.get_text("blocks") - text_blocks = [ - block - for block in blocks - if block[6] == 0 - and (block[2] - block[0]) > 20 - and (block[3] - block[1]) > 10 - ] - - is_multi_col_page = False - if len(text_blocks) >= MULTI_COL_MIN_BLOCKS: - min_x_gap = page.rect.width * MULTI_COL_GAP_RATIO - side_by_side_count = 0 - - for i in range(len(text_blocks)): - for j in range(i + 1, len(text_blocks)): - block_i = text_blocks[i] - block_j = text_blocks[j] - y_overlap = min(block_i[3], block_j[3]) - max( - block_i[1], block_j[1] - ) - if y_overlap <= 0: - continue - x_gap = max(block_j[0] - block_i[2], block_i[0] - block_j[2]) - if x_gap > min_x_gap: - side_by_side_count += 1 - if side_by_side_count >= 3: - is_multi_col_page = True - break - if is_multi_col_page: - break - - if is_multi_col_page: - multi_col_pages += 1 - - is_scan_page = ( - text_len < SCAN_TEXT_THRESHOLD and img_coverage > SCAN_IMAGE_COVERAGE_MIN - ) - if is_scan_page: - scanned_pages += 1 - - page_has_significant_images = ( - page_significant_image_count > 0 or page_medium_image_coverage >= 0.18 - ) - if page_has_significant_images: - pages_with_significant_images += 1 - significant_image_count += page_significant_image_count or 1 - - page_has_large_image = ( - page_max_rect_ratio >= LARGE_IMAGE_PAGE_RATIO or img_coverage >= 0.35 - ) - if page_has_large_image: - large_image_pages += 1 - - max_image_coverage_on_page = max( - max_image_coverage_on_page, page_max_rect_ratio - ) - - is_complex_page = ( - table_hit - or page_has_large_image - or is_multi_col_page - or (page_has_significant_images and text_len < FAST_TEXT_THRESHOLD) - or (drawing_count >= 25 and text_len < FAST_TEXT_THRESHOLD) - ) - if is_complex_page: - complex_pages += 1 - - page_details.append( - { - "page": idx + 1, - "text_len": text_len, - "image_count": len(images), - "img_coverage": round(img_coverage, 3), - "font_count": len(fonts), - "drawing_count": drawing_count, - "line_like_items": line_like_items, - "horizontal_line_items": horizontal_line_items, - "vertical_line_items": vertical_line_items, - "table_hit": table_hit, - "detected_table_count": detected_table_count, - "stroked_rect_count": rect_items, - "fill_only_rect_count": fill_only_rect_items, - "significant_image_count": page_significant_image_count, - "max_image_coverage": round(page_max_rect_ratio, 3), - "is_multi_col_page": is_multi_col_page, - "is_scan_page": is_scan_page, - "is_complex_page": is_complex_page, - "text_block_count": len(text_blocks), - } - ) - - del text_blocks - del blocks - del drawings - del fonts - del images - del page - - doc.close() - del doc - - n_sampled = len(sample_indices) - profile.avg_text_density = total_text_len / n_sampled if n_sampled > 0 else 0.0 - profile.avg_image_coverage = ( - total_image_coverage / n_sampled if n_sampled > 0 else 0.0 - ) - profile.has_embedded_fonts = has_any_fonts - profile.has_tables = has_any_tables - profile.is_multi_column = multi_col_pages > (n_sampled * 0.3) - profile.is_degraded_electronic = degraded_pages > (n_sampled * DEGRADED_PAGE_RATIO) - profile.sample_text = " ".join(all_text_parts)[:500] - profile.page_details = page_details - - profile.has_significant_images = pages_with_significant_images > 0 - profile.significant_image_count = significant_image_count - profile.max_image_coverage_on_page = max_image_coverage_on_page - profile.pages_with_significant_images = pages_with_significant_images - profile.large_image_page_ratio = ( - large_image_pages / n_sampled if n_sampled > 0 else 0.0 - ) - - profile.table_signal_pages = table_signal_pages - profile.table_signal_strength = ( - total_table_signal_strength / n_sampled if n_sampled > 0 else 0.0 - ) - - profile.complex_pages = complex_pages - profile.complex_page_ratio = complex_pages / n_sampled if n_sampled > 0 else 0.0 - profile.max_drawing_count = max_drawing_count - profile.min_text_density_page = min(text_lengths) if text_lengths else 0.0 - profile.text_density_std = _stddev(text_lengths) - - scan_ratio = scanned_pages / n_sampled if n_sampled > 0 else 0.0 - if scan_ratio >= SCAN_PAGE_RATIO: - profile.scan_type = "scanned" - reasons.append( - f"scanned: {scanned_pages}/{n_sampled} sampled pages are scanned ({scan_ratio:.0%})" - ) - elif scanned_pages > 0: - profile.scan_type = "mixed" - reasons.append(f"mixed: {scanned_pages}/{n_sampled} sampled pages are scanned") - else: - profile.scan_type = "electronic" - reasons.append( - f"electronic: sampled pages contain extractable text (avg={profile.avg_text_density:.0f})" - ) - - landscape_ratio = landscape_pages / n_sampled if n_sampled > 0 else 0.0 - - # ── Linear atlas gate: VLM always makes the final call ── - # Any document meeting all 4 conditions is sent for VLM visual confirmation. - # We do NOT heuristically commit here — VLM decides in parse_service. - is_atlas_candidate = ( - profile.avg_text_density - < ATLAS_TEXT_THRESHOLD # text-sparse (< 200 chars/page) - and profile.avg_image_coverage - > ATLAS_CANDIDATE_IMAGE_COVERAGE_MIN # image-heavy (> 30%) - and landscape_ratio >= ATLAS_MIN_LANDSCAPE_RATIO # mostly landscape (>= 50%) - and profile.page_count >= ATLAS_MIN_PAGES # multi-page (>= 2) - ) - if is_atlas_candidate: - profile.doc_category = ( - "generic" # provisional — VLM will promote to "atlas" if confirmed - ) - profile.atlas_candidate = True - reasons.append( - f"atlas_candidate: text={profile.avg_text_density:.0f}<{ATLAS_TEXT_THRESHOLD}, " - f"img={profile.avg_image_coverage:.1%}>{ATLAS_CANDIDATE_IMAGE_COVERAGE_MIN:.0%}, " - f"landscape={landscape_ratio:.0%}>={ATLAS_MIN_LANDSCAPE_RATIO:.0%}, " - f"pages={profile.page_count}>={ATLAS_MIN_PAGES} → VLM confirmation required" - ) - else: - profile.doc_category = "generic" - - if landscape_ratio >= 0.8 and profile.doc_category == "generic": - slide_ratios = [1.333, 1.778, 1.600] - tolerance = 0.05 - ref_page = doc_page_sizes[0] if doc_page_sizes else None - if ref_page: - page_ratio = ref_page[0] / ref_page[1] if ref_page[1] > 0 else 0.0 - is_slide_ratio = any( - abs(page_ratio - ratio) < tolerance for ratio in slide_ratios - ) - if is_slide_ratio: - profile.doc_category = "ppt_converted" - reasons.append( - f"ppt_converted: {landscape_pages}/{n_sampled} landscape, ratio={page_ratio:.2f}" - ) - - route, decision_band, benefit, risk, route_reasons = _classify_route(profile) - profile.route = route - profile.decision_band = decision_band - profile.estimated_fast_benefit = benefit - profile.estimated_risk_score = risk - reasons.extend(route_reasons) - - profile.reasoning = " | ".join(reasons) - _publish_profile_result(queue, profile) - - -def _profile_pdf(file_path: str) -> DocProfile: - """Profile a PDF by running PyMuPDF analysis in a spawned child process.""" - result = run_in_child_process(_profile_pdf_worker, file_path, timeout=300) - profile = DocProfile(**result["profile"]) - logger.info( - f"[doc-profiler] route={profile.route} band={profile.decision_band} " - f"category={profile.doc_category} scan={profile.scan_type} " - f"pages={profile.page_count} text_density={profile.avg_text_density:.0f} " - f"img_coverage={profile.avg_image_coverage:.1%} risk={profile.estimated_risk_score:.2f} " - f"gain={profile.estimated_fast_benefit:.2f}" - ) - return profile +from app.services.document_parser.doc_profile_model import DocProfile +from app.services.document_parser.doc_profile_pdf import profile_pdf def profile_document(file_path: str, filename: str = "") -> DocProfile: @@ -786,7 +31,7 @@ def profile_document(file_path: str, filename: str = "") -> DocProfile: ext = os.path.splitext(filename)[1].lower() if ext == ".pdf": - return _profile_pdf(file_path) + return profile_pdf(file_path) return DocProfile( file_type=ext.lstrip("."), @@ -794,11 +39,3 @@ def profile_document(file_path: str, filename: str = "") -> DocProfile: decision_band="safe_standard", reasoning=f"Non-PDF format ({ext}), using default route", ) - - -def save_profile_metadata(profile: DocProfile, output_dir: str): - """Save profile to output_dir/profile.json.""" - profile_path = os.path.join(output_dir, "profile.json") - with open(profile_path, "w", encoding="utf-8") as file_obj: - json.dump(profile.to_dict(), file_obj, ensure_ascii=False, indent=2) - logger.debug(f"Profile metadata saved to {profile_path}") diff --git a/apps/worker/app/services/document_parser/docx_asset_accumulator.py b/apps/worker/app/services/document_parser/docx_asset_accumulator.py new file mode 100644 index 000000000..ab0d67509 --- /dev/null +++ b/apps/worker/app/services/document_parser/docx_asset_accumulator.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from app.services.document_parser.docx_asset_store import DocxAssetStore + + +ImageHandler = Callable[ + [ + list[list[object]], + dict[str, Any], + DocxAssetStore, + list[dict[str, Any]], + str, + int, + bool, + dict[str, dict[str, str]], + ], + tuple[list[dict[str, Any]], list[list[object]], bool], +] + +TableHandler = Callable[..., tuple[list[dict[str, Any]], list[list[object]], int]] + + +@dataclass +class DocxAssetAccumulator: + asset_store: DocxAssetStore + should_summary_image: bool + should_summary_table: bool + image_handler: ImageHandler + table_handler: TableHandler + _rows: list[list[object]] = field(default_factory=list) + _image_count: int = 0 + _table_count: int = 0 + _seen_images: dict[str, dict[str, str]] = field(default_factory=dict) + + def append_image( + self, + image_meta: dict[str, Any], + headings_stack: list[dict[str, Any]], + current_heading: str, + ) -> list[dict[str, Any]]: + headings_stack, self._rows, is_new_image = self.image_handler( + self._rows, + image_meta, + self.asset_store, + headings_stack, + current_heading, + self._image_count, + self.should_summary_image, + self._seen_images, + ) + if is_new_image: + self._image_count += 1 + return headings_stack + + def append_table( + self, + block: Any, + headings_stack: list[dict[str, Any]], + current_heading: str, + cell_images: Any, + ) -> list[dict[str, Any]]: + headings_stack, self._rows, self._image_count = self.table_handler( + self._rows, + block, + self.asset_store, + headings_stack, + current_heading, + self._table_count, + summary_table=self.should_summary_table, + summary_image=self.should_summary_image, + cell_images=cell_images, + img_count=self._image_count, + seen_images=self._seen_images, + ) + self._table_count += 1 + return headings_stack + + def rows(self) -> list[list[object]]: + return self._rows diff --git a/apps/worker/app/services/document_parser/docx_asset_store.py b/apps/worker/app/services/document_parser/docx_asset_store.py new file mode 100644 index 000000000..f97f8b0a8 --- /dev/null +++ b/apps/worker/app/services/document_parser/docx_asset_store.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import os +import shutil +from dataclasses import dataclass + + +@dataclass(frozen=True) +class StoredDocxAsset: + absolute_path: str + relative_path: str + name: str + extension: str + + +class DocxAssetStore: + def __init__(self, output_dir: str) -> None: + self.output_dir = output_dir + self.image_dir = os.path.join(output_dir, "images") + self.table_dir = os.path.join(output_dir, "tables") + + def reset(self) -> None: + self._reset_dir(self.table_dir) + self._reset_dir(self.image_dir) + + def write_image(self, name: str, extension: str, data: bytes) -> StoredDocxAsset: + absolute_path = os.path.join(self.image_dir, f"{name}{extension}") + with open(absolute_path, "wb") as image_file: + image_file.write(data) + return StoredDocxAsset( + absolute_path=absolute_path, + relative_path=f"images/{name}{extension}", + name=name, + extension=extension, + ) + + def rename_image(self, asset: StoredDocxAsset, new_name: str) -> StoredDocxAsset: + new_absolute_path = os.path.join(self.image_dir, f"{new_name}{asset.extension}") + if asset.absolute_path != new_absolute_path: + os.rename(asset.absolute_path, new_absolute_path) + return StoredDocxAsset( + absolute_path=new_absolute_path, + relative_path=f"images/{new_name}{asset.extension}", + name=new_name, + extension=asset.extension, + ) + + def write_table(self, name: str, html: str) -> StoredDocxAsset: + absolute_path = os.path.join(self.table_dir, f"{name}.html") + with open(absolute_path, "w", encoding="utf-8") as table_file: + table_file.write(html) + return StoredDocxAsset( + absolute_path=absolute_path, + relative_path=f"tables/{name}.html", + name=name, + extension=".html", + ) + + @staticmethod + def _reset_dir(directory: str) -> None: + if os.path.isdir(directory): + shutil.rmtree(directory) + os.makedirs(directory, exist_ok=True) diff --git a/apps/worker/app/services/document_parser/docx_block_stream.py b/apps/worker/app/services/document_parser/docx_block_stream.py new file mode 100644 index 000000000..9832e0e54 --- /dev/null +++ b/apps/worker/app/services/document_parser/docx_block_stream.py @@ -0,0 +1,315 @@ +# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false +from __future__ import annotations + +import io +import zipfile + +from app.services.document_parser.toc_docx import detect_doc_tocs, detect_sdt_toc +from docx import Document +from docx.oxml.table import CT_Tbl +from docx.oxml.text.paragraph import CT_P +from docx.table import Table +from docx.text.paragraph import Paragraph +from loguru import logger +from lxml import etree + + +def iter_block_items(doc_data): + doc_stream = io.BytesIO(doc_data) + doc = Document(doc_stream) + + # python-docx mapping + p_tbl_map = [] + for child in doc.element.body: + if isinstance(child, CT_P): + p_tbl_map.append(("p", child)) + elif isinstance(child, CT_Tbl): + p_tbl_map.append(("tbl", child)) + + with zipfile.ZipFile(io.BytesIO(doc_data), "r") as docx: + xml = docx.read("word/document.xml") + rels = etree.fromstring(docx.read("word/_rels/document.xml.rels")) + rel_map = { + r.get("Id"): r.get("Target") for r in rels.findall(".//{*}Relationship") + } + ns = { + "a": "http://schemas.openxmlformats.org/drawingml/2006/main", + "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", + "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main", + "v": "urn:schemas-microsoft-com:vml", + "o": "urn:schemas-microsoft-com:office:office", + } + r_ns = "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}" + + root = etree.fromstring(xml) + body = root.find(".//w:body", namespaces=ns) + + ele_num = 1 + map_index = 0 # point to p_tbl_map + toc_field_active = False + + for elem in body.iterchildren(): + if not isinstance(elem.tag, str): + continue + + tag = etree.QName(elem.tag).localname + + # --- SDT (Structured Document Tag) container --- + # TOC generated by MS Word is usually in sdt + if tag == "sdt": + sdt_toc_info = detect_sdt_toc(elem, ns) + is_toc_sdt = sdt_toc_info["is_toc_sdt"] + + sdt_content = elem.find(".//w:sdtContent", namespaces=ns) + if sdt_content is not None: + for p_elem in sdt_content.findall(".//w:p", namespaces=ns): + texts = p_elem.xpath(".//w:t/text()", namespaces=ns) + text = "".join(texts).strip() + + if is_toc_sdt: + label = "TOC-AREA" + toc_info = detect_doc_tocs(p_elem, ns) + else: + toc_info = detect_doc_tocs(p_elem, ns) + if toc_info["is_style"] or toc_info["is_field_start"]: + label = "TOC-AREA" + else: + label = "PTXT" + + if text: + meta = None + if "TOC" in label: + meta = { + "toc_level": toc_info.get("toc_level"), + "toc_outline_level": toc_info.get("outline_level"), + "toc_left_indent": toc_info.get("left_indent"), + "toc_style_name": toc_info.get("style_name"), + "toc_source": "sdt", + } + yield ele_num, text, label, meta + ele_num += 1 + continue + + # --- text paras --- + if tag == "p": + texts = elem.xpath(".//w:t/text()", namespaces=ns) + text = "".join(texts).strip() + + if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "p": + p_obj = Paragraph(p_tbl_map[map_index][1], doc) + else: + p_obj = None + + toc_info = detect_doc_tocs(elem, ns) + if toc_info["is_field_start"]: + toc_field_active = True + + if toc_info["is_style"] or toc_field_active: + label = "TOC-AREA" + else: + label = "PTXT" + + if text or p_obj is not None: + meta = None + if "TOC" in label: + meta = { + "toc_level": toc_info.get("toc_level"), + "toc_outline_level": toc_info.get("outline_level"), + "toc_left_indent": toc_info.get("left_indent"), + "toc_style_name": toc_info.get("style_name"), + "toc_source": "paragraph", + } + yield ele_num, p_obj or text, label, meta + ele_num += 1 + + # images (DrawingML: ) + seen_rids = set() + blips = elem.xpath(".//a:blip", namespaces=ns) + for b in blips: + rid = b.get(f"{r_ns}embed") + if not rid or rid in seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + seen_rids.add(rid) + data = docx.read("word/" + target) + yield ( + ele_num, + None, + "IMAGE", + { + "image_name": target.split("/")[-1], + "from": "paragraph", + "size": len(data), + "data": data, + }, + ) + ele_num += 1 + + # TODO: Re-evaluate VML group extraction strategy. + # Complex VML composite images () are currently skipped because extracting + # piece-by-piece loses textual overlay and positioning. + # Future plan: Use LibreOffice headless conversion to render the entire document + # and map the perfectly rendered images back to the layout via text anchors. + # + # Temporary: detect VML-only paragraphs and inject a placeholder so the + # paragraph isn't silently swallowed, leaving its parent section empty. + if not text and not seen_rids: + # No text and no DrawingML images — check for VML content + vml_groups = elem.xpath(".//v:group", namespaces=ns) + vml_images_check = elem.xpath(".//v:imagedata", namespaces=ns) + if vml_groups or vml_images_check: + vml_placeholder = "[VML graphic \u2014 extraction not yet supported]" + yield ele_num, vml_placeholder, "PTXT", None + ele_num += 1 + logger.debug( + f"Injected VML placeholder for paragraph with " + f"{len(vml_groups)} v:group, {len(vml_images_check)} v:imagedata" + ) + """ + # images (VML: ) — convert to PNG + from PIL import Image as PILImage + + vml_images = elem.xpath(".//v:imagedata", namespaces=ns) + for v in vml_images: + rid = v.get(f"{r_ns}id") + if not rid or rid in seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + seen_rids.add(rid) + raw_data = docx.read("word/" + target) + # Convert to PNG for uniform downstream handling + try: + pil_img = PILImage.open(io.BytesIO(raw_data)) + png_buf = io.BytesIO() + pil_img.save(png_buf, format="PNG") + png_data = png_buf.getvalue() + except Exception as e: + logger.warning(f"Failed to convert VML image to PNG: {e}") + continue + orig_name = target.split("/")[-1] + png_name = os.path.splitext(orig_name)[0] + ".png" + yield ( + ele_num, + None, + "IMAGE", + { + "image_name": png_name, + "from": "paragraph_vml", + "size": len(png_data), + "data": png_data, + }, + ) + ele_num += 1 + """ + map_index += 1 + + if toc_info["is_field_end"]: + toc_field_active = False + + # --- tables --- + elif tag == "tbl": + if map_index < len(p_tbl_map) and p_tbl_map[map_index][0] == "tbl": + tbl = Table(p_tbl_map[map_index][1], doc) + else: + tbl = Table(elem, doc) + + # Extract images from each cell, keyed by (row_idx, col_idx) + cell_images = {} # {(row_idx, col_idx): [{'image_name', 'data', 'size'}]} + for row_idx, tr in enumerate(elem.findall(".//w:tr", namespaces=ns)): + for col_idx, tc in enumerate(tr.findall(".//w:tc", namespaces=ns)): + cell_seen_rids = set() + imgs_in_cell = [] + # DrawingML images in cell + blips = tc.xpath(".//a:blip", namespaces=ns) + for b in blips: + rid = b.get(f"{r_ns}embed") + if not rid or rid in cell_seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + cell_seen_rids.add(rid) + data = docx.read("word/" + target) + if ( + len(data) < 10 * 1024 + ): # Skip small images (<10KB, likely icons) + continue + imgs_in_cell.append( + { + "image_name": target.split("/")[-1], + "data": data, + "size": len(data), + } + ) + # TODO: VML in tables is temporarily skipped to avoid extracting + # fragmented textless background images. (Same as paragraph VML logic) + """ + # VML images in cell — convert to PNG + from PIL import Image as PILImage + + vml_in_cell = tc.xpath(".//v:imagedata", namespaces=ns) + for v in vml_in_cell: + rid = v.get(f"{r_ns}id") + if not rid or rid in cell_seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + cell_seen_rids.add(rid) + raw_data = docx.read("word/" + target) + try: + pil_img = PILImage.open(io.BytesIO(raw_data)) + png_buf = io.BytesIO() + pil_img.save(png_buf, format="PNG") + png_data = png_buf.getvalue() + except Exception as e: + logger.warning( + f"Failed to convert VML cell image to PNG: {e}" + ) + continue + if len(png_data) < 10 * 1024: + continue + orig_name = target.split("/")[-1] + png_name = os.path.splitext(orig_name)[0] + ".png" + imgs_in_cell.append( + { + "image_name": png_name, + "data": png_data, + "size": len(png_data), + } + ) + """ + if imgs_in_cell: + cell_images[(row_idx, col_idx)] = imgs_in_cell + + yield ele_num, tbl, "TABLE", cell_images if cell_images else None + ele_num += 1 + map_index += 1 + else: + continue + + # --- handle p_tbl_map at the end --- + while map_index < len(p_tbl_map): + tag, node = p_tbl_map[map_index] + if tag == "p": + toc_info = detect_doc_tocs(node, ns) + label = "TOC-AREA" if toc_info["is_style"] else "PTXT" + meta = None + if "TOC" in label: + meta = { + "toc_level": toc_info.get("toc_level"), + "toc_outline_level": toc_info.get("outline_level"), + "toc_left_indent": toc_info.get("left_indent"), + "toc_style_name": toc_info.get("style_name"), + "toc_source": "tail-map", + } + yield ele_num, Paragraph(node, doc), label, meta + elif tag == "tbl": + yield ele_num, Table(node, doc), "TABLE", None + ele_num += 1 + map_index += 1 + diff --git a/apps/worker/app/services/document_parser/docx_table_html.py b/apps/worker/app/services/document_parser/docx_table_html.py new file mode 100644 index 000000000..6c60f46fa --- /dev/null +++ b/apps/worker/app/services/document_parser/docx_table_html.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from typing import Any + +from docx.table import Table as DocxTable + + +def table2html(table: DocxTable, cell_image_map: dict | None = None) -> str: + """Convert a DOCX table to HTML with colspan, rowspan, and nested tables.""" + + namespace = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"} + + def get_cell_vmerge(cell: Any) -> str | None: + tc = cell._tc + tc_pr = tc.find(".//w:tcPr", namespaces=namespace) + if tc_pr is None: + return None + + v_merge = tc_pr.find(".//w:vMerge", namespaces=namespace) + if v_merge is None: + return None + + val = v_merge.get( + "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val" + ) + return val if val else "continue" + + row_count = len(table.rows) + if row_count == 0: + return "
" + + grid = [] + for row in table.rows: + row_data = [] + previous_tc_id = None + for cell in row.cells: + tc_id = id(cell._tc) + is_new = tc_id != previous_tc_id + row_data.append((tc_id, cell, is_new)) + previous_tc_id = tc_id + grid.append(row_data) + + column_count = max(len(row) for row in grid) if grid else 0 + colspan_grid = [[0] * column_count for _ in range(row_count)] + + for row_idx in range(row_count): + row_len = len(grid[row_idx]) + col_idx = 0 + while col_idx < row_len: + tc_id = grid[row_idx][col_idx][0] + span = 1 + while ( + col_idx + span < row_len + and grid[row_idx][col_idx + span][0] == tc_id + ): + span += 1 + colspan_grid[row_idx][col_idx] = span + col_idx += span + + rowspan_grid = [[1] * column_count for _ in range(row_count)] + + for col_idx in range(column_count): + row_idx = 0 + while row_idx < row_count: + if col_idx >= len(grid[row_idx]): + row_idx += 1 + continue + + cell = grid[row_idx][col_idx][1] + vmerge = get_cell_vmerge(cell) + + if vmerge == "restart": + span = 1 + while row_idx + span < row_count: + if col_idx >= len(grid[row_idx + span]): + break + next_cell = grid[row_idx + span][col_idx][1] + next_vmerge = get_cell_vmerge(next_cell) + if next_vmerge == "continue": + span += 1 + else: + break + rowspan_grid[row_idx][col_idx] = span + row_idx += span + elif vmerge == "continue": + rowspan_grid[row_idx][col_idx] = 0 + row_idx += 1 + else: + row_idx += 1 + + html_parts = [""] + + for row_idx in range(row_count): + html_parts.append("") + col_idx = 0 + unique_col_idx = 0 + + while col_idx < len(grid[row_idx]): + _, cell, is_new = grid[row_idx][col_idx] + + if not is_new: + col_idx += 1 + continue + + rowspan = rowspan_grid[row_idx][col_idx] + if rowspan == 0: + unique_col_idx += 1 + col_idx += 1 + continue + + colspan = colspan_grid[row_idx][col_idx] + + if cell.tables: + content = "".join(table2html(nested_table) for nested_table in cell.tables) + else: + content = cell.text.strip().replace("\n", "
") + + if cell_image_map: + image_description = cell_image_map.get((row_idx, unique_col_idx)) + if image_description: + content += f"
{image_description}" + + attrs = [] + if colspan > 1: + attrs.append(f'colspan="{colspan}"') + if rowspan > 1: + attrs.append(f'rowspan="{rowspan}"') + + attr_str = " " + " ".join(attrs) if attrs else "" + html_parts.append(f"{content}") + + unique_col_idx += 1 + col_idx += colspan + + html_parts.append("") + + html_parts.append("
") + return "".join(html_parts) diff --git a/apps/worker/app/services/document_parser/excel_structure_parser.py b/apps/worker/app/services/document_parser/excel_structure_parser.py new file mode 100644 index 000000000..3ff1ea6b8 --- /dev/null +++ b/apps/worker/app/services/document_parser/excel_structure_parser.py @@ -0,0 +1,745 @@ +from __future__ import annotations + +import datetime +import io +from typing import List, Optional, Tuple, Union +from typing import cast + +import openpyxl +import pandas as pd +from pandas._typing import Axes +from loguru import logger + +from shared.core.exceptions.domain_exceptions import TableParsingException + + +def parse_excel_structure( + file_source: Union[str, io.BytesIO], + sheet_name: Optional[str] = None, + split_subtables: bool = True, + include_hidden_sheets: bool = False, +) -> dict[str, pd.DataFrame]: + try: + if isinstance(file_source, str): + workbook = openpyxl.load_workbook(file_source, data_only=True) + else: + file_source.seek(0) + workbook = openpyxl.load_workbook(file_source, data_only=True) + + results: dict[str, pd.DataFrame] = {} + sheets_to_parse = [sheet_name] if sheet_name else workbook.sheetnames + + for selected_sheet_name in sheets_to_parse: + if selected_sheet_name not in workbook.sheetnames: + logger.warning( + f"Sheet '{selected_sheet_name}' not found in workbook, skipping" + ) + continue + + worksheet = workbook[selected_sheet_name] + + if not include_hidden_sheets and worksheet.sheet_state != "visible": + logger.info( + f"Sheet '{selected_sheet_name}' is hidden " + f"(state={worksheet.sheet_state}), skipping" + ) + continue + + if worksheet.max_row is None or worksheet.max_row == 0: + logger.debug(f"Sheet '{selected_sheet_name}' is empty, skipping") + continue + + merged_ranges = list(worksheet.merged_cells.ranges) + logger.debug( + f"Sheet '{selected_sheet_name}': found {len(merged_ranges)} merged cell ranges" + ) + + if split_subtables: + subtable_regions = _split_sheet_recursive( + worksheet, + (1, worksheet.max_row), + (1, worksheet.max_column or 1), + merged_ranges, + ) + before_count = len(subtable_regions) + subtable_regions = _merge_small_subtables(worksheet, subtable_regions) + if len(subtable_regions) != before_count: + logger.info( + f"Sheet '{selected_sheet_name}': merged {before_count} subtables → " + f"{len(subtable_regions)} " + f"(absorbed {before_count - len(subtable_regions)} small fragments)" + ) + logger.debug( + f"Sheet '{selected_sheet_name}': {len(subtable_regions)} subtables after merge" + ) + + for index, (row_range, col_range) in enumerate(subtable_regions): + result = _parse_subtable( + worksheet, + row_range, + col_range, + merged_ranges, + ) + dataframe = result["df"] + dataframe.attrs["row_header_cols"] = len(result["header_cols"]) + key = selected_sheet_name if index == 0 else f"{selected_sheet_name}_{index + 1}" + logger.debug( + f"Subtable '{key}': rows={row_range}, cols={col_range}, " + f"header_rows={result['header_rows']}, header_cols={result['header_cols']}" + ) + results[key] = dataframe + else: + row_range = (1, worksheet.max_row) + col_range = (1, worksheet.max_column or 1) + result = _parse_subtable(worksheet, row_range, col_range, merged_ranges) + dataframe = result["df"] + dataframe.attrs["row_header_cols"] = len(result["header_cols"]) + logger.debug( + f"Sheet '{selected_sheet_name}': header_rows={result['header_rows']}, " + f"header_cols={result['header_cols']}, " + f"fallback_col={result['fallback_col_header']}, " + f"fallback_row={result['fallback_row_header']}" + ) + results[selected_sheet_name] = dataframe + + workbook.close() + return results + except Exception as exc: + logger.error(f"Error parsing Excel with precision mode: {exc}") + raise TableParsingException( + user_message="Failed to parse Excel file headers", + reason="EXCEL_PRECISION_PARSE_FAILED", + internal_message=str(exc), + original_exception=exc, + ) from exc + + +DATA_TYPES_TO_EXCLUDE = (int, float, datetime.datetime) + + +def _get_merged_cell_value(ws, row: int, col: int, merged_ranges: list): + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + return ws.cell(merged_range.min_row, merged_range.min_col).value + return ws.cell(row, col).value + + +def _get_unique_cells_in_row( + ws, + row: int, + col_range: Tuple[int, int], + merged_ranges: list, +) -> List[dict]: + col_start, col_end = col_range + cells = [] + visited_cols = set() + + for col in range(col_start, col_end + 1): + if col in visited_cols: + continue + + in_merge = False + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + value = ws.cell(merged_range.min_row, merged_range.min_col).value + merge_col_end = min(merged_range.max_col, col_end) + + for merged_col in range(merged_range.min_col, merge_col_end + 1): + visited_cols.add(merged_col) + + cells.append( + { + "col_start": merged_range.min_col, + "col_end": merge_col_end, + "value": value, + "is_merged": True, + } + ) + in_merge = True + break + + if not in_merge: + value = ws.cell(row, col).value + cells.append( + {"col_start": col, "col_end": col, "value": value, "is_merged": False} + ) + visited_cols.add(col) + + return cells + + +def _get_unique_cells_in_col( + ws, + col: int, + row_range: Tuple[int, int], + merged_ranges: list, +) -> List[dict]: + row_start, row_end = row_range + cells = [] + visited_rows = set() + + for row in range(row_start, row_end + 1): + if row in visited_rows: + continue + + in_merge = False + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + value = ws.cell(merged_range.min_row, merged_range.min_col).value + merge_row_end = min(merged_range.max_row, row_end) + + for merged_row in range(merged_range.min_row, merge_row_end + 1): + visited_rows.add(merged_row) + + cells.append( + { + "row_start": merged_range.min_row, + "row_end": merge_row_end, + "value": value, + "is_merged": True, + } + ) + in_merge = True + break + + if not in_merge: + value = ws.cell(row, col).value + cells.append( + {"row_start": row, "row_end": row, "value": value, "is_merged": False} + ) + visited_rows.add(row) + + return cells + + +def _is_candidate_header_row( + ws, + row: int, + col_range: Tuple[int, int], + merged_ranges: list, + exclude_types: tuple = DATA_TYPES_TO_EXCLUDE, +) -> bool: + cells = _get_unique_cells_in_row(ws, row, col_range, merged_ranges) + + has_any_value = False + for cell in cells: + value = cell["value"] + if value is None: + continue + has_any_value = True + + if isinstance(value, bool): + continue + if isinstance(value, exclude_types): + return False + + return has_any_value + + +def _is_candidate_header_col( + ws, + col: int, + row_range: Tuple[int, int], + merged_ranges: list, + exclude_types: tuple = DATA_TYPES_TO_EXCLUDE, +) -> bool: + cells = _get_unique_cells_in_col(ws, col, row_range, merged_ranges) + + has_any_value = False + for cell in cells: + value = cell["value"] + if value is None: + continue + has_any_value = True + + if isinstance(value, bool): + continue + if isinstance(value, exclude_types): + return False + + return has_any_value + + +def _detect_header_regions( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], + merged_ranges: list, +) -> Tuple[List[int], List[int]]: + row_start, row_end = row_range + col_start, col_end = col_range + + header_rows = [] + for row in range(row_start, row_end + 1): + if _is_candidate_header_row(ws, row, col_range, merged_ranges): + header_rows.append(row) + else: + break + + data_row_start = header_rows[-1] + 1 if header_rows else row_start + if data_row_start > row_end: + return header_rows, [] + + header_cols = [] + data_row_range = (data_row_start, row_end) + for col in range(col_start, col_end + 1): + if _is_candidate_header_col(ws, col, data_row_range, merged_ranges): + header_cols.append(col) + else: + break + + return header_rows, header_cols + + +def _build_column_multiindex( + ws, + header_rows: List[int], + col_range: Tuple[int, int], + merged_ranges: list, +) -> Union[pd.Index, pd.MultiIndex]: + col_start, col_end = col_range + levels = [] + + for row in header_rows: + row_values = [] + for col in range(col_start, col_end + 1): + value = _get_merged_cell_value(ws, row, col, merged_ranges) + row_values.append(str(value).strip() if value else "") + levels.append(row_values) + + for index, level in enumerate(levels): + filled = [] + last = "" + for value in level: + if value: + last = value + filled.append(last if last else value) + levels[index] = filled + + if len(levels) == 1: + return pd.Index(levels[0]) + return pd.MultiIndex.from_arrays(levels) + + +def _build_row_multiindex( + ws, + header_cols: List[int], + row_range: Tuple[int, int], + merged_ranges: list, + header_rows: List[int] | None = None, +) -> Union[pd.Index, pd.MultiIndex]: + row_start, row_end = row_range + levels = [] + names = [] + + for col in header_cols: + col_values = [] + for row in range(row_start, row_end + 1): + value = _get_merged_cell_value(ws, row, col, merged_ranges) + col_values.append(str(value).strip() if value else "") + levels.append(col_values) + + if header_rows: + name_row = header_rows[-1] + name_value = _get_merged_cell_value(ws, name_row, col, merged_ranges) + names.append(str(name_value).strip() if name_value else None) + else: + names.append(None) + + for index, level in enumerate(levels): + filled = [] + last = "" + for value in level: + if value: + last = value + filled.append(last if last else value) + levels[index] = filled + + if len(levels) == 1: + row_index = pd.Index(levels[0]) + row_index.name = names[0] if names else None + return row_index + return pd.MultiIndex.from_arrays(levels, names=names) + + +def _parse_subtable( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], + merged_ranges: list, +) -> dict: + row_start, row_end = row_range + col_start, col_end = col_range + + header_rows, header_cols = _detect_header_regions( + ws, row_range, col_range, merged_ranges + ) + + total_rows = row_end - row_start + 1 + total_cols = col_end - col_start + 1 + fallback_col_header = len(header_rows) == total_rows + fallback_row_header = len(header_cols) == total_cols + + if fallback_col_header: + data_row_start = row_start + columns = None + else: + data_row_start = header_rows[-1] + 1 if header_rows else row_start + columns = ( + _build_column_multiindex(ws, header_rows, col_range, merged_ranges) + if header_rows + else None + ) + + if fallback_row_header: + data_col_start = col_start + row_index = None + else: + data_col_start = header_cols[-1] + 1 if header_cols else col_start + row_index = ( + _build_row_multiindex( + ws, header_cols, (data_row_start, row_end), merged_ranges, header_rows + ) + if header_cols + else None + ) + + data = [] + for row in range(data_row_start, row_end + 1): + row_data = [] + for col in range(data_col_start, col_end + 1): + value = _get_merged_cell_value(ws, row, col, merged_ranges) + row_data.append(value) + data.append(row_data) + + if columns is not None and header_cols and not fallback_row_header: + columns = cast(Axes, columns[len(header_cols) :]) + + dataframe = pd.DataFrame(data, columns=columns, index=row_index) + + excel_row_numbers = list(range(data_row_start, row_end + 1)) + if isinstance(dataframe.columns, pd.MultiIndex): + level_count = dataframe.columns.nlevels + src_row_key = tuple(["_src_row"] + [""] * (level_count - 1)) + dataframe[src_row_key] = excel_row_numbers + else: + dataframe["_src_row"] = excel_row_numbers + + return { + "df": dataframe, + "header_rows": header_rows if not fallback_col_header else [], + "header_cols": header_cols if not fallback_row_header else [], + "fallback_col_header": fallback_col_header, + "fallback_row_header": fallback_row_header, + } + + +def _find_effective_range( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], +) -> Tuple[Tuple[int, int], Tuple[int, int]]: + row_start, row_end = row_range + col_start, col_end = col_range + + effective_row_start = None + effective_row_end = None + effective_col_start = None + effective_col_end = None + + for row in range(row_start, row_end + 1): + for col in range(col_start, col_end + 1): + if ws.cell(row, col).value is not None: + if effective_row_start is None: + effective_row_start = row + effective_row_end = row + if effective_col_start is None or col < effective_col_start: + effective_col_start = col + if effective_col_end is None or col > effective_col_end: + effective_col_end = col + + if effective_row_start is None: + return ((row_start, row_start), (col_start, col_start)) + + if ( + effective_row_end is None + or effective_col_start is None + or effective_col_end is None + ): + return ((row_start, row_start), (col_start, col_start)) + + return ( + (effective_row_start, effective_row_end), + (effective_col_start, effective_col_end), + ) + + +def _is_true_separator_row( + ws, + row: int, + effective_col_range: Tuple[int, int], + merged_ranges: list | None = None, +) -> bool: + col_start, col_end = effective_col_range + merged_ranges = merged_ranges or [] + + for col in range(col_start, col_end + 1): + if ws.cell(row, col).value is not None: + return False + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + return False + return True + + +def _is_true_separator_col( + ws, + col: int, + effective_row_range: Tuple[int, int], + merged_ranges: list | None = None, +) -> bool: + row_start, row_end = effective_row_range + merged_ranges = merged_ranges or [] + + for row in range(row_start, row_end + 1): + if ws.cell(row, col).value is not None: + return False + for merged_range in merged_ranges: + if ( + merged_range.min_row <= row <= merged_range.max_row + and merged_range.min_col <= col <= merged_range.max_col + ): + return False + return True + + +def _find_separator_groups(items: List[int]) -> List[List[int]]: + if not items: + return [] + + groups = [] + current_group = [items[0]] + + for index in range(1, len(items)): + if items[index] == items[index - 1] + 1: + current_group.append(items[index]) + else: + groups.append(current_group) + current_group = [items[index]] + + groups.append(current_group) + return groups + + +def _split_sheet_recursive( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], + merged_ranges: list | None = None, + min_rows: int = 2, + min_cols: int = 2, +) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: + row_start, row_end = row_range + col_start, col_end = col_range + merged_ranges = merged_ranges or [] + + (effective_row_start, effective_row_end), ( + effective_col_start, + effective_col_end, + ) = _find_effective_range(ws, row_range, col_range) + + if ( + effective_row_end - effective_row_start + 1 < min_rows + or effective_col_end - effective_col_start + 1 < min_cols + ): + if effective_row_start is not None: + return [ + ( + (effective_row_start, effective_row_end), + (effective_col_start, effective_col_end), + ) + ] + return [] + + separator_rows = [] + for row in range(effective_row_start + 1, effective_row_end): + if _is_true_separator_row( + ws, row, (effective_col_start, effective_col_end), merged_ranges + ): + separator_rows.append(row) + + separator_cols = [] + for col in range(effective_col_start + 1, effective_col_end): + if _is_true_separator_col( + ws, col, (effective_row_start, effective_row_end), merged_ranges + ): + separator_cols.append(col) + + row_groups = _find_separator_groups(separator_rows) + col_groups = _find_separator_groups(separator_cols) + + should_split_rows = len(row_groups) > 0 and ( + len(col_groups) == 0 or len(row_groups) <= len(col_groups) + ) + should_split_cols = len(col_groups) > 0 and not should_split_rows + + if should_split_rows: + subtables = [] + previous_end = effective_row_start + for group in row_groups: + if group[0] > previous_end: + sub_result = _split_sheet_recursive( + ws, + (previous_end, group[0] - 1), + (effective_col_start, effective_col_end), + merged_ranges, + min_rows, + min_cols, + ) + subtables.extend(sub_result) + previous_end = group[-1] + 1 + if previous_end <= effective_row_end: + sub_result = _split_sheet_recursive( + ws, + (previous_end, effective_row_end), + (effective_col_start, effective_col_end), + merged_ranges, + min_rows, + min_cols, + ) + subtables.extend(sub_result) + return subtables + + if should_split_cols: + subtables = [] + previous_end = effective_col_start + for group in col_groups: + if group[0] > previous_end: + sub_result = _split_sheet_recursive( + ws, + (effective_row_start, effective_row_end), + (previous_end, group[0] - 1), + merged_ranges, + min_rows, + min_cols, + ) + subtables.extend(sub_result) + previous_end = group[-1] + 1 + if previous_end <= effective_col_end: + sub_result = _split_sheet_recursive( + ws, + (effective_row_start, effective_row_end), + (previous_end, effective_col_end), + merged_ranges, + min_rows, + min_cols, + ) + subtables.extend(sub_result) + return subtables + + return [((effective_row_start, effective_row_end), (effective_col_start, effective_col_end))] + + +def _count_non_empty_cells( + ws, + row_range: Tuple[int, int], + col_range: Tuple[int, int], +) -> int: + count = 0 + for row in range(row_range[0], row_range[1] + 1): + for col in range(col_range[0], col_range[1] + 1): + if ws.cell(row, col).value is not None: + count += 1 + return count + + +def _merge_small_subtables( + ws, + subtables: List[Tuple[Tuple[int, int], Tuple[int, int]]], + min_cells: int = 4, +) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: + if len(subtables) <= 1: + return subtables + + items = [] + for row_range, col_range in subtables: + count = _count_non_empty_cells(ws, row_range, col_range) + items.append({"rr": row_range, "cr": col_range, "cells": count}) + + changed = True + while changed and len(items) > 1: + changed = False + + min_index = None + for index, item in enumerate(items): + if item["cells"] < min_cells: + if min_index is None or item["cells"] < items[min_index]["cells"]: + min_index = index + + if min_index is None: + break + + source = items[min_index] + best_index = None + best_distance = float("inf") + for index, target in enumerate(items): + if index == min_index: + continue + row_gap = max( + 0, + target["rr"][0] - source["rr"][1] - 1, + source["rr"][0] - target["rr"][1] - 1, + ) + col_gap = max( + 0, + target["cr"][0] - source["cr"][1] - 1, + source["cr"][0] - target["cr"][1] - 1, + ) + distance = row_gap + col_gap + if distance < best_distance or ( + best_index is not None + and distance == best_distance + and target["cells"] > items[best_index]["cells"] + ): + best_distance = distance + best_index = index + + if best_index is None: + break + + target = items[best_index] + merged_row_range = ( + min(source["rr"][0], target["rr"][0]), + max(source["rr"][1], target["rr"][1]), + ) + merged_col_range = ( + min(source["cr"][0], target["cr"][0]), + max(source["cr"][1], target["cr"][1]), + ) + items[best_index] = { + "rr": merged_row_range, + "cr": merged_col_range, + "cells": source["cells"] + target["cells"], + } + + logger.debug( + f"Merged small fragment (rows={source['rr']}, cols={source['cr']}, " + f"cells={source['cells']}) into neighbor (rows={target['rr']}, cols={target['cr']})" + ) + + del items[min_index] + changed = True + + return [(item["rr"], item["cr"]) for item in items] diff --git a/apps/worker/app/services/document_parser/heading_candidates.py b/apps/worker/app/services/document_parser/heading_candidates.py new file mode 100644 index 000000000..ccf293e16 --- /dev/null +++ b/apps/worker/app/services/document_parser/heading_candidates.py @@ -0,0 +1,488 @@ +from __future__ import annotations + +import re +import unicodedata +from collections import defaultdict +from typing import Any + +import pandas as pd +from docx.oxml.ns import qn +from loguru import logger +from pandas import Index + +from app.services.document_parser.text_helpers import count_cn_en + +HEADING_COLUMNS = Index(["id", "heading", "level", "reason"]) + + +def get_max_lvl(code_str: str): + match = re.search(r"\[([^]]+)]", code_str) + if not match: + return "Sure" + + nums = [int(item.strip()) for item in match.group(1).split(",")] + max_value = int(max(nums)) + return max_value if max_value > 1 else -2 + + +def judge_by_conditions( + text, + scope: int = 20, + return_detail: bool = False, + cn_special_index: int = 12, + **legacy_options: Any, +): + legacy_cn_special_index = legacy_options.pop("CN_SPECIAL_IDX", None) + if legacy_cn_special_index is not None: + cn_special_index = int(legacy_cn_special_index) + if legacy_options: + unknown_options = ", ".join(sorted(legacy_options)) + raise TypeError(f"Unknown heading condition option(s): {unknown_options}") + + text = text.replace("\u3000", " ") + text = unicodedata.normalize("NFKC", text)[:scope] + + pos_regex_conditions = [ + r"^\d+(?:\s*\.\s*\d+)+(?![、，。！？；：])(?=\s|$|\w|[一-龥])", + r"^\d、\s{0,4}(?=\S|$)", + r"^\d+\.(?!\d)\s{0,4}(?=\S)", + r"^[0-9]{1,2}\s{1,8}(?=\S)", + r"^\d+(?:\.\d+)*、\s*(?=[A-Za-z一-龥])", + r"^[一二三四五六七八九十百千万]+、\s{0,4}(?=\S|$)", + r"^[一二三四五六七八九十百千万]+(?:\s*\.[一二三四五六七八九十百千万\d]+)+", + r"^[一二三四五六七八九十百千万]+(?=\s|$)", + r"^[\(\（]\s*\d+(?:\.\d+)*(?!\.0)\s*[\)\）]", + r"^\d+(?:\.\d+)*(?!\.0)\s*[\)\）]", + r"^[\(\（]\s*[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]", + r"^[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]", + r"^第[一二三四五六七八九十百千万\d]+(?:\.[一二三四五六七八九十百千万\d]+)*(章|节|条|部分|款|目|项|编|篇|卷|辑)?(?=$|\s|[A-Za-z0-9\u4e00-\u9fa5])", + r"^[A-Za-z](?:\.\d+)*[\.、](?=\s*\S)", + r"^[\(\（]\s*[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]", + r"^[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]", + r"^((附件|附录|附表|附图)|(?i:appendix))[\s_\-—]{0,4}(?:\[)?[一二三四五六七八九十A-Za-z\d]", + ] + + pos_triggered_code = [] + reason_suffix_parts = [] + + for index, regex in enumerate(pos_regex_conditions): + match = re.match(regex, text) + if match: + matched_text = match.group(0) + count = sum(matched_text.count(symbol) for symbol in ".-") + 1 + if index == cn_special_index and return_detail: + unit_match = re.search(r"(章|节|条|部分|款|目|项|编|篇|卷|辑)", matched_text) + if unit_match: + reason_suffix_parts.append(f"[CN:{unit_match.group(1)}]") + pos_triggered_code.append(count) + else: + pos_triggered_code.append(0) + + if return_detail: + reason_suffix = " ".join(reason_suffix_parts) if reason_suffix_parts else "" + return pos_triggered_code, { + "reason_suffix": f" {reason_suffix}" if reason_suffix else "" + } + return pos_triggered_code + + +def remove_by_conditions(text, include_punc: bool = False): + neg_conditions = [ + r"^\d{3,}", + r"(?i)(^https?://\S+|^www\.\S+|^P\.S|^\b\d{0,2}\s*(?:a\.m|p\.m)\b)", + ( + r"(?:" + r"\$[^$]*\\[A-Za-z]+(?:\s*\{[^{}]*\})?[^$]*\$" + r"|" + r"\\(?:times|div|cdot|pm|mp|leq|geq|neq|approx|equiv|sim|infty" + r"|sum|prod|int|sqrt|frac|mathrm|mathbf|mathit|mathcal" + r"|text(?:bf|it|rm)?|alpha|beta|gamma|delta|epsilon|theta" + r"|lambda|mu|sigma|pi|omega|partial|nabla" + r"|left|right|begin|end|overline|underline|hat|vec|tilde)\b" + r")" + ), + r"^0\.\d+[\u4e00-\u9fa5A-Za-z\S]*", + r"^\d*\.\d+$", + r"[。！；].+", + ( + r"^\d+\.?\d*\s{0,2}" + r"(?:mm|cm|km|nm|μm|inch(?:es)?|ft|yd|mi" + r"|kg|mg|μg|lb|oz" + r"|kPa|MPa|GPa|Pa|psi|bar" + r"|°[CFK]" + r"|Hz|kHz|MHz|GHz" + r"|mol|mL|dL|dB|Nm|kN|MN|kW|MW|GW|hp|rpm|cc|cal|kcal)\b" + ), + ] + + neg_triggered_code = [] + for regex in neg_conditions: + neg_triggered_code.append(1 if re.search(regex, text) else 0) + + if include_punc: + neg_triggered_code.append(1 if re.search(r"[.,;，。；]$", text) else 0) + else: + neg_triggered_code.append(0) + + return neg_triggered_code + + +def md_heading_match(line, as_is: bool = True): + match = re.match(r"^\s*(#+)\s*(.*)$", line) + if not match: + return line, -1 + + level = len(match.group(1)) + return (line, level) if as_is else (line.lstrip("#").strip(), level) + + +def filter_markdown_headings( + md_lines: list[str], + num_pos: int = 17, + num_neg: int = 7, + layout_json_path: str | None = None, +) -> pd.DataFrame: + meta_ctx = None + if layout_json_path: + try: + from app.services.document_parser.metadata_extractor import MetadataContext + + meta_ctx = MetadataContext(md_lines, layout_json_path) + except Exception as exc: + logger.warning(f"Failed to create MetadataContext: {exc}") + + raw_candidates = [] + for line_index, line in enumerate(md_lines): + line = line.strip() + if not line: + continue + + if _is_non_heading_markdown_line(line): + est_level = -1 + zero_pos_code = [0] * num_pos + zero_neg_code = [0] * num_neg + reason = f"POS {zero_pos_code} NEG {zero_neg_code}" + if meta_ctx: + reason += " META [0, 0, 0]" + line = "Figure/Image" + else: + est_level, reason, line = _estimate_markdown_heading_level(line, meta_ctx) + + raw_candidates.append((line_index, line, est_level, reason)) + + return pd.DataFrame(raw_candidates, columns=HEADING_COLUMNS, index=None) + + +def filter_document_headings( + heading_infos: list[tuple[Any, Any, str]], + *, + enable_regex: bool = True, +) -> pd.DataFrame: + raw_candidates = [] + logger.debug("Filtering docx heading candidates... total_items={}", len(heading_infos)) + + for element_id, paragraph, text in heading_infos: + reason = "" + est_level = None + style_level = _find_docx_style_level(paragraph) + setting_level = _find_docx_outline_level(paragraph) + + if style_level is not None: + est_level = style_level + reason = f"style-{style_level}" + elif setting_level is not None: + est_level = setting_level + reason = f"outline-{setting_level}" + + is_bold = 1 if _is_bold_docx_paragraph(paragraph) else 0 + + if enable_regex: + pos_code, detail_info = judge_by_conditions(text, return_detail=True) + neg_code = remove_by_conditions(text) + + if any(value > 0 for value in neg_code): + code_level = -1 + code_reason = ( + f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" + ) + elif any(value > 0 for value in pos_code) and all( + value == 0 for value in neg_code + ): + code_level = get_max_lvl(str(pos_code)) + code_reason = ( + f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" + ) + else: + code_level = -1 + code_reason = f"POS {pos_code} NEG {neg_code}" + + if is_bold: + code_reason += f" META [0, 0, {is_bold}]" + + if est_level is None: + est_level = code_level + reason = code_reason + else: + reason = f"{reason} AND {code_reason}" + + raw_candidates.append((element_id, text, est_level, reason)) + + candidates = pd.DataFrame(raw_candidates, columns=HEADING_COLUMNS, index=None) + if candidates.empty: + return pd.DataFrame(columns=HEADING_COLUMNS) + + candidates = postprocess_headings(candidates, task="merge_continuous") + return postprocess_headings(candidates, task="merge_short") + + +def postprocess_headings(df: pd.DataFrame, task: str, max_depth: int = -1) -> pd.DataFrame: + if task == "judge_negs": + return _judge_negative_headings(df) + + if task == "merge_continuous": + return _merge_continuous_non_headings(df) + + if task == "merge_short" or task == "collapse": + return _collapse_heading_groups(df, task) + + return df + + +def _is_non_heading_markdown_line(line: str) -> bool: + return ( + ("" in line) + or line.startswith("|") + or line.startswith("") + or ("![" in line and "](" in line) + ) + + +def _estimate_markdown_heading_level(line: str, meta_ctx: Any | None): + from app.services.document_parser.metadata_extractor import detect_and_strip_md_bold + + line_clean, hash_level = md_heading_match(line, as_is=False) + stripped_line, is_full_bold = detect_and_strip_md_bold(line_clean) + pos_code, detail_info = judge_by_conditions(stripped_line, return_detail=True) + neg_code = remove_by_conditions(stripped_line) + + if any(value > 0 for value in neg_code): + code_level = -1 + code_reason = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" + elif any(value > 0 for value in pos_code) and all(value == 0 for value in neg_code): + code_level = get_max_lvl(str(pos_code)) + code_reason = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" + else: + code_level = -1 + code_reason = f"POS {pos_code} NEG {neg_code}" + + if meta_ctx: + size_rank, occurrence = meta_ctx.get_meta_for_line(line_clean) + bold_value = 1 if is_full_bold else 0 + code_reason += meta_ctx.format_meta_suffix(size_rank, occurrence, bold_value) + elif is_full_bold: + code_reason += " META [0, 0, 1]" + + if hash_level <= 0: + return code_level, code_reason, line_clean + + if isinstance(code_level, int): + est_level = max(hash_level, code_level) + else: + est_level = code_level + return est_level, f"{hash_level}# AND {code_reason}", line_clean + + +def _find_docx_style_level(paragraph: Any): + try: + style_name = paragraph.style.name + except Exception: + style_name = "normal" + + if not (style_name.startswith("Heading") or style_name.startswith("标题")): + return None + + try: + return int(style_name.split(" ")[1]) + except Exception: + return -2 + + +def _find_docx_outline_level(paragraph: Any): + paragraph_properties = paragraph._element.find(qn("w:pPr")) + if paragraph_properties is None: + return None + + outline_level = paragraph_properties.find(qn("w:outlineLvl")) + if outline_level is None: + return None + + return int(outline_level.get(qn("w:val"))) + 1 + + +def _is_bold_docx_paragraph(paragraph: Any): + if paragraph.runs and all(run.bold for run in paragraph.runs if run.text.strip()): + return True + return None + + +def _judge_negative_headings(df: pd.DataFrame) -> pd.DataFrame: + for index, row in df.iterrows(): + neg_code = remove_by_conditions(row["heading"], include_punc=True) + if any(value > 0 for value in neg_code): + current_code = str(df.loc[index, "reason"]) + + neg_match = re.search(r"(.*NEG\s*)\[[^\]]*\](.*)", current_code) + if neg_match: + updated_code = f"{neg_match.group(1)}{neg_code}{neg_match.group(2)}" + else: + updated_code = f"{current_code} NEG {neg_code}" + + df.loc[index, "level"] = -1 + df.loc[index, "reason"] = updated_code + return df + + +def _merge_continuous_non_headings(df: pd.DataFrame) -> pd.DataFrame: + denoised_rows = [] + punc_pattern = re.compile(r'[.,!?;:，。！？；：）】〕｝〉》’”"]$') + + index = 0 + while index < len(df): + row = df.iloc[index] + current_content = str(row["heading"]).strip() + current_level = row["level"] + + next_index = index + 1 + while next_index < len(df): + next_row = df.iloc[next_index] + next_content = str(next_row["heading"]).strip() + next_level = next_row["level"] + + expected_id = row["id"] + (next_index - index) + if next_row["id"] != expected_id: + break + + current_not_punc = not punc_pattern.search(current_content[-2:]) + if (current_level == -1 and next_level == -1) and current_not_punc: + current_content += " " + next_content + next_index += 1 + else: + break + + merged_row = row.copy() + merged_row["heading"] = current_content + denoised_rows.append(tuple(merged_row)) + index = next_index + + return pd.DataFrame(denoised_rows, columns=HEADING_COLUMNS) + + +def _collapse_heading_groups(df: pd.DataFrame, task: str) -> pd.DataFrame: + group_to_indices = defaultdict(list) + for index, row in df.iterrows(): + level = row["level"] + reason = row["reason"] + if level != -1: + group_to_indices[(level, reason)].append(index) + + checked_pairs = set() + for _, indices in group_to_indices.items(): + _collapse_recursive(df, task, indices, merge_threshold=3, checked_pairs=checked_pairs) + + if task == "merge_short": + drop_between = df.index[ + df["reason"].astype(str).str.startswith("Merged into", na=False) + ].tolist() + if drop_between: + logger.debug( + f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows" + ) + df.drop(drop_between, inplace=True) + df.reset_index(drop=True, inplace=True) + return df + + +def _collapse_recursive( + df: pd.DataFrame, + task: str, + indices: list[int], + merge_threshold: int = 3, + checked_pairs: set[tuple[int, int]] | None = None, +) -> None: + if checked_pairs is None: + checked_pairs = set() + + if len(indices) < 2: + return + + for group_index in range(len(indices) - 1): + index, next_index = indices[group_index], indices[group_index + 1] + if (index, next_index) in checked_pairs: + continue + checked_pairs.add((index, next_index)) + + between = df.loc[index + 1 : next_index - 1] + current_text = df.at[index, "heading"].strip() + next_text = df.at[next_index, "heading"].strip() + + if task == "merge_short" and len(between) > 0: + _merge_short_between_headings( + df, between, index, next_index, current_text, merge_threshold + ) + elif task == "collapse" and len(between) == 0: + logger.debug( + f"⚠️ Empty between i={current_text[:15]}, j={next_text[:15]} => set i.level=-1, j.level=Not Sure" + ) + df.at[index, "level"] = -2 + df.at[next_index, "level"] = -2 + + sub_between = between[between["level"] != -1] + code_to_sub = defaultdict(list) + for row_index, row in sub_between.iterrows(): + level = row["level"] + reason = row["reason"] + if level != -1: + code_to_sub[(level, reason)].append(row_index) + + for _, sub_indices in code_to_sub.items(): + _collapse_recursive(df, task, sub_indices, merge_threshold, checked_pairs) + + +def _merge_short_between_headings( + df: pd.DataFrame, + between: pd.DataFrame, + index: int, + next_index: int, + current_text: str, + merge_threshold: int, +) -> None: + between_lengths = [count_cn_en(content) for content in between["heading"].tolist()] + between_levels = [level for level in between["level"].tolist()] + half_current_length = int(count_cn_en(current_text) / 2) + too_short = ( + sum(between_lengths) <= merge_threshold + or sum(between_lengths) < half_current_length + ) + + if not too_short or not all(level == -1 for level in between_levels): + return + + next_text = df.at[next_index, "heading"].strip() + logger.debug( + f"⚠️ too short between {index}=>{current_text[:15]} and {next_index}=>{next_text[:15]} => merge to {index}" + ) + between_texts = [ + heading_text.strip() + for _, row in between.iterrows() + if isinstance(heading_text := row["heading"], str) and heading_text.strip() + ] + + joined_text = "" + if between_texts: + joined_text = "\n".join(between_texts) + df.at[index, "heading"] = f"{current_text} {joined_text}" + + for row_index in between.index: + df.at[row_index, "level"] = -1 + df.at[row_index, "reason"] = f"Merged into {index}" + logger.debug(f"\tmerged texts: {joined_text[:50]}...") diff --git a/apps/worker/app/services/document_parser/heading_hierarchy.py b/apps/worker/app/services/document_parser/heading_hierarchy.py new file mode 100644 index 000000000..e53075b62 --- /dev/null +++ b/apps/worker/app/services/document_parser/heading_hierarchy.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + +import pandas as pd + +from app.services.document_parser.layout_parser import pred_titles + + +@dataclass(frozen=True) +class HeadingHierarchyInput: + infos: Any + doc_type: Literal["pptx", "md", "docx"] + toc_hierarchies: Any | None = None + prompt_limit: int = 4000 + enable_regex: bool = True + smart_parse: bool = False + model_name: str | None = None + output_dir: str | None = None + layout_json_path: str | None = None + first_toc_ele_num: int | None = None + + +def predict_heading_hierarchy(heading_input: HeadingHierarchyInput) -> pd.DataFrame: + return pred_titles( + heading_input.infos, + doc_type=heading_input.doc_type, + toc_hierarchies=heading_input.toc_hierarchies, + prompt_limt=heading_input.prompt_limit, + enable_regx=heading_input.enable_regex, + smart_parse=heading_input.smart_parse, + model_name=heading_input.model_name, + output_dir=heading_input.output_dir, + layout_json_path=heading_input.layout_json_path, + first_toc_ele_num=heading_input.first_toc_ele_num, + ) diff --git a/apps/worker/app/services/document_parser/heading_llm_executor.py b/apps/worker/app/services/document_parser/heading_llm_executor.py new file mode 100644 index 000000000..2f72789e6 --- /dev/null +++ b/apps/worker/app/services/document_parser/heading_llm_executor.py @@ -0,0 +1,610 @@ +# pyright: reportArgumentType=false, reportAssignmentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalSubscript=false, reportReturnType=false +from __future__ import annotations + +import os +import re +from collections import Counter +from collections.abc import Callable +from typing import Any + +import pandas as pd +from app.services.document_parser.metadata_extractor import clean_md_text_for_llm +from app.services.document_parser.stage_profiler import stage_timer +from app.services.document_parser.text_helpers import count_cn_en, truncate_text_by_tokens +from loguru import logger + +from shared.core.exceptions.domain_exceptions import WorkerHandlingException + +PLACEHOLDER_REASON = "__PLACEHOLDER__" + +HierarchyJudge = Callable[..., list[dict[str, Any]]] +FallbackHierarchy = Callable[[pd.DataFrame], pd.DataFrame] +SaveIntermediateCsv = Callable[[pd.DataFrame, str | None, str], None] + + +def build_level_mapping( + df: pd.DataFrame, origin_lvls: list[int], mode: str = "max" +) -> tuple[pd.DataFrame, dict[str, dict[str, Any]]]: + mapped_df = df.copy() + mapped_df["origin_level"] = origin_lvls + + mapping = mapped_df.groupby("reason")["level"].apply(list).to_dict() + + processed_mapping: dict[str, dict[str, Any]] = {} + for reason, lvls in mapping.items(): + positive_lvls = [lvl for lvl in lvls if lvl > -1] + counts = Counter(lvls) + + if not positive_lvls: + mapped_lvl = -1 + elif mode == "max": + mapped_lvl = max(positive_lvls) + elif mode == "freq": + mapped_lvl = counts.most_common(1)[0][0] + else: + raise WorkerHandlingException( + internal_message=f"wrong input mode: {mode}. Must be 'max' or 'freq'" + ) + + processed_mapping[reason] = { + "lvls": lvls, + "positive_lvls": positive_lvls, + "freqs": dict(counts), + "mapped_lvl": mapped_lvl, + } + return mapped_df, processed_mapping + + +def execute_level_mapping( + df: pd.DataFrame, mapping: dict[str, dict[str, Any]] +) -> pd.DataFrame: + def map_row(row: pd.Series) -> int: + reason = row["reason"] + if reason in mapping: + return int(mapping[reason]["mapped_lvl"]) + return int(row["level"]) + + mapped_df = df.copy() + origin_est_lvls = mapped_df["level"].tolist() + mapped_df["level"] = mapped_df.apply(map_row, axis=1) + mapped_df["origin_level"] = origin_est_lvls + return mapped_df + + +def extract_non_neg_code(reason_str: str) -> str: + """Extract the non-NEG code from a heading reason string.""" + if not reason_str or not isinstance(reason_str, str): + return "" + neg_match = re.search(r"\s*NEG\s*\[[^\]]*\]", reason_str) + if neg_match: + before_neg = reason_str[: neg_match.start()] + after_neg = reason_str[neg_match.end() :] + return (before_neg + after_neg).strip() + return reason_str.strip() + + +def build_non_neg_mapping(lvl_mapping: dict[str, dict[str, Any]]) -> dict[str, int]: + non_neg_levels: dict[str, list[int]] = {} + for reason, info in lvl_mapping.items(): + non_neg_code = extract_non_neg_code(reason) + mapped_lvl = int(info.get("mapped_lvl", -1)) + if non_neg_code: + non_neg_levels.setdefault(non_neg_code, []).append(mapped_lvl) + + non_neg_mapping: dict[str, int] = {} + for non_neg_code, levels in non_neg_levels.items(): + positive_levels = [lvl for lvl in levels if lvl > -1] + if positive_levels: + level_counts = Counter(positive_levels) + non_neg_mapping[non_neg_code] = level_counts.most_common(1)[0][0] + else: + non_neg_mapping[non_neg_code] = -1 + + return non_neg_mapping + + +def handle_unseen_codes( + df: pd.DataFrame, + level_dfs: list[pd.DataFrame], + lvl_mapping: dict[str, dict[str, Any]], + output_dir: str | None = None, + window_half_size: int = 10, + strategy: str = "double_mapping", +) -> dict[str, dict[str, Any]]: + """Extend first-chunk reason mapping to reason codes only seen in later chunks.""" + + def extract_reason_signature(reason: str) -> str: + return reason.strip() if reason else "" + + def has_neg_signal(reason_str: str) -> bool: + if not reason_str or not isinstance(reason_str, str): + return False + neg_match = re.search(r"NEG\s*\[([^\]]*)\]", reason_str) + if not neg_match: + return False + neg_content = neg_match.group(1) + try: + nums = [int(x.strip()) for x in neg_content.split(",") if x.strip()] + return any(x >= 1 for x in nums) + except Exception: + return False + + def build_context_window( + target_idx: int, known_codes_set: set[str], total_rows: int, half_size: int = 10 + ) -> dict[str, Any]: + min_start = max(0, target_idx - half_size) + min_end = min(total_rows - 1, target_idx + half_size) + + start_idx = min_start + end_idx = min_end + + found_known_above = False + found_known_below = False + known_positions: list[int] = [] + + for index in range(start_idx, target_idx): + reason = df.iloc[index].get("reason", "") + sig = extract_reason_signature(reason) + if sig in known_codes_set: + found_known_above = True + known_positions.append(index) + + for index in range(target_idx + 1, end_idx + 1): + reason = df.iloc[index].get("reason", "") + sig = extract_reason_signature(reason) + if sig in known_codes_set: + found_known_below = True + known_positions.append(index) + + if not found_known_above and min_start > 0: + search_idx = min_start - 1 + while search_idx >= 0: + reason = df.iloc[search_idx].get("reason", "") + sig = extract_reason_signature(reason) + if sig in known_codes_set: + found_known_above = True + known_positions.append(search_idx) + start_idx = search_idx + break + search_idx -= 1 + + if not found_known_below and min_end < total_rows - 1: + search_idx = min_end + 1 + while search_idx < total_rows: + reason = df.iloc[search_idx].get("reason", "") + sig = extract_reason_signature(reason) + if sig in known_codes_set: + found_known_below = True + known_positions.append(search_idx) + end_idx = search_idx + break + search_idx += 1 + + return { + "start": start_idx, + "end": end_idx, + "found_known": found_known_above or found_known_below, + "known_positions": known_positions, + } + + non_neg_mapping = build_non_neg_mapping(lvl_mapping) + known_codes = set(lvl_mapping.keys()) + + all_codes_in_full: dict[str, dict[str, Any]] = {} + for seg_idx, seg_df in enumerate(level_dfs): + for _, row in seg_df.iterrows(): + reason = row.get("reason", "") + sig = extract_reason_signature(reason) + if not sig or sig == PLACEHOLDER_REASON: + continue + if sig not in all_codes_in_full: + all_codes_in_full[sig] = { + "first_seg": seg_idx, + "first_id": row.get("id", 0), + "reason": reason, + } + + unseen_codes: dict[str, dict[str, Any]] = {} + unseen_neg_filtered: dict[str, dict[str, Any]] = {} + for sig, info in all_codes_in_full.items(): + if sig in known_codes: + continue + if has_neg_signal(info["reason"]): + unseen_neg_filtered[sig] = info + else: + unseen_codes[sig] = info + + logger.info( + f"Unseen codes total: {len(unseen_codes) + len(unseen_neg_filtered)}, " + f"NEG filtered: {len(unseen_neg_filtered)}, to process: {len(unseen_codes)}" + ) + + for sig in unseen_neg_filtered: + lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NEG_FILTERED"} + + if unseen_codes: + if strategy == "double_mapping": + fallback_success = 0 + fallback_failed = 0 + failed_codes = [] + for sig in unseen_codes: + non_neg_code = extract_non_neg_code(sig) + if non_neg_code in non_neg_mapping: + mapped_level = non_neg_mapping[non_neg_code] + lvl_mapping[sig] = { + "mapped_lvl": mapped_level, + "note": f"NON_NEG_FALLBACK from '{non_neg_code}'", + } + fallback_success += 1 + else: + lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NO_MATCH_FALLBACK"} + fallback_failed += 1 + failed_codes.append( + f"'{non_neg_code}' (from '{sig[:60]}...')" + if len(sig) > 60 + else f"'{non_neg_code}' (from '{sig}')" + ) + + logger.info( + f"Double mapping result: success={fallback_success}, failed={fallback_failed}" + ) + if failed_codes: + logger.warning( + f"Failed codes (non_neg not in mapping): {failed_codes[:5]}" + f"{'...' if len(failed_codes) > 5 else ''}" + ) + + elif strategy == "window_llm" and output_dir: + total_rows = len(df) + windows: list[dict[str, Any]] = [] + for sig, info in unseen_codes.items(): + first_id = info["first_id"] + first_seg = info["first_seg"] + df_indices = df.index[df["id"] == first_id].tolist() + if df_indices: + first_df_idx = df_indices[0] + window_info = build_context_window( + first_df_idx, known_codes, total_rows, window_half_size + ) + windows.append( + { + "code": sig, + "first_id": first_id, + "first_seg": first_seg, + "start": window_info["start"], + "end": window_info["end"], + "found_known": window_info["found_known"], + } + ) + + sorted_windows = sorted(windows, key=lambda window: window["start"]) + merged_windows: list[dict[str, Any]] = [] + current_window: dict[str, Any] | None = None + + for window in sorted_windows: + if current_window is None: + current_window = { + "start": window["start"], + "end": window["end"], + "codes": [window["code"]], + "segments": [window["first_seg"]], + } + elif window["start"] <= current_window["end"]: + current_window["end"] = max(current_window["end"], window["end"]) + current_window["codes"].append(window["code"]) + current_window["segments"].append(window["first_seg"]) + else: + merged_windows.append(current_window) + current_window = { + "start": window["start"], + "end": window["end"], + "codes": [window["code"]], + "segments": [window["first_seg"]], + } + + if current_window: + merged_windows.append(current_window) + + windows_dir = os.path.join(output_dir, "merged_windows") + os.makedirs(windows_dir, exist_ok=True) + + unseen_codes_set = set(unseen_codes.keys()) + unseen_neg_set = set(unseen_neg_filtered.keys()) + + for index, merged_window in enumerate(merged_windows): + window_df = df.iloc[ + merged_window["start"] : merged_window["end"] + 1 + ].copy() + + def get_code_status(row: pd.Series) -> str: + reason = row.get("reason", "") + sig = extract_reason_signature(reason) + if not sig: + return "" + if sig in unseen_codes_set: + return "UNSEEN_TARGET" + if sig in unseen_neg_set: + return "NEG_TO_NEGATIVE_ONE" + if sig in known_codes: + return "KNOWN" + return "" + + window_df["code_status"] = window_df.apply(get_code_status, axis=1) + window_path = os.path.join( + windows_dir, + f"window_{index + 1:02d}_rows_" + f"{merged_window['start']}-{merged_window['end']}.csv", + ) + window_df.to_csv(window_path, index=False, encoding="utf-8-sig") + + logger.debug( + f"Window LLM: {len(merged_windows)} windows created in {windows_dir}" + ) + + return lvl_mapping + + +def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: + """Collapse consecutive body rows into placeholder rows before LLM chunking.""" + if df is None or len(df) == 0: + return pd.DataFrame(columns=["id", "heading", "level", "reason"]) + + rows: list[dict[str, Any]] = [] + index = 0 + row_count = len(df) + while index < row_count: + lvl_raw = df.iloc[index]["level"] + try: + lvl_int = int(lvl_raw) + except (TypeError, ValueError): + lvl_int = None + + if lvl_int == -1: + end_index = index + while end_index < row_count: + try: + next_level = int(df.iloc[end_index]["level"]) + except (TypeError, ValueError): + break + if next_level != -1: + break + end_index += 1 + start_id = int(df.iloc[index]["id"]) + end_id = int(df.iloc[end_index - 1]["id"]) + run_length = end_index - index + rows.append( + { + "id": f"{start_id}-{end_id}", + "heading": f"[{run_length} BODY LINES]", + "level": "-", + "reason": PLACEHOLDER_REASON, + } + ) + index = end_index + else: + row = df.iloc[index] + rows.append( + { + "id": int(row["id"]), + "heading": str(row["heading"]), + "level": ( + int(lvl_int) + if lvl_int is not None and lvl_int != -2 + else "Not Sure" + ), + "reason": str(row.get("reason", "") or ""), + } + ) + index += 1 + + return pd.DataFrame(rows, columns=["id", "heading", "level", "reason"]) + + +def split_heading_table( + df: pd.DataFrame, threshold: int = 3000, max_start: int = 50, max_end: int = 10 +) -> tuple[list[pd.DataFrame], list[str]]: + raw_headings = df["heading"].tolist() + working_df = df.copy() + working_df["heading"] = working_df["heading"].apply( + lambda heading: truncate_text_by_tokens(heading, max_start, max_end) + ) + + sub_dfs: list[pd.DataFrame] = [] + current_rows: list[list[Any]] = [] + current_len = 0 + for _, row in working_df.iterrows(): + row_filtered = row.drop(labels=["reason"], errors="ignore") + row_len = sum(count_cn_en(str(value)) for value in row_filtered.values) + + if current_len + row_len > threshold and current_rows: + sub_dfs.append(pd.DataFrame(current_rows, columns=working_df.columns)) + current_rows = [row.tolist()] + current_len = row_len + else: + current_rows.append(row.tolist()) + current_len += row_len + + if current_rows: + sub_dfs.append(pd.DataFrame(current_rows, columns=working_df.columns)) + return sub_dfs, raw_headings + + +def execute_llm_heading_hierarchy( + raw_preds: pd.DataFrame, + prompt_limt: int, + hierarchy_judge: HierarchyJudge, + fallback_hierarchy: FallbackHierarchy, + save_intermediate_csv: SaveIntermediateCsv, + toc_hierarchies: Any | None = None, + max_len: int = 30, + max_depth: int = 6, + model_name: str | None = None, + output_dir: str | None = None, + csv_suffix: str = "", +) -> pd.DataFrame: + if len(raw_preds) == 0: + return pd.DataFrame(columns=["id", "heading", "level", "reason"]) + + compact_enabled = os.environ.get( + "KB_LAYOUT_LLM_COMPACT_INPUT", "true" + ).strip().lower() in ("true", "1", "yes", "on") + preds_for_llm = compact_for_llm(raw_preds) if compact_enabled else raw_preds.copy() + if compact_enabled: + placeholder_count = int(preds_for_llm["reason"].eq(PLACEHOLDER_REASON).sum()) + logger.info( + f"smart parse => compact input: {len(raw_preds)} -> {len(preds_for_llm)} rows " + f"({placeholder_count} placeholder groups)" + ) + + non_placeholder = ( + preds_for_llm[preds_for_llm["reason"].astype(str) != PLACEHOLDER_REASON] + if compact_enabled + else preds_for_llm + ) + if len(non_placeholder) == 0: + logger.info( + "smart parse => no heading candidates, skipping LLM hierarchy detection" + ) + fallback = raw_preds.copy()[["id", "heading", "level", "reason"]] + fallback["level"] = -1 + return fallback.sort_values("id").reset_index(drop=True) + + level_dfs, _raw_headings = split_heading_table( + preds_for_llm, threshold=prompt_limt, max_start=max_len, max_end=5 + ) + chunk_sizes = [len(dataframe) for dataframe in level_dfs] + logger.info( + f"smart parse => {len(level_dfs)} chunk(s) | rows per chunk: {chunk_sizes} | " + f"threshold={prompt_limt} | max_start={max_len}" + ) + + basic_idx = 0 + for idx, chunk in enumerate(level_dfs): + if (chunk["reason"].astype(str) != PLACEHOLDER_REASON).any(): + basic_idx = idx + break + basic_df = level_dfs[basic_idx] + if basic_idx != 0: + logger.info( + f"smart parse => promoted chunk {basic_idx} as basic_df " + f"(chunks 0..{basic_idx - 1} contain only placeholders)" + ) + + full_preds: pd.DataFrame | None = None + try: + with stage_timer( + "heading.hierarchy_llm", + chunk_count=len(level_dfs), + base_chunk_rows=len(basic_df), + compact_enabled=compact_enabled, + source_row_count=len(raw_preds), + model_name=model_name, + ): + logger.debug("smart parse => interpreting hierarchy patterns...") + df4llm = basic_df.drop(columns=["reason"]).copy() + df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) + logger.debug(f"DataFrame transformation completed, rows: {len(df4llm)}") + + layout_res = hierarchy_judge( + df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings" + ) + + layout_level_by_id: dict[Any, Any] = {} + if isinstance(layout_res, list): + for item in layout_res: + if isinstance(item, dict) and "id" in item and "level" in item: + layout_level_by_id[item["id"]] = item["level"] + + def level_for(row_id: Any) -> Any: + if row_id in layout_level_by_id: + return layout_level_by_id[row_id] + try: + return layout_level_by_id.get(int(row_id), -1) + except (TypeError, ValueError): + return -1 + + base_preds = ( + basic_df[["id", "heading", "reason"]].copy().reset_index(drop=True) + ) + base_preds.insert(2, "level", base_preds["id"].map(level_for)) + save_intermediate_csv( + base_preds, output_dir, f"preds_3_llm_base{csv_suffix}" + ) + + llm_levels: dict[int, Any] = {} + for _, row in base_preds.iterrows(): + row_id = row["id"] + if isinstance(row_id, bool): + continue + if isinstance(row_id, int): + llm_levels[row_id] = row["level"] + + if len(level_dfs) > 1: + placeholder_mask_base = base_preds["reason"].eq(PLACEHOLDER_REASON) + figure_mask_base = base_preds["heading"].eq("Figure/Image") + exclude_mask_base = placeholder_mask_base | figure_mask_base + base_preds_for_mapping = base_preds[~exclude_mask_base].copy() + base_origin_for_mapping = basic_df.loc[ + ~exclude_mask_base.values, "level" + ].tolist() + + base_preds_for_mapping, lvl_mapping = build_level_mapping( + base_preds_for_mapping, base_origin_for_mapping, mode="freq" + ) + logger.debug( + f"mapping development finished: {len(lvl_mapping)} rules " + f"(placeholders and Figure/Image excluded)" + ) + + logger.debug( + f"mapping dataframe to levels across {len(level_dfs)} chunks..." + ) + lvl_mapping = handle_unseen_codes( + preds_for_llm, level_dfs, lvl_mapping, output_dir + ) + + for level_df in level_dfs: + placeholder_mask_chunk = level_df["reason"].eq(PLACEHOLDER_REASON) + figure_mask_chunk = level_df["heading"].eq("Figure/Image") + exclude_mask_chunk = placeholder_mask_chunk | figure_mask_chunk + non_excluded = level_df[~exclude_mask_chunk].copy() + if not non_excluded.empty: + non_excluded = execute_level_mapping(non_excluded, lvl_mapping) + for _, row in non_excluded.iterrows(): + row_id = row["id"] + if isinstance(row_id, bool): + continue + if isinstance(row_id, int): + llm_levels[row_id] = row["level"] + logger.info( + f"multi-chunk mapping produced {len(llm_levels)} id->level entries" + ) + else: + logger.info( + "single chunk - skipping reason-code mapping, using LLM output directly" + ) + + full_preds = raw_preds.copy()[["id", "heading", "level", "reason"]] + + def resolve_level(row_id: Any) -> int: + try: + int_id = int(row_id) + except (TypeError, ValueError): + return -1 + level = llm_levels.get(int_id, -1) + try: + return int(level) + except (TypeError, ValueError): + return -1 + + full_preds["level"] = full_preds["id"].map(resolve_level).astype(int) + save_intermediate_csv( + full_preds, output_dir, f"preds_4_llm_final{csv_suffix}" + ) + + except Exception as exc: + logger.warning( + f"LLM-based parsing fails due to {exc}, using non-llm pipeline..." + ) + full_preds = fallback_hierarchy(raw_preds.copy()) + return full_preds diff --git a/apps/worker/app/services/document_parser/heading_tree.py b/apps/worker/app/services/document_parser/heading_tree.py new file mode 100644 index 000000000..3f5c39dcc --- /dev/null +++ b/apps/worker/app/services/document_parser/heading_tree.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import pandas as pd +from loguru import logger + + +def build_tree_from_dataframe( + heading_preds: pd.DataFrame, +) -> tuple[dict[str, dict], dict[tuple[str, str], int], dict[int, dict]]: + headings = heading_preds[heading_preds["level"] > -1].copy() + + node_to_id: dict[tuple[str, str], int] = {} + id_to_row: dict[int, dict] = {} + root: dict[str, dict] = {} + stack: list[tuple[int, dict, str, str]] = [(0, root, "ROOT", "")] + + for _, row in headings.iterrows(): + heading_text = str(row["heading"]) + row_id = int(row["id"]) + level = int(row["level"]) + + id_to_row[row_id] = row.to_dict() + + while len(stack) > 1 and stack[-1][0] >= level: + stack.pop() + + _, parent_dict, _, parent_path = stack[-1] + tree_node_key = heading_text + if tree_node_key in parent_dict: + tree_node_key = f"{heading_text}#{row_id}" + + node_key = (tree_node_key, parent_path) + node_to_id[node_key] = row_id + + parent_dict[tree_node_key] = {} + current_path = ( + f"{parent_path}/{tree_node_key}" if parent_path else tree_node_key + ) + stack.append((level, parent_dict[tree_node_key], tree_node_key, current_path)) + + return root, node_to_id, id_to_row + + +def tree_to_dataframe( + tree: dict[str, dict], + node_to_id: dict[tuple[str, str], int], + original_df: pd.DataFrame, +) -> pd.DataFrame: + preserved_headings = _extract_headings_from_tree(tree, node_to_id) + preserved_ids = {heading["id"] for heading in preserved_headings} + + updated_df = original_df.copy() + removed_count = 0 + level_changed_count = 0 + + for index, row in original_df.iterrows(): + row_id = int(row["id"]) + old_level = int(row["level"]) if row["level"] not in [-2, "nan", -1] else -1 + + if old_level <= -1: + continue + + if row_id in preserved_ids: + new_level = next( + ( + heading["level"] + for heading in preserved_headings + if heading["id"] == row_id + ), + old_level, + ) + updated_df.at[index, "level"] = new_level + if new_level != old_level: + level_changed_count += 1 + else: + updated_df.at[index, "level"] = -1 + removed_count += 1 + + logger.debug( + f"Tree changed: removed headings={removed_count}, " + f"level changed={level_changed_count}, preserved headings={len(preserved_ids)}" + ) + return updated_df + + +def remove_isolated_nodes(tree: dict[str, dict]) -> dict[str, dict]: + return _remove_isolated_nodes_recursive(tree) + + +def cleanup_heading_tree(heading_preds: pd.DataFrame) -> pd.DataFrame: + if heading_preds.empty: + return heading_preds + + tree, node_to_id, _ = build_tree_from_dataframe(heading_preds) + processed_tree = remove_isolated_nodes(tree) + return tree_to_dataframe(processed_tree, node_to_id, heading_preds) + + +def _extract_headings_from_tree( + node_dict: dict[str, dict], + node_to_id: dict[tuple[str, str], int], + *, + current_level: int = 1, + parent_path: str = "", +) -> list[dict[str, object]]: + results: list[dict[str, object]] = [] + for tree_node_key, children in node_dict.items(): + node_key = (tree_node_key, parent_path) + row_id = node_to_id.get(node_key, -1) + + if row_id >= 0: + original_heading = ( + tree_node_key.split("#")[0] if "#" in tree_node_key else tree_node_key + ) + results.append( + { + "id": row_id, + "heading": original_heading, + "level": current_level, + "tree_key": tree_node_key, + "parent_path": parent_path, + } + ) + + if isinstance(children, dict) and children: + current_path = ( + f"{parent_path}/{tree_node_key}" if parent_path else tree_node_key + ) + results.extend( + _extract_headings_from_tree( + children, + node_to_id, + current_level=current_level + 1, + parent_path=current_path, + ) + ) + return results + + +def _remove_isolated_nodes_recursive( + node_dict: dict[str, dict], + *, + parent_path: str = "", +) -> dict[str, dict]: + result_dict: dict[str, dict] = {} + + for heading, children in node_dict.items(): + if isinstance(children, dict) and len(children) == 1: + child_heading = list(children.keys())[0] + grandchildren = children[child_heading] + + if not grandchildren or ( + isinstance(grandchildren, dict) and len(grandchildren) == 0 + ): + result_dict[heading] = {} + logger.debug( + f"remove isolated heading: {parent_path}/{heading}/{child_heading}" + ) + else: + result_dict[heading] = _remove_isolated_nodes_recursive( + children, + parent_path=f"{parent_path}/{heading}" if parent_path else heading, + ) + elif isinstance(children, dict) and children: + result_dict[heading] = _remove_isolated_nodes_recursive( + children, + parent_path=f"{parent_path}/{heading}" if parent_path else heading, + ) + else: + result_dict[heading] = children + + return result_dict diff --git a/apps/worker/app/services/document_parser/html_parser.py b/apps/worker/app/services/document_parser/html_parser.py index 8c65a4ec7..f412067d6 100644 --- a/apps/worker/app/services/document_parser/html_parser.py +++ b/apps/worker/app/services/document_parser/html_parser.py @@ -1,21 +1,18 @@ # pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportReturnType=false """ -HTML rendering utilities for DataFrame to HTML conversion. +HTML table parsing and extraction utilities. This module provides functions for converting pandas DataFrames to HTML tables with support for: -- MultiIndex columns with colspan/rowspan merging -- Row headers (semantic structure where: - - Horizontally adjacent identical values are merged with colspan - - Vertically repeated values are merged with rowspan - - Args: - columns: pandas MultiIndex representing the column headers - escape: Whether to HTML-escape the cell values - - Returns: - HTML string for the element - """ - import html as html_lib - - n_levels = columns.nlevels - n_cols = len(columns) - - # Build a 2D grid of values [level][col] - grid = [] - for level in range(n_levels): - row = [columns.get_level_values(level)[col] for col in range(n_cols)] - grid.append(row) - - # Calculate colspan for each cell (horizontal merging) - # colspan[level][col] = number of columns this cell spans - colspan = [[1] * n_cols for _ in range(n_levels)] - - for level in range(n_levels): - col = 0 - while col < n_cols: - span = 1 - while col + span < n_cols and grid[level][col] == grid[level][col + span]: - # Check if the parent cells also match (for correct hierarchical merging) - parent_match = True - for parent_level in range(level): - if grid[parent_level][col] != grid[parent_level][col + span]: - parent_match = False - break - if parent_match: - span += 1 - else: - break - colspan[level][col] = span - col += span - - # Calculate rowspan for each cell (vertical merging) - # A cell has rowspan > 1 if all cells in the same column below have the same value - # AND if they would have the same colspan - rowspan = [[1] * n_cols for _ in range(n_levels)] - - for col in range(n_cols): - level = 0 - while level < n_levels: - span = 1 - # Check if cells below have the same value AND same colspan - while level + span < n_levels: - if ( - grid[level][col] == grid[level + span][col] - and colspan[level][col] == colspan[level + span][col] - ): - span += 1 - else: - break - rowspan[level][col] = span - level += span - - # Build HTML rows - # Track which cells are "covered" by rowspan from above - covered = [[False] * n_cols for _ in range(n_levels)] - - html_parts = [""] - - for level in range(n_levels): - html_parts.append('') - col = 0 - while col < n_cols: - if covered[level][col]: - # This cell is covered by a rowspan from above, skip it - col += 1 - continue - - # Get cell value - val = grid[level][col] - val_str = str(val) if val is not None else "" - if escape: - val_str = html_lib.escape(val_str) - - # Get spans - cs = colspan[level][col] - rs = rowspan[level][col] - - # Mark covered cells - for r_offset in range(rs): - for c_offset in range(cs): - if r_offset > 0 or c_offset > 0: - if level + r_offset < n_levels and col + c_offset < n_cols: - covered[level + r_offset][col + c_offset] = True - - # Build th element with attributes - attrs = [] - if cs > 1: - attrs.append(f'colspan="{cs}"') - if rs > 1: - attrs.append(f'rowspan="{rs}"') - - attr_str = " " + " ".join(attrs) if attrs else "" - html_parts.append(f"{val_str}") - - col += cs - - html_parts.append("") - - html_parts.append("") - return "".join(html_parts) - - -def render_tbody_with_row_headers( - tb_df: pd.DataFrame, - row_header_cols: int = 0, - na_rep: str = "—", - escape: bool = False, -) -> str: - """ - Render DataFrame body with support for row headers and cell merging. - - This function generates a proper structure where: - - Row header columns use element - """ - import html as html_lib - - if row_header_cols <= 0: - # No row headers - simple rendering without merging - html_parts = [""] - for _, row in tb_df.iterrows(): - html_parts.append("") - for val in row: - if pd.isna(val): - val_str = na_rep - else: - val_str = str(val) - if escape: - val_str = html_lib.escape(val_str) - html_parts.append(f"") - html_parts.append("") - html_parts.append("") - return "".join(html_parts) - - n_rows = len(tb_df) - n_cols = len(tb_df.columns) - - if n_rows == 0: - return "" - - # Build 2D grid of values for row header columns - # grid[row_idx][col_idx] = value - grid = [] - for row_idx in range(n_rows): - row_values = [] - for col_idx in range(row_header_cols): - val = tb_df.iloc[row_idx, col_idx] - if pd.isna(val): - val = na_rep - else: - val = str(val) - row_values.append(val) - grid.append(row_values) - - # Calculate colspan for each cell (horizontal merging within same row) - # colspan[row_idx][col_idx] = number of columns this cell spans - colspan = [[1] * row_header_cols for _ in range(n_rows)] - - for row_idx in range(n_rows): - col_idx = 0 - while col_idx < row_header_cols: - span = 1 - while ( - col_idx + span < row_header_cols - and grid[row_idx][col_idx] == grid[row_idx][col_idx + span] - ): - span += 1 - colspan[row_idx][col_idx] = span - col_idx += span - - # Calculate rowspan for each cell (vertical merging) - # Only calculate rowspan for cells that start a colspan group - # rowspan[row_idx][col_idx] = number of rows this cell spans - rowspan = [[1] * row_header_cols for _ in range(n_rows)] - - col_idx = 0 - while col_idx < row_header_cols: - row_idx = 0 - while row_idx < n_rows: - # Only process cells that start a colspan group (not covered by colspan from left) - if col_idx > 0 and grid[row_idx][col_idx] == grid[row_idx][col_idx - 1]: - row_idx += 1 - continue - - current_colspan = colspan[row_idx][col_idx] - span = 1 - - while row_idx + span < n_rows: - # Check if the value matches - if grid[row_idx][col_idx] != grid[row_idx + span][col_idx]: - break - # Check if colspan in the next row also matches - if colspan[row_idx + span][col_idx] != current_colspan: - break - # Check if all parent columns (to the left) also have same rowspan behavior - parent_match = True - for parent_col in range(col_idx): - if grid[row_idx][parent_col] != grid[row_idx + span][parent_col]: - parent_match = False - break - if parent_match: - span += 1 - else: - break - - rowspan[row_idx][col_idx] = span - row_idx += span - col_idx += 1 - - # Track which cells are covered by rowspan from above or colspan from left - covered = [[False] * row_header_cols for _ in range(n_rows)] - - # Mark cells covered by colspan (horizontal) - for row_idx in range(n_rows): - col_idx = 0 - while col_idx < row_header_cols: - cs = colspan[row_idx][col_idx] - for offset in range(1, cs): - if col_idx + offset < row_header_cols: - covered[row_idx][col_idx + offset] = True - col_idx += cs - - # Mark cells covered by rowspan (vertical) - for row_idx in range(n_rows): - for col_idx in range(row_header_cols): - if covered[row_idx][col_idx]: - continue # Skip cells already covered by colspan - rs = rowspan[row_idx][col_idx] - for offset in range(1, rs): - if row_idx + offset < n_rows: - # Mark all cells in the rowspan as covered - cs = colspan[row_idx][col_idx] - for c_offset in range(cs): - if col_idx + c_offset < row_header_cols: - covered[row_idx + offset][col_idx + c_offset] = True - - # Build HTML - html_parts = [""] - - for row_idx in range(n_rows): - html_parts.append("") - - # Render row header columns with rowspan/colspan - for col_idx in range(row_header_cols): - if covered[row_idx][col_idx]: - # This cell is covered by a rowspan/colspan, skip it - continue - - val_str = grid[row_idx][col_idx] - if escape: - val_str = html_lib.escape(val_str) - - rs = rowspan[row_idx][col_idx] - cs = colspan[row_idx][col_idx] - - attrs = [] - if rs > 1: - attrs.append(f'rowspan="{rs}"') - if cs > 1: - attrs.append(f'colspan="{cs}"') - - attr_str = " " + " ".join(attrs) if attrs else "" - html_parts.append(f'') - - # Render data columns - for col_idx in range(row_header_cols, n_cols): - val = tb_df.iloc[row_idx, col_idx] - if pd.isna(val): - val_str = na_rep - else: - val_str = str(val) - if escape: - val_str = html_lib.escape(val_str) - html_parts.append(f"") - - html_parts.append("") - - html_parts.append("") - return "".join(html_parts) - - -def df2html( - tb_df: pd.DataFrame, - *, - index: bool = False, - classes: Union[str, List[str], None] = "table table-striped", - na_rep: str = "—", - escape: bool = False, - row_header_cols: int = 0, -) -> str: - """Convert DataFrame to HTML table. - - Supports: - - MultiIndex columns with proper colspan/rowspan merging - - Row headers (leftmost columns rendered as ' in dataframe_html + assert "" in dataframe_html + + +def test_mineru_modules_separate_client_and_task_polling( + worker_contract_environment: None, +) -> None: + from app.services.document_parser.mineru_client import get_mineru_headers + from app.services.document_parser.mineru_task_polling import ( + get_batch_status, + get_polling_interval_for_state, + ) + + assert get_mineru_headers("token") == { + "Content-Type": "application/json", + "Authorization": "Bearer token", + } + assert get_batch_status({"data": {"extract_result": [{"state": "done"}]}}) == { + "state": "done" + } + assert get_batch_status({"data": {"extract_result": {"state": "failed"}}}) == { + "state": "failed" + } + assert get_polling_interval_for_state("pending", 2) == 8.0 + assert get_polling_interval_for_state("running", 2) == 10.0 + assert get_polling_interval_for_state("waiting-file", 2) == 15.0 + + +def test_doc_profile_model_owns_profile_contract_and_metadata( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + import json + + from app.services.document_parser.doc_profile_model import ( + DocProfile, + save_profile_metadata, + ) + + profile = DocProfile( + file_type="pdf", + route="fast", + decision_band="safe_fast", + page_count=3, + avg_text_density=123.4, + avg_image_coverage=0.05, + page_details=[{"page": 1}], + sample_text="hidden", + ) + + save_profile_metadata(profile, str(tmp_path)) + saved_profile = json.loads((tmp_path / "profile.json").read_text(encoding="utf-8")) + + assert "page_details" not in saved_profile + assert "sample_text" not in saved_profile + assert saved_profile["file_type"] == "pdf" + assert "route=fast" in profile.summary() + + +def test_doc_profiler_dispatches_pdf_to_pdf_profile_module( + worker_contract_environment: None, + monkeypatch: Any, +) -> None: + from app.services.document_parser.doc_profile_model import DocProfile + from app.services.document_parser.doc_profiler import profile_document + + called_paths: list[str] = [] + + def fake_profile_pdf(path: str) -> DocProfile: + called_paths.append(path) + return DocProfile(file_type="pdf", page_count=2) + + monkeypatch.setattr( + "app.services.document_parser.doc_profiler.profile_pdf", + fake_profile_pdf, + ) + + pdf_profile = profile_document("/tmp/input.bin", filename="report.pdf") + docx_profile = profile_document("/tmp/input.bin", filename="report.docx") + + assert called_paths == ["/tmp/input.bin"] + assert pdf_profile.file_type == "pdf" + assert docx_profile.file_type == "docx" + assert docx_profile.route == "standard" + + +def test_excel_structure_parser_is_table_structure_seam( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + import openpyxl + + from app.services.document_parser.excel_structure_parser import ( + parse_excel_structure, + ) + + workbook = openpyxl.Workbook() + worksheet = workbook.active + worksheet.title = "Budget" + worksheet["A1"] = "Region" + worksheet["B1"] = "Value" + worksheet["A2"] = "North" + worksheet["B2"] = 10 + workbook_path = tmp_path / "budget.xlsx" + workbook.save(workbook_path) + + parsed_sheets = parse_excel_structure(str(workbook_path), split_subtables=False) + + assert list(parsed_sheets.keys()) == ["Budget"] + assert parsed_sheets["Budget"].attrs["row_header_cols"] >= 0 + assert "North" in parsed_sheets["Budget"].astype(str).to_string() + + +def test_heading_hierarchy_exposes_candidate_and_tree_modules( + worker_contract_environment: None, +) -> None: + from app.services.document_parser.heading_candidates import filter_markdown_headings + from app.services.document_parser.heading_tree import cleanup_heading_tree + + candidates = filter_markdown_headings(["# Intro", "body", "## Detail"]) + cleaned = cleanup_heading_tree( + pd.DataFrame( + [ + {"id": 0, "heading": "Intro", "level": 1, "reason": ""}, + {"id": 2, "heading": "Detail", "level": 2, "reason": ""}, + ] + ) + ) + + assert candidates[["id", "heading", "level"]].to_dict("records")[0] == { + "id": 0, + "heading": "Intro", + "level": 1, + } + assert cleaned["heading"].tolist() == ["Intro", "Detail"] + + +def test_heading_llm_executor_owns_prompt_execution_and_fallback( + worker_contract_environment: None, + monkeypatch: Any, +) -> None: + from app.services.document_parser.heading_llm_executor import ( + execute_llm_heading_hierarchy, + ) + + monkeypatch.setenv("KB_LAYOUT_LLM_COMPACT_INPUT", "true") + raw_preds = pd.DataFrame( + [ + {"id": 0, "heading": "body", "level": -1, "reason": ""}, + {"id": 1, "heading": "Intro", "level": -2, "reason": "POS [1] NEG [0]"}, + {"id": 2, "heading": "body", "level": -1, "reason": ""}, + ] + ) + judged_prompts: list[pd.DataFrame] = [] + saved_files: list[str] = [] + + def fake_hierarchy_judge( + df: pd.DataFrame, + *_args: Any, + **_kwargs: Any, + ) -> list[dict[str, int]]: + judged_prompts.append(df.copy()) + return [{"id": 1, "level": 1}] + + def unexpected_fallback(_df: pd.DataFrame) -> pd.DataFrame: + raise AssertionError("fallback should not run") + + actual_df = execute_llm_heading_hierarchy( + raw_preds=raw_preds, + prompt_limt=4000, + hierarchy_judge=fake_hierarchy_judge, + fallback_hierarchy=unexpected_fallback, + save_intermediate_csv=lambda _df, _output_dir, filename: saved_files.append( + filename + ), + model_name="hierarchy-model", + ) + + assert actual_df["level"].tolist() == [-1, 1, -1] + assert "Intro" in judged_prompts[0]["heading"].tolist() + assert judged_prompts[0]["heading"].tolist().count("[1 BODY LINES]") == 2 + assert saved_files == ["preds_3_llm_base", "preds_4_llm_final"] + + body_only = pd.DataFrame( + [{"id": 0, "heading": "body", "level": -1, "reason": ""}] + ) + skipped_df = execute_llm_heading_hierarchy( + raw_preds=body_only, + prompt_limt=4000, + hierarchy_judge=lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("LLM should not run without heading candidates") + ), + fallback_hierarchy=unexpected_fallback, + save_intermediate_csv=lambda *_args: None, + ) + + assert skipped_df["level"].tolist() == [-1] + + +def test_markdown_deferred_summary_module_updates_rows_and_refs( + worker_contract_environment: None, + monkeypatch: Any, + tmp_path: Path, +) -> None: + import app.services.document_parser.markdown_deferred_summary as deferred_summary + from app.services.document_parser.markdown_deferred_summary import ( + MarkdownDeferredSummaryInput, + apply_markdown_deferred_summaries, + ) + + image_dir = tmp_path / "images" + table_dir = tmp_path / "tables" + image_dir.mkdir() + table_dir.mkdir() + (image_dir / "image-3-old.png").write_bytes(b"image") + (table_dir / "table-0 old.html").write_text("
elements) -- Proper HTML escaping and formatting -- DOCX table to HTML conversion - HTML to DataFrame/Markdown conversion +- HTML header expansion +- Nested HTML table parsing """ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional import pandas as pd from bs4 import BeautifulSoup -from docx.table import Table as DocxTable from shared.core.exceptions.domain_exceptions import TableParsingException from shared.utils.text_utils import remove_duplicates_orderkept @@ -449,557 +446,3 @@ def first_cols_rows_html(html_str, max_items=10, max_chars=20): first_col_text = " | ".join(unique_col) if unique_col else "" return first_row_text, first_col_text - - -def table2html(table: DocxTable, cell_image_map: dict = None) -> str: - """Convert a DOCX table to HTML string with proper colspan/rowspan handling. - - Handles merged cells by: - - Detecting horizontal merges via comparing cell._tc objects - - Detecting vertical merges via vMerge XML attribute - - Generating proper colspan and rowspan attributes - - Args: - table: python-docx Table object - cell_image_map: Optional dict mapping (row_idx, col_idx) to image description - strings. col_idx corresponds to the unique tc index in each row - (matching XML ordering, not expanded python-docx cells). - - Returns: - HTML string representation of the table with merged cells - """ - - NS = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"} - - def get_cell_vmerge(cell): - """Get vMerge status: 'restart', 'continue', or None""" - tc = cell._tc - tcPr = tc.find(".//w:tcPr", namespaces=NS) - if tcPr is not None: - vMerge = tcPr.find(".//w:vMerge", namespaces=NS) - if vMerge is not None: - val = vMerge.get( - "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val" - ) - return ( - val if val else "continue" - ) # If no val attribute, it's a continuation - return None - - n_rows = len(table.rows) - if n_rows == 0: - return "
" - - # Build grid: track unique cells and their positions - # grid[row][col] = (cell_tc_id, cell, is_new_cell) - # We use id(cell._tc) as unique identifier for cells - - grid = [] - for row_idx, row in enumerate(table.rows): - row_data = [] - prev_tc_id = None - for cell in row.cells: - tc_id = id(cell._tc) - is_new = tc_id != prev_tc_id - row_data.append((tc_id, cell, is_new)) - prev_tc_id = tc_id - grid.append(row_data) - - # Rows may have different cell counts due to complex merges; - # use the maximum for grid allocation, per-row length for access. - n_cols = max(len(r) for r in grid) if grid else 0 - - # Calculate colspan for each cell (count consecutive cells with same _tc) - colspan_grid = [[0] * n_cols for _ in range(n_rows)] - - for row_idx in range(n_rows): - row_len = len(grid[row_idx]) - col_idx = 0 - while col_idx < row_len: - tc_id = grid[row_idx][col_idx][0] - span = 1 - while ( - col_idx + span < row_len and grid[row_idx][col_idx + span][0] == tc_id - ): - span += 1 - colspan_grid[row_idx][col_idx] = span - col_idx += span - - # Calculate rowspan for cells with vMerge='restart' - rowspan_grid = [[1] * n_cols for _ in range(n_rows)] - - for col_idx in range(n_cols): - row_idx = 0 - while row_idx < n_rows: - if col_idx >= len(grid[row_idx]): - row_idx += 1 - continue - cell = grid[row_idx][col_idx][1] - vmerge = get_cell_vmerge(cell) - - if vmerge == "restart": - # Count how many 'continue' cells follow - span = 1 - while row_idx + span < n_rows: - if col_idx >= len(grid[row_idx + span]): - break - next_cell = grid[row_idx + span][col_idx][1] - next_vmerge = get_cell_vmerge(next_cell) - if next_vmerge == "continue": - span += 1 - else: - break - rowspan_grid[row_idx][col_idx] = span - row_idx += span - elif vmerge == "continue": - # This cell is part of a vertical merge, mark as 0 (skip) - rowspan_grid[row_idx][col_idx] = 0 - row_idx += 1 - else: - row_idx += 1 - - # Build HTML - html_parts = [""] - - for row_idx in range(n_rows): - html_parts.append("") - col_idx = 0 - unique_col_idx = 0 # Tracks unique tc index per row (matches XML order) - - while col_idx < len(grid[row_idx]): - tc_id, cell, is_new = grid[row_idx][col_idx] - - # Skip if this cell is a horizontal continuation - if not is_new: - col_idx += 1 - continue - - # Skip if this cell is a vertical continuation - rowspan = rowspan_grid[row_idx][col_idx] - if rowspan == 0: - unique_col_idx += 1 - col_idx += 1 - continue - - colspan = colspan_grid[row_idx][col_idx] - - # Build cell content - if cell.tables: - # Nested table - content = "".join( - table2html(nested_table) for nested_table in cell.tables - ) - else: - content = cell.text.strip().replace("\n", "
") - - # Append image descriptions if available - if cell_image_map: - img_desc = cell_image_map.get((row_idx, unique_col_idx)) - if img_desc: - content += f"
{img_desc}" - - # Build attributes - attrs = [] - if colspan > 1: - attrs.append(f'colspan="{colspan}"') - if rowspan > 1: - attrs.append(f'rowspan="{rowspan}"') - - attr_str = " " + " ".join(attrs) if attrs else "" - html_parts.append(f"{content}") - - unique_col_idx += 1 - col_idx += colspan - - html_parts.append("") - - html_parts.append("
") - return "".join(html_parts) - - -def render_multiindex_thead(columns: pd.MultiIndex, escape: bool = False) -> str: - """ - Convert MultiIndex columns to HTML thead with colspan/rowspan. - - This function generates a proper multi-row
instead of - - Horizontally adjacent identical values in row headers are merged with colspan - - Vertically adjacent identical values in row headers are merged with rowspan - - Merging respects hierarchical structure - - Args: - tb_df: DataFrame to render - row_header_cols: Number of leftmost columns to render as - na_rep: String representation for NaN values - escape: Whether to HTML-escape values - - Returns: - HTML string for the
{val_str}
{val_str}{val_str}
) - - Custom CSS classes and NA representation - - Args: - tb_df: DataFrame to convert - index: Whether to include the DataFrame index (not commonly used) - classes: CSS classes to add to the table - na_rep: String representation for NaN values - escape: Whether to HTML-escape values - row_header_cols: Number of leftmost columns to render as row headers (). - These columns will use instead of . - - Returns: - HTML table string - """ - class_str = ( - classes if isinstance(classes, str) else " ".join(classes) if classes else "" - ) - - # Check if columns are MultiIndex - use advanced rendering - if isinstance(tb_df.columns, pd.MultiIndex): - # Use specialized rendering for MultiIndex columns - thead_html = render_multiindex_thead(tb_df.columns, escape=escape) - tbody_html = render_tbody_with_row_headers( - tb_df, row_header_cols, na_rep, escape - ) - return f'{thead_html}{tbody_html}
' - - # Simple columns case - if row_header_cols <= 0: - # Use default pandas to_html for simple case without row headers - table_html = tb_df.to_html( - index=index, - na_rep=na_rep, - classes=classes, - escape=escape, - border=0, - justify="center", - ) - return table_html.replace("\n", "") - - # Simple columns with row headers - custom rendering - import html as html_lib - - html_parts = [f''] - - # Build thead - html_parts.append("") - html_parts.append('') - for col in tb_df.columns: - col_str = str(col) if col is not None else "" - if escape: - col_str = html_lib.escape(col_str) - html_parts.append(f"") - html_parts.append("") - html_parts.append("") - - # Build tbody with row headers - tbody_html = render_tbody_with_row_headers(tb_df, row_header_cols, na_rep, escape) - html_parts.append(tbody_html) - - html_parts.append("
{col_str}
") - - return "".join(html_parts) diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/image_parser.py index 80723ed4e..cc9be0ab1 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/image_parser.py @@ -10,6 +10,7 @@ import pandas as pd from app.services.document_parser.dataframe_helpers import process_dup_paths_df from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder from loguru import logger from PIL import Image @@ -224,7 +225,6 @@ def parse_image( relative_root=None, ): split_char = settings.SPLIT_CHAR or "/" - df_list = [] time_stamp = get_str_time() os.makedirs(output_dir, exist_ok=True) img_dir = os.path.join(output_dir, "images") @@ -355,23 +355,19 @@ def parse_image( ) img_ref = build_chunk_ref(relative_img_path) img_bottom_content = f"{img_ref}\nImage Content:\n{image_content}" - df_list.append( - [ - img_bottom_content, - relative_img_path, - "image", - len(img_bottom_content), - "", - image_summary, - temp_uid, - "", - "", - time_stamp, - "", - ] + rows_builder = ParsedRowsBuilder() + rows_builder.append( + ParsedRow( + content=img_bottom_content, + path=relative_img_path, + type="image", + summary=image_summary, + know_id=temp_uid, + addtime=time_stamp, + ) ) - img_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(",")) + img_df = rows_builder.to_dataframe() img_df = process_dup_paths_df(img_df) return img_df diff --git a/apps/worker/app/services/document_parser/inline_asset.py b/apps/worker/app/services/document_parser/inline_asset.py new file mode 100644 index 000000000..4cb16b531 --- /dev/null +++ b/apps/worker/app/services/document_parser/inline_asset.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from app.services.document_parser.parser_rows import ParsedRow + + +def build_image_asset_row( + *, + content: str, + relative_path: str, + summary: str, + know_id: str, + addtime: str, + page_nums: str = "", +) -> ParsedRow: + return ParsedRow( + content=content, + path=relative_path, + type="image", + keywords="", + summary=summary, + know_id=know_id, + tokens="", + connectto="", + addtime=addtime, + page_nums=page_nums, + ) + + +def build_table_asset_row( + *, + content: str, + relative_path: str, + summary: str, + keywords: str, + know_id: str, + addtime: str, + page_nums: str = "", +) -> ParsedRow: + return ParsedRow( + content=content, + path=relative_path, + type="table", + keywords=keywords, + summary=summary, + know_id=know_id, + tokens="", + connectto="", + addtime=addtime, + page_nums=page_nums, + ) diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py index 6aa83ba63..6a3b3a8bf 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/layout_parser.py @@ -1,15 +1,30 @@ # pyright: reportArgumentType=false, reportAssignmentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalSubscript=false import os -import re -import unicodedata -from collections import Counter, defaultdict import gevent import pandas as pd -from app.services.document_parser.text_helpers import count_cn_en, truncate_text_by_tokens +from app.services.document_parser.heading_candidates import ( + filter_document_headings, + filter_markdown_headings, + judge_by_conditions, + postprocess_headings, +) +from app.services.document_parser.heading_llm_executor import ( + build_level_mapping, + execute_level_mapping, + execute_llm_heading_hierarchy, +) +from app.services.document_parser.heading_tree import ( + build_tree_from_dataframe as build_heading_tree_from_dataframe, +) +from app.services.document_parser.heading_tree import ( + remove_isolated_nodes as remove_isolated_heading_nodes, +) +from app.services.document_parser.heading_tree import ( + tree_to_dataframe as heading_tree_to_dataframe, +) from app.services.document_parser.stage_profiler import stage_timer from app.services.document_parser.table_parser import df2md -from docx.oxml.ns import qn from gevent.pool import Pool as GeventPool try: @@ -24,7 +39,6 @@ def convert(self, content): from loguru import logger from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import WorkerHandlingException # TaskRedis dependency is removed, use Redis directly to track from shared.services.ai.prompt_service import build_prompt @@ -68,187 +82,15 @@ def save_intermediate_csv(df: pd.DataFrame, output_dir: str, filename: str): def build_tree_from_dataframe(df): - """ - develop json tree from dataframe - - Args: - df: DataFrame, including id, heading, level columns - - Returns: - tree: pure nested dict structure - node_to_id: map from tree node to id (use unique node key) - id_to_row: map from id to original row data - """ - headings = df[df["level"] > -1].copy() - - node_to_id = {} # {(tree_node_key, parent_path): id} - id_to_node_info = {} # {id: (tree_node_key, parent_path)} - id_to_row = {} - root = {} - stack = [(0, root, "ROOT", "")] - - for _, row in headings.iterrows(): - heading_txt = row["heading"] - row_id = int(row["id"]) - level = int(row["level"]) - - # record id to row mapping - id_to_row[row_id] = row.to_dict() - - # find suitable parent node - while len(stack) > 1 and stack[-1][0] >= level: - stack.pop() - - # get parent node info - parent_level, parent_dict, parent_heading, parent_path = stack[-1] - - # create unique key for tree node: if there are duplicate headings under the same parent, add ID suffix - tree_node_key = heading_txt - - if tree_node_key in parent_dict: - tree_node_key = f"{heading_txt}#{row_id}" - - # build mapping: use (tree_node_key, parent_path) as key - node_key = (tree_node_key, parent_path) - node_to_id[node_key] = row_id - id_to_node_info[row_id] = node_key - - parent_dict[tree_node_key] = {} - current_path = ( - f"{parent_path}/{tree_node_key}" if parent_path else tree_node_key - ) - stack.append((level, parent_dict[tree_node_key], tree_node_key, current_path)) - return root, node_to_id, id_to_row + return build_heading_tree_from_dataframe(df) def tree_to_dataframe(tree, node_to_id, original_df): - """ - convert processed tree structure back to dataframe - - Args: - tree: processed pure nested dict structure - node_to_id: map from node to id {(tree_node_key, parent_path): id} - original_df: original dataframe - - Returns: - updated_df: updated dataframe - """ - - # extract all retained headings from tree - def extract_headings(node_dict, current_level=1, parent_path=""): - """recursively extract all headings and their new levels""" - results = [] - for tree_node_key, children in node_dict.items(): - # use (tree_node_key, parent_path) as key to find ID - node_key = (tree_node_key, parent_path) - row_id = node_to_id.get(node_key, -1) - - if row_id >= 0: - # extract original heading from tree_node_key (remove possible ID suffix) - original_heading = ( - tree_node_key.split("#")[0] - if "#" in tree_node_key - else tree_node_key - ) - - results.append( - { - "id": row_id, - "heading": original_heading, - "level": current_level, - "tree_key": tree_node_key, - "parent_path": parent_path, - } - ) - # recursively process child nodes - if isinstance(children, dict) and children: - current_path = ( - f"{parent_path}/{tree_node_key}" - if parent_path - else tree_node_key - ) - results.extend( - extract_headings(children, current_level + 1, current_path) - ) - return results - - preserved_headings = extract_headings(tree) - preserved_ids = set([h["id"] for h in preserved_headings]) - - updated_df = original_df.copy() - removed_count = 0 - level_changed_count = 0 - - for idx, row in original_df.iterrows(): - row_id = int(row["id"]) - old_level = int(row["level"]) if row["level"] not in [-2, "nan", -1] else -1 - - if old_level > -1: - if row_id in preserved_ids: - new_level = next( - (h["level"] for h in preserved_headings if h["id"] == row_id), - old_level, - ) - updated_df.at[idx, "level"] = new_level - if new_level != old_level: - level_changed_count += 1 - else: - updated_df.at[idx, "level"] = -1 - removed_count += 1 - - logger.debug( - f"Tree changed: removed headings={removed_count}, level changed={level_changed_count}, preserved headings={len(preserved_ids)}" - ) - return updated_df + return heading_tree_to_dataframe(tree, node_to_id, original_df) def remove_isolated_nodes(tree): - """ - rules: if a heading has only one child heading, and the child heading has no further child headings, - then delete this isolated child heading - - Args: - tree: pure nested dict structure, format as {heading: {child_heading: {...}}} - - Returns: - processed_tree: processed tree structure - """ - - def recursive_check_and_remove(node_dict, parent_path=""): - if not isinstance(node_dict, dict): - return node_dict - - result_dict = {} - - for heading, children in node_dict.items(): - if isinstance(children, dict) and len(children) == 1: - child_heading = list(children.keys())[0] - grandchildren = children[child_heading] - - if not grandchildren or ( - isinstance(grandchildren, dict) and len(grandchildren) == 0 - ): - result_dict[heading] = {} - logger.debug( - f"remove isolated heading: {parent_path}/{heading}/{child_heading}" - ) - else: - processed_children = recursive_check_and_remove( - children, f"{parent_path}/{heading}" if parent_path else heading - ) - result_dict[heading] = processed_children - elif isinstance(children, dict) and children: - processed_children = recursive_check_and_remove( - children, f"{parent_path}/{heading}" if parent_path else heading - ) - result_dict[heading] = processed_children - else: - result_dict[heading] = children - - return result_dict - - processed_tree = recursive_check_and_remove(tree) - return processed_tree + return remove_isolated_heading_nodes(tree) # def if_no_pos_code(reason_str: str) -> bool: @@ -270,869 +112,12 @@ def recursive_check_and_remove(node_dict, parent_path=""): # return True -# ==================== Level Mapping Functions ==================== - - -def build_level_mapping(df, origin_lvls, mode="max"): - df = df.copy() - df["origin_level"] = origin_lvls - - mapping = df.groupby("reason")["level"].apply(list).to_dict() - - processed_mapping = {} - for reason, lvls in mapping.items(): - positive_lvls = [lvl for lvl in lvls if lvl > -1] - counts = Counter(lvls) - - if not positive_lvls: - mapped_lvl = -1 - elif mode == "max": - mapped_lvl = max(positive_lvls) - elif mode == "freq": - mapped_lvl = counts.most_common(1)[0][0] - else: - raise WorkerHandlingException( - internal_message=f"wrong input mode: {mode}. Must be 'max' or 'freq'" - ) - - processed_mapping[reason] = { - "lvls": lvls, - "positive_lvls": positive_lvls, - "freqs": dict(counts), - "mapped_lvl": mapped_lvl, - } - return df, processed_mapping - - -def execute_level_mapping(df: pd.DataFrame, mapping: dict) -> pd.DataFrame: - def map_row(row): - reason = row["reason"] - if reason in mapping: - return mapping[reason]["mapped_lvl"] - return row["level"] - - df = df.copy() - origin_est_lvls = df["level"].tolist() - df["level"] = df.apply(map_row, axis=1) - df["origin_level"] = origin_est_lvls - return df - - -def extract_non_neg_code(reason_str: str) -> str: - """ - Extract the non-NEG code from reason_str (strip only NEG part, preserve META) - - Example: "POS [1, 0, 0] NEG [0, 0, 0] META [1, 2, 1]" -> "POS [1, 0, 0] META [1, 2, 1]" - Example: "3# AND POS [1, 0] NEG [0, 0]" -> "3# AND POS [1, 0]" - Example: "3# AND POS [1, 0] NEG [0, 0] META [1, 1, 0]" -> "3# AND POS [1, 0] META [1, 1, 0]" - """ - if not reason_str or not isinstance(reason_str, str): - return "" - neg_match = re.search(r"\s*NEG\s*\[[^\]]*\]", reason_str) - if neg_match: - # Remove only the NEG [...] part, keep everything before and after - before_neg = reason_str[: neg_match.start()] - after_neg = reason_str[neg_match.end() :] - return (before_neg + after_neg).strip() - return reason_str.strip() - - -def build_non_neg_mapping(lvl_mapping: dict) -> dict: - """ - Build non-NEG code mapping from complete lvl_mapping, select by highest frequency - - Args: - lvl_mapping: reason -> level mapping - - Returns: - non_neg_mapping: {non_neg_code: mapped_lvl} - """ - # collect all levels for each non_neg_code - non_neg_levels = {} - for reason, info in lvl_mapping.items(): - non_neg_code = extract_non_neg_code(reason) - mapped_lvl = info.get("mapped_lvl", -1) - if non_neg_code: - if non_neg_code not in non_neg_levels: - non_neg_levels[non_neg_code] = [] - non_neg_levels[non_neg_code].append(mapped_lvl) - - # select by highest frequency - non_neg_mapping = {} - for non_neg_code, levels in non_neg_levels.items(): - positive_levels = [lvl for lvl in levels if lvl > -1] - if positive_levels: - level_counts = Counter(positive_levels) - most_common_level = level_counts.most_common(1)[0][0] - non_neg_mapping[non_neg_code] = most_common_level - else: - non_neg_mapping[non_neg_code] = -1 - - return non_neg_mapping - - -def handle_unseen_codes( - df: pd.DataFrame, - level_dfs: list, - lvl_mapping: dict, - output_dir: str = None, - window_half_size: int = 10, - strategy: str = "double_mapping", -) -> dict: - """ - Handle unseen codes with configurable strategy - - Args: - df: original complete DataFrame - level_dfs: segment DataFrames - lvl_mapping: existing level mapping - output_dir: output directory (optional, only used for window_llm strategy) - window_half_size: window half size (how many rows above and below) - strategy: "double_mapping" or "window_llm" - - double_mapping: use non-neg code fallback (fast, no LLM call) - - window_llm: create windows for LLM to judge (slower, more accurate) - - Returns: - updated lvl_mapping - """ - - def extract_reason_signature(reason: str) -> str: - """Extract reason signature""" - return reason.strip() if reason else "" - - def has_neg_signal(reason_str: str) -> bool: - """Check if NEG signal exists (any value >= 1)""" - if not reason_str or not isinstance(reason_str, str): - return False - neg_match = re.search(r"NEG\s*\[([^\]]*)\]", reason_str) - if not neg_match: - return False - neg_content = neg_match.group(1) - try: - nums = [int(x.strip()) for x in neg_content.split(",") if x.strip()] - return any(x >= 1 for x in nums) - except Exception: - return False - - def build_context_window( - target_idx: int, known_codes_set: set, total_rows: int, half_size: int = 10 - ) -> dict: - """ - Build context window for unseen codes - 1. window size: half_size - 2. window should contain at least one known code - """ - min_start = max(0, target_idx - half_size) - min_end = min(total_rows - 1, target_idx + half_size) - - start_idx = min_start - end_idx = min_end - - found_known_above = False - found_known_below = False - known_positions = [] - - # check above - for i in range(start_idx, target_idx): - reason = df.iloc[i].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_above = True - known_positions.append(i) - - # check below - for i in range(target_idx + 1, end_idx + 1): - reason = df.iloc[i].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_below = True - known_positions.append(i) - - # expand above if needed - if not found_known_above and min_start > 0: - search_idx = min_start - 1 - while search_idx >= 0: - reason = df.iloc[search_idx].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_above = True - known_positions.append(search_idx) - start_idx = search_idx - break - search_idx -= 1 - - # expand below if needed - if not found_known_below and min_end < total_rows - 1: - search_idx = min_end + 1 - while search_idx < total_rows: - reason = df.iloc[search_idx].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_below = True - known_positions.append(search_idx) - end_idx = search_idx - break - search_idx += 1 - - return { - "start": start_idx, - "end": end_idx, - "found_known": found_known_above or found_known_below, - "known_positions": known_positions, - } - - # build non-neg mapping - non_neg_mapping = build_non_neg_mapping(lvl_mapping) - - # get known codes - known_codes = set(lvl_mapping.keys()) - - # record all codes from all segments. Placeholder rows (reason == - # PLACEHOLDER_REASON) are injected by _compact_for_llm and are never real - # heading candidates, so they must be skipped here — otherwise they would - # show up as an "unseen code" and fall through to NO_MATCH_FALLBACK, adding - # harmless but noisy warnings to the log. - all_codes_in_full = {} - for seg_idx, seg_df in enumerate(level_dfs): - for _, row in seg_df.iterrows(): - reason = row.get("reason", "") - sig = extract_reason_signature(reason) - if not sig or sig == PLACEHOLDER_REASON: - continue - if sig not in all_codes_in_full: - all_codes_in_full[sig] = { - "first_seg": seg_idx, - "first_id": row.get("id", 0), - "reason": reason, - } - - # find unseen codes - unseen_codes = {} - unseen_neg_filtered = {} - for sig, info in all_codes_in_full.items(): - if sig in known_codes: - continue - if has_neg_signal(info["reason"]): - unseen_neg_filtered[sig] = info - else: - unseen_codes[sig] = info - - logger.info( - f"Unseen codes total: {len(unseen_codes) + len(unseen_neg_filtered)}, NEG filtered: {len(unseen_neg_filtered)}, to process: {len(unseen_codes)}" - ) - - # if neg signal, map to -1 - for sig in unseen_neg_filtered: - lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NEG_FILTERED"} - - # handle remaining unseen_codes based on strategy - if unseen_codes: - if strategy == "double_mapping": - # Strategy 1: use non-neg code fallback - fallback_success = 0 - fallback_failed = 0 - failed_codes = [] - for sig, info in unseen_codes.items(): - non_neg_code = extract_non_neg_code(sig) - if non_neg_code in non_neg_mapping: - mapped_level = non_neg_mapping[non_neg_code] - lvl_mapping[sig] = { - "mapped_lvl": mapped_level, - "note": f"NON_NEG_FALLBACK from '{non_neg_code}'", - } - fallback_success += 1 - else: - lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NO_MATCH_FALLBACK"} - fallback_failed += 1 - failed_codes.append( - f"'{non_neg_code}' (from '{sig[:60]}...')" - if len(sig) > 60 - else f"'{non_neg_code}' (from '{sig}')" - ) - - logger.info( - f"Double mapping result: success={fallback_success}, failed={fallback_failed}" - ) - if failed_codes: - logger.warning( - f"Failed codes (non_neg not in mapping): {failed_codes[:5]}{'...' if len(failed_codes) > 5 else ''}" - ) - - elif strategy == "window_llm" and output_dir: - # Strategy 2: create windows for LLM to judge - total_rows = len(df) - windows = [] - for sig, info in unseen_codes.items(): - first_id = info["first_id"] - first_seg = info["first_seg"] - df_indices = df.index[df["id"] == first_id].tolist() - if df_indices: - first_df_idx = df_indices[0] - window_info = build_context_window( - first_df_idx, known_codes, total_rows, window_half_size - ) - windows.append( - { - "code": sig, - "first_id": first_id, - "first_seg": first_seg, - "start": window_info["start"], - "end": window_info["end"], - "found_known": window_info["found_known"], - } - ) - - # merge windows - sorted_windows = sorted(windows, key=lambda x: x["start"]) - merged_windows = [] - current_window = None - - for w in sorted_windows: - if current_window is None: - current_window = { - "start": w["start"], - "end": w["end"], - "codes": [w["code"]], - "segments": [w["first_seg"]], - } - elif w["start"] <= current_window["end"]: - current_window["end"] = max(current_window["end"], w["end"]) - current_window["codes"].append(w["code"]) - current_window["segments"].append(w["first_seg"]) - else: - merged_windows.append(current_window) - current_window = { - "start": w["start"], - "end": w["end"], - "codes": [w["code"]], - "segments": [w["first_seg"]], - } - - if current_window: - merged_windows.append(current_window) - - # save windows - windows_dir = os.path.join(output_dir, "merged_windows") - os.makedirs(windows_dir, exist_ok=True) - - unseen_codes_set = set(unseen_codes.keys()) - unseen_neg_set = set(unseen_neg_filtered.keys()) - - for i, mw in enumerate(merged_windows): - window_df = df.iloc[mw["start"] : mw["end"] + 1].copy() - - def get_code_status(row): - reason = row.get("reason", "") - sig = extract_reason_signature(reason) - if not sig: - return "" - if sig in unseen_codes_set: - return "★ UNSEEN_TARGET" - elif sig in unseen_neg_set: - return "NEG→-1" - elif sig in known_codes: - return "KNOWN" - else: - return "" - - window_df["code_status"] = window_df.apply(get_code_status, axis=1) - window_path = os.path.join( - windows_dir, - f"window_{i + 1:02d}_rows_{mw['start']}-{mw['end']}.csv", - ) - window_df.to_csv(window_path, index=False, encoding="utf-8-sig") - - logger.debug( - f"Window LLM: {len(merged_windows)} windows created in {windows_dir}" - ) - # TODO: use llm to assign level based on window data - - return lvl_mapping - def detect_outlines_md(line): pos_code = judge_by_conditions(line) any(x > 0 for x in pos_code) -def get_max_lvl(code_str: str): - match = re.search(r"\[([^]]+)]", code_str) - if not match: - return "Sure" - - nums = [int(x.strip()) for x in match.group(1).split(",")] - max_val = int(max(nums)) - return max_val if max_val > 1 else -2 # -2 = "Not Sure" sentinel (int-safe) - - -PLACEHOLDER_REASON = "__PLACEHOLDER__" - - -def _compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: - """Collapse consecutive ``level == -1`` rows into a single placeholder row. - - Rows with ``level >= 1`` (heading candidates) and ``level == -2`` ("Not Sure") - are preserved verbatim so the LLM can still judge them. Each run of - consecutive ``-1`` rows becomes one placeholder row whose: - - id = "start-end" (always a range; "N-N" when the run is one row) - heading = "[N BODY LINES]" where N is the run length - level = "-" - reason = ``PLACEHOLDER_REASON`` - - The id is ALWAYS a hyphenated string, even for single-row runs, so that - ``int(id)`` fails for every placeholder. This lets downstream code identify - placeholders structurally (non-integer id) without depending on ``reason`` - or length heuristics. - """ - if df is None or len(df) == 0: - return pd.DataFrame(columns=["id", "heading", "level", "reason"]) - - rows = [] - i = 0 - n = len(df) - while i < n: - lvl_raw = df.iloc[i]["level"] - try: - lvl_int = int(lvl_raw) - except (TypeError, ValueError): - lvl_int = None - - if lvl_int == -1: - j = i - while j < n: - try: - nxt_lvl = int(df.iloc[j]["level"]) - except (TypeError, ValueError): - break - if nxt_lvl != -1: - break - j += 1 - start_id = int(df.iloc[i]["id"]) - end_id = int(df.iloc[j - 1]["id"]) - run = j - i - rows.append( - { - "id": f"{start_id}-{end_id}", - "heading": f"[{run} BODY LINES]", - "level": "-", - "reason": PLACEHOLDER_REASON, - } - ) - i = j - else: - r = df.iloc[i] - rows.append( - { - "id": int(r["id"]), - "heading": str(r["heading"]), - "level": ( - int(lvl_int) - if lvl_int is not None and lvl_int != -2 - else "Not Sure" - ), - "reason": str(r.get("reason", "") or ""), - } - ) - i += 1 - - return pd.DataFrame(rows, columns=["id", "heading", "level", "reason"]) - - -def heading_tb_transfer(df, threshold=3000, max_start=50, max_end=10): - raw_headings = df["heading"].tolist() - df["heading"] = df["heading"].apply( - lambda x: truncate_text_by_tokens(x, max_start, max_end) - ) - - sub_dfs = [] - current_rows = [] - current_len = 0 - for _, row in df.iterrows(): - row_filtered = row.drop(labels=["reason"], errors="ignore") - row_len = sum(count_cn_en(str(v)) for v in row_filtered.values) - - if current_len + row_len > threshold and current_rows: - sub_dfs.append(pd.DataFrame(current_rows, columns=df.columns)) - current_rows = [row.tolist()] - current_len = row_len - else: - current_rows.append(row.tolist()) - current_len += row_len - - if current_rows: - sub_dfs.append(pd.DataFrame(current_rows, columns=df.columns)) - return sub_dfs, raw_headings - - -def judge_by_conditions(text, scope=20, return_detail=False, CN_SPECIAL_IDX=12): - """ - judge level features as one-hot embeddings for texts - - Args: - text: input text - scope: text scope for judging - return_detail: whether to return detailed information (including unit type) - CN_SPECIAL_IDX: index of special Chinese number - - Returns: - if return_detail=False: return pos_triggered_code list - if return_detail=True: return (pos_triggered_code, detail_info) tuple - where detail_info is a dictionary containing additional information, such as Chinese unit type - """ - text = text.replace("\u3000", " ") - text = unicodedata.normalize("NFKC", text)[:scope] - - # ========== English Numbering ========== - regex_en_num_dots = r"^\d+(?:\s*\.\s*\d+)+(?![、，。！？；：])(?=\s|$|\w|[一-龥])" - regex_en_num_dun = r"^\d、\s{0,4}(?=\S|$)" # 1、xxx - regex_en_num_dots_dun = r"^\d+(?:\.\d+)*、\s*(?=[A-Za-z一-龥])" - regex_en_num_single_dot = r"^\d+\.(?!\d)\s{0,4}(?=\S)" # 1.xxx - regex_en_num_space = r"^[0-9]{1,2}\s{1,8}(?=\S)" # 1 xxx - # ========== Chinese Numbering ========== - regex_cn_num_dun = r"^[一二三四五六七八九十百千万]+、\s{0,4}(?=\S|$)" - regex_cn_num_mix = ( - r"^[一二三四五六七八九十百千万]+(?:\s*\.[一二三四五六七八九十百千万\d]+)+" - ) - regex_cn_num_plain = r"^[一二三四五六七八九十百千万]+(?=\s|$)" - # ========== English Bracketing ========== - regex_en_brac_paren = r"^[\(\（]\s*\d+(?:\.\d+)*(?!\.0)\s*[\)\）]" - regex_en_brac_right = r"^\d+(?:\.\d+)*(?!\.0)\s*[\)\）]" - # ========== Chinese Bracketing ========== - regex_cn_brac_paren = r"^[\(\（]\s*[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]" - regex_cn_brac_right = r"^[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]" - # ========== Chinese Special ========== - regex_cn_special = r"^第[一二三四五六七八九十百千万\d]+(?:\.[一二三四五六七八九十百千万\d]+)*(章|节|条|部分|款|目|项|编|篇|卷|辑)?(?=$|\s|[A-Za-z0-9\u4e00-\u9fa5])" - # ========== English Letter Numbering ========== - regex_letter_dot = r"^[A-Za-z](?:\.\d+)*[\.、](?=\s*\S)" - regex_letter_brac_paren = r"^[\(\（]\s*[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]" - regex_letter_brac_right = r"^[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]" - # ========== Appendix ========== - regex_appendix = r"^((附件|附录|附表|附图)|(?i:appendix))[\s_\-—]{0,4}(?:\[)?[一二三四五六七八九十A-Za-z\d]" - - pos_regex_conditions = [ - # English Numbering - regex_en_num_dots, - regex_en_num_dun, - regex_en_num_single_dot, - regex_en_num_space, - regex_en_num_dots_dun, - # Chinese Numbering - regex_cn_num_dun, - regex_cn_num_mix, - regex_cn_num_plain, - # English Bracketing - regex_en_brac_paren, - regex_en_brac_right, - # Chinese Bracketing - regex_cn_brac_paren, - regex_cn_brac_right, - # Chinese Special - regex_cn_special, - # English Letter Numbering - regex_letter_dot, - regex_letter_brac_paren, - regex_letter_brac_right, - # Appendix - regex_appendix, - ] - - pos_triggered_code = [] - reason_suffix_parts = [] - - for idx, regex in enumerate(pos_regex_conditions): - match = re.match(regex, text) - if match: - matched_text = match.group(0) - symbols = ".-" - count_ = sum(matched_text.count(s) for s in symbols) + 1 - - # Special handling for Chinese chapter/section/item markers. - if idx == CN_SPECIAL_IDX and return_detail: - unit_match = re.search( - r"(章|节|条|部分|款|目|项|编|篇|卷|辑)", matched_text - ) - if unit_match: - unit = unit_match.group(1) - reason_suffix_parts.append(f"[CN:{unit}]") - pos_triggered_code.append(count_) - else: - pos_triggered_code.append(0) - - if return_detail: - detail_info = { - "reason_suffix": ( - " ".join(reason_suffix_parts) if reason_suffix_parts else "" - ) - } - if detail_info["reason_suffix"]: - detail_info["reason_suffix"] = " " + detail_info["reason_suffix"] - return pos_triggered_code, detail_info - return pos_triggered_code - - -def remove_by_conditions(text, include_punc=False): - neg_condition_num = r"^\d{3,}" - neg_condition_zero = r"^0\.\d+[\u4e00-\u9fa5A-Za-z\S]*" # 0.2xxx - neg_decimal_only = r"^\d*\.\d+$" # 0.2 .23 - neg_condition_http = ( - r"(?i)(^https?://\S+|^www\.\S+|^P\.S|^\b\d{0,2}\s*(?:a\.m|p\.m)\b)" - ) - # LaTeX: wrapped ($..\cmd..$) OR bare commands (\times, \mathrm, etc.) - neg_condition_latex = ( - r"(?:" - r"\$[^$]*\\[A-Za-z]+(?:\s*\{[^{}]*\})?[^$]*\$" # wrapped: $...\cmd...$ - r"|" - r"\\(?:times|div|cdot|pm|mp|leq|geq|neq|approx|equiv|sim|infty" - r"|sum|prod|int|sqrt|frac|mathrm|mathbf|mathit|mathcal" - r"|text(?:bf|it|rm)?|alpha|beta|gamma|delta|epsilon|theta" - r"|lambda|mu|sigma|pi|omega|partial|nabla" - r"|left|right|begin|end|overline|underline|hat|vec|tilde)\b" - r")" - ) - # Number immediately followed by measurement unit (e.g. 25.40mm, 100kPa) - neg_condition_unit = ( - r"^\d+\.?\d*\s{0,2}" - r"(?:mm|cm|km|nm|μm|inch(?:es)?|ft|yd|mi" - r"|kg|mg|μg|lb|oz" - r"|kPa|MPa|GPa|Pa|psi|bar" - r"|°[CFK]" - r"|Hz|kHz|MHz|GHz" - r"|mol|mL|dL|dB|Nm|kN|MN|kW|MW|GW|hp|rpm|cc|cal|kcal)\b" - ) - neg_condition_punc_mid = r"[。！；].+" - neg_condition_punc_end = r"[.,;，。；]$" - - neg_conditions = [ - neg_condition_num, - neg_condition_http, - neg_condition_latex, - neg_condition_zero, - neg_decimal_only, - neg_condition_punc_mid, - neg_condition_unit, - ] - - neg_triggered_code = [] - for regex in neg_conditions: - match = re.search(regex, text) - neg_triggered_code.append(1 if match else 0) - - if include_punc: - match = re.search(neg_condition_punc_end, text) - neg_triggered_code.append(1 if match else 0) - else: - neg_triggered_code.append(0) - - return neg_triggered_code - - -def md_heading_match(line, as_is=True): - """handle markdown headings, considering # < ! [....""" - match = re.match(r"^\s*(#+)\s*(.*)$", line) - if match: - level = len(match.group(1)) # count the number of '#' - if as_is: # determine if remove the '#' - return line, level - else: - return line.lstrip("#").strip(), level - else: - return line, -1 - - -def filter_md_headings(md_lines, num_pos=17, num_neg=7, layout_json_path=None): - """filter candidate headings for .md - - Args: - md_lines: list of markdown lines - num_pos: number of positive conditions - num_neg: number of negative conditions - layout_json_path: optional path to layout.json for META features (size ranking) - """ - # Create MetadataContext if layout_json_path is provided - meta_ctx = None - if layout_json_path: - try: - from .metadata_extractor import MetadataContext - - meta_ctx = MetadataContext(md_lines, layout_json_path) - except Exception as e: - logger.warning(f"Failed to create MetadataContext: {e}") - - raw_candidates = [] - for i, line in enumerate(md_lines): - line = line.strip() - if not line: - continue - - if ( - ("" in line) # annotation line - or line.startswith("|") # table line - or line.startswith("") - or "![" in line - and "](" in line # image line - ): - est_lvl = -1 - zero_pos_code = [0] * num_pos - zero_neg_code = [0] * num_neg - str_lvl = f"POS {zero_pos_code} NEG {zero_neg_code}" - if meta_ctx: - str_lvl += " META [0, 0, 0]" - line = "Figure/Image" - else: - line_clean, hash_lvl = md_heading_match( - line, as_is=False - ) # detect "#" in .md lines - - # NEW: detect and strip full-line bold markers (e.g. **3.4 Title** -> 3.4 Title) - from .metadata_extractor import detect_and_strip_md_bold - - line_clean_stripped, is_full_bold = detect_and_strip_md_bold(line_clean) - - # Use stripped text for POS/NEG analysis (fixes '**3.4' -> '3.4' issue) - pos_code, detail_info = judge_by_conditions( - line_clean_stripped, return_detail=True - ) - neg_code = remove_by_conditions(line_clean_stripped) - - if any(x > 0 for x in neg_code): - code_lvl = -1 - code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" - - elif any(x > 0 for x in pos_code) and all(x == 0 for x in neg_code): - code_lvl = get_max_lvl(str(pos_code)) - code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" - - else: - code_lvl = -1 - code_str = f"POS {pos_code} NEG {neg_code}" - - # Add META suffix with bold dimension - if meta_ctx: - size_rank, occurrence = meta_ctx.get_meta_for_line(line_clean) - is_bold_int = 1 if is_full_bold else 0 - code_str += meta_ctx.format_meta_suffix( - size_rank, occurrence, is_bold_int - ) - else: - # Even without layout.json, output bold info in META - if is_full_bold: - code_str += " META [0, 0, 1]" - - if hash_lvl <= 0: - est_lvl = code_lvl - str_lvl = code_str - else: - if isinstance(code_lvl, int): - est_lvl = max( - hash_lvl, code_lvl - ) # current miner tend to produce fewer #s - else: - est_lvl = code_lvl # code_lvl could be not sure - str_lvl = f"{hash_lvl}# AND {code_str}" - raw_candidates.append((i, line, est_lvl, str_lvl)) - - preds_df = pd.DataFrame( - raw_candidates, columns=["id", "heading", "level", "reason"], index=None - ) - return preds_df - - -def filter_doc_headings(titles_material, enable_regx=True, enable_style_check=False): - """filter candidate headings for docx""" - - def find_docstyle(para_): - try: - style_name = para_.style.name - except Exception: - style_name = "normal" - if style_name.startswith("Heading") or style_name.startswith("标题"): - try: - outline_level = int(style_name.split(" ")[1]) - except Exception: - outline_level = -2 # "Not Sure" sentinel - return outline_level - else: - return None - - def find_otsetting(para_): - ppr = para_._element.find(qn("w:pPr")) - if ppr is not None: - plvl = ppr.find(qn("w:outlineLvl")) - else: - return None - - if plvl is not None: - outline_level = int(plvl.get(qn("w:val"))) + 1 - return outline_level - else: - return None - - def find_bold(para_): - if para_.runs and all(run.bold for run in para_.runs if run.text.strip()): - return True - else: - return None - - raw_candidates = [] - logger.debug( - "Filtering docx heading candidates... total_items={}", len(titles_material) - ) - for ele_id, para, text in titles_material: - str_lvl = "" - est_lvl = None - style_lvl = find_docstyle(para) - setting_lvl = find_otsetting(para) - - # 1. check .docx style settings - if style_lvl is not None: - est_lvl = style_lvl - str_lvl = f"style-{style_lvl}" - - # 2. check .docx paragraph numbering settings - elif setting_lvl is not None: - est_lvl = setting_lvl - str_lvl = f"outline-{setting_lvl}" - - # 3. detect bold (unconditionally, encode as META dimension) - is_bold = 1 if find_bold(para) else 0 - - # 4. proceed condition judge - if enable_regx: - pos_code, detail_info = judge_by_conditions(text, return_detail=True) - neg_code = remove_by_conditions(text) - - if any(x > 0 for x in neg_code): - code_lvl = -1 - code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" - elif any(x > 0 for x in pos_code) and all(x == 0 for x in neg_code): - code_lvl = get_max_lvl(str(pos_code)) - code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}" - else: - code_lvl = -1 - code_str = f"POS {pos_code} NEG {neg_code}" - - # Append bold as META dimension (DOCX has no layout.json, so only bold) - if is_bold: - code_str += f" META [0, 0, {is_bold}]" - - if est_lvl is None: - est_lvl = code_lvl - str_lvl = code_str - else: - str_lvl = f"{str_lvl} AND {code_str}" - raw_candidates.append((ele_id, text, est_lvl, str_lvl)) - - preds_df = pd.DataFrame( - raw_candidates, columns=["id", "heading", "level", "reason"], index=None - ) - - # initial merge isolated and short texts - preds_df = postprocess_headings(preds_df, task="merge_continuous") - preds_df = postprocess_headings(preds_df, task="merge_short") - return preds_df - - def format_toc_context_for_llm(toc_context) -> str: """Convert TOC hierarchy or structured payloads into compact LLM-friendly plain text.""" if not toc_context: @@ -1434,11 +419,11 @@ def pred_titles( ) if doc_type == "pptx": - raw_preds = filter_md_headings(infos) + raw_preds = filter_markdown_headings(infos) elif doc_type == "md": - raw_preds = filter_md_headings(infos, layout_json_path=layout_json_path) + raw_preds = filter_markdown_headings(infos, layout_json_path=layout_json_path) elif doc_type == "docx": - raw_preds = filter_doc_headings(infos, enable_regx) + raw_preds = filter_document_headings(infos, enable_regex=enable_regx) else: raw_preds = pd.DataFrame(columns=["id", "heading", "level", "reason"]) @@ -1706,358 +691,26 @@ def est_hierarchies_llm( raw_preds: raw data prompt_limt: prompt character limit toc_hierarchies: TOC hierarchies - max_len: maximum heading length (passed through to heading_tb_transfer) + max_len: maximum heading length for executor chunk preparation max_depth: maximum hierarchy depth model_name: LLM model name output_dir: output directory, used to save intermediate results CSV csv_suffix: suffix for intermediate CSV filenames """ model_name = _resolve_hierarchy_model_name(model_name) - if len(raw_preds) == 0: - return pd.DataFrame(columns=["id", "heading", "level", "reason"]) - - compact_enabled = os.environ.get( - "KB_LAYOUT_LLM_COMPACT_INPUT", "true" - ).strip().lower() in ("true", "1", "yes", "on") - preds_for_llm = _compact_for_llm(raw_preds) if compact_enabled else raw_preds.copy() - if compact_enabled: - placeholder_count = int(preds_for_llm["reason"].eq(PLACEHOLDER_REASON).sum()) - logger.info( - f"smart parse => compact input: {len(raw_preds)} → {len(preds_for_llm)} rows " - f"({placeholder_count} placeholder groups)" - ) - - # Short-circuit: if there are no heading candidates to judge (all rows were - # collapsed into placeholders, or raw_preds contains only level==-1 rows - # with compaction disabled), skip the LLM entirely and return raw_preds - # with all levels set to -1. - non_placeholder = ( - preds_for_llm[preds_for_llm["reason"].astype(str) != PLACEHOLDER_REASON] - if compact_enabled - else preds_for_llm + return execute_llm_heading_hierarchy( + raw_preds=raw_preds, + prompt_limt=prompt_limt, + hierarchy_judge=hiearchy_llm, + fallback_hierarchy=est_hierarchies_naive, + save_intermediate_csv=save_intermediate_csv, + toc_hierarchies=toc_hierarchies, + max_len=max_len, + max_depth=max_depth, + model_name=model_name, + output_dir=output_dir, + csv_suffix=csv_suffix, ) - if len(non_placeholder) == 0: - logger.info( - "smart parse => no heading candidates, skipping LLM hierarchy detection" - ) - fallback = raw_preds.copy()[["id", "heading", "level", "reason"]] - fallback["level"] = -1 - return fallback.sort_values("id").reset_index(drop=True) - - level_dfs, _raw_headings = heading_tb_transfer( - preds_for_llm, threshold=prompt_limt, max_start=max_len, max_end=5 - ) - chunk_sizes = [len(d) for d in level_dfs] - logger.info( - f"smart parse => {len(level_dfs)} chunk(s) | rows per chunk: {chunk_sizes} | " - f"threshold={prompt_limt} | max_start={max_len}" - ) - - # Pick the first chunk that actually contains heading candidates. When - # compaction is enabled a small prompt_limt may push a placeholder-only - # chunk to index 0 — using it would waste an LLM call and produce an empty - # mapping. Placeholder chunks that precede the chosen one contribute no - # reason-code signal (their ids map to -1 anyway). - basic_idx = 0 - for idx, chunk in enumerate(level_dfs): - if (chunk["reason"].astype(str) != PLACEHOLDER_REASON).any(): - basic_idx = idx - break - basic_df = level_dfs[basic_idx] - if basic_idx != 0: - logger.info( - f"smart parse => promoted chunk {basic_idx} as basic_df " - f"(chunks 0..{basic_idx - 1} contain only placeholders)" - ) - full_preds = None - try: - with stage_timer( - "heading.hierarchy_llm", - chunk_count=len(level_dfs), - base_chunk_rows=len(basic_df), - compact_enabled=compact_enabled, - source_row_count=len(raw_preds), - model_name=model_name, - ): - logger.debug("🚀 smart parse => interpreting hierarchy patterns...") - df4llm = basic_df.drop(columns=["reason"]).copy() - from .metadata_extractor import clean_md_text_for_llm - - # Keep formatting signals in `reason` / preliminary `level`, but let the LLM - # judge hierarchy from the semantic heading text instead of raw markdown markers. - df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) - logger.debug(f"DataFrame transformation completed, rows: {len(df4llm)}") - - layout_res = hiearchy_llm( - df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings" - ) - - # Build base_preds by aligning on basic_df["id"]: we always have one - # row per chunk-0 row in the rendered output, regardless of how many - # entries the LLM actually returned. Missing ids -> level=-1. - layout_level_by_id = {} - if isinstance(layout_res, list): - for item in layout_res: - if isinstance(item, dict) and "id" in item and "level" in item: - layout_level_by_id[item["id"]] = item["level"] - - def _level_for(rid): - if rid in layout_level_by_id: - return layout_level_by_id[rid] - try: - return layout_level_by_id.get(int(rid), -1) - except (TypeError, ValueError): - return -1 - - base_preds = ( - basic_df[["id", "heading", "reason"]].copy().reset_index(drop=True) - ) - base_preds.insert(2, "level", base_preds["id"].map(_level_for)) - - # Save base_preds as preds_3 (reflects what the LLM saw, compact or not) - save_intermediate_csv( - base_preds, output_dir, f"preds_3_llm_base{csv_suffix}" - ) - - # Collect {int_id -> level} from chunk-0 LLM output. Placeholder rows - # have non-integer ids and are skipped. - llm_levels = {} - for _, row in base_preds.iterrows(): - rid = row["id"] - if isinstance(rid, bool): - continue - if isinstance(rid, int): - llm_levels[rid] = row["level"] - - if len(level_dfs) > 1: - # Multi-chunk: build reason-code mapping from chunk-0 candidates and - # apply it to chunks 1..N to infer levels for headings beyond chunk 0. - # Placeholder rows are excluded from both the mapping source and the - # per-chunk application — they always map to level=-1 in the final df. - placeholder_mask_base = base_preds["reason"].eq(PLACEHOLDER_REASON) - figure_mask_base = base_preds["heading"].eq("Figure/Image") - exclude_mask_base = placeholder_mask_base | figure_mask_base - base_preds_for_mapping = base_preds[~exclude_mask_base].copy() - base_origin_for_mapping = basic_df.loc[ - ~exclude_mask_base.values, "level" - ].tolist() - - base_preds_for_mapping, lvl_mapping = build_level_mapping( - base_preds_for_mapping, base_origin_for_mapping, mode="freq" - ) - logger.debug( - f"mapping development finished: {len(lvl_mapping)} rules " - f"(placeholders and Figure/Image excluded)" - ) - - logger.debug( - f"mapping dataframe to levels across {len(level_dfs)} chunks..." - ) - lvl_mapping = handle_unseen_codes( - preds_for_llm, level_dfs, lvl_mapping, output_dir - ) - - for level_df in level_dfs: - placeholder_mask_chunk = level_df["reason"].eq(PLACEHOLDER_REASON) - figure_mask_chunk = level_df["heading"].eq("Figure/Image") - exclude_mask_chunk = placeholder_mask_chunk | figure_mask_chunk - non_excluded = level_df[~exclude_mask_chunk].copy() - if not non_excluded.empty: - non_excluded = execute_level_mapping(non_excluded, lvl_mapping) - for _, row in non_excluded.iterrows(): - rid = row["id"] - if isinstance(rid, bool): - continue - if isinstance(rid, int): - # Mapping may override chunk-0 LLM decisions when - # two rows share the same reason; accept that (the - # mapping is by construction the "representative" - # level for each reason-code). - llm_levels[rid] = row["level"] - logger.info( - f"multi-chunk mapping produced {len(llm_levels)} id→level entries" - ) - else: - logger.info( - "single chunk — skipping reason-code mapping, using LLM output directly" - ) - - # Expand back onto the original raw_preds: heading candidates take - # the LLM/mapping-assigned level; everything else is body text (-1). - full_preds = raw_preds.copy() - full_preds = full_preds[["id", "heading", "level", "reason"]] - - def _resolve_level(rid): - try: - int_id = int(rid) - except (TypeError, ValueError): - return -1 - lvl = llm_levels.get(int_id, -1) - try: - return int(lvl) - except (TypeError, ValueError): - return -1 - - full_preds["level"] = full_preds["id"].map(_resolve_level).astype(int) - - save_intermediate_csv( - full_preds, output_dir, f"preds_4_llm_final{csv_suffix}" - ) - - except Exception as e: - logger.warning(f"LLM-based parsing fails due to {e}, using non-llm pipeline...") - full_preds = est_hierarchies_naive(raw_preds.copy()) - return full_preds - - -def collapse_recursive(df, task, indices, merge_th=3, checked_pairs=None, depth=0): - """recursive collapse""" - if checked_pairs is None: - checked_pairs = set() - - if len(indices) < 2: - return - - for k in range(len(indices) - 1): - i, j = indices[k], indices[k + 1] - if (i, j) in checked_pairs: - continue - checked_pairs.add((i, j)) - - between = df.loc[i + 1 : j - 1] - i_txt = df.at[i, "heading"].strip() - j_txt = df.at[j, "heading"].strip() - - if task == "merge_short" and len(between) > 0: - between_lens = [count_cn_en(c) for c in between["heading"].tolist()] - between_lvls = [bl for bl in between["level"].tolist()] - i_half_len = int(count_cn_en(i_txt) / 2) - too_short = sum(between_lens) <= merge_th or sum(between_lens) < i_half_len - - if too_short and all( - bl == -1 for bl in between_lvls - ): # only non-headings can be merged - logger.debug( - f"⚠️ too short between {i}=>{i_txt[:15]} and {j}=>{j_txt[:15]} => merge to {i}" - ) - between_txts = [ - str(r["heading"]).strip() - for _, r in between.iterrows() - if isinstance(r.get("heading"), str) and r["heading"].strip() - ] - - if between_txts: - joined_txt = "\n".join(between_txts) - df.at[i, "heading"] = f"{i_txt} {joined_txt}" - - for idx in between.index: - df.at[idx, "level"] = -1 - df.at[idx, "reason"] = f"Merged into {i}" - logger.debug(f"\tmerged texts: {joined_txt[:50]}...") - - elif task == "collapse" and len(between) == 0: - logger.debug( - f"⚠️ Empty between i={i_txt[:15]}, j={j_txt[:15]} => set i.level=-1, j.level=Not Sure" - ) - df.at[i, "level"] = -2 # "Not Sure" sentinel (int-safe) - df.at[j, "level"] = -2 # "Not Sure" sentinel (int-safe) - - # ========== get subgroups for recursive tasks ========== - sub_between = between[between["level"] != -1] - code2sub = defaultdict(list) - for idx, row in sub_between.iterrows(): - level = row["level"] - reason = row["reason"] - if level != -1: - code2sub[(level, reason)].append(idx) - - for _, sub_indices in code2sub.items(): - collapse_recursive( - df, task, sub_indices, merge_th, checked_pairs, depth + 1 - ) - - -def postprocess_headings(df, task, max_depth=-1): - """postprocess headings""" - if task == "judge_negs": - for i, row in df.iterrows(): - neg_code = remove_by_conditions(row["heading"], include_punc=True) - if any(x > 0 for x in neg_code): - current_code = str(df.loc[i, "reason"]) - - neg_match = re.search(r"(.*NEG\s*)\[[^\]]*\](.*)", current_code) - if neg_match: - update_code = f"{neg_match.group(1)}{neg_code}{neg_match.group(2)}" - else: - update_code = f"{current_code} NEG {neg_code}" - - df.loc[i, "level"] = -1 - df.loc[i, "reason"] = update_code - return df - - elif task == "merge_continuous": - denoised_rows = [] - punc_pattern = re.compile(r'[.,!?;:，。！？；：）】〕｝〉》’”"]$') - - i = 0 - while i < len(df): - row = df.iloc[i] - current_content = str(row["heading"]).strip() - current_level = row["level"] - - j = i + 1 - while j < len(df): - next_row = df.iloc[j] - next_content = str(next_row["heading"]).strip() - next_level = next_row["level"] - - # Skip merge if ID is not continuous (indicates table/image was skipped in between) - expected_id = row["id"] + (j - i) - if next_row["id"] != expected_id: - break - - # both current and next rows are not heading & current row has no punctuation -> merge - current_not_punc = not punc_pattern.search(current_content[-2:]) - if (current_level == -1 and next_level == -1) and current_not_punc: - current_content += " " + next_content - j += 1 - else: - break - - merge_row = row.copy() - merge_row["heading"] = current_content - denoised_rows.append(tuple(merge_row)) - i = j - return pd.DataFrame(denoised_rows, columns=["id", "heading", "level", "reason"]) - - elif task == "merge_short" or task == "collapse": - group2indices = defaultdict(list) - for idx, row in df.iterrows(): - level = row["level"] - reason = row["reason"] - if level != -1: - group2indices[(level, reason)].append(idx) - - checked_pairs = set() - for _, indices in group2indices.items(): - collapse_recursive( - df, task, indices, merge_th=3, checked_pairs=checked_pairs, depth=0 - ) - - if task == "merge_short": - drop_between = df.index[ - df["reason"].astype(str).str.startswith("Merged into", na=False) - ].tolist() - if drop_between: - logger.debug( - f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows" - ) - df.drop(drop_between, inplace=True) - df.reset_index(drop=True, inplace=True) - return df - - else: - return None # def parse_outline_hier(markdown_text): diff --git a/apps/worker/app/services/document_parser/markdown_deferred_summary.py b/apps/worker/app/services/document_parser/markdown_deferred_summary.py new file mode 100644 index 000000000..b11da7c80 --- /dev/null +++ b/apps/worker/app/services/document_parser/markdown_deferred_summary.py @@ -0,0 +1,292 @@ +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from typing import Any, TypeGuard + +import gevent +from app.services.document_parser.image_parser import _get_vision_client, ask_image +from app.services.document_parser.stage_profiler import stage_timer +from app.services.document_parser.table_parser import sanitize_table_name_from_header +from app.services.document_parser.txt_parser import ( + extract_title_keywords_summary, + split_title_summary, +) +from gevent.pool import Pool as GeventPool +from loguru import logger + +from shared.core.config import settings +from shared.utils.chunk_refs import build_chunk_ref +from shared.utils.file_utils import path_handle + +DeferredResult = ( + tuple[ + int, + str, + tuple[str | None, str | None] | tuple[str, str, str] | tuple[str, str], + ] +) +ImageSummaryResult = tuple[str | None, str | None] +TableSummaryResult = tuple[str, str, str] +TextSummaryResult = tuple[str, str] + + +@dataclass(frozen=True) +class MarkdownDeferredSummaryInput: + rows: list[list[str | int]] + tasks: list[tuple[Any, ...]] + output_dir: str + summary_len: int = 1500 + + +def apply_markdown_deferred_summaries( + deferred_input: MarkdownDeferredSummaryInput, +) -> None: + if not deferred_input.tasks: + return + + image_task_count = sum(1 for task in deferred_input.tasks if task[0] == "image") + table_task_count = sum(1 for task in deferred_input.tasks if task[0] == "table") + text_task_count = sum(1 for task in deferred_input.tasks if task[0] == "text") + logger.info( + f"Running {len(deferred_input.tasks)} deferred summary LLM calls in parallel" + ) + max_concurrent = getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 8) + + with stage_timer( + "md.deferred_summaries", + total_tasks=len(deferred_input.tasks), + image_tasks=image_task_count, + table_tasks=table_task_count, + text_tasks=text_task_count, + max_concurrent=min(max_concurrent, len(deferred_input.tasks)), + ): + results = _run_deferred_summary_tasks(deferred_input, max_concurrent) + _apply_deferred_summary_results(deferred_input, results) + + logger.info(f"Completed {len(deferred_input.tasks)} deferred summary LLM calls") + + +def replace_chunk_ref_in_rows( + rows: list[list[str | int]], old_path: str, new_path: str +) -> None: + old_ref = build_chunk_ref(old_path) + new_ref = build_chunk_ref(new_path) + if not old_ref or old_ref == new_ref: + return + + for row in rows: + if len(row) > 0 and isinstance(row[0], str): + row[0] = row[0].replace(old_ref, new_ref) + if len(row) > 1 and row[1] == old_path: + row[1] = new_path + if len(row) > 2 and isinstance(row[2], str): + row[2] = row[2].replace(old_ref, new_ref) + if len(row) > 8 and isinstance(row[8], str): + row[8] = row[8].replace(old_ref, new_ref) + + +def _run_deferred_summary_tasks( + deferred_input: MarkdownDeferredSummaryInput, + max_concurrent: int, +) -> list[DeferredResult | None]: + pool = GeventPool(size=min(max_concurrent, len(deferred_input.tasks))) + greenlets = [ + pool.spawn(_run_deferred_summary_task, task, deferred_input) + for task in deferred_input.tasks + ] + gevent.joinall(greenlets) + return [greenlet.value for greenlet in greenlets] + + +def _run_deferred_summary_task( + task: tuple[Any, ...], + deferred_input: MarkdownDeferredSummaryInput, +) -> DeferredResult | None: + task_type, row_index = task[0], task[1] + try: + if task_type == "image": + relative_path = task[2] + client = _get_vision_client() + # TODO: Risk of missing text content if MinerU outputted a pure text image. + # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. + llm_resp = ask_image( + client, deferred_input.output_dir, paths_=[relative_path] + ) + if llm_resp: + img_title, img_summary = split_title_summary(llm_resp) + else: + img_title, img_summary = None, None + return row_index, task_type, (img_title, img_summary) + + if task_type == "table": + table_html = task[2] + title, keywords, summary = extract_title_keywords_summary( + table_html, max_keywords=3 + ) + return row_index, task_type, (title, keywords, summary) + + if task_type == "text": + text_content = task[2] + _, keywords, summary = extract_title_keywords_summary( + text_content, + max_keywords=3, + summary_len=deferred_input.summary_len, + ) + return row_index, task_type, (keywords, summary) + except Exception as exc: + logger.warning( + f"Deferred {task_type} LLM call failed for idx={row_index}: {exc}" + ) + return None + + logger.warning(f"Unknown deferred markdown summary task type: {task_type}") + return None + + +def _apply_deferred_summary_results( + deferred_input: MarkdownDeferredSummaryInput, + results: list[DeferredResult | None], +) -> None: + deferred_by_index = {task[1]: task for task in deferred_input.tasks} + + for result in results: + if result is None: + continue + + row_index, task_type, task_result = result + if task_type == "image": + if not _is_image_summary_result(task_result): + logger.warning(f"Invalid image deferred result for idx={row_index}") + continue + _apply_image_summary_result( + deferred_input.rows, + deferred_by_index[row_index], + row_index, + task_result, + ) + elif task_type == "table": + if not _is_table_summary_result(task_result): + logger.warning(f"Invalid table deferred result for idx={row_index}") + continue + _apply_table_summary_result( + deferred_input.rows, + deferred_by_index[row_index], + row_index, + task_result, + ) + elif task_type == "text": + if not _is_text_summary_result(task_result): + logger.warning(f"Invalid text deferred result for idx={row_index}") + continue + _apply_text_summary_result(deferred_input.rows, row_index, task_result) + + +def _is_image_summary_result(result: object) -> TypeGuard[ImageSummaryResult]: + return ( + isinstance(result, tuple) + and len(result) == 2 + and all(isinstance(value, (str, type(None))) for value in result) + ) + + +def _is_table_summary_result(result: object) -> TypeGuard[TableSummaryResult]: + return ( + isinstance(result, tuple) + and len(result) == 3 + and all(isinstance(value, str) for value in result) + ) + + +def _is_text_summary_result(result: object) -> TypeGuard[TextSummaryResult]: + return ( + isinstance(result, tuple) + and len(result) == 2 + and all(isinstance(value, str) for value in result) + ) + + +def _apply_image_summary_result( + rows: list[list[str | int]], + original_task: tuple[Any, ...], + row_index: int, + result: ImageSummaryResult, +) -> None: + img_title, img_summary = result + row = rows[row_index] + if img_summary: + image_index = str(row[5]).split("\n")[0] if row[5] else "image" + row[5] = f"{image_index}\n{img_summary}" + + if not img_title: + return + + image_dir, old_img_name, image_suffix = ( + original_task[3], + original_task[4], + original_task[5], + ) + safe_title = path_handle(str(img_title), mode="clean_single") + img_num_match = re.match(r"image-(\d+)", str(old_img_name)) + img_num = ( + img_num_match.group(1) + if img_num_match + else str(old_img_name).split("-")[1] + if "-" in str(old_img_name) + else "0" + ) + new_img_name = path_handle(f"image-{img_num}-{safe_title}", mode="clean_single") + old_path = os.path.join(image_dir, f"{old_img_name}{image_suffix}") + new_path = os.path.join(image_dir, f"{new_img_name}{image_suffix}") + if old_path == new_path or not os.path.exists(old_path): + return + + os.rename(old_path, new_path) + new_relative_path = f"images/{new_img_name}{image_suffix}" + replace_chunk_ref_in_rows(rows, str(row[1]), new_relative_path) + row[1] = new_relative_path + + +def _apply_table_summary_result( + rows: list[list[str | int]], + original_task: tuple[Any, ...], + row_index: int, + result: TableSummaryResult, +) -> None: + title, keywords, summary = result + row = rows[row_index] + row[4] = keywords if isinstance(keywords, str) else "" + if summary: + table_index = str(row[5]) if "\n" not in str(row[5]) else str(row[5]).split("\n")[0] + row[5] = f"{table_index}\n{summary}" + + if not title: + return + + table_dir, old_table_name, table_count = ( + original_task[3], + original_task[4], + original_task[5], + ) + safe_title = sanitize_table_name_from_header(str(title)) + new_table_name = path_handle( + f"table-{table_count} {safe_title}", mode="clean_single" + ) + old_path = os.path.join(table_dir, f"{old_table_name}.html") + new_path = os.path.join(table_dir, f"{new_table_name}.html") + if old_path == new_path or not os.path.exists(old_path): + return + + os.rename(old_path, new_path) + new_relative_path = f"tables/{new_table_name}.html" + replace_chunk_ref_in_rows(rows, str(row[1]), new_relative_path) + row[1] = new_relative_path + + +def _apply_text_summary_result( + rows: list[list[str | int]], row_index: int, result: TextSummaryResult +) -> None: + keywords, summary = result + rows[row_index][4] = keywords if isinstance(keywords, str) else "" + rows[row_index][5] = summary if isinstance(summary, str) else "" diff --git a/apps/worker/app/services/document_parser/markdown_parse_state.py b/apps/worker/app/services/document_parser/markdown_parse_state.py new file mode 100644 index 000000000..eb8837f0f --- /dev/null +++ b/apps/worker/app/services/document_parser/markdown_parse_state.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +import re +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import pandas as pd + +from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder + +ParserRowValues = list[str | int] + +RowUpdater = Callable[ + [list[ParserRowValues], list[str], str, dict[str, Any], str, str, int, bool], + list[ParserRowValues], +] + + +@dataclass +class MarkdownParseState: + relative_root: str + split_char: str + llm_parameters: dict[str, Any] + timestamp: str + row_updater: RowUpdater + rows: list[ParserRowValues] = field(default_factory=list) + content_items: list[str] = field(default_factory=list) + path_stack: list[tuple[str, int]] = field(default_factory=list) + inner_paths: list[str] = field(default_factory=list) + error_line_numbers: list[int] = field(default_factory=list) + table_lines: list[str] = field(default_factory=list) + current_page_number: int = 0 + chunk_pages: set[int] = field(default_factory=set) + base_level: int | None = None + path: str = "" + path_counter: dict[str, int] = field(default_factory=dict) + deferred_llm_tasks: list[tuple[Any, ...]] = field(default_factory=list) + seen_images: dict[str, dict[str, str]] = field(default_factory=dict) + image_count: int = 1 + table_count: int = 1 + + def __post_init__(self) -> None: + if not self.path: + self.path = self.relative_root + + def record_page_marker(self, line: str) -> bool: + if "" not in line: + return False + if "page" not in line and "Slide number" not in line: + return False + + page_match = re.search(r"page\s+(\d+)", line) + if page_match: + self.current_page_number = int(page_match.group(1)) + else: + self.current_page_number += 1 + self.chunk_pages.add(self.current_page_number) + return True + + def flush_current_content(self) -> None: + page_numbers = self._format_chunk_pages() + self.rows = self.row_updater( + self.rows, + self.content_items, + self.path, + self.llm_parameters, + self.timestamp, + page_numbers, + 1500, + True, + ) + self.content_items = [] + self.chunk_pages = set() + if self.current_page_number > 0: + self.chunk_pages.add(self.current_page_number) + + def flush_placeholder_chunk(self) -> None: + page_numbers = self._format_chunk_pages() + self.rows = self.row_updater( + self.rows, + [], + self.path, + self.llm_parameters, + self.timestamp, + page_numbers, + 1500, + True, + ) + + def enter_heading(self, heading: str, level: int) -> None: + if self.base_level is None: + self.base_level = level + elif level < self.base_level: + self.base_level = level + + adjusted_level = level - self.base_level + 1 + self.path_stack = [ + (item_heading, item_level) + for item_heading, item_level in self.path_stack + if item_level < adjusted_level + ] + + current_heading = ( + heading.replace(self.split_char, "∕") + if self.split_char in heading + else heading + ) + tentative_names = [item_heading for item_heading, _ in self.path_stack] + tentative_names.append(current_heading) + tentative_path_parts = [self.relative_root] if self.relative_root else [] + tentative_path_parts.extend(tentative_names) + tentative_path = self.split_char.join(tentative_path_parts) + + if tentative_path in self.path_counter: + self.path_counter[tentative_path] += 1 + current_heading = f"{current_heading}_{self.path_counter[tentative_path]}" + else: + self.path_counter[tentative_path] = 1 + + self.path_stack.append((current_heading, adjusted_level)) + heading_names = [item_heading for item_heading, _ in self.path_stack] + path_parts = [self.relative_root] if self.relative_root else [] + path_parts.extend(heading_names) + self.inner_paths.append(self.split_char.join(heading_names)) + self.path = self.split_char.join(path_parts) + + def append_content_item(self, item: str) -> None: + self.content_items.append(item) + + def append_plain_text(self, text: str) -> None: + self.content_items.append(text.strip()) + if self.current_page_number > 0: + self.chunk_pages.add(self.current_page_number) + + def append_row(self, row: ParserRowValues) -> None: + self.rows.append(row) + + def schedule_deferred_task(self, task: tuple[Any, ...]) -> None: + self.deferred_llm_tasks.append(task) + + def collect_text_summary_tasks(self, summary_len: int) -> None: + if not self.llm_parameters.get("summary_txt"): + return + + for index, entry in enumerate(self.rows): + marker = entry[2] + if isinstance(marker, str) and marker.strip().split("\n", 1)[0].lower() in { + "image", + "table", + }: + continue + content = str(entry[0]) + if len(content) > summary_len and not entry[4] and not entry[5]: + self.deferred_llm_tasks.append(("text", index, content)) + + def to_dataframe(self) -> pd.DataFrame: + rows_builder = ParsedRowsBuilder() + for row_values in self.rows: + rows_builder.append( + ParsedRow( + content=str(row_values[0]), + path=str(row_values[1]), + type=str(row_values[2]), + length=int(row_values[3]), + keywords=str(row_values[4]), + summary=str(row_values[5]), + know_id=str(row_values[6]), + tokens=str(row_values[7]), + connectto=str(row_values[8]), + addtime=str(row_values[9]), + page_nums=str(row_values[10]), + ) + ) + return process_dup_paths_df(rows_builder.to_dataframe()) + + def _format_chunk_pages(self) -> str: + return ",".join(str(page) for page in sorted(self.chunk_pages)) diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py index 80c1a99bd..83f8a3e58 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/md_parser.py @@ -5,10 +5,17 @@ import shutil from pathlib import Path -import gevent -import pandas as pd -from app.services.document_parser.dataframe_helpers import process_dup_paths_df from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.inline_asset import ( + build_image_asset_row, + build_table_asset_row, +) +from app.services.document_parser.markdown_deferred_summary import ( + MarkdownDeferredSummaryInput, + apply_markdown_deferred_summaries, +) +from app.services.document_parser.markdown_parse_state import MarkdownParseState +from app.services.document_parser.parser_rows import ParsedRow from app.services.document_parser.path_helpers import find_matches_parsing from app.services.document_parser.html_parser import ( first_cols_rows_html, @@ -16,12 +23,14 @@ ) from app.services.document_parser.image_parser import ( MD_IMAGE_PATTERN, - _get_vision_client, - ask_image, detect_summary_img_md, perceptual_hash, ) -from app.services.document_parser.layout_parser import md_heading_match, pred_titles +from app.services.document_parser.heading_hierarchy import ( + HeadingHierarchyInput, + predict_heading_hierarchy, +) +from app.services.document_parser.heading_candidates import md_heading_match from app.services.document_parser.stage_profiler import stage_timer from app.services.document_parser.table_parser import ( extract_tables_by_forms, @@ -29,11 +38,7 @@ sanitize_table_name_from_header, ) from app.services.document_parser.toc_parser import detect_tocs_in_texts -from app.services.document_parser.txt_parser import ( - extract_title_keywords_summary, - split_title_summary, -) -from gevent.pool import Pool as GeventPool +from app.services.document_parser.txt_parser import extract_title_keywords_summary from loguru import logger from shared.core.config import settings @@ -136,15 +141,17 @@ def eval_md_headings( layout_json_path=None, ): """Evaluate markdown headings with optional TOC hierarchies context""" - heading_preds = pred_titles( - md_lines, - source_type, - toc_hierarchies=toc_hierarchies, - enable_regx=True, - smart_parse=smart_parse, - model_name=model_name, - output_dir=output_dir, - layout_json_path=layout_json_path, + heading_preds = predict_heading_hierarchy( + HeadingHierarchyInput( + infos=md_lines, + doc_type=source_type, + toc_hierarchies=toc_hierarchies, + enable_regex=True, + smart_parse=smart_parse, + model_name=model_name, + output_dir=output_dir, + layout_json_path=layout_json_path, + ) ) if len(heading_preds) == 0: @@ -181,24 +188,6 @@ def clean_md_table_lines(table_lines, start_line_num): return cleaned_lines, error_lines -def replace_chunk_ref_in_rows(df_list, old_path: str, new_path: str) -> None: - """Rewrite readable chunk refs after deferred image/table renames.""" - old_ref = build_chunk_ref(old_path) - new_ref = build_chunk_ref(new_path) - if not old_ref or old_ref == new_ref: - return - - for row in df_list: - if len(row) > 0 and isinstance(row[0], str): - row[0] = row[0].replace(old_ref, new_ref) - if len(row) > 1 and row[1] == old_path: - row[1] = new_path - if len(row) > 2 and isinstance(row[2], str): - row[2] = row[2].replace(old_ref, new_ref) - if len(row) > 8 and isinstance(row[8], str): - row[8] = row[8].replace(old_ref, new_ref) - - def update_df_list( df_list, content_items, @@ -238,19 +227,17 @@ def update_df_list( ) df_list.append( - [ - bottom_content, - path, - match_type, - len(bottom_content), - keywords, - summary, - know_id, - bottom_tokens, - "", - time_stamp, - page_nums, - ] + ParsedRow( + content=bottom_content, + path=path, + type=match_type, + keywords=keywords, + summary=summary, + know_id=know_id, + tokens=bottom_tokens, + addtime=time_stamp, + page_nums=page_nums, + ).to_list() ) return df_list @@ -323,22 +310,13 @@ def parse_md( # initialize vars split_char = settings.SPLIT_CHAR or "/" - df_list = [] - path_stack = [] # sxjg: uses (heading, level) tuples - inner_paths = [] - error_line_numbers = [] - table_lines = [] - current_pg_num = 0 - chunk_pages = set() # collect all page numbers seen during current chunk - base_level = None - content_items = [] - # Use relative_root as initial path (not absolute output_dir) - path = relative_root if relative_root else "" - table_count = 1 - img_count = 1 - path_counter = {} # Track path occurrences for deduplication - deferred_llm_tasks = [] # Collected during loop, executed in parallel after - _seen_images: dict[str, dict] = {} # sha256_hex -> cached image info for dedup + parser_state = MarkdownParseState( + relative_root=relative_root or "", + split_char=split_char, + llm_parameters=base_llm_paras, + timestamp=get_str_time(), + row_updater=update_df_list, + ) # Find layout.json path layout_json_path = os.path.join(output_dir, "layout.json") @@ -363,19 +341,10 @@ def parse_md( layout_json_path=layout_json_path, ) - time_stamp = get_str_time() logger.debug("Parsing md data... total_lines={}", len(lines_with_heading)) for i, line in enumerate(lines_with_heading): - if "" in line: - if "page" in line or "Slide number" in line: - # Parse actual page number from marker: - pg_match = re.search(r"page\s+(\d+)", line) - if pg_match: - current_pg_num = int(pg_match.group(1)) - else: - current_pg_num += 1 # fallback for Slide number or old format - chunk_pages.add(current_pg_num) - continue + if parser_state.record_page_marker(line): + continue last_context = find_surround_context( lines_with_heading, i @@ -385,88 +354,19 @@ def parse_md( if ( not current_heading_level == -1 ): # indicate a new path should be evaluated or added - if content_items: # record contents of the last path and reset content - # Build page_nums from collected pages during this chunk - chunk_page_str = ( - ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else "" - ) - df_list = update_df_list( - df_list, - content_items, - path, - base_llm_paras, - time_stamp, - page_nums=chunk_page_str, - skip_llm=True, - ) - content_items = [] - chunk_pages = set() # reset for next chunk - if current_pg_num > 0: - chunk_pages.add( - current_pg_num - ) # carry current page into next chunk - elif path and path != (relative_root or ""): + if parser_state.content_items: + parser_state.flush_current_content() + elif parser_state.path and parser_state.path != (relative_root or ""): # Consecutive headings with no body text between them: # Create a placeholder chunk so the previous heading's path - chunk_page_str = ( - ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else "" - ) - df_list = update_df_list( - df_list, - [], - path, - base_llm_paras, - time_stamp, - page_nums=chunk_page_str, - skip_llm=True, - ) - - # update path based on path name and level - if base_level is None: - base_level = current_heading_level - elif current_heading_level < base_level: - base_level = current_heading_level - - adjusted_level = current_heading_level - base_level + 1 - path_stack = [(h, lvl) for h, lvl in path_stack if lvl < adjusted_level] - - # Build tentative path to check for duplicates - # Sanitize heading: replace split_char in heading text to prevent path corruption - current_heading = ( - current_heading.replace(split_char, "∕") - if split_char in current_heading - else current_heading - ) - tentative_heading = current_heading - tentative_names = [h for h, lvl in path_stack] + [tentative_heading] - tentative_path_parts = [relative_root] if relative_root else [] - tentative_path_parts.extend(tentative_names) - tentative_path = split_char.join(tentative_path_parts) - - # Deduplicate: if path already exists, add suffix - if tentative_path in path_counter: - path_counter[tentative_path] += 1 - suffix = path_counter[tentative_path] - current_heading = ( - f"{current_heading}_{suffix}" # Modify heading with suffix - ) - else: - path_counter[tentative_path] = 1 + parser_state.flush_placeholder_chunk() - path_stack.append((current_heading, adjusted_level)) - - # Extract pure heading names for path construction - heading_names = [h for h, lvl in path_stack] - # Use relative_root as prefix - path_parts = [relative_root] if relative_root else [] - path_parts.extend(heading_names) - inner_paths.append(split_char.join(heading_names)) - path = split_char.join(path_parts) # path with relative root + parser_state.enter_heading(current_heading, current_heading_level) else: # no path change, remain in the same hierarchy # a. handle lines containing images (LLM deferred to post-loop parallel batch) img_name_context = path_handle(last_context[:10], mode="clean_single") - img_name = f"image-{str(img_count)}-{img_name_context}" + img_name = f"image-{str(parser_state.image_count)}-{img_name_context}" # Always skip inline LLM — vision calls are deferred to parallel batch imgs = detect_summary_img_md(line, last_context, output_dir, mode=False) @@ -478,30 +378,27 @@ def parse_md( source_path = resolve_markdown_image_source_path(output_dir, img_path) if source_path is None or not source_path.exists(): logger.warning(f"Image file not found, skipping rename: {img_path}") - img_count += 1 + parser_state.image_count += 1 continue # Document-level dedup: perceptual hash for visual duplicates with open(source_path, "rb") as f: img_binary_hash = perceptual_hash(f.read()) - if img_binary_hash in _seen_images: - cached = _seen_images[img_binary_hash] - content_items.append(cached["img_content"]) - df_list.append( - [ - cached["img_content"], - cached["relative_img_path"], - "image", - len(cached["img_content"]), - "", - cached["img_summary_field"], - cached["temp_uid"], - "", - "", - time_stamp, - str(current_pg_num) if current_pg_num > 0 else "", - ] + if img_binary_hash in parser_state.seen_images: + cached = parser_state.seen_images[img_binary_hash] + parser_state.append_content_item(cached["img_content"]) + parser_state.append_row( + build_image_asset_row( + content=cached["img_content"], + relative_path=cached["relative_img_path"], + summary=cached["img_summary_field"], + know_id=cached["temp_uid"], + addtime=parser_state.timestamp, + page_nums=str(parser_state.current_page_number) + if parser_state.current_page_number > 0 + else "", + ).to_list() ) logger.debug( f"Skipped duplicate image (hash={img_binary_hash[:12]}...)" @@ -516,7 +413,7 @@ def parse_md( os.rename(source_path, update_img_path) # Image index (always present) - image_index = f"image-{img_count}" + image_index = f"image-{parser_state.image_count}" # Fallback: LLM summary -> last_context -> None effective_summary = img_summary or last_context or None @@ -538,26 +435,23 @@ def parse_md( else: img_content = f"\n{img_ref}\n" - content_items.append(img_content) - - df_list.append( - [ - img_content, - relative_img_path, - "image", - len(img_content), - "", - img_summary_field, - temp_uid, - "", - "", - time_stamp, - str(current_pg_num) if current_pg_num > 0 else "", - ] + parser_state.append_content_item(img_content) + + parser_state.append_row( + build_image_asset_row( + content=img_content, + relative_path=relative_img_path, + summary=img_summary_field, + know_id=temp_uid, + addtime=parser_state.timestamp, + page_nums=str(parser_state.current_page_number) + if parser_state.current_page_number > 0 + else "", + ).to_list() ) # Cache result for document-level dedup - _seen_images[img_binary_hash] = { + parser_state.seen_images[img_binary_hash] = { "relative_img_path": relative_img_path, "img_content": img_content, "img_summary_field": img_summary_field, @@ -566,17 +460,17 @@ def parse_md( if base_llm_paras["summary_image"]: # Store img_dir, img_name, img_suffix for post-loop rename (mirrors table deferred task) - deferred_llm_tasks.append( + parser_state.schedule_deferred_task( ( "image", - len(df_list) - 1, + len(parser_state.rows) - 1, relative_img_path, img_dir, img_name, img_suffix, ) ) - img_count += 1 + parser_state.image_count += 1 # TODO for large and dense tables, such as "Epstein flight logs", # integrate tabula-py as an independent extraction path to solve VLM hallucinations and misplacement @@ -588,7 +482,7 @@ def parse_md( tb_str = line elif form == "md": # For MD tables, accumulate lines until table ends - table_lines.append(line) + parser_state.table_lines.append(line) if i + 1 >= len(lines_with_heading): tb_bool_next = False else: @@ -598,10 +492,10 @@ def parse_md( if not tb_bool_next or i == len(lines_with_heading) - 1: cleaned_table_lines, error_lines = clean_md_table_lines( - table_lines, start_line_num=i + parser_state.table_lines, start_line_num=i ) tb_str = "\n".join(cleaned_table_lines) - error_line_numbers.extend(error_lines) + parser_state.error_line_numbers.extend(error_lines) tb_str = extract_tables_by_forms(tb_str, form="md") else: continue # Keep accumulating MD table lines @@ -612,7 +506,7 @@ def parse_md( first_row_text, first_col_text = first_cols_rows_html(tb_str) # Table index (always present) - table_index = f"table-{table_count}" + table_index = f"table-{parser_state.table_count}" # LLM title + keywords + summary deferred to post-loop parallel batch llm_title = None @@ -633,18 +527,19 @@ def parse_md( # Use LLM title for filename when available, fallback to sanitized header effective_name = llm_title if llm_title else raw_tb_name tb_name = path_handle( - f"table-{str(table_count)} {effective_name}", mode="clean_single" + f"table-{str(parser_state.table_count)} {effective_name}", + mode="clean_single", ) - temp_uid = gen_str_codes((tb_str + str(table_count))) + temp_uid = gen_str_codes((tb_str + str(parser_state.table_count))) relative_tb_path = f"tables/{tb_name}.html" tb_ref = build_chunk_ref(relative_tb_path) # Build table_ref for content: optional LLM summary + table path ref if llm_summary: - content_items.append(f"\n{llm_summary}\n{tb_ref}\n") + parser_state.append_content_item(f"\n{llm_summary}\n{tb_ref}\n") else: - content_items.append(f"\n{tb_ref}\n") + parser_state.append_content_item(f"\n{tb_ref}\n") tb_path = os.path.join(tb_dir, f"{tb_name}.html") # Add border to HTML tables for consistent display tb_str_with_border = tb_str.replace( @@ -653,215 +548,53 @@ def parse_md( with open(tb_path, "w", encoding="utf-8") as f: f.write(tb_str_with_border) - df_list.append( - [ - tb_str, - relative_tb_path, - "table", - len(tb_str), - tb_keywords, - tb_summary, - temp_uid, - "", - "", - time_stamp, - str(current_pg_num) if current_pg_num > 0 else "", - ] + parser_state.append_row( + build_table_asset_row( + content=tb_str, + relative_path=relative_tb_path, + summary=tb_summary, + keywords=tb_keywords, + know_id=temp_uid, + addtime=parser_state.timestamp, + page_nums=str(parser_state.current_page_number) + if parser_state.current_page_number > 0 + else "", + ).to_list() ) if base_llm_paras["summary_table"]: - deferred_llm_tasks.append( + parser_state.schedule_deferred_task( ( "table", - len(df_list) - 1, + len(parser_state.rows) - 1, tb_str, tb_dir, tb_name, - table_count - 1, + parser_state.table_count - 1, ) ) - table_lines = [] # Reset table_lines after storing the DataFrame - table_count += 1 + parser_state.table_lines = [] + parser_state.table_count += 1 # c. handle plain texts if len(imgs) == 0 and not tb_bool: - content_items.append(line.strip()) - if current_pg_num > 0: - chunk_pages.add(current_pg_num) # track page for this content line + parser_state.append_plain_text(line) - if content_items: # handle the remaining contents, append them to the last section - chunk_page_str = ( - ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else "" - ) - df_list = update_df_list( - df_list, - content_items, - path, - base_llm_paras, - time_stamp, - page_nums=chunk_page_str, - skip_llm=True, - ) + if parser_state.content_items: + parser_state.flush_current_content() # Collect text chunk deferred tasks (entries needing summary/keywords) summary_len = 1500 - if base_llm_paras.get("summary_txt"): - for idx, entry in enumerate(df_list): - marker = entry[2] # col 2: match_type / img_id / table_id - if isinstance(marker, str) and marker.strip().split("\n", 1)[0].lower() in { - "image", - "table", - }: - continue - if len(entry[0]) > summary_len and not entry[4] and not entry[5]: - deferred_llm_tasks.append(("text", idx, entry[0])) - - # ── Post-loop: execute all deferred LLM calls in parallel via gevent ── - if deferred_llm_tasks: - image_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "image") - table_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "table") - text_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "text") - logger.info( - f"Running {len(deferred_llm_tasks)} deferred summary LLM calls in parallel" + parser_state.collect_text_summary_tasks(summary_len) + apply_markdown_deferred_summaries( + MarkdownDeferredSummaryInput( + rows=parser_state.rows, + tasks=parser_state.deferred_llm_tasks, + output_dir=output_dir, + summary_len=summary_len, ) - max_concurrent = getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 8) - - with stage_timer( - "md.deferred_summaries", - total_tasks=len(deferred_llm_tasks), - image_tasks=image_task_count, - table_tasks=table_task_count, - text_tasks=text_task_count, - max_concurrent=min(max_concurrent, len(deferred_llm_tasks)), - ): - - def _run_deferred(task): - task_type, idx = task[0], task[1] - try: - if task_type == "image": - relative_path = task[2] - client = _get_vision_client() - # TODO: Risk of missing text content if MinerU outputted a pure text image. - # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. - llm_resp = ask_image(client, output_dir, paths_=[relative_path]) - if llm_resp: - img_title, img_summary = split_title_summary(llm_resp) - else: - img_title, img_summary = None, None - return idx, task_type, (img_title, img_summary) - elif task_type == "table": - tb_html = task[2] - title, kw, summary = extract_title_keywords_summary( - tb_html, max_keywords=3 - ) - return idx, task_type, (title, kw, summary) - elif task_type == "text": - text_content = task[2] - _, kw, summary = extract_title_keywords_summary( - text_content, max_keywords=3, summary_len=summary_len - ) - return idx, task_type, (kw, summary) - except Exception as e: - logger.warning( - f"Deferred {task_type} LLM call failed for idx={idx}: {e}" - ) - return idx, task_type, None - - pool = GeventPool(size=min(max_concurrent, len(deferred_llm_tasks))) - greenlets = [pool.spawn(_run_deferred, task) for task in deferred_llm_tasks] - gevent.joinall(greenlets) - - # Build a lookup from deferred task list: idx -> original task tuple - deferred_by_idx = {task[1]: task for task in deferred_llm_tasks} + ) - for g in greenlets: - if g.value is None: - continue - idx, task_type, result = g.value - if result is None: - continue - if task_type == "image": - img_title, img_summary = result - entry = df_list[idx] - if img_summary: - image_index = entry[5].split("\n")[0] if entry[5] else "image" - entry[5] = f"{image_index}\n{img_summary}" - # Rename image file if LLM provided a better title (mirrors table rename logic) - if img_title: - orig_task = deferred_by_idx[idx] - i_dir, old_img_name, i_suffix = ( - orig_task[3], - orig_task[4], - orig_task[5], - ) - safe_title = path_handle(img_title, mode="clean_single") - # Derive image index number from old_img_name (e.g. "image-3-xxx" -> "3") - img_num_match = re.match(r"image-(\d+)", old_img_name) - img_num = ( - img_num_match.group(1) - if img_num_match - else ( - old_img_name.split("-")[1] - if "-" in old_img_name - else "0" - ) - ) - new_img_name = path_handle( - f"image-{img_num}-{safe_title}", mode="clean_single" - ) - old_path = os.path.join(i_dir, f"{old_img_name}{i_suffix}") - new_path = os.path.join(i_dir, f"{new_img_name}{i_suffix}") - if old_path != new_path and os.path.exists(old_path): - os.rename(old_path, new_path) - new_relative_path = f"images/{new_img_name}{i_suffix}" - replace_chunk_ref_in_rows( - df_list, entry[1], new_relative_path - ) - entry[1] = new_relative_path - elif task_type == "table": - title, kw, summary = result - entry = df_list[idx] - entry[4] = kw if isinstance(kw, str) else "" - if summary: - table_index = ( - entry[5] - if "\n" not in entry[5] - else entry[5].split("\n")[0] - ) - entry[5] = f"{table_index}\n{summary}" - # Rename table file if LLM provided a better title - if title: - orig_task = deferred_by_idx[idx] - t_dir, old_tb_name, t_count = ( - orig_task[3], - orig_task[4], - orig_task[5], - ) - safe_title = ( - sanitize_table_name_from_header(title) if title else "" - ) - new_tb_name = path_handle( - f"table-{t_count} {safe_title}", mode="clean_single" - ) - old_path = os.path.join(t_dir, f"{old_tb_name}.html") - new_path = os.path.join(t_dir, f"{new_tb_name}.html") - if old_path != new_path and os.path.exists(old_path): - os.rename(old_path, new_path) - new_relative_path = f"tables/{new_tb_name}.html" - replace_chunk_ref_in_rows( - df_list, entry[1], new_relative_path - ) - entry[1] = new_relative_path - elif task_type == "text": - kw, summary = result - df_list[idx][4] = kw if isinstance(kw, str) else "" - df_list[idx][5] = summary if isinstance(summary, str) else "" - - logger.info( - f"Completed {len(deferred_llm_tasks)} deferred summary LLM calls" - ) - - with stage_timer("md.build_dataframe", row_count=len(df_list)): - doc_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(",")) - doc_df = process_dup_paths_df(doc_df) + with stage_timer("md.build_dataframe", row_count=len(parser_state.rows)): + doc_df = parser_state.to_dataframe() return doc_df diff --git a/apps/worker/app/services/document_parser/mineru_client.py b/apps/worker/app/services/document_parser/mineru_client.py new file mode 100644 index 000000000..1cf5e3b97 --- /dev/null +++ b/apps/worker/app/services/document_parser/mineru_client.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from typing import Any, Optional + +import requests +from app.services.document_parser.mineru_quota_manager import get_mineru_quota_manager +from loguru import logger +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import UnavailableException + + +def build_mineru_session() -> requests.Session: + session = requests.Session() + retry_strategy = Retry( + total=settings.MINERU_UPLOAD_RETRY_TOTAL, + backoff_factor=settings.MINERU_UPLOAD_RETRY_BACKOFF_FACTOR, + status_forcelist=[429, 502, 503, 504], + allowed_methods=["GET", "POST", "PUT"], + raise_on_status=False, + ) + adapter = HTTPAdapter( + max_retries=retry_strategy, + pool_connections=1, + pool_maxsize=settings.MINERU_POOL_MAXSIZE, + ) + session.mount("https://", adapter) + session.mount("http://", adapter) + return session + + +_mineru_session: Optional[requests.Session] = None + + +def get_mineru_session() -> requests.Session: + global _mineru_session + if _mineru_session is None: + _mineru_session = build_mineru_session() + return _mineru_session + + +def get_mineru_headers(api_key: str) -> dict[str, str]: + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + + +def mineru_logger(step: str, **fields: Any): + return logger.bind(service="mineru", step=step, **fields) + + +def get_retry_after_seconds( + response: requests.Response, default_retry_after: int +) -> int: + retry_after_header = response.headers.get("Retry-After") + if retry_after_header: + try: + return max( + 1, + min( + int(retry_after_header), settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER + ), + ) + except ValueError: + logger.debug(f"Invalid MinerU Retry-After header: {retry_after_header}") + + return max(1, min(default_retry_after, settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER)) + + +def raise_mineru_unavailable( + token_id: str, response: requests.Response, operation: str +) -> None: + retry_after = get_retry_after_seconds( + response, settings.MINERU_TOKEN_COOLDOWN_SECONDS + ) + quota_manager = get_mineru_quota_manager() + quota_manager.mark_rate_limited(token_id, retry_after) + mineru_logger( + "rate_limited", + operation=operation, + token_id=token_id, + status_code=response.status_code, + retry_after=retry_after, + ).warning("MinerU request rate-limited") + raise UnavailableException( + internal_message=f"MinerU rate limited during {operation}", + retry_after=retry_after, + limit=settings.MINERU_TOKEN_RPM_LIMIT, + period="minute", + user_message="Document processing is busy right now. Please retry shortly.", + ) diff --git a/apps/worker/app/services/document_parser/mineru_pdf_service.py b/apps/worker/app/services/document_parser/mineru_pdf_service.py index d3cde60e4..2025609ff 100644 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ b/apps/worker/app/services/document_parser/mineru_pdf_service.py @@ -1,28 +1,29 @@ -# pyright: reportUnusedExpression=false import os -import time -from typing import Any, Callable, Optional +from typing import Optional import requests +from app.services.document_parser.mineru_client import ( + get_mineru_headers, + get_mineru_session, + mineru_logger, + raise_mineru_unavailable, +) from app.services.document_parser.mineru_quota_manager import get_mineru_quota_manager +from app.services.document_parser.mineru_task_polling import ( + get_batch_status, + poll_mineru_task, +) from app.services.document_parser.parser_log_utils import truncate_log_value -from loguru import logger -from requests.adapters import HTTPAdapter -from urllib3.util.retry import Retry from shared.core.config import settings from shared.core.constants import APIConstants from shared.core.exceptions.domain_exceptions import ( MinerUServiceException, - PDFParsingException, StorageServiceException, - TimeoutException, UnavailableException, ) -from shared.core.exceptions.knowhere_exception import KnowhereException from shared.services.storage.job_file_storage import JobFileStorage from shared.utils.file_loading import is_remote -from shared.utils.zip_download import download_and_extract_zip MINERU_UPLOAD_TIMEOUT = ( settings.MINERU_UPLOAD_CONNECT_TIMEOUT, @@ -30,49 +31,6 @@ ) -def _build_mineru_session() -> requests.Session: - session = requests.Session() - # Hybrid rate-limit control: urllib3 handles transient 429s with backoff - # (respects Retry-After header); application-level handling in callers - # covers persistent rate limits with Redis token marking and pool rotation. - retry_strategy = Retry( - total=settings.MINERU_UPLOAD_RETRY_TOTAL, - backoff_factor=settings.MINERU_UPLOAD_RETRY_BACKOFF_FACTOR, - status_forcelist=[429, 502, 503, 504], - allowed_methods=["GET", "POST", "PUT"], - raise_on_status=False, - ) - adapter = HTTPAdapter( - max_retries=retry_strategy, - pool_connections=1, - pool_maxsize=settings.MINERU_POOL_MAXSIZE, - ) - session.mount("https://", adapter) - session.mount("http://", adapter) - return session - - -_mineru_session: Optional[requests.Session] = None - - -def get_mineru_session() -> requests.Session: - global _mineru_session - if _mineru_session is None: - _mineru_session = _build_mineru_session() - return _mineru_session - - -def get_mineru_headers(api_key: str) -> dict[str, str]: - return { - "Content-Type": "application/json", - "Authorization": f"Bearer {api_key}", - } - - -def _mineru_logger(step: str, **fields: Any): - return logger.bind(service="mineru", step=step, **fields) - - def _should_use_mineru_s3_url_mode(s3_key: Optional[str]) -> bool: if settings.FORCE_MINERU_UPLOAD_ENABLED: return False @@ -86,7 +44,7 @@ def _log_mineru_url_mode_storage_fallback( local_file_path: Optional[str], exc: Exception, ) -> None: - _mineru_logger( + mineru_logger( "url_mode_storage_fallback", operation=operation, source_s3_key=s3_key, @@ -104,7 +62,7 @@ def _log_mineru_url_mode_ingestion_fallback( pdf_url: str, exc: Exception, ) -> None: - _mineru_logger( + mineru_logger( "url_mode_ingestion_fallback", operation=operation, source_s3_key=s3_key, @@ -133,7 +91,7 @@ def _inspect_mineru_source_s3_key(s3_key: Optional[str]) -> tuple[Optional[str], return None, False if existing_file.get("exists"): - _mineru_logger( + mineru_logger( "url_mode_source_reused", source_s3_key=s3_key, ).info("Reusing existing S3 source for MinerU URL mode") @@ -175,7 +133,7 @@ def resolve_mineru_source_s3_key( ) return None - _mineru_logger( + mineru_logger( "url_mode_source_uploaded", source_s3_key=s3_key, local_file_path=local_file_path, @@ -183,293 +141,11 @@ def resolve_mineru_source_s3_key( return s3_key -def _get_retry_after_seconds( - response: requests.Response, default_retry_after: int -) -> int: - """Parse Retry-After header with sane bounds for worker backoff.""" - retry_after_header = response.headers.get("Retry-After") - if retry_after_header: - try: - return max( - 1, - min( - int(retry_after_header), settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER - ), - ) - except ValueError: - logger.debug(f"Invalid MinerU Retry-After header: {retry_after_header}") - - return max(1, min(default_retry_after, settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER)) - - -def _raise_mineru_unavailable( - token_id: str, response: requests.Response, operation: str -) -> None: - retry_after = _get_retry_after_seconds( - response, settings.MINERU_TOKEN_COOLDOWN_SECONDS - ) - quota_manager = get_mineru_quota_manager() - quota_manager.mark_rate_limited(token_id, retry_after) - _mineru_logger( - "rate_limited", - operation=operation, - token_id=token_id, - status_code=response.status_code, - retry_after=retry_after, - ).warning("MinerU request rate-limited") - raise UnavailableException( - internal_message=f"MinerU rate limited during {operation}", - retry_after=retry_after, - limit=settings.MINERU_TOKEN_RPM_LIMIT, - period="minute", - user_message="Document processing is busy right now. Please retry shortly.", - ) - - -def _polling_interval_for_state(state: str, attempt: int) -> float: - """Return seconds to sleep before the next poll. - - Tuned so that ``WORKER_CONCURRENCY`` scales safely with MinerU limits. - With a 4-token pool (300 RPM each), total budget is 1200 req/min: - - 150 tasks × (60 s / 15 s) = 600 req/min total - 600 / 4 tokens = 150 req/min per token ← leaves headroom - - Observed data (Logfire, 2026-03-08 dev batch): - - 99 % of tasks never enter ``running``; lifecycle is - ``waiting-file`` → ``done`` in 2-4 s on MinerU's side. - - Longest observed task: 21 s (5 poll attempts). - - Peak burst: 351 concurrent tasks, 308 req/min → rate-limited. - """ - if state == "pending": - return min(20.0, 5.0 + attempt * 1.5) - if state == "running": - return 10.0 - # waiting-file, converting, unknown, etc. - return 15.0 - - -def _get_batch_status(data: dict[str, Any]) -> Optional[dict[str, Any]]: - extract_result = data.get("data", {}).get("extract_result") - if isinstance(extract_result, list): - return extract_result[0] if extract_result else None - return extract_result - - -def poll_mineru_task( - status_url: str, - task_id: str, - output_dir: str, - get_status: Callable[[dict[str, Any]], Optional[dict[str, Any]]], - preferred_token_id: Optional[str] = None, -) -> None: - quota_manager = get_mineru_quota_manager() - polling_logger = _mineru_logger( - "poll_status", - operation="poll_status", - task_id=task_id, - preferred_token_id=preferred_token_id, - ) - - max_polling_attempts = 120 - polling_interval = 5.0 - max_wait_time = 6000 - - start_time = time.time() - attempt = 0 - last_token_id: Optional[str] = None - last_state: Optional[str] = None - - polling_logger.info("Starting MinerU polling") - - while attempt < max_polling_attempts: - if time.time() - start_time > max_wait_time: - polling_logger.bind( - attempt=attempt + 1, - max_polling_attempts=max_polling_attempts, - max_wait_time=max_wait_time, - ).warning("MinerU polling timed out") - raise TimeoutException( - internal_message=f"PDF parsing timed out, exceeded {max_wait_time} seconds", - retry_after=60, - user_message="PDF parsing timed out. Please try again.", - ) - - try: - logger.debug( - f"parse_pdfs status_url: {status_url} " - f"(attempt {attempt + 1}/{max_polling_attempts})" - ) - lease = quota_manager.acquire_request( - operation="poll_status", - preferred_token_id=preferred_token_id, - ) - if lease.token_id != last_token_id: - polling_logger.bind( - token_id=lease.token_id, - attempt=attempt + 1, - ).info("Acquired MinerU token for polling") - last_token_id = lease.token_id - - response = get_mineru_session().get( - status_url, - headers=get_mineru_headers(lease.api_key), - timeout=settings.MINERU_API_TIMEOUT, - ) - - if response.status_code == 429: - # urllib3 already retried with backoff — if we still see 429, - # mark the token and let Celery handle task-level retry. - _raise_mineru_unavailable( - lease.token_id, response, operation="poll_status" - ) - - if response.status_code == 200: - response_json = response.json() - if response_json.get("code") != 0: - response_message = str(response_json.get("msg") or "Unknown error") - if "rate limit" in response_message.lower(): - quota_manager.mark_rate_limited( - lease.token_id, - settings.MINERU_TOKEN_COOLDOWN_SECONDS, - ) - raise UnavailableException( - internal_message=( - f"MinerU rate limited during poll_status: {response_message}" - ), - retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, - limit=lease.rpm_limit, - period="minute", - user_message="Document processing is busy right now. Please retry shortly.", - ) - raise MinerUServiceException( - internal_message=f"MinerU API Error: {response_message}" - ) - - status = get_status(response_json) - if not status: - polling_logger.bind( - token_id=lease.token_id, - attempt=attempt + 1, - ).warning("Received empty MinerU status payload") - time.sleep(polling_interval) - attempt += 1 - continue - - state = status.get("state", "unknown") - if state != last_state: - # polling_logger.bind( - # token_id=lease.token_id, - # attempt=attempt + 1, - # state=state, - # ).info("MinerU status changed") - last_state = state - - if state == "done": - download_and_extract_zip( - status["full_zip_url"], - dest_dir=output_dir, - keep_exts=(".md", ".jpg", ".jpeg", ".png", ".gif", ".json"), - exclude_patterns=("content_list", "middle.json", "model.json"), - ) - polling_logger.bind(token_id=lease.token_id).info( - "MinerU parsing completed" - ) - break - - if state == "running": - if "extract_progress" in status: - try: - ( - status["extract_progress"]["extracted_pages"] - / status["extract_progress"]["total_pages"] - ) - # polling_logger.bind( - # token_id=lease.token_id, - # progress=progress, - # ).info("MinerU parsing progress updated") - except (KeyError, ZeroDivisionError): - polling_logger.bind(token_id=lease.token_id).info( - "MinerU parsing in progress" - ) - else: - polling_logger.bind(token_id=lease.token_id).info( - "MinerU parsing in progress" - ) - elif state == "failed": - error_message = status.get("err_msg", "Unknown error") - polling_logger.bind( - token_id=lease.token_id, - error_message=error_message, - ).error("MinerU parsing reported failed state") - raise PDFParsingException( - user_message="Failed to parse the PDF file", - internal_message=f"MinerU failed with state 'failed': {error_message}", - ) - elif state == "pending": - polling_logger.bind(token_id=lease.token_id).debug( - "MinerU parsing pending" - ) - elif state == "waiting-file": - polling_logger.bind(token_id=lease.token_id).debug( - "MinerU waiting for file queueing" - ) - elif state == "converting": - polling_logger.bind(token_id=lease.token_id).debug( - "MinerU converting file" - ) - else: - polling_logger.bind( - token_id=lease.token_id, - state=state, - ).warning("MinerU returned unknown state") - - time.sleep(_polling_interval_for_state(state, attempt)) - attempt += 1 - else: - polling_logger.bind( - token_id=lease.token_id, - attempt=attempt + 1, - status_code=response.status_code, - ).warning("MinerU status query failed") - time.sleep(polling_interval * 2) - attempt += 1 - - except requests.RequestException as exc: - polling_logger.bind( - attempt=attempt + 1, - error_message=str(exc), - ).warning("MinerU polling network request failed") - time.sleep(polling_interval * 2) - attempt += 1 - except KnowhereException: - raise - except Exception as exc: - polling_logger.bind( - attempt=attempt + 1, - error_message=str(exc), - ).error("Unexpected error during MinerU polling") - raise PDFParsingException( - user_message="An unexpected error occurred while parsing the PDF", - internal_message=str(exc), - original_exception=exc, - ) - - if attempt >= max_polling_attempts: - raise TimeoutException( - internal_message=( - f"minerU PDF parsing timed out after {max_polling_attempts} attempts, " - f"Task ID: {task_id}" - ), - retry_after=60, - user_message="PDF parsing timed out. Please try again.", - ) - def _request_upload_target(pdf_url: str, filename: str) -> tuple[str, str, str]: base_url = settings.MINERU_URL quota_manager = get_mineru_quota_manager() - upload_logger = _mineru_logger( + upload_logger = mineru_logger( "upload_url", operation="upload_url", filename=filename, @@ -501,7 +177,7 @@ def _request_upload_target(pdf_url: str, filename: str) -> tuple[str, str, str]: timeout=settings.MINERU_API_TIMEOUT, ) if response.status_code == 429: - _raise_mineru_unavailable(lease.token_id, response, operation="upload_url") + raise_mineru_unavailable(lease.token_id, response, operation="upload_url") if response.status_code != 200: upload_logger.bind( token_id=lease.token_id, @@ -551,7 +227,7 @@ def _request_upload_target(pdf_url: str, filename: str) -> tuple[str, str, str]: def _upload_file_to_mineru( pdf_url: str, filename: str, upload_url: str, token_id: str ) -> None: - upload_logger = _mineru_logger( + upload_logger = mineru_logger( "file_upload", operation="file_upload", filename=filename, @@ -645,7 +321,7 @@ def _submit_url_task(presigned_url: str, filename: str) -> tuple[str, str]: """ base_url = settings.MINERU_URL quota_manager = get_mineru_quota_manager() - submit_logger = _mineru_logger( + submit_logger = mineru_logger( "submit_url_task", operation="submit_url_task", filename=filename, @@ -675,7 +351,7 @@ def _submit_url_task(presigned_url: str, filename: str) -> tuple[str, str]: ) if response.status_code == 429: - _raise_mineru_unavailable(lease.token_id, response, operation="submit_url_task") + raise_mineru_unavailable(lease.token_id, response, operation="submit_url_task") if response.status_code != 200: submit_logger.bind( @@ -730,7 +406,7 @@ def parse_via_full( resolved_s3_key, expires_in=settings.MINERU_URL_MODE_PRESIGN_EXPIRY ) presigned_url = presigned["download_url"] - _mineru_logger("ingestion_mode", mode="s3_url").info( + mineru_logger("ingestion_mode", mode="s3_url").info( "Using S3 URL mode for MinerU ingestion" ) batch_id, token_id = _submit_url_task(presigned_url, filename) @@ -744,7 +420,7 @@ def parse_via_full( resolved_s3_key = None if resolved_s3_key is None: - _mineru_logger("ingestion_mode", mode="direct_upload").info( + mineru_logger("ingestion_mode", mode="direct_upload").info( "Using direct upload mode for MinerU ingestion" ) batch_id, upload_url, token_id = _request_upload_target(pdf_url, filename) @@ -754,6 +430,6 @@ def parse_via_full( status_url=f"{settings.MINERU_URL}/extract-results/batch/{batch_id}", task_id=batch_id, output_dir=output_dir, - get_status=_get_batch_status, + get_status=get_batch_status, preferred_token_id=token_id, ) diff --git a/apps/worker/app/services/document_parser/mineru_task_polling.py b/apps/worker/app/services/document_parser/mineru_task_polling.py new file mode 100644 index 000000000..c09c1ea86 --- /dev/null +++ b/apps/worker/app/services/document_parser/mineru_task_polling.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any, Optional + +import requests +from app.services.document_parser.mineru_client import ( + get_mineru_headers, + get_mineru_session, + mineru_logger, + raise_mineru_unavailable, +) +from app.services.document_parser.mineru_quota_manager import get_mineru_quota_manager +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + MinerUServiceException, + PDFParsingException, + TimeoutException, + UnavailableException, +) +from shared.core.exceptions.knowhere_exception import KnowhereException +from shared.utils.zip_download import download_and_extract_zip + + +def get_batch_status(data: dict[str, Any]) -> Optional[dict[str, Any]]: + extract_result = data.get("data", {}).get("extract_result") + if isinstance(extract_result, list): + return extract_result[0] if extract_result else None + return extract_result + + +def get_polling_interval_for_state(state: str, attempt: int) -> float: + """Return seconds to sleep before the next MinerU status poll.""" + if state == "pending": + return min(20.0, 5.0 + attempt * 1.5) + if state == "running": + return 10.0 + return 15.0 + + +def poll_mineru_task( + status_url: str, + task_id: str, + output_dir: str, + get_status: Callable[[dict[str, Any]], Optional[dict[str, Any]]], + preferred_token_id: Optional[str] = None, +) -> None: + quota_manager = get_mineru_quota_manager() + polling_logger = mineru_logger( + "poll_status", + operation="poll_status", + task_id=task_id, + preferred_token_id=preferred_token_id, + ) + + max_polling_attempts = 120 + polling_interval = 5.0 + max_wait_time = 6000 + + start_time = time.time() + attempt = 0 + last_token_id: Optional[str] = None + last_state: Optional[str] = None + + polling_logger.info("Starting MinerU polling") + + while attempt < max_polling_attempts: + if time.time() - start_time > max_wait_time: + polling_logger.bind( + attempt=attempt + 1, + max_polling_attempts=max_polling_attempts, + max_wait_time=max_wait_time, + ).warning("MinerU polling timed out") + raise TimeoutException( + internal_message=f"PDF parsing timed out, exceeded {max_wait_time} seconds", + retry_after=60, + user_message="PDF parsing timed out. Please try again.", + ) + + try: + logger.debug( + f"parse_pdfs status_url: {status_url} " + f"(attempt {attempt + 1}/{max_polling_attempts})" + ) + lease = quota_manager.acquire_request( + operation="poll_status", + preferred_token_id=preferred_token_id, + ) + if lease.token_id != last_token_id: + polling_logger.bind( + token_id=lease.token_id, + attempt=attempt + 1, + ).info("Acquired MinerU token for polling") + last_token_id = lease.token_id + + response = get_mineru_session().get( + status_url, + headers=get_mineru_headers(lease.api_key), + timeout=settings.MINERU_API_TIMEOUT, + ) + + if response.status_code == 429: + raise_mineru_unavailable( + lease.token_id, response, operation="poll_status" + ) + + if response.status_code == 200: + response_json = response.json() + if response_json.get("code") != 0: + response_message = str(response_json.get("msg") or "Unknown error") + if "rate limit" in response_message.lower(): + quota_manager.mark_rate_limited( + lease.token_id, + settings.MINERU_TOKEN_COOLDOWN_SECONDS, + ) + raise UnavailableException( + internal_message=( + f"MinerU rate limited during poll_status: {response_message}" + ), + retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS, + limit=lease.rpm_limit, + period="minute", + user_message="Document processing is busy right now. Please retry shortly.", + ) + raise MinerUServiceException( + internal_message=f"MinerU API Error: {response_message}" + ) + + status = get_status(response_json) + if not status: + polling_logger.bind( + token_id=lease.token_id, + attempt=attempt + 1, + ).warning("Received empty MinerU status payload") + time.sleep(polling_interval) + attempt += 1 + continue + + state = status.get("state", "unknown") + if state != last_state: + last_state = state + + if state == "done": + download_and_extract_zip( + status["full_zip_url"], + dest_dir=output_dir, + keep_exts=(".md", ".jpg", ".jpeg", ".png", ".gif", ".json"), + exclude_patterns=("content_list", "middle.json", "model.json"), + ) + polling_logger.bind(token_id=lease.token_id).info( + "MinerU parsing completed" + ) + break + + if state == "running": + if "extract_progress" in status: + try: + _progress = ( + status["extract_progress"]["extracted_pages"] + / status["extract_progress"]["total_pages"] + ) + except (KeyError, ZeroDivisionError): + polling_logger.bind(token_id=lease.token_id).info( + "MinerU parsing in progress" + ) + else: + polling_logger.bind(token_id=lease.token_id).info( + "MinerU parsing in progress" + ) + elif state == "failed": + error_message = status.get("err_msg", "Unknown error") + polling_logger.bind( + token_id=lease.token_id, + error_message=error_message, + ).error("MinerU parsing reported failed state") + raise PDFParsingException( + user_message="Failed to parse the PDF file", + internal_message=f"MinerU failed with state 'failed': {error_message}", + ) + elif state == "pending": + polling_logger.bind(token_id=lease.token_id).debug( + "MinerU parsing pending" + ) + elif state == "waiting-file": + polling_logger.bind(token_id=lease.token_id).debug( + "MinerU waiting for file queueing" + ) + elif state == "converting": + polling_logger.bind(token_id=lease.token_id).debug( + "MinerU converting file" + ) + else: + polling_logger.bind( + token_id=lease.token_id, + state=state, + ).warning("MinerU returned unknown state") + + time.sleep(get_polling_interval_for_state(state, attempt)) + attempt += 1 + else: + polling_logger.bind( + token_id=lease.token_id, + attempt=attempt + 1, + status_code=response.status_code, + ).warning("MinerU status query failed") + time.sleep(polling_interval * 2) + attempt += 1 + + except requests.RequestException as exc: + polling_logger.bind( + attempt=attempt + 1, + error_message=str(exc), + ).warning("MinerU polling network request failed") + time.sleep(polling_interval * 2) + attempt += 1 + except KnowhereException: + raise + except Exception as exc: + polling_logger.bind( + attempt=attempt + 1, + error_message=str(exc), + ).error("Unexpected error during MinerU polling") + raise PDFParsingException( + user_message="An unexpected error occurred while parsing the PDF", + internal_message=str(exc), + original_exception=exc, + ) + + if attempt >= max_polling_attempts: + raise TimeoutException( + internal_message=( + f"minerU PDF parsing timed out after {max_polling_attempts} attempts, " + f"Task ID: {task_id}" + ), + retry_after=60, + user_message="PDF parsing timed out. Please try again.", + ) diff --git a/apps/worker/app/services/document_parser/orchestration/format_adapters.py b/apps/worker/app/services/document_parser/orchestration/format_adapters.py new file mode 100644 index 000000000..882efbaab --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -0,0 +1,221 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +import pandas as pd + +from app.services.document_parser.orchestration.parse_session import ParseSession + + +class DocumentParseAdapter(Protocol): + @property + def document_format(self) -> object: + """Document format handled by this adapter.""" + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + """Parse a document session into the parser output directory and DataFrame.""" + ... + + +@dataclass(frozen=True) +class FragmentParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.fragment_parser import parse_fragment + + full_output_dir, _relative_root, parsed_df = parse_fragment( + session.fragment_content, + filename=session.filename, + output_dir=session.output_dir, + kb_dir=session.kb_dir, + base_llm_paras=session.base_llm_paras, + ) + return full_output_dir, parsed_df + + +@dataclass(frozen=True) +class TextParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.md_parser import parse_md + from app.services.document_parser.txt_parser import parse_texts + + text_lines = parse_texts(file_path=session.file_full_path, baseurl=session.base_url) + parsed_df = parse_md( + session.full_output_dir, + source_type="md", + md_lines=text_lines, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + +@dataclass(frozen=True) +class ImageParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.image_parser import parse_image + + parsed_df = parse_image( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + baseurl=session.base_url, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + +@dataclass(frozen=True) +class PdfParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.pdf_parser import parse_pdfs + + parsed_df = parse_pdfs( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + base_llm_paras=session.base_llm_paras, + profile=session.profile, + relative_root=session.relative_root, + s3_key=session.s3_key, + ) + return session.full_output_dir, parsed_df + + +@dataclass(frozen=True) +class DocParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.legacy_converter import doc_to_docx + + converted_docx_path, _ = doc_to_docx( + session.file_full_path, + outdir=session.full_output_dir, + ) + return _parse_docx_path(converted_docx_path, session) + + +@dataclass(frozen=True) +class DocxParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + return _parse_docx_path(session.file_full_path, session) + + +@dataclass(frozen=True) +class XlsParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.legacy_converter import xls_to_xlsx + + converted_xlsx_path, _ = xls_to_xlsx( + session.file_full_path, + outdir=session.full_output_dir, + ) + return _parse_xlsx_path(converted_xlsx_path, session) + + +@dataclass(frozen=True) +class XlsxParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + return _parse_xlsx_path(session.file_full_path, session) + + +@dataclass(frozen=True) +class PptxParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.pptx_parser import parse_pptx + + parsed_df = parse_pptx( + session.file_full_path, + filename=session.filename, + output_dir=session.full_output_dir, + base_llm_paras=session.base_llm_paras, + strategy="to_pdf_api", + job_id=session.job_id, + relative_root=session.relative_root, + baseurl=session.base_url, + ) + return session.full_output_dir, parsed_df + + +@dataclass(frozen=True) +class MarkdownParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.md_parser import parse_md + + parsed_df = parse_md( + session.full_output_dir, + source_type="md", + file_path=session.file_full_path, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + +@dataclass(frozen=True) +class JsonParseAdapter: + document_format: object + + def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + return session.full_output_dir, None + + +def _parse_docx_path( + docx_path: str, + session: ParseSession, +) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx + + parsed_structure, dataframe_list = parse_docx( + docx_path, + session.base_llm_paras, + session.full_output_dir, + session.filename, + session.base_url, + relative_root=session.relative_root, + ) + parsed_df = convert_doc2dics( + parsed_structure, + dataframe_list, + session.full_output_dir, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df + + +def _parse_xlsx_path( + xlsx_path: str, + session: ParseSession, +) -> tuple[str, pd.DataFrame | None]: + from app.services.document_parser.table_parser import parse_xlsx + + parsed_df = parse_xlsx( + xlsx_path, + session.filename, + session.full_output_dir, + session.base_url, + base_llm_paras=session.base_llm_paras, + relative_root=session.relative_root, + ) + return session.full_output_dir, parsed_df diff --git a/apps/worker/app/services/document_parser/orchestration/format_router.py b/apps/worker/app/services/document_parser/orchestration/format_router.py new file mode 100644 index 000000000..b94cf2787 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/format_router.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import os +from enum import Enum + +from app.services.document_parser.orchestration import format_adapters +from app.services.document_parser.orchestration.format_adapters import ( + DocumentParseAdapter, +) +from shared.core.exceptions.domain_exceptions import ValidationException + + +class DocumentFormat(str, Enum): + TEXT = "text" + FRAGMENT = "fragment" + IMAGE = "image" + PDF = "pdf" + DOC = "doc" + DOCX = "docx" + XLS = "xls" + XLSX = "xlsx" + PPTX = "pptx" + MARKDOWN = "markdown" + JSON = "json" + + +SUPPORTED_FILE_TYPES: tuple[str, ...] = ( + ".txt", + ".fragment", + ".png", + ".jpg", + ".jpeg", + ".pdf", + ".doc", + ".docx", + ".xls", + ".xlsx", + ".pptx", + ".md", + ".json", +) + + +def resolve_document_format(file_path: str) -> DocumentFormat: + extension = os.path.splitext(file_path)[1].lower() + if extension == ".fragment": + return DocumentFormat.FRAGMENT + if extension == ".txt": + return DocumentFormat.TEXT + if extension in (".png", ".jpg", ".jpeg"): + return DocumentFormat.IMAGE + if extension == ".pdf": + return DocumentFormat.PDF + if extension == ".doc": + return DocumentFormat.DOC + if extension == ".docx": + return DocumentFormat.DOCX + if extension == ".xls": + return DocumentFormat.XLS + if extension == ".xlsx": + return DocumentFormat.XLSX + if extension == ".pptx": + return DocumentFormat.PPTX + if extension == ".md": + return DocumentFormat.MARKDOWN + if extension == ".json": + return DocumentFormat.JSON + + raise ValidationException( + user_message=f"Unsupported file type: {extension}", + violations=[ + { + "field": "file_type", + "description": f"Must be one of: {', '.join(SUPPORTED_FILE_TYPES)}", + } + ], + ) + + +def get_document_parse_adapter(document_format: DocumentFormat) -> DocumentParseAdapter: + adapter_by_format: dict[DocumentFormat, DocumentParseAdapter] = { + DocumentFormat.FRAGMENT: format_adapters.FragmentParseAdapter(document_format), + DocumentFormat.TEXT: format_adapters.TextParseAdapter(document_format), + DocumentFormat.IMAGE: format_adapters.ImageParseAdapter(document_format), + DocumentFormat.PDF: format_adapters.PdfParseAdapter(document_format), + DocumentFormat.DOC: format_adapters.DocParseAdapter(document_format), + DocumentFormat.DOCX: format_adapters.DocxParseAdapter(document_format), + DocumentFormat.XLS: format_adapters.XlsParseAdapter(document_format), + DocumentFormat.XLSX: format_adapters.XlsxParseAdapter(document_format), + DocumentFormat.PPTX: format_adapters.PptxParseAdapter(document_format), + DocumentFormat.MARKDOWN: format_adapters.MarkdownParseAdapter(document_format), + DocumentFormat.JSON: format_adapters.JsonParseAdapter(document_format), + } + return adapter_by_format[document_format] diff --git a/apps/worker/app/services/document_parser/orchestration/parse_input.py b/apps/worker/app/services/document_parser/orchestration/parse_input.py new file mode 100644 index 000000000..89249b3d0 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/parse_input.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class ParseOptions: + llm_histories: int = 5 + smart_title_parse: bool = True + summary_image: bool = True + summary_table: bool = True + summary_txt: bool = True + stopwords: list[str] | None = None + doc_type: str = "auto" + add_frag_desc: str = "" + + +@dataclass(frozen=True) +class ParseInput: + file_full_path: str + filename: str + output_dir: str + internal_output_filename: str + job_id: str | None = None + kb_dir: str = "Default_Root" + options: ParseOptions = field(default_factory=ParseOptions) + base_url: str = "" + fragment_content: str = "" + s3_key: str | None = None diff --git a/apps/worker/app/services/document_parser/orchestration/parse_session.py b/apps/worker/app/services/document_parser/orchestration/parse_session.py index 44a58e343..832e13969 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -6,6 +6,7 @@ from app.services.document_parser.atlas_classifier import classify_atlas_with_vlm from app.services.document_parser.doc_profiler import profile_document +from app.services.document_parser.orchestration.parse_input import ParseInput from app.services.document_parser.stage_profiler import stage_timer from loguru import logger @@ -32,69 +33,84 @@ class ParseSession: relative_root: str s3_key: str | None + @classmethod + def from_input( + cls, + *, + parse_input: ParseInput, + base_llm_paras: dict[str, object], + full_output_dir: str, + profile: Any, + relative_root: str, + ) -> "ParseSession": + return cls( + base_llm_paras=base_llm_paras, + base_url=parse_input.base_url, + file_full_path=parse_input.file_full_path, + filename=parse_input.filename, + fragment_content=parse_input.fragment_content, + full_output_dir=full_output_dir, + internal_output_filename=parse_input.internal_output_filename, + job_id=parse_input.job_id, + kb_dir=parse_input.kb_dir, + output_dir=parse_input.output_dir, + profile=profile, + relative_root=relative_root, + s3_key=parse_input.s3_key, + ) -def build_parse_session( - *, - add_frag_desc: str, - base_url: str, - doc_type: str, - file_full_path: str, - filename: str, - fragment_content: str, - internal_output_filename: str, - job_id: str | None, - kb_dir: str, - llm_histories: int, - output_dir: str, - s3_key: str | None, - smart_title_parse: bool, - stopwords: list[str] | None, - summary_image: bool, - summary_table: bool, - summary_txt: bool, -) -> ParseSession: + +def build_parse_session(parse_input: ParseInput) -> ParseSession: """Build the parser routing session from explicit parse inputs.""" + parse_options = parse_input.options base_llm_paras = { - "llm_histories": llm_histories, - "smart_title_parse": smart_title_parse, - "summary_image": summary_image, - "summary_table": summary_table, - "summary_txt": summary_txt, - "stopwords": stopwords, - "doc_type": doc_type, - "frag_desc": add_frag_desc, + "llm_histories": parse_options.llm_histories, + "smart_title_parse": parse_options.smart_title_parse, + "summary_image": parse_options.summary_image, + "summary_table": parse_options.summary_table, + "summary_txt": parse_options.summary_txt, + "stopwords": parse_options.stopwords, + "doc_type": parse_options.doc_type, + "frag_desc": parse_options.add_frag_desc, "model_name": settings.NORMOL_MODEL, "hierarchy_model_name": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, } - logger.debug(f"baseurl: {base_url}") - logger.debug(f"file_full_path: {file_full_path}") + logger.debug(f"baseurl: {parse_input.base_url}") + logger.debug(f"file_full_path: {parse_input.file_full_path}") relative_root, full_output_dir = _resolve_output_paths( - filename=filename, - internal_output_filename=internal_output_filename, - kb_dir=kb_dir, - output_dir=output_dir, + filename=parse_input.filename, + internal_output_filename=parse_input.internal_output_filename, + kb_dir=parse_input.kb_dir, + output_dir=parse_input.output_dir, ) logger.debug(f"relative_root: {relative_root}") logger.debug(f"full_output_dir: {full_output_dir}") - with stage_timer("document.profile", filename=filename): - profile = profile_document(file_full_path, internal_output_filename) + with stage_timer("document.profile", filename=parse_input.filename): + profile = profile_document( + parse_input.file_full_path, + parse_input.internal_output_filename, + ) logger.info(f"📋 DocProfile: {profile.summary()}") logger.debug(f"📋 Reasoning: {profile.reasoning}") if profile.atlas_candidate and profile.doc_category not in ("atlas", "ppt_converted"): - logger.info(f"🔍 Atlas candidate detected, running VLM visual check for {filename}") - with stage_timer("document.atlas_vlm_check", filename=filename): - vlm_is_atlas = classify_atlas_with_vlm(file_full_path) + logger.info( + f"🔍 Atlas candidate detected, running VLM visual check for {parse_input.filename}" + ) + with stage_timer("document.atlas_vlm_check", filename=parse_input.filename): + vlm_is_atlas = classify_atlas_with_vlm(parse_input.file_full_path) if vlm_is_atlas: profile.doc_category = "atlas" profile.reasoning += " | vlm_confirmed_atlas=True" - logger.info(f"✅ VLM confirmed atlas for {filename}") + logger.info(f"✅ VLM confirmed atlas for {parse_input.filename}") else: profile.reasoning += " | vlm_confirmed_atlas=False" - logger.info(f"ℹ️ VLM rejected atlas for {filename}, routing as generic") + logger.info( + f"ℹ️ VLM rejected atlas for {parse_input.filename}, routing as generic" + ) if profile.file_type == "pdf" and profile.page_count > PDF_PAGE_LIMIT: raise ValidationException( @@ -111,28 +127,34 @@ def build_parse_session( ) if profile.doc_category == "atlas": - filename, internal_output_filename, relative_root, full_output_dir = _rename_atlas_output( + filename, internal_output_filename, relative_root, full_output_dir = ( + _rename_atlas_output( + filename=parse_input.filename, + internal_output_filename=parse_input.internal_output_filename, + kb_dir=parse_input.kb_dir, + output_dir=parse_input.output_dir, + ) + ) + logger.info(f"📐 Atlas output renamed: {filename}") + parse_input = ParseInput( + file_full_path=parse_input.file_full_path, filename=filename, + output_dir=parse_input.output_dir, internal_output_filename=internal_output_filename, - kb_dir=kb_dir, - output_dir=output_dir, + job_id=parse_input.job_id, + kb_dir=parse_input.kb_dir, + options=parse_input.options, + base_url=parse_input.base_url, + fragment_content=parse_input.fragment_content, + s3_key=parse_input.s3_key, ) - logger.info(f"📐 Atlas output renamed: {filename}") - return ParseSession( + return ParseSession.from_input( + parse_input=parse_input, base_llm_paras=base_llm_paras, - base_url=base_url, - file_full_path=file_full_path, - filename=filename, - fragment_content=fragment_content, full_output_dir=full_output_dir, - internal_output_filename=internal_output_filename, - job_id=job_id, - kb_dir=kb_dir, - output_dir=output_dir, profile=profile, relative_root=relative_root, - s3_key=s3_key, ) diff --git a/apps/worker/app/services/document_parser/orchestration/route_parse.py b/apps/worker/app/services/document_parser/orchestration/route_parse.py index faeda3dee..85402339e 100644 --- a/apps/worker/app/services/document_parser/orchestration/route_parse.py +++ b/apps/worker/app/services/document_parser/orchestration/route_parse.py @@ -1,199 +1,14 @@ -from __future__ import annotations - -import os - import pandas as pd -from app.services.document_parser.orchestration.parse_session import ParseSession -from shared.core.exceptions.domain_exceptions import ValidationException - -SUPPORTED_FILE_TYPES: tuple[str, ...] = ( - ".txt", - ".fragment", - ".png", - ".jpg", - ".jpeg", - ".pdf", - ".doc", - ".docx", - ".xls", - ".xlsx", - ".pptx", - ".md", - ".json", +from app.services.document_parser.orchestration.format_router import ( + get_document_parse_adapter, + resolve_document_format, ) +from app.services.document_parser.orchestration.parse_session import ParseSession def route_document_parse(session: ParseSession) -> tuple[str, pd.DataFrame | None]: """Route a parser session to the correct adapter and return its output.""" - file_path_lower = session.file_full_path.lower() - - if ".fragment" in file_path_lower: - from app.services.document_parser.fragment_parser import parse_fragment - - full_output_dir, _relative_root, parsed_df = parse_fragment( - session.fragment_content, - filename=session.filename, - output_dir=session.output_dir, - kb_dir=session.kb_dir, - base_llm_paras=session.base_llm_paras, - ) - return full_output_dir, parsed_df - - if ".txt" in file_path_lower: - from app.services.document_parser.md_parser import parse_md - from app.services.document_parser.txt_parser import parse_texts - - text_lines = parse_texts(file_path=session.file_full_path, baseurl=session.base_url) - parsed_df = parse_md( - session.full_output_dir, - source_type="md", - md_lines=text_lines, - base_llm_paras=session.base_llm_paras, - relative_root=session.relative_root, - ) - return session.full_output_dir, parsed_df - - if any(extension in file_path_lower for extension in (".png", ".jpg", ".jpeg")): - from app.services.document_parser.image_parser import parse_image - - parsed_df = parse_image( - session.file_full_path, - filename=session.filename, - output_dir=session.full_output_dir, - baseurl=session.base_url, - base_llm_paras=session.base_llm_paras, - relative_root=session.relative_root, - ) - return session.full_output_dir, parsed_df - - if ".pdf" in file_path_lower: - from app.services.document_parser.pdf_parser import parse_pdfs - - parsed_df = parse_pdfs( - session.file_full_path, - filename=session.filename, - output_dir=session.full_output_dir, - base_llm_paras=session.base_llm_paras, - profile=session.profile, - relative_root=session.relative_root, - s3_key=session.s3_key, - ) - return session.full_output_dir, parsed_df - - if ".doc" in file_path_lower and ".docx" not in file_path_lower: - from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx - from app.services.document_parser.legacy_converter import doc_to_docx - - converted_docx_path, _ = doc_to_docx( - session.file_full_path, - outdir=session.full_output_dir, - ) - parsed_structure, dataframe_list = parse_docx( - converted_docx_path, - session.base_llm_paras, - session.full_output_dir, - session.filename, - session.base_url, - relative_root=session.relative_root, - ) - parsed_df = convert_doc2dics( - parsed_structure, - dataframe_list, - session.full_output_dir, - base_llm_paras=session.base_llm_paras, - relative_root=session.relative_root, - ) - return session.full_output_dir, parsed_df - - if ".docx" in file_path_lower: - from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx - - parsed_structure, dataframe_list = parse_docx( - session.file_full_path, - session.base_llm_paras, - session.full_output_dir, - session.filename, - session.base_url, - relative_root=session.relative_root, - ) - parsed_df = convert_doc2dics( - parsed_structure, - dataframe_list, - session.full_output_dir, - base_llm_paras=session.base_llm_paras, - relative_root=session.relative_root, - ) - return session.full_output_dir, parsed_df - - if ".xls" in file_path_lower and ".xlsx" not in file_path_lower: - from app.services.document_parser.legacy_converter import xls_to_xlsx - from app.services.document_parser.table_parser import parse_xlsx - - converted_xlsx_path, _ = xls_to_xlsx( - session.file_full_path, - outdir=session.full_output_dir, - ) - parsed_df = parse_xlsx( - converted_xlsx_path, - session.filename, - session.full_output_dir, - session.base_url, - base_llm_paras=session.base_llm_paras, - relative_root=session.relative_root, - ) - return session.full_output_dir, parsed_df - - if ".xlsx" in file_path_lower: - from app.services.document_parser.table_parser import parse_xlsx - - parsed_df = parse_xlsx( - session.file_full_path, - session.filename, - session.full_output_dir, - session.base_url, - base_llm_paras=session.base_llm_paras, - relative_root=session.relative_root, - ) - return session.full_output_dir, parsed_df - - if ".pptx" in file_path_lower: - from app.services.document_parser.pptx_parser import parse_pptx - - parsed_df = parse_pptx( - session.file_full_path, - filename=session.filename, - output_dir=session.full_output_dir, - base_llm_paras=session.base_llm_paras, - strategy="to_pdf_api", - job_id=session.job_id, - relative_root=session.relative_root, - baseurl=session.base_url, - ) - return session.full_output_dir, parsed_df - - if ".md" in file_path_lower: - from app.services.document_parser.md_parser import parse_md - - parsed_df = parse_md( - session.full_output_dir, - source_type="md", - file_path=session.file_full_path, - base_llm_paras=session.base_llm_paras, - relative_root=session.relative_root, - ) - return session.full_output_dir, parsed_df - - if ".json" in file_path_lower: - return session.full_output_dir, None - - file_ext = os.path.splitext(session.file_full_path)[1].lower() - raise ValidationException( - user_message=f"Unsupported file type: {file_ext}", - violations=[ - { - "field": "file_type", - "description": f"Must be one of: {', '.join(SUPPORTED_FILE_TYPES)}", - } - ], - ) + document_format = resolve_document_format(session.file_full_path) + adapter = get_document_parse_adapter(document_format) + return adapter.parse(session) diff --git a/apps/worker/app/services/document_parser/parse_service.py b/apps/worker/app/services/document_parser/parse_service.py index 482147099..000cef7bf 100644 --- a/apps/worker/app/services/document_parser/parse_service.py +++ b/apps/worker/app/services/document_parser/parse_service.py @@ -2,6 +2,7 @@ import pandas as pd +from app.services.document_parser.orchestration.parse_input import ParseInput, ParseOptions from app.services.document_parser.orchestration.parse_session import build_parse_session from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess from app.services.document_parser.orchestration.route_parse import route_document_parse @@ -27,25 +28,28 @@ def checkerboard_inject_parse( s3_key: str | None = None, ) -> tuple[str, pd.DataFrame | None]: """Run the stable parser seam using dedicated orchestration modules.""" - session = build_parse_session( - add_frag_desc=add_frag_desc, - base_url=base_url, - doc_type=doc_type, + parse_input = ParseInput( file_full_path=file_full_path, filename=filename, - fragment_content=fragment_content, internal_output_filename=internal_output_filename, job_id=job_id, kb_dir=kb_dir, - llm_histories=llm_histories, output_dir=output_dir, + options=ParseOptions( + add_frag_desc=add_frag_desc, + doc_type=doc_type, + llm_histories=llm_histories, + smart_title_parse=smart_title_parse, + stopwords=stopwords, + summary_image=summary_image, + summary_table=summary_table, + summary_txt=summary_txt, + ), + base_url=base_url, + fragment_content=fragment_content, s3_key=s3_key, - smart_title_parse=smart_title_parse, - stopwords=stopwords, - summary_image=summary_image, - summary_table=summary_table, - summary_txt=summary_txt, ) + session = build_parse_session(parse_input) full_output_dir, parsed_df = route_document_parse(session) parsed_df = apply_parse_postprocess(full_output_dir, parsed_df) return full_output_dir, parsed_df diff --git a/apps/worker/app/services/document_parser/parser_rows.py b/apps/worker/app/services/document_parser/parser_rows.py new file mode 100644 index 000000000..839ddd9f6 --- /dev/null +++ b/apps/worker/app/services/document_parser/parser_rows.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pandas as pd +from pandas import Index + +from shared.core.config import settings + +PARSER_ROW_COLUMNS: tuple[str, ...] = tuple(settings.ALL_DF_COLS.split(",")) + + +@dataclass(frozen=True) +class ParsedRow: + content: str + path: str + type: str + know_id: str + addtime: str + keywords: str = "" + summary: str = "" + tokens: str = "" + connectto: str = "" + page_nums: str = "" + length: int | None = None + + def to_list(self) -> list[object]: + content_length = self.length if self.length is not None else len(self.content) + return [ + self.content, + self.path, + self.type, + content_length, + self.keywords, + self.summary, + self.know_id, + self.tokens, + self.connectto, + self.addtime, + self.page_nums, + ] + + def to_dict(self) -> dict[str, object]: + return dict(zip(PARSER_ROW_COLUMNS, self.to_list())) + + +class ParsedRowsBuilder: + def __init__(self) -> None: + self._rows: list[ParsedRow] = [] + + def append(self, row: ParsedRow) -> None: + self._rows.append(row) + + def extend(self, rows: list[ParsedRow]) -> None: + self._rows.extend(rows) + + def to_dataframe(self) -> pd.DataFrame: + return pd.DataFrame( + [row.to_list() for row in self._rows], + columns=Index(PARSER_ROW_COLUMNS), + ) diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/pptx_parser.py index 8782be796..19063937a 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/pptx_parser.py @@ -11,13 +11,11 @@ _convert_with_libreoffice, ) from app.services.document_parser.md_parser import parse_md -from app.services.document_parser.mineru_pdf_service import ( - get_existing_mineru_source_s3_key, -) from app.services.document_parser.parser_log_utils import truncate_log_value -from app.services.document_parser.pdf_parser import parse_pdfs -from app.services.document_parser.pptx_pdf_rendering import ( - render_pdf_to_image_pdf as _render_pdf_to_image_pdf, +from app.services.document_parser.rendered_pdf_transform import ( + build_rendered_pdf_s3_key, + parse_cached_rendered_pdf, + parse_rendered_pdf_bytes, ) from loguru import logger from markitdown import MarkItDown @@ -28,7 +26,6 @@ FileSystemException, ) from shared.core.logging import LogEvent -from shared.services.storage.job_file_storage import JobFileStorage from shared.utils.file_loading import load_file_bytes from shared.utils.file_utils import path_handle @@ -306,50 +303,6 @@ class _ILoveApiConcurrencyExceeded(Exception): pass -def _build_rendered_pdf_s3_key(job_id: str | None) -> str | None: - """Store rendered parser artifacts under a stable transform/ prefix.""" - if settings.ENVIRONMENT == "development" or not job_id: - return None - return f"transform/{job_id}.rendered.pdf" - - -def _parse_cached_rendered_pdf( - rendered_pdf_s3_key: str | None, - filename: str, - output_dir: str, - base_llm_paras, - relative_root, -): - """Parse a previously rendered PPTX PDF from S3 without re-reading the source deck.""" - if rendered_pdf_s3_key is None: - return None - - cached_rendered_pdf_s3_key = get_existing_mineru_source_s3_key(rendered_pdf_s3_key) - if cached_rendered_pdf_s3_key is None: - return None - - logger.info( - f"[parse_pptx] Reusing rendered PDF for MinerU URL mode: {rendered_pdf_s3_key}" - ) - cached_rendered_pdf_path = JobFileStorage().download_upload_to_temp( - cached_rendered_pdf_s3_key, - suffix=".pdf", - temp_dir=output_dir, - ) - try: - return parse_pdfs( - cached_rendered_pdf_path, - filename, - output_dir, - base_llm_paras, - relative_root=relative_root, - s3_key=cached_rendered_pdf_s3_key, - ) - finally: - if os.path.exists(cached_rendered_pdf_path): - os.remove(cached_rendered_pdf_path) - - # ==================== main parsing entrance ==================== @@ -372,12 +325,12 @@ def parse_pptx( - "to_pdf_api": use iLoveAPI to convert to PDF, then parse via MinerU (recommended) """ rendered_pdf_s3_key = ( - _build_rendered_pdf_s3_key(job_id) + build_rendered_pdf_s3_key(job_id) if strategy in {"to_pdf_api", "to_pdf"} else None ) if strategy in {"to_pdf_api", "to_pdf"}: - cached_result = _parse_cached_rendered_pdf( + cached_result = parse_cached_rendered_pdf( rendered_pdf_s3_key=rendered_pdf_s3_key, filename=filename, output_dir=output_dir, @@ -491,27 +444,14 @@ def _parse_pptx_via_api( # Step 1: PPTX → PDF (in memory) pdf_bytes = _pptx_bytes_to_pdf_bytes(pptx_data, filename) - # Step 2: PDF → image-only PDF (in memory) - img_pdf_bytes = _render_pdf_to_image_pdf(pdf_bytes) - - # Step 3: Write to output_dir for MinerU upload, then clean up - tmp_path = os.path.join(output_dir, "_pptx_tmp.pdf") - with open(tmp_path, "wb") as f: - f.write(img_pdf_bytes) - - try: - parsed_df = parse_pdfs( - tmp_path, - filename=filename, - output_dir=output_dir, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - s3_key=rendered_pdf_s3_key, - ) - return parsed_df - finally: - if os.path.exists(tmp_path): - os.remove(tmp_path) + return parse_rendered_pdf_bytes( + pdf_bytes=pdf_bytes, + filename=filename, + output_dir=output_dir, + base_llm_paras=base_llm_paras, + relative_root=relative_root, + rendered_pdf_s3_key=rendered_pdf_s3_key, + ) def _parse_pptx_via_libreoffice( @@ -552,25 +492,14 @@ def _parse_pptx_via_libreoffice( finally: shutil.rmtree(tmp_dir, ignore_errors=True) - img_pdf_bytes = _render_pdf_to_image_pdf(pdf_bytes) - - tmp_path = os.path.join(output_dir, "_pptx_tmp.pdf") - with open(tmp_path, "wb") as f: - f.write(img_pdf_bytes) - - try: - parsed_df = parse_pdfs( - tmp_path, - filename=filename, - output_dir=output_dir, - base_llm_paras=base_llm_paras, - relative_root=relative_root, - s3_key=rendered_pdf_s3_key, - ) - return parsed_df - finally: - if os.path.exists(tmp_path): - os.remove(tmp_path) + return parse_rendered_pdf_bytes( + pdf_bytes=pdf_bytes, + filename=filename, + output_dir=output_dir, + base_llm_paras=base_llm_paras, + relative_root=relative_root, + rendered_pdf_s3_key=rendered_pdf_s3_key, + ) def _parse_pptx_to_md(pptx_data, filename, output_dir, base_llm_paras, relative_root): diff --git a/apps/worker/app/services/document_parser/rendered_pdf_transform.py b/apps/worker/app/services/document_parser/rendered_pdf_transform.py new file mode 100644 index 000000000..cab7ca822 --- /dev/null +++ b/apps/worker/app/services/document_parser/rendered_pdf_transform.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import os + +import pandas as pd + +from app.services.document_parser.mineru_pdf_service import ( + get_existing_mineru_source_s3_key, +) +from app.services.document_parser.pdf_parser import parse_pdfs +from app.services.document_parser.pptx_pdf_rendering import render_pdf_to_image_pdf +from loguru import logger + +from shared.core.config import settings +from shared.services.storage.job_file_storage import JobFileStorage + +RENDERED_PDF_TEMP_FILENAME = "_pptx_tmp.pdf" + + +def build_rendered_pdf_s3_key(job_id: str | None) -> str | None: + """Store rendered parser artifacts under a stable transform/ prefix.""" + if settings.ENVIRONMENT == "development" or not job_id: + return None + return f"transform/{job_id}.rendered.pdf" + + +def parse_cached_rendered_pdf( + *, + rendered_pdf_s3_key: str | None, + filename: str, + output_dir: str, + base_llm_paras: dict[str, object], + relative_root: str | None, +) -> pd.DataFrame | None: + """Parse a previously rendered PDF from S3 without re-reading the source deck.""" + if rendered_pdf_s3_key is None: + return None + + cached_rendered_pdf_s3_key = get_existing_mineru_source_s3_key(rendered_pdf_s3_key) + if cached_rendered_pdf_s3_key is None: + return None + + logger.info( + f"[rendered_pdf_transform] Reusing rendered PDF for MinerU URL mode: {rendered_pdf_s3_key}" + ) + cached_rendered_pdf_path = JobFileStorage().download_upload_to_temp( + cached_rendered_pdf_s3_key, + suffix=".pdf", + temp_dir=output_dir, + ) + try: + return parse_pdfs( + cached_rendered_pdf_path, + filename, + output_dir, + base_llm_paras, + relative_root=relative_root, + s3_key=cached_rendered_pdf_s3_key, + ) + finally: + if os.path.exists(cached_rendered_pdf_path): + os.remove(cached_rendered_pdf_path) + + +def parse_rendered_pdf_bytes( + *, + pdf_bytes: bytes, + filename: str, + output_dir: str, + base_llm_paras: dict[str, object], + relative_root: str | None, + rendered_pdf_s3_key: str | None = None, +) -> pd.DataFrame: + image_only_pdf_bytes = render_pdf_to_image_pdf(pdf_bytes) + temporary_pdf_path = os.path.join(output_dir, RENDERED_PDF_TEMP_FILENAME) + with open(temporary_pdf_path, "wb") as temporary_pdf_file: + temporary_pdf_file.write(image_only_pdf_bytes) + + try: + return parse_pdfs( + temporary_pdf_path, + filename=filename, + output_dir=output_dir, + base_llm_paras=base_llm_paras, + relative_root=relative_root, + s3_key=rendered_pdf_s3_key, + ) + finally: + if os.path.exists(temporary_pdf_path): + os.remove(temporary_pdf_path) diff --git a/apps/worker/app/services/document_parser/table_asset_writer.py b/apps/worker/app/services/document_parser/table_asset_writer.py new file mode 100644 index 000000000..91599f14d --- /dev/null +++ b/apps/worker/app/services/document_parser/table_asset_writer.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass + +from app.services.document_parser.parser_rows import ParsedRow + + +@dataclass(frozen=True) +class TableAssetInput: + html: str + output_dir: str + table_name: str + summary: str + keywords: str + know_id: str + addtime: str + page_nums: str = "" + content: str | None = None + tokens: str = "" + length: int | None = None + + +def write_table_asset(table_input: TableAssetInput) -> ParsedRow: + table_dir = os.path.join(table_input.output_dir, "tables") + os.makedirs(table_dir, exist_ok=True) + table_filename = _ensure_html_extension(table_input.table_name) + table_path = os.path.join(table_dir, table_filename) + with open(table_path, "w", encoding="utf-8") as table_file: + table_file.write(table_input.html) + row_content = table_input.content if table_input.content is not None else table_input.html + return ParsedRow( + content=row_content, + path=f"tables/{table_filename}", + type="table", + keywords=table_input.keywords, + summary=table_input.summary, + know_id=table_input.know_id, + tokens=table_input.tokens, + connectto="", + addtime=table_input.addtime, + page_nums=table_input.page_nums, + length=table_input.length, + ) + + +def _ensure_html_extension(table_name: str) -> str: + return table_name if table_name.endswith(".html") else f"{table_name}.html" diff --git a/apps/worker/app/services/document_parser/table_parser.py b/apps/worker/app/services/document_parser/table_parser.py index 44b951baa..7ff58a16e 100755 --- a/apps/worker/app/services/document_parser/table_parser.py +++ b/apps/worker/app/services/document_parser/table_parser.py @@ -6,19 +6,22 @@ import threading import uuid from collections import OrderedDict -from typing import Dict, List, Optional, Tuple, Union import numpy as np -import openpyxl import pandas as pd from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.excel_structure_parser import parse_excel_structure from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder from app.services.document_parser.path_helpers import flatten_dic2paths, remove_spaces -from app.services.document_parser.html_parser import df2html +from app.services.document_parser.table_asset_writer import ( + TableAssetInput, + write_table_asset, +) +from app.services.document_parser.dataframe_html_renderer import df2html from bs4 import BeautifulSoup from loguru import logger -from shared.core.config import settings from shared.core.exceptions.domain_exceptions import TableParsingException from shared.core.exceptions.knowhere_exception import KnowhereException from shared.services.ai.prompt_service import build_prompt @@ -82,817 +85,6 @@ def sanitize_table_name_from_header(raw_header_text: str) -> str: g_tbl_lock = threading.Lock() -# ============================================================================ -# PRECISION MODE: Excel Header Detection with Merge Cell Metadata -# ============================================================================ - - -def _get_merged_cell_value(ws, row: int, col: int, merged_ranges: list): - """ - Get the value of a cell, accounting for merged cell regions. - For merged cells, returns the value from the top-left corner of the merge range. - - Args: - ws: openpyxl worksheet - row: 1-indexed row number - col: 1-indexed column number - merged_ranges: list of merged cell ranges from ws.merged_cells.ranges - - Returns: - The cell value (from merge origin if applicable) - """ - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - # This cell is part of a merged region, get value from top-left - return ws.cell(mr.min_row, mr.min_col).value - # Not a merged cell, return direct value - return ws.cell(row, col).value - - -# ============================================================================ -# NEW: Enhanced Header Detection with Row/Column MultiIndex Support -# ============================================================================ - -# Data types that indicate a cell is data, not header (parameterized for future extension) -DATA_TYPES_TO_EXCLUDE = (int, float, datetime.datetime) - - -def _get_unique_cells_in_row( - ws, row: int, col_range: Tuple[int, int], merged_ranges: list -) -> List[dict]: - """Get all unique cells in a row, treating merged cells as single cells. - - Returns: List of {col_start, col_end, value, is_merged} - """ - c_start, c_end = col_range - cells = [] - visited_cols = set() - - for col in range(c_start, c_end + 1): - if col in visited_cols: - continue - - in_merge = False - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - val = ws.cell(mr.min_row, mr.min_col).value - merge_col_end = min(mr.max_col, c_end) - - for mc in range(mr.min_col, merge_col_end + 1): - visited_cols.add(mc) - - cells.append( - { - "col_start": mr.min_col, - "col_end": merge_col_end, - "value": val, - "is_merged": True, - } - ) - in_merge = True - break - - if not in_merge: - val = ws.cell(row, col).value - cells.append( - {"col_start": col, "col_end": col, "value": val, "is_merged": False} - ) - visited_cols.add(col) - - return cells - - -def _get_unique_cells_in_col( - ws, col: int, row_range: Tuple[int, int], merged_ranges: list -) -> List[dict]: - """Get all unique cells in a column, treating merged cells as single cells.""" - r_start, r_end = row_range - cells = [] - visited_rows = set() - - for row in range(r_start, r_end + 1): - if row in visited_rows: - continue - - in_merge = False - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - val = ws.cell(mr.min_row, mr.min_col).value - merge_row_end = min(mr.max_row, r_end) - - for mr_row in range(mr.min_row, merge_row_end + 1): - visited_rows.add(mr_row) - - cells.append( - { - "row_start": mr.min_row, - "row_end": merge_row_end, - "value": val, - "is_merged": True, - } - ) - in_merge = True - break - - if not in_merge: - val = ws.cell(row, col).value - cells.append( - {"row_start": row, "row_end": row, "value": val, "is_merged": False} - ) - visited_rows.add(row) - - return cells - - -def _is_candidate_header_row( - ws, - row: int, - col_range: Tuple[int, int], - merged_ranges: list, - exclude_types: tuple = DATA_TYPES_TO_EXCLUDE, -) -> bool: - """Check if a row is a candidate header row. - - Logic: Row is a candidate if all cells are text (no numbers/dates). - Merged cells are treated as single cells. - """ - cells = _get_unique_cells_in_row(ws, row, col_range, merged_ranges) - - has_any_value = False - for cell in cells: - val = cell["value"] - if val is None: - continue - has_any_value = True - - if isinstance(val, bool): - continue - if isinstance(val, exclude_types): - return False - - return has_any_value - - -def _is_candidate_header_col( - ws, - col: int, - row_range: Tuple[int, int], - merged_ranges: list, - exclude_types: tuple = DATA_TYPES_TO_EXCLUDE, -) -> bool: - """Check if a column is a candidate header column (for row index).""" - cells = _get_unique_cells_in_col(ws, col, row_range, merged_ranges) - - has_any_value = False - for cell in cells: - val = cell["value"] - if val is None: - continue - has_any_value = True - - if isinstance(val, bool): - continue - if isinstance(val, exclude_types): - return False - - return has_any_value - - -def _detect_header_regions( - ws, row_range: Tuple[int, int], col_range: Tuple[int, int], merged_ranges: list -) -> Tuple[List[int], List[int]]: - """Detect header rows and columns. - - Scans rows first, then scans columns only in the data region (excluding header rows). - This prevents header row content from influencing column header detection. - - Returns: - header_rows: List of candidate header row numbers (1-indexed) - header_cols: List of candidate header column numbers (1-indexed) - """ - r_start, r_end = row_range - c_start, c_end = col_range - - # Scan for candidate header rows (top to bottom) - header_rows = [] - for row in range(r_start, r_end + 1): - if _is_candidate_header_row(ws, row, col_range, merged_ranges): - header_rows.append(row) - else: - break - - # Determine data region (excluding header rows) - data_row_start = header_rows[-1] + 1 if header_rows else r_start - - # Skip column scanning if no data rows remain - if data_row_start > r_end: - return header_rows, [] - - # Scan for candidate header columns (left to right) - only in data region - header_cols = [] - data_row_range = (data_row_start, r_end) - for col in range(c_start, c_end + 1): - if _is_candidate_header_col(ws, col, data_row_range, merged_ranges): - header_cols.append(col) - else: - break - - return header_rows, header_cols - - -def _build_column_multiindex( - ws, header_rows: List[int], col_range: Tuple[int, int], merged_ranges: list -) -> Union[pd.Index, pd.MultiIndex]: - """Build column MultiIndex from header rows.""" - c_start, c_end = col_range - levels = [] - - for row in header_rows: - row_values = [] - for col in range(c_start, c_end + 1): - val = _get_merged_cell_value(ws, row, col, merged_ranges) - row_values.append(str(val).strip() if val else "") - levels.append(row_values) - - # Forward fill for merged cells - for idx, level in enumerate(levels): - filled = [] - last = "" - for val in level: - if val: - last = val - filled.append(last if last else val) - levels[idx] = filled - - if len(levels) == 1: - return pd.Index(levels[0]) - return pd.MultiIndex.from_arrays(levels) - - -def _build_row_multiindex( - ws, - header_cols: List[int], - row_range: Tuple[int, int], - merged_ranges: list, - header_rows: List[int] = None, -) -> Union[pd.Index, pd.MultiIndex]: - """Build row MultiIndex from header columns. - - Args: - header_rows: If provided, use the last header row's values as index names - """ - r_start, r_end = row_range - levels = [] - names = [] - - for col in header_cols: - col_values = [] - for row in range(r_start, r_end + 1): - val = _get_merged_cell_value(ws, row, col, merged_ranges) - col_values.append(str(val).strip() if val else "") - levels.append(col_values) - - # Get the column name from the last header row - if header_rows: - name_row = header_rows[-1] - name_val = _get_merged_cell_value(ws, name_row, col, merged_ranges) - names.append(str(name_val).strip() if name_val else None) - else: - names.append(None) - - # Forward fill for merged cells - for idx, level in enumerate(levels): - filled = [] - last = "" - for val in level: - if val: - last = val - filled.append(last if last else val) - levels[idx] = filled - - if len(levels) == 1: - idx = pd.Index(levels[0]) - idx.name = names[0] if names else None - return idx - return pd.MultiIndex.from_arrays(levels, names=names) - - -def _parse_subtable( - ws, row_range: Tuple[int, int], col_range: Tuple[int, int], merged_ranges: list -) -> dict: - """Parse a subtable with new header detection logic. - - Returns: - dict with keys: df, header_rows, header_cols, fallback_col_header, fallback_row_header - """ - r_start, r_end = row_range - c_start, c_end = col_range - - header_rows, header_cols = _detect_header_regions( - ws, row_range, col_range, merged_ranges - ) - - total_rows = r_end - r_start + 1 - total_cols = c_end - c_start + 1 - - # Fall-back check: if all rows/cols are headers, treat as no-header - fallback_col_header = len(header_rows) == total_rows - fallback_row_header = len(header_cols) == total_cols - - # Determine data region - if fallback_col_header: - data_row_start = r_start - columns = None - else: - data_row_start = header_rows[-1] + 1 if header_rows else r_start - columns = ( - _build_column_multiindex(ws, header_rows, col_range, merged_ranges) - if header_rows - else None - ) - - if fallback_row_header: - data_col_start = c_start - row_index = None - else: - data_col_start = header_cols[-1] + 1 if header_cols else c_start - row_index = ( - _build_row_multiindex( - ws, header_cols, (data_row_start, r_end), merged_ranges, header_rows - ) - if header_cols - else None - ) - - # Read data - data = [] - for row in range(data_row_start, r_end + 1): - row_data = [] - for col in range(data_col_start, c_end + 1): - val = _get_merged_cell_value(ws, row, col, merged_ranges) - row_data.append(val) - data.append(row_data) - - # Adjust column index if there are row index columns - if columns is not None and header_cols and not fallback_row_header: - if isinstance(columns, pd.MultiIndex): - columns = columns[len(header_cols) :] - else: - columns = columns[len(header_cols) :] - - df = pd.DataFrame(data, columns=columns, index=row_index) - - # Append original Excel row numbers as the last column for cross-referencing - excel_row_numbers = list(range(data_row_start, r_end + 1)) - if isinstance(df.columns, pd.MultiIndex): - n_levels = df.columns.nlevels - src_row_key = tuple(["_src_row"] + [""] * (n_levels - 1)) - df[src_row_key] = excel_row_numbers - else: - df["_src_row"] = excel_row_numbers - - return { - "df": df, - "header_rows": header_rows if not fallback_col_header else [], - "header_cols": header_cols if not fallback_row_header else [], - "fallback_col_header": fallback_col_header, - "fallback_row_header": fallback_row_header, - } - - -# ============================================================================ -# Sheet Splitting: Detect true separators and split into subtables -# ============================================================================ - - -def _find_effective_range( - ws, row_range: Tuple[int, int], col_range: Tuple[int, int] -) -> Tuple[Tuple[int, int], Tuple[int, int]]: - """Find the effective (non-empty) row and column ranges within a region.""" - r_start, r_end = row_range - c_start, c_end = col_range - - eff_r_start, eff_r_end = None, None - eff_c_start, eff_c_end = None, None - - for row in range(r_start, r_end + 1): - for col in range(c_start, c_end + 1): - if ws.cell(row, col).value is not None: - if eff_r_start is None: - eff_r_start = row - eff_r_end = row - if eff_c_start is None or col < eff_c_start: - eff_c_start = col - if eff_c_end is None or col > eff_c_end: - eff_c_end = col - - if eff_r_start is None: - return ((r_start, r_start), (c_start, c_start)) - - return ((eff_r_start, eff_r_end), (eff_c_start, eff_c_end)) - - -def _is_true_separator_row( - ws, row: int, effective_col_range: Tuple[int, int], merged_ranges: list = None -) -> bool: - """Check if a row is a true separator (all empty within effective column range). - - Considers merged cells - a cell is not empty if it's part of any merged range. - """ - c_start, c_end = effective_col_range - merged_ranges = merged_ranges or [] - - for col in range(c_start, c_end + 1): - # Check if cell has a value - if ws.cell(row, col).value is not None: - return False - # Check if cell is part of a merged range - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - return False # Part of a merge, not truly empty - return True - - -def _is_true_separator_col( - ws, col: int, effective_row_range: Tuple[int, int], merged_ranges: list = None -) -> bool: - """Check if a column is a true separator (all empty within effective row range). - - Considers merged cells - a cell is not empty if it's part of any merged range. - """ - r_start, r_end = effective_row_range - merged_ranges = merged_ranges or [] - - for row in range(r_start, r_end + 1): - # Check if cell has a value - if ws.cell(row, col).value is not None: - return False - # Check if cell is part of a merged range - for mr in merged_ranges: - if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col: - return False # Part of a merge, not truly empty - return True - - -def _find_separator_groups(items: List[int]) -> List[List[int]]: - """Group consecutive separator items together.""" - if not items: - return [] - - groups = [] - current_group = [items[0]] - - for i in range(1, len(items)): - if items[i] == items[i - 1] + 1: - current_group.append(items[i]) - else: - groups.append(current_group) - current_group = [items[i]] - - groups.append(current_group) - return groups - - -def _split_sheet_recursive( - ws, - row_range: Tuple[int, int], - col_range: Tuple[int, int], - merged_ranges: list = None, - min_rows: int = 2, - min_cols: int = 2, -) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: - """ - Recursively split a sheet region into subtables based on true separators. - - Args: - merged_ranges: List of merged cell ranges to consider when detecting separators - - Returns list of (row_range, col_range) tuples for each subtable. - """ - r_start, r_end = row_range - c_start, c_end = col_range - merged_ranges = merged_ranges or [] - - # Find effective range (trim empty edges) - (eff_r_start, eff_r_end), (eff_c_start, eff_c_end) = _find_effective_range( - ws, row_range, col_range - ) - - # If region is too small or empty, return as-is or empty - if eff_r_end - eff_r_start + 1 < min_rows or eff_c_end - eff_c_start + 1 < min_cols: - if eff_r_start is not None: - return [((eff_r_start, eff_r_end), (eff_c_start, eff_c_end))] - return [] - - # Find true separator rows (considering merged cells) - separator_rows = [] - for row in range(eff_r_start + 1, eff_r_end): - if _is_true_separator_row(ws, row, (eff_c_start, eff_c_end), merged_ranges): - separator_rows.append(row) - - # Find true separator columns (considering merged cells) - separator_cols = [] - for col in range(eff_c_start + 1, eff_c_end): - if _is_true_separator_col(ws, col, (eff_r_start, eff_r_end), merged_ranges): - separator_cols.append(col) - - # Group consecutive separators - row_groups = _find_separator_groups(separator_rows) - col_groups = _find_separator_groups(separator_cols) - - # Choose split direction - do_row_split = len(row_groups) > 0 and ( - len(col_groups) == 0 or len(row_groups) <= len(col_groups) - ) - do_col_split = len(col_groups) > 0 and not do_row_split - - if do_row_split: - subtables = [] - prev_end = eff_r_start - for group in row_groups: - if group[0] > prev_end: - sub_result = _split_sheet_recursive( - ws, - (prev_end, group[0] - 1), - (eff_c_start, eff_c_end), - merged_ranges, - min_rows, - min_cols, - ) - subtables.extend(sub_result) - prev_end = group[-1] + 1 - if prev_end <= eff_r_end: - sub_result = _split_sheet_recursive( - ws, - (prev_end, eff_r_end), - (eff_c_start, eff_c_end), - merged_ranges, - min_rows, - min_cols, - ) - subtables.extend(sub_result) - return subtables - - elif do_col_split: - subtables = [] - prev_end = eff_c_start - for group in col_groups: - if group[0] > prev_end: - sub_result = _split_sheet_recursive( - ws, - (eff_r_start, eff_r_end), - (prev_end, group[0] - 1), - merged_ranges, - min_rows, - min_cols, - ) - subtables.extend(sub_result) - prev_end = group[-1] + 1 - if prev_end <= eff_c_end: - sub_result = _split_sheet_recursive( - ws, - (eff_r_start, eff_r_end), - (prev_end, eff_c_end), - merged_ranges, - min_rows, - min_cols, - ) - subtables.extend(sub_result) - return subtables - - else: - return [((eff_r_start, eff_r_end), (eff_c_start, eff_c_end))] - - -# ============================================================================ -# Post-split Merge: Absorb small fragments into nearest neighbor -# ============================================================================ - - -def _count_non_empty_cells( - ws, row_range: Tuple[int, int], col_range: Tuple[int, int] -) -> int: - """Count non-empty cells in a region.""" - count = 0 - for r in range(row_range[0], row_range[1] + 1): - for c in range(col_range[0], col_range[1] + 1): - if ws.cell(r, c).value is not None: - count += 1 - return count - - -def _merge_small_subtables( - ws, subtables: List[Tuple[Tuple[int, int], Tuple[int, int]]], min_cells: int = 4 -) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: - """Merge subtables that have too few non-empty cells into their nearest neighbor. - - This is a post-processing step after _split_sheet_recursive to prevent - over-fragmentation. Fragments with fewer than min_cells non-empty cells - are iteratively absorbed into the nearest neighbor subtable (by bounding-box - distance), expanding the neighbor's bounding box to encompass both regions. - - Args: - ws: openpyxl worksheet - subtables: list of (row_range, col_range) tuples from _split_sheet_recursive - min_cells: minimum non-empty cells for a subtable to be kept standalone - - Returns: - Merged list of (row_range, col_range) tuples - """ - if len(subtables) <= 1: - return subtables - - # Build working list with cell counts - items = [] - for rr, cr in subtables: - count = _count_non_empty_cells(ws, rr, cr) - items.append({"rr": rr, "cr": cr, "cells": count}) - - # Iteratively merge the smallest sub-threshold fragment - changed = True - while changed and len(items) > 1: - changed = False - - # Find the smallest fragment below threshold - min_idx = None - for i, item in enumerate(items): - if item["cells"] < min_cells: - if min_idx is None or item["cells"] < items[min_idx]["cells"]: - min_idx = i - - if min_idx is None: - break # All subtables are above threshold - - # Find nearest neighbor by bounding-box gap distance - src = items[min_idx] - best_j = None - best_dist = float("inf") - for j, tgt in enumerate(items): - if j == min_idx: - continue - row_gap = max( - 0, tgt["rr"][0] - src["rr"][1] - 1, src["rr"][0] - tgt["rr"][1] - 1 - ) - col_gap = max( - 0, tgt["cr"][0] - src["cr"][1] - 1, src["cr"][0] - tgt["cr"][1] - 1 - ) - dist = row_gap + col_gap - if dist < best_dist or ( - dist == best_dist and tgt["cells"] > items[best_j]["cells"] - ): - best_dist = dist - best_j = j - - if best_j is None: - break # Should not happen when len(items) > 1 - - # Merge: expand neighbor's bounding box to encompass both - tgt = items[best_j] - merged_rr = (min(src["rr"][0], tgt["rr"][0]), max(src["rr"][1], tgt["rr"][1])) - merged_cr = (min(src["cr"][0], tgt["cr"][0]), max(src["cr"][1], tgt["cr"][1])) - items[best_j] = { - "rr": merged_rr, - "cr": merged_cr, - "cells": src["cells"] + tgt["cells"], - } - - logger.debug( - f"Merged small fragment (rows={src['rr']}, cols={src['cr']}, " - f"cells={src['cells']}) into neighbor (rows={tgt['rr']}, cols={tgt['cr']})" - ) - - del items[min_idx] - changed = True - - return [(item["rr"], item["cr"]) for item in items] - - -def parse_headers_from_excel( - file_source: Union[str, io.BytesIO], - sheet_name: Optional[str] = None, - split_subtables: bool = True, - include_hidden_sheets: bool = False, -) -> Dict[str, pd.DataFrame]: - """ - Parse Excel file using openpyxl to accurately detect headers via merged cell metadata. - - This is the PRECISION MODE for Excel parsing - it uses the actual merge cell - information from the Excel file to build correct MultiIndex headers without - relying on LLM or heuristics. - - Args: - file_source: Path to Excel file or BytesIO stream - sheet_name: Specific sheet to parse (None = all sheets) - split_subtables: If True, split sheets into subtables based on empty row/column separators (default: True) - include_hidden_sheets: If True, parse hidden/very-hidden sheets. Default False (skip them). - - Returns: - Dictionary mapping sheet/subtable names to DataFrames with correctly set headers - When split_subtables=True, keys are like 'SheetName', 'SheetName_2', 'SheetName_3' etc. - """ - try: - # Load workbook with data_only=True to get calculated values - if isinstance(file_source, str): - wb = openpyxl.load_workbook(file_source, data_only=True) - else: - # BytesIO stream - file_source.seek(0) # Ensure we're at the start - wb = openpyxl.load_workbook(file_source, data_only=True) - - results = {} - sheets_to_parse = [sheet_name] if sheet_name else wb.sheetnames - - for sn in sheets_to_parse: - if sn not in wb.sheetnames: - logger.warning(f"Sheet '{sn}' not found in workbook, skipping") - continue - - ws = wb[sn] - - # Skip hidden sheets unless explicitly included - if not include_hidden_sheets and ws.sheet_state != "visible": - logger.info( - f"Sheet '{sn}' is hidden (state={ws.sheet_state}), skipping" - ) - continue - - # Skip empty sheets - if ws.max_row is None or ws.max_row == 0: - logger.debug(f"Sheet '{sn}' is empty, skipping") - continue - - # Get merged cell ranges - merged_ranges = list(ws.merged_cells.ranges) - logger.debug(f"Sheet '{sn}': found {len(merged_ranges)} merged cell ranges") - - if split_subtables: - # Split sheet into subtables (considers merged cells) - subtable_regions = _split_sheet_recursive( - ws, (1, ws.max_row), (1, ws.max_column or 1), merged_ranges - ) - # Merge back small fragments to prevent over-fragmentation - before_count = len(subtable_regions) - subtable_regions = _merge_small_subtables(ws, subtable_regions) - if len(subtable_regions) != before_count: - logger.info( - f"Sheet '{sn}': merged {before_count} subtables → {len(subtable_regions)} " - f"(absorbed {before_count - len(subtable_regions)} small fragments)" - ) - logger.debug( - f"Sheet '{sn}': {len(subtable_regions)} subtables after merge" - ) - - for idx, (row_range, col_range) in enumerate(subtable_regions): - result = _parse_subtable(ws, row_range, col_range, merged_ranges) - df = result["df"] - - # Store header_cols count in DataFrame attrs for later use in HTML rendering - df.attrs["row_header_cols"] = len(result["header_cols"]) - - # Generate unique key for each subtable - if idx == 0: - key = sn - else: - key = f"{sn}_{idx + 1}" - - logger.debug( - f"Subtable '{key}': rows={row_range}, cols={col_range}, " - f"header_rows={result['header_rows']}, header_cols={result['header_cols']}" - ) - - results[key] = df - else: - # Treat entire sheet as one subtable - row_range = (1, ws.max_row) - col_range = (1, ws.max_column or 1) - - result = _parse_subtable(ws, row_range, col_range, merged_ranges) - df = result["df"] - - # Store header_cols count in DataFrame attrs for later use in HTML rendering - df.attrs["row_header_cols"] = len(result["header_cols"]) - - logger.debug( - f"Sheet '{sn}': header_rows={result['header_rows']}, " - f"header_cols={result['header_cols']}, " - f"fallback_col={result['fallback_col_header']}, " - f"fallback_row={result['fallback_row_header']}" - ) - - results[sn] = df - - wb.close() - return results - - except Exception as e: - logger.error(f"Error parsing Excel with precision mode: {e}") - raise TableParsingException( - user_message="Failed to parse Excel file headers", - reason="EXCEL_PRECISION_PARSE_FAILED", - internal_message=str(e), - original_exception=e, - ) - def identify_tables(line): """Identify if a line contains a table. @@ -1490,7 +682,7 @@ def parse_xlsx( # PRECISION MODE: Use openpyxl metadata for accurate header detection logger.info("Using precision mode for Excel header detection") try: - sheets_dict = parse_headers_from_excel( + sheets_dict = parse_excel_structure( table_stream, include_hidden_sheets=include_hidden_sheets ) precision_mode_active = True @@ -1520,13 +712,13 @@ def parse_xlsx( if len(tb) == 0 or tb.empty or tb.isna().all().all(): continue - # In precision mode, headers are already correctly set by parse_headers_from_excel + # In precision mode, headers are already set by parse_excel_structure # In legacy mode, use LLM/heuristic header parsing if not precision_mode_active: tb = parse_headers(tb, paras=base_llm_paras) # Drop _src_row column before converting to HTML/keywords - # (_src_row is a debug column added by _parse_subtable for cross-referencing) + # (_src_row is a debug column added by Excel structure parsing for cross-referencing) src_row_cols = [ c for c in tb.columns @@ -1535,7 +727,7 @@ def parse_xlsx( if src_row_cols: tb = tb.drop(columns=src_row_cols) - # Get row header column count from DataFrame attrs (set in parse_headers_from_excel) + # Get row header column count from DataFrame attrs (set in parse_excel_structure) row_header_cols = tb.attrs.get("row_header_cols", 0) tb_paths, tb_strs = parse_tb_contents( @@ -1579,16 +771,12 @@ def parse_xlsx( ) + ".html" ) - tb_path = os.path.join(tb_dir, tb_name) soup = BeautifulSoup(tb_strs, features="html.parser") tb_html_str = soup.prettify() - with open(tb_path, "w", encoding="utf-8") as f: - f.write(tb_html_str) # Use same temp_uid for both marker and know_id (aligned with doc_parser/md_parser) temp_uid = gen_str_codes(tb_strs + str(sheet_name)) - relative_tb_path = f"tables/{tb_name}" - tb_ref = build_chunk_ref(relative_tb_path) + tb_ref = build_chunk_ref(f"tables/{tb_name}") tb_bottom_content = f"{tb_ref}\nTable summary:\n{tb_summary}\nMain columns:\n{tb_keywords}" bottom_tokens = tokenize2stw_remove( @@ -1596,22 +784,21 @@ def parse_xlsx( ) all_tb_paths.extend(tb_paths) - # Use relative path for tables: "tables/xxx.html" - df_list.append( - [ - tb_bottom_content, - relative_tb_path, - "table", - len(tb_strs), - tb_keywords, - tb_summary, - temp_uid, - bottom_tokens, - "", - time_stamp, - "", - ] + table_row = write_table_asset( + TableAssetInput( + html=tb_html_str, + output_dir=output_dir, + table_name=tb_name, + summary=tb_summary, + keywords=tb_keywords, + know_id=temp_uid, + addtime=time_stamp, + content=tb_bottom_content, + tokens=bottom_tokens, + length=len(tb_strs), + ) ) + df_list.append(table_row.to_list()) except KnowhereException: raise @@ -1624,6 +811,23 @@ def parse_xlsx( original_exception=e, ) - table_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(",")) + rows_builder = ParsedRowsBuilder() + for row_values in df_list: + rows_builder.append( + ParsedRow( + content=str(row_values[0]), + path=str(row_values[1]), + type=str(row_values[2]), + length=int(row_values[3]), + keywords=str(row_values[4]), + summary=str(row_values[5]), + know_id=str(row_values[6]), + tokens=str(row_values[7]), + connectto=str(row_values[8]), + addtime=str(row_values[9]), + page_nums=str(row_values[10]), + ) + ) + table_df = rows_builder.to_dataframe() table_df = process_dup_paths_df(table_df) return table_df diff --git a/apps/worker/app/services/document_parser/toc_docx.py b/apps/worker/app/services/document_parser/toc_docx.py new file mode 100644 index 000000000..86d612add --- /dev/null +++ b/apps/worker/app/services/document_parser/toc_docx.py @@ -0,0 +1,291 @@ +from __future__ import annotations + +import re + +from app.services.document_parser.heading_candidates import ( + judge_by_conditions, + remove_by_conditions, +) +from app.services.document_parser.toc_hierarchy import build_toc_hierarchy_payload +import lxml.etree as etree + +TOC_TITLE_KEYWORDS = {"目录", "目次", "contents", "table of contents"} + + +def parse_w_int_attr(elem, ns, attr_names): + if elem is None: + return None + + for attr_name in attr_names: + raw_val = elem.get("{%s}%s" % (ns["w"], attr_name)) + if raw_val is None: + continue + try: + return int(raw_val) + except (TypeError, ValueError): + continue + return None + + +def get_docx_toc_layout_hints(elem, ns): + ppr = elem.find("./w:pPr", namespaces=ns) + if ppr is None: + ppr = elem.find(".//w:pPr", namespaces=ns) + + if ppr is None: + return { + "outline_level": None, + "left_indent": None, + } + + outline_elem = ppr.find("./w:outlineLvl", namespaces=ns) + outline_level = None + if outline_elem is not None: + outline_val = parse_w_int_attr(outline_elem, ns, ["val"]) + if outline_val is not None: + outline_level = outline_val + 1 + + indent_elem = ppr.find("./w:ind", namespaces=ns) + left_indent = parse_w_int_attr( + indent_elem, ns, ["left", "start", "leftChars", "startChars"] + ) + + return { + "outline_level": outline_level, + "left_indent": left_indent, + } + + +def infer_toc_level_from_text(text: str): + text_clean = str(text).strip() + if not text_clean: + return None + + normalized = re.sub(r"\s+", " ", text_clean).lower() + if normalized in TOC_TITLE_KEYWORDS: + return None + + raw_pos_code = judge_by_conditions(text_clean) + if not isinstance(raw_pos_code, list): + return None + positive_codes = [ + int(value) + for value in raw_pos_code + if isinstance(value, int) and value > 0 + ] + neg_code = remove_by_conditions(text_clean) + if any(value > 0 for value in neg_code) or not positive_codes: + return None + + return max(positive_codes) + + +def is_toc_title_text(text: str) -> bool: + normalized = re.sub(r"\s+", " ", str(text).strip()).lower() + return normalized in TOC_TITLE_KEYWORDS + + +def infer_toc_levels_from_indentation(entries: list) -> None: + indent_values = sorted( + { + entry["left_indent"] + for entry in entries + if entry.get("level") is None + and entry.get("left_indent") is not None + and not is_toc_title_text(entry.get("heading", "")) + } + ) + + if not indent_values: + return + + indent_to_level = {indent: idx + 1 for idx, indent in enumerate(indent_values)} + for entry in entries: + if entry.get("level") is not None: + continue + if is_toc_title_text(entry.get("heading", "")): + continue + left_indent = entry.get("left_indent") + if left_indent is None: + continue + entry["level"] = indent_to_level.get(left_indent) + + +def get_docx_toc_style_info(elem, ns): + style = elem.find(".//w:pPr/w:pStyle", namespaces=ns) + if style is None: + return { + "is_toc_style": False, + "toc_level": None, + "style_name": None, + } + + val = style.get("{%s}val" % ns["w"]) + if not val: + return { + "is_toc_style": False, + "toc_level": None, + "style_name": None, + } + + val_lower = val.lower().strip() + if "toc" not in val_lower and "目录" not in val: + return { + "is_toc_style": False, + "toc_level": None, + "style_name": val, + } + + level = None + match = re.search(r"(?:toc|目录)\s*[_-]?(\d+)$", val_lower) + if match: + level = int(match.group(1)) + + layout_hints = get_docx_toc_layout_hints(elem, ns) + if level is None: + level = layout_hints["outline_level"] + + return { + "is_toc_style": True, + "toc_level": level, + "style_name": val, + "outline_level": layout_hints["outline_level"], + "left_indent": layout_hints["left_indent"], + } + + +def get_toc_level(elem, ns): + style_info = get_docx_toc_style_info(elem, ns) + if not style_info["is_toc_style"]: + return False + + if style_info["toc_level"] is not None: + return style_info["toc_level"] + return True + + +def detect_sdt_toc(elem, ns): + tag = etree.QName(elem.tag).localname if isinstance(elem.tag, str) else None + + if tag != "sdt": + return {"is_toc_sdt": False, "gallery_type": None} + + is_toc_sdt = False + gallery_type = None + + sdt_pr = elem.find(".//w:sdtPr", namespaces=ns) + if sdt_pr is not None: + doc_part_obj = sdt_pr.find(".//w:docPartObj", namespaces=ns) + if doc_part_obj is not None: + doc_part_gallery = doc_part_obj.find(".//w:docPartGallery", namespaces=ns) + if doc_part_gallery is not None: + gallery_type = doc_part_gallery.get("{%s}val" % ns["w"]) + if gallery_type and "table of contents" in gallery_type.lower(): + is_toc_sdt = True + + return {"is_toc_sdt": is_toc_sdt, "gallery_type": gallery_type} + + +def detect_doc_tocs(elem, ns): + style_info = get_docx_toc_style_info(elem, ns) + is_style = style_info["is_toc_style"] + is_field_start = False + + instrs = elem.findall(".//w:instrText", namespaces=ns) + for instr in instrs: + if instr.text: + instr_text_stripped = instr.text.strip() + instr_text_lower = instr_text_stripped.lower() + if ( + instr_text_lower.startswith("toc") + or "table of contents" in instr_text_lower + or "目录" in instr_text_stripped + ): + is_field_start = True + break + + is_field_end = False + fldchars = elem.findall(".//w:fldChar", namespaces=ns) + for fld in fldchars: + if fld.get("{%s}fldCharType" % ns["w"]) == "end": + is_field_end = True + break + + return { + "is_style": is_style, + "toc_level": style_info["toc_level"], + "style_name": style_info["style_name"], + "outline_level": style_info.get("outline_level"), + "left_indent": style_info.get("left_indent"), + "is_field_start": is_field_start, + "is_field_end": is_field_end, + } + + +def build_docx_toc_hierarchies(block_tuples: list) -> list: + toc_areas = [] + current_area = [] + + for ele_num, block, label, meta in block_tuples: + if "TOC" in label: + current_area.append((ele_num, block, meta or {})) + continue + + if current_area: + toc_areas.append(current_area) + current_area = [] + + if current_area: + toc_areas.append(current_area) + + toc_hierarchies = [] + for area in toc_areas: + toc_entries = [] + for ele_num, block, meta in area: + toc_level = meta.get("toc_level") + try: + toc_level = int(toc_level) if toc_level is not None else None + except (TypeError, ValueError): + toc_level = None + + text = getattr(block, "text", str(block)).strip() + if not text: + continue + + if toc_level is None: + outline_level = meta.get("toc_outline_level") + try: + toc_level = ( + int(outline_level) if outline_level is not None else None + ) + except (TypeError, ValueError): + toc_level = None + + if toc_level is None: + toc_level = infer_toc_level_from_text(text) + + left_indent = meta.get("toc_left_indent") + try: + left_indent = int(left_indent) if left_indent is not None else None + except (TypeError, ValueError): + left_indent = None + + toc_entries.append( + { + "id": ele_num, + "heading": text, + "level": toc_level if toc_level and toc_level > 0 else None, + "left_indent": left_indent, + } + ) + + infer_toc_levels_from_indentation(toc_entries) + payload = build_toc_hierarchy_payload( + toc_entries, + toc_range=(area[0][0], area[-1][0]), + scan_range=(area[0][0], area[-1][0]), + ) + if payload: + toc_hierarchies.append(payload) + + return toc_hierarchies diff --git a/apps/worker/app/services/document_parser/toc_hierarchy.py b/apps/worker/app/services/document_parser/toc_hierarchy.py new file mode 100644 index 000000000..693350a37 --- /dev/null +++ b/apps/worker/app/services/document_parser/toc_hierarchy.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +import pandas as pd +from app.services.document_parser.layout_parser import hiearchy_llm +from app.services.document_parser.stage_profiler import stage_timer +from app.services.document_parser.table_parser import df2md +from app.services.document_parser.text_helpers import normalize_md +from loguru import logger +from pandas import Index + +from shared.core.config import settings + + +def resolve_hierarchy_model_name(model_name: str | None = None) -> str: + return model_name or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL + + +def parse_toc_hierarchy( + toc_df: pd.DataFrame, max_depth: int = 6, model_name: str | None = None +) -> list[dict]: + resolved_model_name = resolve_hierarchy_model_name(model_name) + try: + with stage_timer( + "toc.parse_hierarchy_llm", + model_name=resolved_model_name, + heading_count=len(toc_df), + max_depth=max_depth, + ): + toc_hierarchy = hiearchy_llm( + toc_df, + model_name=resolved_model_name, + max_depth=max_depth, + task="eval-toc-headings", + ) + id_to_level = {item["id"]: item["level"] for item in toc_hierarchy} + + toc_with_level = [] + for _, row in toc_df.iterrows(): + line_id = row["id"] + heading = row["heading"] + level = id_to_level.get(line_id, 1) + toc_with_level.append({"id": line_id, "heading": heading, "level": level}) + return toc_with_level + + except Exception as exc: + logger.error(f"LLM hierarchy analysis failed: {exc}") + return [] + + +def build_tree_tocs(toc_with_level: list[dict]) -> dict: + if not toc_with_level: + return {} + + positive_levels = [item["level"] for item in toc_with_level if item["level"] > 0] + level_for_minus_one = max(positive_levels) + 1 if positive_levels else 1 + + root = {} + stack = [(root, 0)] + + for item in toc_with_level: + heading = item["heading"] + original_level = item["level"] + normalized_level = ( + level_for_minus_one if original_level == -1 else original_level + ) + while len(stack) > 1 and stack[-1][1] >= normalized_level: + stack.pop() + + parent_dict = stack[-1][0] + parent_dict[heading] = {} + stack.append((parent_dict[heading], normalized_level)) + return root + + +def build_toc_hierarchy_payload( + toc_entries: list[dict], + toc_range: tuple | None = None, + scan_range: tuple | None = None, +) -> dict | None: + valid_entries = [] + for entry in toc_entries: + heading = str(entry.get("heading", "")).strip() + level = entry.get("level") + if not heading or not isinstance(level, int) or level <= 0: + continue + + valid_entries.append( + { + "id": entry.get("id"), + "heading": heading, + "level": level, + } + ) + + if not valid_entries: + return None + + toc_df = pd.DataFrame(valid_entries, columns=Index(["id", "heading", "level"])) + payload = { + "toc_range": toc_range or (valid_entries[0]["id"], valid_entries[-1]["id"]), + "toc_with_level": df2md(toc_df, index=False), + "toc_tree": build_tree_tocs(valid_entries), + } + if scan_range is not None: + payload["scan_range"] = scan_range + return payload + + +def eval_toc_levels( + toc_lines: list[str], model_name: str | None = None, max_depth: int = 6 +) -> tuple[str, dict]: + toc_title_keywords = {"目录", "目次", "tableofcontents", "contents"} + valid_data = [] + + for index, line in enumerate(toc_lines): + heading = line.strip() + if not heading: + continue + if normalize_md(heading) in toc_title_keywords: + logger.debug( + f"eval_toc_levels: skipping TOC keyword title line id={index}: {heading[:60]}" + ) + continue + + valid_data.append({"id": index, "heading": heading, "level": "Not Sure"}) + + toc_df = pd.DataFrame(valid_data) + + if toc_df.empty: + logger.info("No valid TOC content, skip hierarchy analysis") + return "", {} + + llm_result = parse_toc_hierarchy(toc_df, max_depth, model_name) + id_to_level = {item["id"]: item["level"] for item in llm_result} + + valid_items_for_tree = [] + for data in valid_data: + line_id = data["id"] + heading = data["heading"] + level = id_to_level.get(line_id, -1) + if level > 0: + valid_items_for_tree.append( + {"id": line_id, "heading": heading, "level": level} + ) + + payload = build_toc_hierarchy_payload(valid_items_for_tree) + if not payload: + return "", {} + return payload["toc_with_level"], payload["toc_tree"] diff --git a/apps/worker/app/services/document_parser/toc_parser.py b/apps/worker/app/services/document_parser/toc_parser.py index c885f8b4c..7a0229859 100644 --- a/apps/worker/app/services/document_parser/toc_parser.py +++ b/apps/worker/app/services/document_parser/toc_parser.py @@ -5,308 +5,25 @@ Provides functionality for: - Detecting TOC (Table of Contents) candidates in markdown documents -- Detecting TOC in DOCX documents (SDT containers, styles, field codes) - Using LLM to determine precise TOC boundaries -- Analyzing TOC hierarchy structure -- Building nested tree structures from TOC """ import re import gevent import pandas as pd +from app.services.document_parser.toc_hierarchy import eval_toc_levels from app.services.document_parser.text_helpers import normalize_md, truncate_text_by_tokens -from app.services.document_parser.layout_parser import ( - hiearchy_llm, - judge_by_conditions, - remove_by_conditions, -) from app.services.document_parser.stage_profiler import stage_timer from app.services.document_parser.table_parser import df2md from gevent.pool import Pool as GeventPool from loguru import logger -from lxml import etree -from shared.core.config import settings from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response from shared.utils.OpenAICompatibleClientSync import get_openai_client -def _resolve_hierarchy_model_name(model_name: str | None = None) -> str: - """Resolve dedicated hierarchy model, falling back to the normal model.""" - return model_name or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL - - -# ==================== DOCX TOC Detection Functions ==================== - -TOC_TITLE_KEYWORDS = {"目录", "目次", "contents", "table of contents"} - - -def _parse_w_int_attr(elem, ns, attr_names): - """Parse the first integer-valued OOXML attribute from a list of names.""" - if elem is None: - return None - - for attr_name in attr_names: - raw_val = elem.get("{%s}%s" % (ns["w"], attr_name)) - if raw_val is None: - continue - try: - return int(raw_val) - except (TypeError, ValueError): - continue - return None - - -def get_docx_toc_layout_hints(elem, ns): - """Extract outline/indent hints for TOC paragraphs without numeric TOC styles.""" - ppr = elem.find("./w:pPr", namespaces=ns) - if ppr is None: - ppr = elem.find(".//w:pPr", namespaces=ns) - - if ppr is None: - return { - "outline_level": None, - "left_indent": None, - } - - outline_elem = ppr.find("./w:outlineLvl", namespaces=ns) - outline_level = None - if outline_elem is not None: - outline_val = _parse_w_int_attr(outline_elem, ns, ["val"]) - if outline_val is not None: - outline_level = outline_val + 1 - - indent_elem = ppr.find("./w:ind", namespaces=ns) - left_indent = _parse_w_int_attr( - indent_elem, ns, ["left", "start", "leftChars", "startChars"] - ) - - return { - "outline_level": outline_level, - "left_indent": left_indent, - } - - -def infer_toc_level_from_text(text: str): - """Fallback TOC level inference from numbering patterns in TOC text.""" - text_clean = str(text).strip() - if not text_clean: - return None - - normalized = re.sub(r"\s+", " ", text_clean).lower() - if normalized in TOC_TITLE_KEYWORDS: - return None - - pos_code = judge_by_conditions(text_clean) - neg_code = remove_by_conditions(text_clean) - if any(x > 0 for x in neg_code) or not any(x > 0 for x in pos_code): - return None - - return max(int(x) for x in pos_code) - - -def is_toc_title_text(text: str) -> bool: - """Return True when the line is likely the standalone TOC heading itself.""" - normalized = re.sub(r"\s+", " ", str(text).strip()).lower() - return normalized in TOC_TITLE_KEYWORDS - - -def infer_toc_levels_from_indentation(entries: list) -> None: - """Populate missing TOC levels by ranking paragraph indentation within one TOC area.""" - indent_values = sorted( - { - entry["left_indent"] - for entry in entries - if entry.get("level") is None - and entry.get("left_indent") is not None - and not is_toc_title_text(entry.get("heading", "")) - } - ) - - if not indent_values: - return - - indent_to_level = {indent: idx + 1 for idx, indent in enumerate(indent_values)} - for entry in entries: - if entry.get("level") is not None: - continue - if is_toc_title_text(entry.get("heading", "")): - continue - left_indent = entry.get("left_indent") - if left_indent is None: - continue - entry["level"] = indent_to_level.get(left_indent) - - -def get_docx_toc_style_info(elem, ns): - """ - Parse TOC style metadata from a DOCX paragraph element. - - Returns: - dict: { - 'is_toc_style': bool, - 'toc_level': Optional[int], - 'style_name': Optional[str] - } - """ - style = elem.find(".//w:pPr/w:pStyle", namespaces=ns) - if style is None: - return { - "is_toc_style": False, - "toc_level": None, - "style_name": None, - } - - val = style.get("{%s}val" % ns["w"]) - if not val: - return { - "is_toc_style": False, - "toc_level": None, - "style_name": None, - } - - val_lower = val.lower().strip() - if "toc" not in val_lower and "目录" not in val: - return { - "is_toc_style": False, - "toc_level": None, - "style_name": val, - } - - level = None - match = re.search(r"(?:toc|目录)\s*[_-]?(\d+)$", val_lower) - if match: - level = int(match.group(1)) - - layout_hints = get_docx_toc_layout_hints(elem, ns) - if level is None: - level = layout_hints["outline_level"] - - return { - "is_toc_style": True, - "toc_level": level, - "style_name": val, - "outline_level": layout_hints["outline_level"], - "left_indent": layout_hints["left_indent"], - } - - -def get_toc_level(elem, ns): - """ - Detect whether a paragraph uses a TOC style. - - Args: - elem: XML paragraph element. - ns: XML namespace map. - - Returns: - bool: True when the paragraph uses a TOC style. - """ - style_info = get_docx_toc_style_info(elem, ns) - if not style_info["is_toc_style"]: - return False - - if style_info["toc_level"] is not None: - return style_info["toc_level"] - return True - - -def detect_sdt_toc(elem, ns): - """ - Detect an SDT (Structured Document Tag) TOC container. - Word-generated TOCs are often wrapped in ``sdt`` elements. - - Args: - elem: SDT element. - ns: XML namespace map. - - Returns: - dict: { - 'is_toc_sdt': bool - whether the element is a TOC SDT, - 'gallery_type': str - docPartGallery type - } - """ - tag = etree.QName(elem.tag).localname if isinstance(elem.tag, str) else None - - if tag != "sdt": - return {"is_toc_sdt": False, "gallery_type": None} - - is_toc_sdt = False - gallery_type = None - - sdt_pr = elem.find(".//w:sdtPr", namespaces=ns) - if sdt_pr is not None: - doc_part_obj = sdt_pr.find(".//w:docPartObj", namespaces=ns) - if doc_part_obj is not None: - doc_part_gallery = doc_part_obj.find(".//w:docPartGallery", namespaces=ns) - if doc_part_gallery is not None: - gallery_type = doc_part_gallery.get("{%s}val" % ns["w"]) - if gallery_type and "table of contents" in gallery_type.lower(): - is_toc_sdt = True - - return {"is_toc_sdt": is_toc_sdt, "gallery_type": gallery_type} - - -def detect_doc_tocs(elem, ns): - """ - Detect TOC regions using two strategies: - 1. paragraph style detection (TOC styles) - 2. field code detection (instrText) - - Note: SDT container detection is handled by ``detect_sdt_toc``. - - Args: - elem: XML paragraph element. - ns: XML namespace map. - - Returns: - dict: { - 'is_style': bool - whether the paragraph uses a TOC style, - 'is_field_start': bool - whether this starts a TOC field, - 'is_field_end': bool - whether this ends a field - } - """ - style_info = get_docx_toc_style_info(elem, ns) - is_style = style_info["is_toc_style"] - is_field_start = False - - instrs = elem.findall(".//w:instrText", namespaces=ns) - for instr in instrs: - if instr.text: - instr_text_stripped = instr.text.strip() - instr_text_lower = instr_text_stripped.lower() - # Match standalone TOC field commands, NOT "PAGEREF _TocXXXX" - # TOC fields start with "TOC" as the command word - if ( - instr_text_lower.startswith("toc") - or "table of contents" in instr_text_lower - or "目录" in instr_text_stripped - ): - is_field_start = True - break - - is_field_end = False - # Always check for fldChar end, even on TOC-styled paragraphs, - # so that the outer TOC field boundary can be properly closed. - fldchars = elem.findall(".//w:fldChar", namespaces=ns) - for fld in fldchars: - if fld.get("{%s}fldCharType" % ns["w"]) == "end": - is_field_end = True - break - - return { - "is_style": is_style, - "toc_level": style_info["toc_level"], - "style_name": style_info["style_name"], - "outline_level": style_info.get("outline_level"), - "left_indent": style_info.get("left_indent"), - "is_field_start": is_field_start, - "is_field_end": is_field_end, - } - - # ==================== Markdown TOC Detection Functions ==================== @@ -655,275 +372,6 @@ def _judge_single_area(idx, lines_, invalid_ids, area_start, area_end): return toc_ranges -def parse_toc_hierarchy(toc_df, max_depth: int = 6, model_name: str = None) -> list: - """ - Parse TOC hierarchy using LLM - - Args: - toc_df: DataFrame with id, heading columns - max_depth: max depth of hierarchy - model_name: model name (optional) - - Returns: - List of dicts with id, heading, level - """ - resolved_model_name = _resolve_hierarchy_model_name(model_name) - try: - with stage_timer( - "toc.parse_hierarchy_llm", - model_name=resolved_model_name, - heading_count=len(toc_df), - max_depth=max_depth, - ): - toc_hierarchy = hiearchy_llm( - toc_df, - model_name=resolved_model_name, - max_depth=max_depth, - task="eval-toc-headings", - ) - id_to_level = {item["id"]: item["level"] for item in toc_hierarchy} - - toc_with_level = [] - for _, row in toc_df.iterrows(): - line_id = row["id"] - heading = row["heading"] - level = id_to_level.get(line_id, 1) - toc_with_level.append({"id": line_id, "heading": heading, "level": level}) - return toc_with_level - - except Exception as e: - logger.error(f"LLM hierarchy analysis failed: {e}") - return [] - - -def build_tree_tocs(toc_with_level: list) -> dict: - """ - Build nested JSON from TOC with level - - Args: - toc_with_level: [{"id": line index, "heading": content, "level": level, "reason": ...}, ...] - level: 1 for h1, 2 for h2..., -1 will be treated as the lowest level title - - Returns: - nested JSON structure - - Notes: - in the TOC scenario, all lines are treated as titles: - - normal levels (1,2,3...) are treated as is - - -1 is treated as a level deeper than all normal levels - """ - if not toc_with_level: - return {} - - # Step 1: collect all levels (exclude -1) - positive_levels = [item["level"] for item in toc_with_level if item["level"] > 0] - - # Step 2: determine the level that -1 should be mapped to - if positive_levels: - # if there are normal levels, -1 is mapped to max + 1 - max_positive_level = max(positive_levels) - level_for_minus_one = max_positive_level + 1 - else: - level_for_minus_one = 1 - - # Step 3: build nested structure - root = {} - stack = [(root, 0)] - - for item in toc_with_level: - heading = item["heading"] - original_level = item["level"] - - # normalize level: -1 -> level_for_minus_one - normalized_level = ( - level_for_minus_one if original_level == -1 else original_level - ) - while len(stack) > 1 and stack[-1][1] >= normalized_level: - stack.pop() - - parent_dict = stack[-1][0] - parent_dict[heading] = {} - stack.append((parent_dict[heading], normalized_level)) - return root - - -def build_toc_hierarchy_payload( - toc_entries: list, - toc_range: tuple | None = None, - scan_range: tuple | None = None, -) -> dict | None: - """ - Build a toc_hierarchies-compatible payload from structured TOC entries. - """ - valid_entries = [] - for entry in toc_entries: - heading = str(entry.get("heading", "")).strip() - level = entry.get("level") - if not heading or not isinstance(level, int) or level <= 0: - continue - - normalized_entry = { - "id": entry.get("id"), - "heading": heading, - "level": level, - } - valid_entries.append(normalized_entry) - - if not valid_entries: - return None - - result_df = pd.DataFrame(valid_entries) - payload = { - "toc_range": toc_range or (valid_entries[0]["id"], valid_entries[-1]["id"]), - "toc_with_level": df2md(result_df[["id", "heading", "level"]], index=False), - "toc_tree": build_tree_tocs(valid_entries), - } - if scan_range is not None: - payload["scan_range"] = scan_range - return payload - - -def build_docx_toc_hierarchies(block_tuples: list) -> list: - """ - Convert DOCX TOC blocks into the same toc_hierarchies structure used by MD/PDF. - """ - toc_areas = [] - current_area = [] - - for ele_num, block, label, meta in block_tuples: - if "TOC" in label: - current_area.append((ele_num, block, meta or {})) - continue - - if current_area: - toc_areas.append(current_area) - current_area = [] - - if current_area: - toc_areas.append(current_area) - - toc_hierarchies = [] - for area in toc_areas: - toc_entries = [] - for ele_num, block, meta in area: - toc_level = meta.get("toc_level") - try: - toc_level = int(toc_level) if toc_level is not None else None - except (TypeError, ValueError): - toc_level = None - - text = getattr(block, "text", str(block)).strip() - if not text: - continue - - if toc_level is None: - outline_level = meta.get("toc_outline_level") - try: - toc_level = ( - int(outline_level) if outline_level is not None else None - ) - except (TypeError, ValueError): - toc_level = None - - if toc_level is None: - toc_level = infer_toc_level_from_text(text) - - left_indent = meta.get("toc_left_indent") - try: - left_indent = int(left_indent) if left_indent is not None else None - except (TypeError, ValueError): - left_indent = None - - toc_entries.append( - { - "id": ele_num, - "heading": text, - "level": toc_level if toc_level and toc_level > 0 else None, - "left_indent": left_indent, - } - ) - - infer_toc_levels_from_indentation(toc_entries) - payload = build_toc_hierarchy_payload( - toc_entries, - toc_range=(area[0][0], area[-1][0]), - scan_range=(area[0][0], area[-1][0]), - ) - if payload: - toc_hierarchies.append(payload) - - return toc_hierarchies - - -def eval_toc_levels( - toc_lines: list, model_name: str = None, max_depth: int = 6 -) -> tuple: - """ - Analyze TOC hierarchy and generate nested JSON - - Args: - toc_lines: list of pre-filtered valid TOC lines (invalid content already removed) - model_name: model name (optional) - max_depth: max depth of hierarchy - - Returns: - (toc_with_level, toc_tree) - - toc_with_level: list with level information - Format: [{"id": int, "heading": str, "level": int, "reason": str}, ...] - - toc_tree: nested JSON structure - """ - # Build data for LLM judgment (all lines are valid, pre-filtered) - # TOC title trigger lines are excluded from - # the LLM input: they stay within toc_range so they are stripped from md_lines, - # but they must not be sent to the hierarchy LLM to avoid a spurious Level=1 entry. - _toc_title_keywords = {"目录", "目次", "tableofcontents", "contents"} - valid_data = [] - - for i, line in enumerate(toc_lines): - heading = line.strip() - if not heading: - continue - if normalize_md(heading) in _toc_title_keywords: - logger.debug( - f"eval_toc_levels: skipping TOC keyword title line id={i}: {heading[:60]}" - ) - continue - - valid_data.append({"id": i, "heading": heading, "level": "Not Sure"}) - - toc_df = pd.DataFrame(valid_data) - - if toc_df.empty: - logger.info("No valid TOC content, skip hierarchy analysis") - return "", {} - - # Evaluate TOC hierarchy with LLM - llm_result = parse_toc_hierarchy(toc_df, max_depth, model_name) - - # Build id -> level mapping from LLM result - id_to_level = {item["id"]: item["level"] for item in llm_result} - - # Build final result - result_data = [] - valid_items_for_tree = [] - - for data in valid_data: - line_id = data["id"] - heading = data["heading"] - level = id_to_level.get(line_id, -1) - - if level > 0: - result_data.append({"id": line_id, "heading": heading, "level": level}) - valid_items_for_tree.append( - {"id": line_id, "heading": heading, "level": level} - ) - - payload = build_toc_hierarchy_payload(valid_items_for_tree) - if not payload: - return "", {} - return payload["toc_with_level"], payload["toc_tree"] - - def detect_tocs_in_texts( md_lines: list, model_name: str = None, diff --git a/apps/worker/tests/contract/test_document_parser_architecture_contract.py b/apps/worker/tests/contract/test_document_parser_architecture_contract.py new file mode 100644 index 000000000..0f3a377ee --- /dev/null +++ b/apps/worker/tests/contract/test_document_parser_architecture_contract.py @@ -0,0 +1,732 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pandas as pd + + +def test_parse_input_builds_typed_llm_parameters( + worker_contract_environment: None, + monkeypatch: Any, + tmp_path: Path, +) -> None: + from app.services.document_parser.orchestration.parse_input import ( + ParseInput, + ParseOptions, + ) + from app.services.document_parser.orchestration.parse_session import ( + build_parse_session, + ) + + monkeypatch.setattr( + "app.services.document_parser.orchestration.parse_session.profile_document", + lambda *_args, **_kwargs: SimpleNamespace( + file_type="pdf", + page_count=3, + atlas_candidate=False, + doc_category="generic", + summary=lambda: "profile", + reasoning="test", + ), + ) + + parse_input = ParseInput( + file_full_path=str(tmp_path / "sample.pdf"), + filename="sample.pdf", + output_dir=str(tmp_path), + internal_output_filename="internal.pdf", + job_id="job-1", + kb_dir="Default_Root", + options=ParseOptions( + doc_type="auto", + llm_histories=7, + smart_title_parse=False, + summary_image=False, + summary_table=True, + summary_txt=False, + stopwords=["the"], + add_frag_desc="fragment", + ), + s3_key="uploads/sample.pdf", + ) + + session = build_parse_session(parse_input) + + assert session.base_llm_paras == { + "llm_histories": 7, + "smart_title_parse": False, + "summary_image": False, + "summary_table": True, + "summary_txt": False, + "stopwords": ["the"], + "doc_type": "auto", + "frag_desc": "fragment", + "model_name": session.base_llm_paras["model_name"], + "hierarchy_model_name": session.base_llm_paras["hierarchy_model_name"], + } + assert session.relative_root == "Default_Root/sample.pdf" + + +def test_document_format_router_uses_adapters( + worker_contract_environment: None, + monkeypatch: Any, + tmp_path: Path, +) -> None: + from app.services.document_parser.orchestration.format_router import ( + DocumentFormat, + get_document_parse_adapter, + resolve_document_format, + ) + from app.services.document_parser.orchestration.parse_input import ParseInput + from app.services.document_parser.orchestration.parse_session import ParseSession + from app.services.document_parser.orchestration.route_parse import route_document_parse + + assert resolve_document_format("/tmp/report.PDF") == DocumentFormat.PDF + assert resolve_document_format("/tmp/report.docx") == DocumentFormat.DOCX + assert get_document_parse_adapter(DocumentFormat.PDF).document_format == DocumentFormat.PDF + + parsed_df = pd.DataFrame([{"content": "ok"}]) + + monkeypatch.setattr( + "app.services.document_parser.pdf_parser.parse_pdfs", + lambda *_args, **_kwargs: parsed_df, + ) + + profile = SimpleNamespace(route="standard", doc_category="generic") + parse_input = ParseInput( + file_full_path=str(tmp_path / "report.pdf"), + filename="report.pdf", + output_dir=str(tmp_path), + internal_output_filename="report.pdf", + ) + session = ParseSession.from_input( + parse_input=parse_input, + base_llm_paras={}, + full_output_dir=str(tmp_path), + profile=profile, + relative_root="Default_Root/report.pdf", + ) + + output_dir, actual_df = route_document_parse(session) + + assert output_dir == str(tmp_path) + assert actual_df is parsed_df + + +def test_rendered_pdf_transform_centralizes_temporary_pdf_cleanup( + worker_contract_environment: None, + monkeypatch: Any, + tmp_path: Path, +) -> None: + from app.services.document_parser.rendered_pdf_transform import ( + parse_rendered_pdf_bytes, + ) + + parsed_df = pd.DataFrame([{"content": "pptx"}]) + seen_pdf_bytes: list[bytes] = [] + + def fake_parse_pdfs(pdf_path: str, **_kwargs: Any) -> pd.DataFrame: + seen_pdf_bytes.append(Path(pdf_path).read_bytes()) + assert Path(pdf_path).exists() + return parsed_df + + monkeypatch.setattr( + "app.services.document_parser.rendered_pdf_transform.parse_pdfs", + fake_parse_pdfs, + ) + monkeypatch.setattr( + "app.services.document_parser.rendered_pdf_transform.render_pdf_to_image_pdf", + lambda pdf_bytes: pdf_bytes, + ) + + actual_df = parse_rendered_pdf_bytes( + pdf_bytes=b"rendered", + filename="slides.pptx", + output_dir=str(tmp_path), + base_llm_paras={}, + relative_root="Default_Root/slides.pptx", + rendered_pdf_s3_key="transforms/job.pdf", + ) + + assert actual_df is parsed_df + assert seen_pdf_bytes == [b"rendered"] + assert not (tmp_path / "_pptx_tmp.pdf").exists() + + +def test_heading_hierarchy_module_wraps_prediction( + worker_contract_environment: None, + monkeypatch: Any, + tmp_path: Path, +) -> None: + from app.services.document_parser.heading_hierarchy import ( + HeadingHierarchyInput, + predict_heading_hierarchy, + ) + + expected_df = pd.DataFrame( + [{"id": 1, "heading": "Intro", "level": 1, "reason": "test"}] + ) + captured: dict[str, Any] = {} + + def fake_pred_titles(*args: Any, **kwargs: Any) -> pd.DataFrame: + captured["args"] = args + captured["kwargs"] = kwargs + return expected_df + + monkeypatch.setattr( + "app.services.document_parser.heading_hierarchy.pred_titles", + fake_pred_titles, + ) + + actual_df = predict_heading_hierarchy( + HeadingHierarchyInput( + infos=[(1, "Intro")], + doc_type="md", + smart_parse=True, + model_name="hierarchy-model", + output_dir=str(tmp_path), + layout_json_path=str(tmp_path / "layout.json"), + ) + ) + + assert actual_df is expected_df + assert captured["kwargs"]["doc_type"] == "md" + assert captured["kwargs"]["smart_parse"] is True + assert captured["kwargs"]["model_name"] == "hierarchy-model" + + +def test_parser_row_builder_owns_dataframe_column_order( + worker_contract_environment: None, +) -> None: + from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder + + builder = ParsedRowsBuilder() + builder.append( + ParsedRow( + content="chunk text", + path="Default_Root/doc/Section", + type="text", + keywords="alpha;beta", + summary="summary", + know_id="chunk-1", + tokens="alpha->beta", + connectto="", + page_nums="1,2", + addtime="now", + ) + ) + + parsed_df = builder.to_dataframe() + + assert list(parsed_df.columns) == [ + "content", + "path", + "type", + "length", + "keywords", + "summary", + "know_id", + "tokens", + "connectto", + "addtime", + "page_nums", + ] + assert parsed_df.iloc[0].to_dict() == { + "content": "chunk text", + "path": "Default_Root/doc/Section", + "type": "text", + "length": len("chunk text"), + "keywords": "alpha;beta", + "summary": "summary", + "know_id": "chunk-1", + "tokens": "alpha->beta", + "connectto": "", + "addtime": "now", + "page_nums": "1,2", + } + + +def test_inline_asset_module_builds_image_and_table_rows( + worker_contract_environment: None, +) -> None: + from app.services.document_parser.inline_asset import ( + build_image_asset_row, + build_table_asset_row, + ) + + image_row = build_image_asset_row( + content="\nImage summary\n[images/image-1.png]\n", + relative_path="images/image-1.png", + summary="image-1\nImage summary", + know_id="image-1", + addtime="now", + page_nums="3", + ) + table_row = build_table_asset_row( + content="
", + relative_path="tables/table-1.html", + summary="table-1\nTable summary", + keywords="column", + know_id="table-1", + addtime="now", + page_nums="4", + ) + + assert image_row.type == "image" + assert image_row.path == "images/image-1.png" + assert image_row.summary == "image-1\nImage summary" + assert table_row.type == "table" + assert table_row.path == "tables/table-1.html" + assert table_row.keywords == "column" + + +def test_table_asset_writer_creates_table_row_and_html_file( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.table_asset_writer import ( + TableAssetInput, + write_table_asset, + ) + + row = write_table_asset( + TableAssetInput( + html="
A
", + output_dir=str(tmp_path), + table_name="table-1", + summary="table-1", + keywords="A", + know_id="table-1", + addtime="now", + ) + ) + + assert (tmp_path / "tables" / "table-1.html").read_text(encoding="utf-8") + assert row.type == "table" + assert row.path == "tables/table-1.html" + assert row.content == "
A
" + + +def test_docx_asset_store_owns_asset_filesystem_lifecycle( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.docx_asset_store import DocxAssetStore + + store = DocxAssetStore(str(tmp_path)) + (tmp_path / "images").mkdir() + (tmp_path / "tables").mkdir() + (tmp_path / "images" / "stale.png").write_bytes(b"stale") + (tmp_path / "tables" / "stale.html").write_text("stale", encoding="utf-8") + + store.reset() + image_asset = store.write_image("image-1 raw", ".png", b"image") + renamed_asset = store.rename_image(image_asset, "image-1 final") + table_asset = store.write_table("table-1 final", "
") + + assert not (tmp_path / "images" / "stale.png").exists() + assert not (tmp_path / "tables" / "stale.html").exists() + assert renamed_asset.relative_path == "images/image-1 final.png" + assert (tmp_path / "images" / "image-1 final.png").read_bytes() == b"image" + assert table_asset.relative_path == "tables/table-1 final.html" + assert (tmp_path / "tables" / "table-1 final.html").read_text( + encoding="utf-8" + ) == "
" + + +def test_docx_block_stream_emits_document_ordered_blocks( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.docx_block_stream import iter_block_items + from docx import Document + from docx.table import Table + from docx.text.paragraph import Paragraph + + docx_path = tmp_path / "sample.docx" + document = Document() + document.add_paragraph("Intro") + table = document.add_table(rows=1, cols=1) + table.cell(0, 0).text = "Cell" + document.save(docx_path) + + block_events = list(iter_block_items(docx_path.read_bytes())) + + assert block_events[0][0] == 1 + assert isinstance(block_events[0][1], Paragraph) + assert block_events[0][1].text == "Intro" + assert block_events[0][2] == "PTXT" + assert isinstance(block_events[1][1], Table) + assert block_events[1][2] == "TABLE" + + +def test_html_table_modules_separate_docx_and_dataframe_rendering( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.dataframe_html_renderer import df2html + from app.services.document_parser.docx_table_html import table2html + from docx import Document + + docx_path = tmp_path / "table.docx" + document = Document() + table = document.add_table(rows=1, cols=2) + table.cell(0, 0).text = "A" + table.cell(0, 1).text = "B" + document.save(docx_path) + + loaded_table = Document(str(docx_path)).tables[0] + docx_html = table2html(loaded_table, cell_image_map={(0, 1): "image summary"}) + + dataframe_html = df2html( + pd.DataFrame([["North", "North", 3]], columns=["Region", "Group", "Value"]), + row_header_cols=2, + ) + + assert docx_html == ( + "" + "
AB
image summary
" + ) + assert '		North3
", encoding="utf-8") + + rows: list[list[str | int]] = [ + [ + "[images/image-3-old.png]", + "images/image-3-old.png", + "image", + 24, + "", + "image-3", + "image-id", + "", + "", + "now", + "", + ], + [ + "[tables/table-0 old.html]", + "tables/table-0 old.html", + "table", + 25, + "", + "table-0", + "table-id", + "", + "", + "now", + "", + ], + [ + "long text", + "Root/Text", + "text", + 9, + "", + "", + "text-id", + "", + "", + "now", + "", + ], + ] + + monkeypatch.setattr(deferred_summary, "_get_vision_client", lambda: object()) + monkeypatch.setattr( + deferred_summary, + "ask_image", + lambda *_args, **_kwargs: "Better Image\nImage summary", + ) + + def fake_extract(text: str, **_kwargs: Any) -> tuple[str, str, str]: + if "", str(table_dir), "table-0 old", 0), + ("text", 2, "long text"), + ], + output_dir=str(tmp_path), + ) + ) + + assert rows[0][1] == "images/image-3-Better Image.png" + assert "[images/image-3-Better Image.png]" in str(rows[0][0]) + assert rows[0][5] == "image-3\nImage summary" + assert (image_dir / "image-3-Better Image.png").exists() + assert rows[1][1] == "tables/table-0 Better Table.html" + assert rows[1][4] == "table-keyword" + assert rows[1][5] == "table-0\ntable summary" + assert (table_dir / "table-0 Better Table.html").exists() + assert rows[2][4] == "text-keyword" + assert rows[2][5] == "text summary" + + +def test_toc_modules_separate_docx_detection_from_hierarchy_payloads( + worker_contract_environment: None, +) -> None: + from app.services.document_parser.toc_docx import infer_toc_level_from_text + from app.services.document_parser.toc_hierarchy import build_toc_hierarchy_payload + + assert infer_toc_level_from_text("1.2 Scope") == 2 + payload = build_toc_hierarchy_payload( + [ + {"id": 3, "heading": "1 Overview", "level": 1}, + {"id": 4, "heading": "1.1 Detail", "level": 2}, + ], + toc_range=(3, 4), + scan_range=(3, 5), + ) + + assert payload is not None + assert payload["toc_range"] == (3, 4) + assert payload["scan_range"] == (3, 5) + assert payload["toc_tree"] == {"1 Overview": {"1.1 Detail": {}}} + + +def test_format_adapters_do_not_expose_lazy_any_wrappers( + worker_contract_environment: None, +) -> None: + import app.services.document_parser.orchestration.format_adapters as format_adapters + + wrapper_names = [ + "parse_fragment", + "parse_texts", + "parse_md", + "parse_image", + "parse_pdfs", + "parse_docx", + "convert_doc2dics", + "doc_to_docx", + "xls_to_xlsx", + "parse_xlsx", + "parse_pptx", + ] + + assert not any(hasattr(format_adapters, wrapper_name) for wrapper_name in wrapper_names) From 21381a959747be5f7949341748fca2676950223e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:02:59 +0000 Subject: [PATCH 37/40] refactor(worker): deepen excel parser contract --- .../test_agentic_answer_policy_contract.py | 3 +- .../contract/test_demo_documents_contract.py | 33 +- .../contract/test_qstash_callback_contract.py | 1 + .../tests/contract/test_retrieval_contract.py | 18 +- .../document_parser/excel_table_parser.py | 286 ++++++++++++++++++ .../orchestration/format_adapters.py | 2 +- .../services/document_parser/table_parser.py | 238 ++------------- .../contract/test_excel_parser_contract.py | 123 ++++++++ .../shared/testing/contract_runtime.py | 18 ++ 9 files changed, 491 insertions(+), 231 deletions(-) create mode 100644 apps/worker/app/services/document_parser/excel_table_parser.py create mode 100644 apps/worker/tests/contract/test_excel_parser_contract.py diff --git a/apps/api/tests/contract/test_agentic_answer_policy_contract.py b/apps/api/tests/contract/test_agentic_answer_policy_contract.py index e0c0cf1d2..6ef4e07cd 100644 --- a/apps/api/tests/contract/test_agentic_answer_policy_contract.py +++ b/apps/api/tests/contract/test_agentic_answer_policy_contract.py @@ -4,9 +4,10 @@ from shared.services.retrieval.agentic.policy import attempt_answer from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState +from shared.services.retrieval.llm_adapter import LLMFnInput -async def _malformed_json_wrapper(_prompt: str) -> str: +async def _malformed_json_wrapper(_prompt: LLMFnInput) -> str: return '{"status": "DONE", "answer": "truncated"' diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index ea21e7de6..a3b63b9ab 100644 --- a/apps/api/tests/contract/test_demo_documents_contract.py +++ b/apps/api/tests/contract/test_demo_documents_contract.py @@ -136,13 +136,16 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( monkeypatch: MonkeyPatch, ) -> None: fake_result_storage = FakeResultStorage() - monkeypatch.setattr( - "shared.services.storage.result_storage.get_result_storage", - lambda: fake_result_storage, - ) monkeypatch.setenv("RETRIEVAL_AGENTIC_ENABLED", "false") async with developer_api_client_factory() as api_client: + import app.services.demo_document_service as demo_document_service + + monkeypatch.setattr( + demo_document_service, + "get_result_storage", + lambda: fake_result_storage, + ) empty_cached_response = await api_client.post( "/api/v1/retrieval/query", json={ @@ -276,12 +279,15 @@ async def test_should_serialize_concurrent_first_demo_materialization( monkeypatch: MonkeyPatch, ) -> None: fake_result_storage = FakeResultStorage() - monkeypatch.setattr( - "shared.services.storage.result_storage.get_result_storage", - lambda: fake_result_storage, - ) async with developer_api_client_factory() as api_client: + import app.services.demo_document_service as demo_document_service + + monkeypatch.setattr( + demo_document_service, + "get_result_storage", + lambda: fake_result_storage, + ) first_response, second_response = await asyncio.gather( api_client.post( "/api/v1/demo/materializations", @@ -369,12 +375,15 @@ async def test_should_reject_mixed_demo_materialization_selection_before_upload( monkeypatch: MonkeyPatch, ) -> None: fake_result_storage = FakeResultStorage() - monkeypatch.setattr( - "shared.services.storage.result_storage.get_result_storage", - lambda: fake_result_storage, - ) async with developer_api_client_factory() as api_client: + import app.services.demo_document_service as demo_document_service + + monkeypatch.setattr( + demo_document_service, + "get_result_storage", + lambda: fake_result_storage, + ) response = await api_client.post( "/api/v1/demo/materializations", json={ diff --git a/apps/api/tests/contract/test_qstash_callback_contract.py b/apps/api/tests/contract/test_qstash_callback_contract.py index e25c3d581..48c78cee3 100644 --- a/apps/api/tests/contract/test_qstash_callback_contract.py +++ b/apps/api/tests/contract/test_qstash_callback_contract.py @@ -344,6 +344,7 @@ async def test_should_return_ok_without_mutating_state_when_the_callback_has_no_ ) assert event_row is not None + assert log_count_row is not None assert event_row["status"] == "pending" assert event_row["attempts"] == 0 assert cast(int, log_count_row["count"]) == 0 diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index cb99fd689..51c4ff48b 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -135,6 +135,10 @@ async def _seed_retrieval_chunk_for_existing_document( } +def _result_source(result: dict[str, object]) -> dict[str, object]: + return cast(dict[str, object], result["source"]) + + @pytest.mark.asyncio async def test_should_return_seeded_retrieval_results_for_the_authenticated_user( developer_api_client_factory: Callable[ @@ -208,7 +212,7 @@ async def test_should_default_the_namespace_to_default_when_it_is_omitted( assert response_json["namespace"] == "default" assert len(results) == 1 - assert results[0]["source"]["document_id"] == seeded_document["document_id"] + assert _result_source(results[0])["document_id"] == seeded_document["document_id"] @pytest.mark.asyncio @@ -435,7 +439,7 @@ async def fake_graph_routing(*_args: object, **_kwargs: object) -> list[dict[str results = cast(list[dict[str, object]], response_json["results"]) assert len(results) == 1 - assert results[0]["source"]["document_id"] == hot_document["document_id"] + assert _result_source(results[0])["document_id"] == hot_document["document_id"] @pytest.mark.asyncio @@ -726,7 +730,7 @@ async def fake_retrieval_run( cast(str, reference["document_id"]) for reference in referenced_chunks } result_document_ids = { - cast(str, result["source"]["document_id"]) for result in results + cast(str, _result_source(result)["document_id"]) for result in results } assert referenced_document_ids == { @@ -837,7 +841,7 @@ async def fake_retrieval_run( cast(str, reference["section_path"]) for reference in referenced_chunks } result_section_paths = { - cast(str, result["source"]["section_path"]) for result in results + cast(str, _result_source(result)["section_path"]) for result in results } assert referenced_section_paths == { @@ -918,7 +922,7 @@ async def test_should_exclude_matching_document_ids_from_the_response( results = cast(list[dict[str, object]], response_json["results"]) assert len(results) == 1 - assert results[0]["source"]["document_id"] == included_document["document_id"] + assert _result_source(results[0])["document_id"] == included_document["document_id"] @pytest.mark.asyncio @@ -963,5 +967,5 @@ async def test_should_exclude_matching_sections_from_the_response( results = cast(list[dict[str, object]], response_json["results"]) assert len(results) == 1 - assert results[0]["source"]["document_id"] == included_document["document_id"] - assert results[0]["source"]["section_path"] == included_document["section_path"] + assert _result_source(results[0])["document_id"] == included_document["document_id"] + assert _result_source(results[0])["section_path"] == included_document["section_path"] diff --git a/apps/worker/app/services/document_parser/excel_table_parser.py b/apps/worker/app/services/document_parser/excel_table_parser.py new file mode 100644 index 000000000..d1d5589c0 --- /dev/null +++ b/apps/worker/app/services/document_parser/excel_table_parser.py @@ -0,0 +1,286 @@ +# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalSubscript=false, reportReturnType=false +from __future__ import annotations + +import io +import os +from dataclasses import dataclass +from typing import Any + +import pandas as pd +from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.excel_structure_parser import parse_excel_structure +from app.services.document_parser.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder +from app.services.document_parser.path_helpers import remove_spaces +from app.services.document_parser.table_asset_writer import ( + TableAssetInput, + write_table_asset, +) +from app.services.document_parser.table_parser import ( + parse_headers, + parse_tb_contents, + parse_tb_keywords, + postprocess_tb, +) +from bs4 import BeautifulSoup +from loguru import logger + +from shared.core.exceptions.domain_exceptions import TableParsingException +from shared.core.exceptions.knowhere_exception import KnowhereException +from shared.utils.chunk_refs import build_chunk_ref +from shared.utils.file_loading import load_file_bytes +from shared.utils.file_utils import path_handle +from shared.utils.text_utils import tokenize2stw_remove + + +@dataclass(frozen=True) +class ExcelWorkbookParseRequest: + file_path: str + file_name: str + output_dir: str + baseurl: str + base_llm_paras: dict[str, Any] + window_h: int + relative_root: str | None + use_precision_mode: bool + include_hidden_sheets: bool + + +def parse_xlsx( + file_path: str, + file_name: str, + output_dir: str, + baseurl: str, + base_llm_paras: dict[str, Any] | None = None, + window_h: int = 10, + relative_root: str | None = None, + use_precision_mode: bool = True, + include_hidden_sheets: bool = False, +) -> pd.DataFrame: + request = ExcelWorkbookParseRequest( + file_path=file_path, + file_name=file_name, + output_dir=output_dir, + baseurl=baseurl, + base_llm_paras=_normalise_llm_parameters(base_llm_paras), + window_h=window_h, + relative_root=relative_root, + use_precision_mode=use_precision_mode, + include_hidden_sheets=include_hidden_sheets, + ) + return parse_excel_workbook(request) + + +def parse_excel_workbook(request: ExcelWorkbookParseRequest) -> pd.DataFrame: + time_stamp = get_str_time() + sheets_dict, precision_mode_active = _load_excel_sheets(request) + parsed_rows: list[ParsedRow] = [] + + for sheet_name, sheet_frame in _iter_unique_sheets(sheets_dict): + table_rows = _parse_excel_sheet( + request=request, + sheet_name=sheet_name, + sheet_frame=sheet_frame, + precision_mode_active=precision_mode_active, + time_stamp=time_stamp, + ) + parsed_rows.extend(table_rows) + + return _rows_to_dataframe(parsed_rows) + + +def _normalise_llm_parameters( + base_llm_paras: dict[str, Any] | None, +) -> dict[str, Any]: + llm_parameters = dict(base_llm_paras or {}) + llm_parameters.setdefault("summary_table", False) + llm_parameters.setdefault("stopwords", []) + return llm_parameters + + +def _load_excel_sheets( + request: ExcelWorkbookParseRequest, +) -> tuple[dict[str, pd.DataFrame], bool]: + table_data = load_file_bytes(request.file_path, file_url=request.baseurl) + table_stream = io.BytesIO(table_data) + + os.makedirs(os.path.join(request.output_dir, "tables"), exist_ok=True) + + if not request.use_precision_mode: + return pd.read_excel(table_stream, sheet_name=None), False + + logger.info("Using precision mode for Excel header detection") + try: + return ( + parse_excel_structure( + table_stream, + include_hidden_sheets=request.include_hidden_sheets, + ), + True, + ) + except Exception as exc: + logger.warning(f"Precision mode failed, falling back to legacy mode: {exc}") + table_stream.seek(0) + return pd.read_excel(table_stream, sheet_name=None), False + + +def _iter_unique_sheets( + sheets_dict: dict[str, pd.DataFrame], +) -> list[tuple[str, pd.DataFrame]]: + used_sheet_names: list[str] = [] + unique_sheets: list[tuple[str, pd.DataFrame]] = [] + + for raw_sheet_name, sheet_content in sheets_dict.items(): + sheet_name = raw_sheet_name.strip() + if sheet_name in used_sheet_names: + sheet_name = sheet_name + str(len(used_sheet_names)) + else: + used_sheet_names.append(sheet_name) + unique_sheets.append((sheet_name, sheet_content)) + + return unique_sheets + + +def _parse_excel_sheet( + *, + request: ExcelWorkbookParseRequest, + sheet_name: str, + sheet_frame: pd.DataFrame, + precision_mode_active: bool, + time_stamp: str, +) -> list[ParsedRow]: + parsed_rows: list[ParsedRow] = [] + + try: + table_frame = postprocess_tb(sheet_frame, drop=True) + if len(table_frame) == 0 or table_frame.empty or table_frame.isna().all().all(): + return parsed_rows + + if not precision_mode_active: + table_frame = parse_headers(table_frame, paras=request.base_llm_paras) + + table_frame = _drop_source_row_columns(table_frame) + row_header_cols = int(table_frame.attrs.get("row_header_cols", 0)) + + _table_paths, table_html = parse_tb_contents( + table_frame, + parent_dic={request.file_name: {sheet_name: {}}}, + file_name=request.file_name, + sheet_name=sheet_name, + row_header_cols=row_header_cols, + ) + + parsed_rows.append( + _write_excel_table_asset( + request=request, + sheet_name=sheet_name, + table_frame=table_frame, + table_html=table_html, + time_stamp=time_stamp, + ) + ) + return parsed_rows + except KnowhereException: + raise + except Exception as exc: + logger.error(f"Table parsing failed: {exc}") + raise TableParsingException( + user_message="Failed to parse Excel table content", + reason="TABLE_PROCESSING_FAILED", + internal_message=str(exc), + original_exception=exc, + ) from exc + + +def _drop_source_row_columns(table_frame: pd.DataFrame) -> pd.DataFrame: + source_row_columns = [ + column + for column in table_frame.columns + if (isinstance(column, tuple) and column[0] == "_src_row") + or column == "_src_row" + ] + if not source_row_columns: + return table_frame + return table_frame.drop(columns=source_row_columns) + + +def _write_excel_table_asset( + *, + request: ExcelWorkbookParseRequest, + sheet_name: str, + table_frame: pd.DataFrame, + table_html: str, + time_stamp: str, +) -> ParsedRow: + title, keywords, summary = _summarize_excel_table( + table_frame=table_frame, + table_html=table_html, + sheet_name=sheet_name, + llm_parameters=request.base_llm_paras, + ) + table_index = f"table-{sheet_name}" + table_summary = f"{table_index}\n{summary}" if summary else table_index + effective_name = title or sheet_name + table_stem = path_handle( + remove_spaces("table-" + effective_name), + mode="clean_single", + ) + if not isinstance(table_stem, str) or not table_stem: + raise ValueError(f"Failed to sanitize Excel table name: {effective_name}") + table_name = table_stem + ".html" + table_html_string = BeautifulSoup(table_html, features="html.parser").prettify() + know_id = gen_str_codes(table_html + str(sheet_name)) + table_ref = build_chunk_ref(f"tables/{table_name}") + table_content = ( + f"{table_ref}\nTable summary:\n{table_summary}\nMain columns:\n{keywords}" + ) + table_tokens = tokenize2stw_remove( + [table_content], + request.base_llm_paras["stopwords"], + ) + + return write_table_asset( + TableAssetInput( + html=table_html_string, + output_dir=request.output_dir, + table_name=table_name, + summary=table_summary, + keywords=keywords, + know_id=know_id, + addtime=time_stamp, + content=table_content, + tokens=table_tokens, + length=len(table_html), + ) + ) + + +def _summarize_excel_table( + *, + table_frame: pd.DataFrame, + table_html: str, + sheet_name: str, + llm_parameters: dict[str, Any], +) -> tuple[str | None, str, str | None]: + mechanical_keywords = parse_tb_keywords(table_frame) + + if llm_parameters["summary_table"]: + from app.services.document_parser.txt_parser import ( + extract_title_keywords_summary, + ) + + title, keywords, summary = extract_title_keywords_summary( + table_html, + max_keywords=3, + ) + return title, keywords or mechanical_keywords, summary + + return None, mechanical_keywords, None + + +def _rows_to_dataframe(parsed_rows: list[ParsedRow]) -> pd.DataFrame: + rows_builder = ParsedRowsBuilder() + for row in parsed_rows: + rows_builder.append(row) + table_df = rows_builder.to_dataframe() + return process_dup_paths_df(table_df) diff --git a/apps/worker/app/services/document_parser/orchestration/format_adapters.py b/apps/worker/app/services/document_parser/orchestration/format_adapters.py index 882efbaab..f10f9edd1 100644 --- a/apps/worker/app/services/document_parser/orchestration/format_adapters.py +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -208,7 +208,7 @@ def _parse_xlsx_path( xlsx_path: str, session: ParseSession, ) -> tuple[str, pd.DataFrame | None]: - from app.services.document_parser.table_parser import parse_xlsx + from app.services.document_parser.excel_table_parser import parse_xlsx parsed_df = parse_xlsx( xlsx_path, diff --git a/apps/worker/app/services/document_parser/table_parser.py b/apps/worker/app/services/document_parser/table_parser.py index 7ff58a16e..342200836 100755 --- a/apps/worker/app/services/document_parser/table_parser.py +++ b/apps/worker/app/services/document_parser/table_parser.py @@ -1,7 +1,5 @@ # pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalOperand=false, reportOptionalSubscript=false, reportReturnType=false import datetime -import io -import os import re import threading import uuid @@ -9,28 +7,16 @@ import numpy as np import pandas as pd -from app.services.document_parser.dataframe_helpers import process_dup_paths_df -from app.services.document_parser.excel_structure_parser import parse_excel_structure -from app.services.document_parser.identifiers import gen_str_codes, get_str_time -from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder -from app.services.document_parser.path_helpers import flatten_dic2paths, remove_spaces -from app.services.document_parser.table_asset_writer import ( - TableAssetInput, - write_table_asset, -) +from app.services.document_parser.identifiers import gen_str_codes +from app.services.document_parser.path_helpers import flatten_dic2paths from app.services.document_parser.dataframe_html_renderer import df2html from bs4 import BeautifulSoup from loguru import logger -from shared.core.exceptions.domain_exceptions import TableParsingException -from shared.core.exceptions.knowhere_exception import KnowhereException from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response -from shared.utils.chunk_refs import build_chunk_ref -from shared.utils.file_loading import load_file_bytes -from shared.utils.file_utils import path_handle from shared.utils.OpenAICompatibleClientSync import get_openai_client -from shared.utils.text_utils import remove_duplicates_orderkept, tokenize2stw_remove +from shared.utils.text_utils import remove_duplicates_orderkept # ── Table filename sanitizer ──────────────────────────── # Max byte-safe filename length. Most filesystems cap at 255 bytes; we leave @@ -638,196 +624,28 @@ def format_tb_scope(df, num): def parse_xlsx( - file_path, - file_name, - output_dir, - baseurl, - base_llm_paras=None, - window_h=10, - relative_root=None, - use_precision_mode=True, - include_hidden_sheets=False, -): - """ - Parse Excel file and extract table content. - - Args: - file_path: Path or URL to the Excel file - file_name: Display name for the file - output_dir: Directory to save extracted tables - baseurl: Base URL for file loading - base_llm_paras: LLM parameters for summarization - window_h: Window size for table scope - relative_root: Root path for relative paths - use_precision_mode: If True, use openpyxl merged cell metadata for accurate - header detection. If False, use LLM/heuristic mode. - Default is True for better accuracy. - include_hidden_sheets: If True, parse hidden/very-hidden sheets. Default False. - - Returns: - DataFrame with parsed table information - """ - time_stamp = get_str_time() - df_list = [] - - table_data = load_file_bytes(file_path, file_url=baseurl) - table_stream = io.BytesIO(table_data) - - tb_dir = os.path.join(output_dir, "tables") - os.makedirs(tb_dir, exist_ok=True) - all_tb_paths = [] - exist_sheets = [] - - if use_precision_mode: - # PRECISION MODE: Use openpyxl metadata for accurate header detection - logger.info("Using precision mode for Excel header detection") - try: - sheets_dict = parse_excel_structure( - table_stream, include_hidden_sheets=include_hidden_sheets - ) - precision_mode_active = True - except Exception as e: - logger.warning(f"Precision mode failed, falling back to legacy mode: {e}") - table_stream.seek(0) # Reset stream position - sheets_dict = pd.read_excel(table_stream, sheet_name=None) - precision_mode_active = False - else: - # LEGACY MODE: Use pandas read_excel + LLM/heuristic header detection - sheets_dict = pd.read_excel(table_stream, sheet_name=None) - precision_mode_active = False - - all_sheets = sheets_dict.items() - - for sheet_name, sheet_content in all_sheets: - sheet_name = sheet_name.strip() - if sheet_name in exist_sheets: - sheet_name = sheet_name + str(len(exist_sheets)) - else: - exist_sheets.append(sheet_name) - - sheet_tbs = [sheet_content] - for tb in sheet_tbs: - try: - tb = postprocess_tb(tb, drop=True) - if len(tb) == 0 or tb.empty or tb.isna().all().all(): - continue - - # In precision mode, headers are already set by parse_excel_structure - # In legacy mode, use LLM/heuristic header parsing - if not precision_mode_active: - tb = parse_headers(tb, paras=base_llm_paras) - - # Drop _src_row column before converting to HTML/keywords - # (_src_row is a debug column added by Excel structure parsing for cross-referencing) - src_row_cols = [ - c - for c in tb.columns - if (isinstance(c, tuple) and c[0] == "_src_row") or c == "_src_row" - ] - if src_row_cols: - tb = tb.drop(columns=src_row_cols) - - # Get row header column count from DataFrame attrs (set in parse_excel_structure) - row_header_cols = tb.attrs.get("row_header_cols", 0) - - tb_paths, tb_strs = parse_tb_contents( - tb, - parent_dic={file_name: {sheet_name: {}}}, - file_name=file_name, - sheet_name=sheet_name, - row_header_cols=row_header_cols, - ) - - # Unified LLM extraction: title + keywords + summary in one call - # (consistent with doc_parser.py and md_parser.py) - llm_title = None - llm_summary = None - tb_keywords = "" - if base_llm_paras["summary_table"]: - from app.services.document_parser.txt_parser import ( - extract_title_keywords_summary, - ) - - llm_title, tb_keywords, llm_summary = ( - extract_title_keywords_summary(tb_strs, max_keywords=3) - ) - - # Build tb_summary: table index + optional LLM summary - table_index = f"table-{sheet_name}" - if llm_summary: - tb_summary = f"{table_index}\n{llm_summary}" - else: - # Fallback: use mechanical column keywords when LLM is off - tb_keywords_fallback = parse_tb_keywords(tb) - tb_summary = table_index - tb_keywords = tb_keywords if tb_keywords else tb_keywords_fallback - - # Use a filesystem-safe filename so LLM titles like "A/B" do not - # accidentally create nested paths under tables/. - effective_name = llm_title if llm_title else sheet_name - tb_name = ( - path_handle( - remove_spaces("table-" + effective_name), mode="clean_single" - ) - + ".html" - ) - soup = BeautifulSoup(tb_strs, features="html.parser") - tb_html_str = soup.prettify() - - # Use same temp_uid for both marker and know_id (aligned with doc_parser/md_parser) - temp_uid = gen_str_codes(tb_strs + str(sheet_name)) - tb_ref = build_chunk_ref(f"tables/{tb_name}") - tb_bottom_content = f"{tb_ref}\nTable summary:\n{tb_summary}\nMain columns:\n{tb_keywords}" - - bottom_tokens = tokenize2stw_remove( - [tb_bottom_content], base_llm_paras["stopwords"] - ) - - all_tb_paths.extend(tb_paths) - table_row = write_table_asset( - TableAssetInput( - html=tb_html_str, - output_dir=output_dir, - table_name=tb_name, - summary=tb_summary, - keywords=tb_keywords, - know_id=temp_uid, - addtime=time_stamp, - content=tb_bottom_content, - tokens=bottom_tokens, - length=len(tb_strs), - ) - ) - df_list.append(table_row.to_list()) - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Table parsing failed: {e}") - raise TableParsingException( - user_message="Failed to parse Excel table content", - reason="TABLE_PROCESSING_FAILED", - internal_message=str(e), - original_exception=e, - ) - - rows_builder = ParsedRowsBuilder() - for row_values in df_list: - rows_builder.append( - ParsedRow( - content=str(row_values[0]), - path=str(row_values[1]), - type=str(row_values[2]), - length=int(row_values[3]), - keywords=str(row_values[4]), - summary=str(row_values[5]), - know_id=str(row_values[6]), - tokens=str(row_values[7]), - connectto=str(row_values[8]), - addtime=str(row_values[9]), - page_nums=str(row_values[10]), - ) - ) - table_df = rows_builder.to_dataframe() - table_df = process_dup_paths_df(table_df) - return table_df + file_path: str, + file_name: str, + output_dir: str, + baseurl: str, + base_llm_paras: dict | None = None, + window_h: int = 10, + relative_root: str | None = None, + use_precision_mode: bool = True, + include_hidden_sheets: bool = False, +) -> pd.DataFrame: + from app.services.document_parser.excel_table_parser import ( + parse_xlsx as parse_excel_xlsx, + ) + + return parse_excel_xlsx( + file_path=file_path, + file_name=file_name, + output_dir=output_dir, + baseurl=baseurl, + base_llm_paras=base_llm_paras, + window_h=window_h, + relative_root=relative_root, + use_precision_mode=use_precision_mode, + include_hidden_sheets=include_hidden_sheets, + ) diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py new file mode 100644 index 000000000..0f86e6e96 --- /dev/null +++ b/apps/worker/tests/contract/test_excel_parser_contract.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from pathlib import Path + +from pytest import MonkeyPatch + + +def _write_contract_workbook(workbook_path: Path) -> None: + import openpyxl + + workbook = openpyxl.Workbook() + visible_sheet = workbook.active + visible_sheet.title = "Visible" + visible_sheet["A1"] = "Region" + visible_sheet["B1"] = "Value" + visible_sheet["A2"] = "North" + visible_sheet["B2"] = 10 + + hidden_sheet = workbook.create_sheet("Hidden") + hidden_sheet.sheet_state = "hidden" + hidden_sheet["A1"] = "Secret" + hidden_sheet["B1"] = "Value" + hidden_sheet["A2"] = "Hidden" + hidden_sheet["B2"] = 99 + + workbook.save(workbook_path) + + +def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.parse_service import checkerboard_inject_parse + + workbook_path = tmp_path / "budget.xlsx" + _write_contract_workbook(workbook_path) + + full_output_dir, parsed_df = checkerboard_inject_parse( + file_full_path=str(workbook_path), + filename="budget.xlsx", + output_dir=str(tmp_path), + internal_output_filename="budget.xlsx", + summary_image=False, + summary_table=False, + summary_txt=False, + smart_title_parse=False, + stopwords=[], + ) + + assert full_output_dir.endswith("Default_Root/budget.xlsx") + assert parsed_df is not None + assert parsed_df["type"].tolist() == ["table"] + assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] + assert parsed_df["summary"].tolist() == ["table-Visible"] + assert "Region" in parsed_df["keywords"].iloc[0] + assert "Value" in parsed_df["keywords"].iloc[0] + + table_html = Path(full_output_dir) / "tables" / "table-Visible.html" + table_html_text = table_html.read_text(encoding="utf-8") + + assert "North" in table_html_text + assert "10" in table_html_text + assert "Secret" not in table_html_text + assert "Hidden" not in table_html_text + + +def test_xlsx_parser_contract_accepts_missing_llm_parameters( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.excel_table_parser import parse_xlsx + + workbook_path = tmp_path / "default-parameters.xlsx" + output_dir = tmp_path / "output" + _write_contract_workbook(workbook_path) + + parsed_df = parse_xlsx( + file_path=str(workbook_path), + file_name="default-parameters.xlsx", + output_dir=str(output_dir), + baseurl="", + base_llm_paras=None, + ) + + assert parsed_df["type"].tolist() == ["table"] + assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] + assert (output_dir / "tables" / "table-Visible.html").exists() + + +def test_xlsx_parser_contract_falls_back_to_column_keywords_when_llm_summary_is_empty( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, + tmp_path: Path, +) -> None: + import app.services.document_parser.txt_parser as txt_parser + from app.services.document_parser.excel_table_parser import parse_xlsx + + workbook_path = tmp_path / "empty-summary.xlsx" + output_dir = tmp_path / "output" + _write_contract_workbook(workbook_path) + + monkeypatch.setattr( + txt_parser, + "extract_title_keywords_summary", + lambda *_args, **_kwargs: (None, "", ""), + ) + + parsed_df = parse_xlsx( + file_path=str(workbook_path), + file_name="empty-summary.xlsx", + output_dir=str(output_dir), + baseurl="", + base_llm_paras={"summary_table": True, "stopwords": []}, + ) + + keywords = str(parsed_df["keywords"].iloc[0]) + content = str(parsed_df["content"].iloc[0]) + + assert "Region" in keywords + assert "Value" in keywords + assert "Main columns:" in content + assert "Region" in content + assert "Value" in content diff --git a/packages/shared-python/shared/testing/contract_runtime.py b/packages/shared-python/shared/testing/contract_runtime.py index 11aa49efb..53c139510 100644 --- a/packages/shared-python/shared/testing/contract_runtime.py +++ b/packages/shared-python/shared/testing/contract_runtime.py @@ -358,6 +358,24 @@ def clear_application_modules() -> None: sys.modules.pop(module_name, None) continue + if module_name == "shared.services.storage" or module_name.startswith( + "shared.services.storage." + ): + sys.modules.pop(module_name, None) + continue + + if module_name == "shared.services.jobs" or module_name.startswith( + "shared.services.jobs." + ): + sys.modules.pop(module_name, None) + continue + + if module_name == "shared.services.webhook" or module_name.startswith( + "shared.services.webhook." + ): + sys.modules.pop(module_name, None) + continue + if module_name == "app" or module_name.startswith("app."): sys.modules.pop(module_name, None) From eaf78c7623ac7e0a9006a95c9f86e970692f4109 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:55:17 +0000 Subject: [PATCH 38/40] refactor(worker): deepen markdown table parser modules --- .../services/document_parser/doc_parser.py | 2 +- .../document_parser/excel_table_parser.py | 2 +- .../services/document_parser/layout_parser.py | 2 +- .../markdown_deferred_summary.py | 92 +-- .../document_parser/markdown_deferred_task.py | 33 + .../document_parser/markdown_image_asset.py | 229 +++++++ .../document_parser/markdown_parse_state.py | 12 +- .../document_parser/markdown_table_asset.py | 101 +++ .../app/services/document_parser/md_parser.py | 278 ++------ .../document_parser/table_frame_parser.py | 422 ++++++++++++ .../services/document_parser/table_parser.py | 642 +----------------- .../document_parser/table_text_parser.py | 149 ++++ .../services/document_parser/toc_hierarchy.py | 2 +- .../services/document_parser/toc_parser.py | 2 +- ...t_document_parser_architecture_contract.py | 229 ++++++- 15 files changed, 1309 insertions(+), 888 deletions(-) create mode 100644 apps/worker/app/services/document_parser/markdown_deferred_task.py create mode 100644 apps/worker/app/services/document_parser/markdown_image_asset.py create mode 100644 apps/worker/app/services/document_parser/markdown_table_asset.py create mode 100644 apps/worker/app/services/document_parser/table_frame_parser.py create mode 100644 apps/worker/app/services/document_parser/table_text_parser.py diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index 956cc3b57..77238675d 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -28,7 +28,7 @@ ask_image, perceptual_hash, ) -from app.services.document_parser.table_parser import sanitize_table_name_from_header +from app.services.document_parser.table_text_parser import sanitize_table_name_from_header from app.services.document_parser.toc_docx import build_docx_toc_hierarchies from app.services.document_parser.txt_parser import postprocess_leaf_dics from docx.text.paragraph import Paragraph diff --git a/apps/worker/app/services/document_parser/excel_table_parser.py b/apps/worker/app/services/document_parser/excel_table_parser.py index d1d5589c0..1770a4163 100644 --- a/apps/worker/app/services/document_parser/excel_table_parser.py +++ b/apps/worker/app/services/document_parser/excel_table_parser.py @@ -16,7 +16,7 @@ TableAssetInput, write_table_asset, ) -from app.services.document_parser.table_parser import ( +from app.services.document_parser.table_frame_parser import ( parse_headers, parse_tb_contents, parse_tb_keywords, diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py index 6a3b3a8bf..8d9e94ed5 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/layout_parser.py @@ -24,7 +24,7 @@ tree_to_dataframe as heading_tree_to_dataframe, ) from app.services.document_parser.stage_profiler import stage_timer -from app.services.document_parser.table_parser import df2md +from app.services.document_parser.table_text_parser import df2md from gevent.pool import Pool as GeventPool try: diff --git a/apps/worker/app/services/document_parser/markdown_deferred_summary.py b/apps/worker/app/services/document_parser/markdown_deferred_summary.py index b11da7c80..270303e90 100644 --- a/apps/worker/app/services/document_parser/markdown_deferred_summary.py +++ b/apps/worker/app/services/document_parser/markdown_deferred_summary.py @@ -3,12 +3,18 @@ import os import re from dataclasses import dataclass -from typing import Any, TypeGuard +from typing import Literal, TypeGuard import gevent +from app.services.document_parser.markdown_deferred_task import ( + ImageDeferredSummaryTask, + MarkdownDeferredSummaryTask, + TableDeferredSummaryTask, + TextDeferredSummaryTask, +) from app.services.document_parser.image_parser import _get_vision_client, ask_image from app.services.document_parser.stage_profiler import stage_timer -from app.services.document_parser.table_parser import sanitize_table_name_from_header +from app.services.document_parser.table_text_parser import sanitize_table_name_from_header from app.services.document_parser.txt_parser import ( extract_title_keywords_summary, split_title_summary, @@ -23,7 +29,7 @@ DeferredResult = ( tuple[ int, - str, + Literal["image", "table", "text"], tuple[str | None, str | None] | tuple[str, str, str] | tuple[str, str], ] ) @@ -35,7 +41,7 @@ @dataclass(frozen=True) class MarkdownDeferredSummaryInput: rows: list[list[str | int]] - tasks: list[tuple[Any, ...]] + tasks: list[MarkdownDeferredSummaryTask] output_dir: str summary_len: int = 1500 @@ -46,9 +52,15 @@ def apply_markdown_deferred_summaries( if not deferred_input.tasks: return - image_task_count = sum(1 for task in deferred_input.tasks if task[0] == "image") - table_task_count = sum(1 for task in deferred_input.tasks if task[0] == "table") - text_task_count = sum(1 for task in deferred_input.tasks if task[0] == "text") + image_task_count = sum( + 1 for task in deferred_input.tasks if isinstance(task, ImageDeferredSummaryTask) + ) + table_task_count = sum( + 1 for task in deferred_input.tasks if isinstance(task, TableDeferredSummaryTask) + ) + text_task_count = sum( + 1 for task in deferred_input.tasks if isinstance(task, TextDeferredSummaryTask) + ) logger.info( f"Running {len(deferred_input.tasks)} deferred summary LLM calls in parallel" ) @@ -101,47 +113,43 @@ def _run_deferred_summary_tasks( def _run_deferred_summary_task( - task: tuple[Any, ...], + task: MarkdownDeferredSummaryTask, deferred_input: MarkdownDeferredSummaryInput, ) -> DeferredResult | None: - task_type, row_index = task[0], task[1] try: - if task_type == "image": - relative_path = task[2] + if isinstance(task, ImageDeferredSummaryTask): client = _get_vision_client() # TODO: Risk of missing text content if MinerU outputted a pure text image. # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. llm_resp = ask_image( - client, deferred_input.output_dir, paths_=[relative_path] + client, deferred_input.output_dir, paths_=[task.relative_path] ) if llm_resp: img_title, img_summary = split_title_summary(llm_resp) else: img_title, img_summary = None, None - return row_index, task_type, (img_title, img_summary) + return task.row_index, "image", (img_title, img_summary) - if task_type == "table": - table_html = task[2] + if isinstance(task, TableDeferredSummaryTask): title, keywords, summary = extract_title_keywords_summary( - table_html, max_keywords=3 + task.table_html, max_keywords=3 ) - return row_index, task_type, (title, keywords, summary) + return task.row_index, "table", (title, keywords, summary) - if task_type == "text": - text_content = task[2] + if isinstance(task, TextDeferredSummaryTask): _, keywords, summary = extract_title_keywords_summary( - text_content, + task.content, max_keywords=3, summary_len=deferred_input.summary_len, ) - return row_index, task_type, (keywords, summary) + return task.row_index, "text", (keywords, summary) except Exception as exc: logger.warning( - f"Deferred {task_type} LLM call failed for idx={row_index}: {exc}" + f"Deferred summary LLM call failed for idx={task.row_index}: {exc}" ) return None - logger.warning(f"Unknown deferred markdown summary task type: {task_type}") + logger.warning(f"Unknown deferred markdown summary task type: {type(task).__name__}") return None @@ -149,7 +157,7 @@ def _apply_deferred_summary_results( deferred_input: MarkdownDeferredSummaryInput, results: list[DeferredResult | None], ) -> None: - deferred_by_index = {task[1]: task for task in deferred_input.tasks} + deferred_by_index = {task.row_index: task for task in deferred_input.tasks} for result in results: if result is None: @@ -162,7 +170,7 @@ def _apply_deferred_summary_results( continue _apply_image_summary_result( deferred_input.rows, - deferred_by_index[row_index], + _get_image_task(deferred_by_index[row_index]), row_index, task_result, ) @@ -172,7 +180,7 @@ def _apply_deferred_summary_results( continue _apply_table_summary_result( deferred_input.rows, - deferred_by_index[row_index], + _get_table_task(deferred_by_index[row_index]), row_index, task_result, ) @@ -207,9 +215,21 @@ def _is_text_summary_result(result: object) -> TypeGuard[TextSummaryResult]: ) +def _get_image_task(task: MarkdownDeferredSummaryTask) -> ImageDeferredSummaryTask: + if isinstance(task, ImageDeferredSummaryTask): + return task + raise TypeError(f"Expected image deferred task, got {type(task).__name__}") + + +def _get_table_task(task: MarkdownDeferredSummaryTask) -> TableDeferredSummaryTask: + if isinstance(task, TableDeferredSummaryTask): + return task + raise TypeError(f"Expected table deferred task, got {type(task).__name__}") + + def _apply_image_summary_result( rows: list[list[str | int]], - original_task: tuple[Any, ...], + original_task: ImageDeferredSummaryTask, row_index: int, result: ImageSummaryResult, ) -> None: @@ -222,11 +242,9 @@ def _apply_image_summary_result( if not img_title: return - image_dir, old_img_name, image_suffix = ( - original_task[3], - original_task[4], - original_task[5], - ) + image_dir = original_task.image_dir + old_img_name = original_task.image_name + image_suffix = original_task.image_suffix safe_title = path_handle(str(img_title), mode="clean_single") img_num_match = re.match(r"image-(\d+)", str(old_img_name)) img_num = ( @@ -250,7 +268,7 @@ def _apply_image_summary_result( def _apply_table_summary_result( rows: list[list[str | int]], - original_task: tuple[Any, ...], + original_task: TableDeferredSummaryTask, row_index: int, result: TableSummaryResult, ) -> None: @@ -264,11 +282,9 @@ def _apply_table_summary_result( if not title: return - table_dir, old_table_name, table_count = ( - original_task[3], - original_task[4], - original_task[5], - ) + table_dir = original_task.table_dir + old_table_name = original_task.table_name + table_count = original_task.table_count safe_title = sanitize_table_name_from_header(str(title)) new_table_name = path_handle( f"table-{table_count} {safe_title}", mode="clean_single" diff --git a/apps/worker/app/services/document_parser/markdown_deferred_task.py b/apps/worker/app/services/document_parser/markdown_deferred_task.py new file mode 100644 index 000000000..da935abaa --- /dev/null +++ b/apps/worker/app/services/document_parser/markdown_deferred_task.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TypeAlias + + +@dataclass(frozen=True) +class ImageDeferredSummaryTask: + row_index: int + relative_path: str + image_dir: str + image_name: str + image_suffix: str + + +@dataclass(frozen=True) +class TableDeferredSummaryTask: + row_index: int + table_html: str + table_dir: str + table_name: str + table_count: int + + +@dataclass(frozen=True) +class TextDeferredSummaryTask: + row_index: int + content: str + + +MarkdownDeferredSummaryTask: TypeAlias = ( + ImageDeferredSummaryTask | TableDeferredSummaryTask | TextDeferredSummaryTask +) diff --git a/apps/worker/app/services/document_parser/markdown_image_asset.py b/apps/worker/app/services/document_parser/markdown_image_asset.py new file mode 100644 index 000000000..10c93b5dc --- /dev/null +++ b/apps/worker/app/services/document_parser/markdown_image_asset.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import cast + +from app.services.document_parser.identifiers import gen_str_codes +from app.services.document_parser.image_parser import perceptual_hash +from app.services.document_parser.inline_asset import build_image_asset_row +from app.services.document_parser.markdown_deferred_task import ( + ImageDeferredSummaryTask, + MarkdownDeferredSummaryTask, +) +from app.services.document_parser.markdown_parse_state import ParserRowValues +from loguru import logger + +from shared.utils.chunk_refs import build_chunk_ref +from shared.utils.file_utils import path_handle + + +@dataclass(frozen=True) +class MarkdownImageAsset: + content_item: str | None + row_values: ParserRowValues | None + cache_key: str | None + cache_entry: dict[str, str] | None + deferred_task: MarkdownDeferredSummaryTask | None + should_advance_image_count: bool + + +@dataclass(frozen=True) +class MarkdownImageAssetRequest: + output_dir: str + image_dir: str + image_path: str + image_name: str + image_count: int + last_context: str + image_summary: str | None + timestamp: str + current_page_number: int + seen_images: dict[str, dict[str, str]] + summary_image: bool + row_index: int + + +def build_markdown_image_asset( + request: MarkdownImageAssetRequest, +) -> MarkdownImageAsset: + image_suffix = os.path.splitext(request.image_path)[-1] + source_path = resolve_markdown_image_source_path( + request.output_dir, + request.image_path, + ) + if source_path is None or not source_path.exists(): + logger.warning(f"Image file not found, skipping rename: {request.image_path}") + return _empty_asset(should_advance_image_count=True) + + with open(source_path, "rb") as image_file: + image_binary_hash = perceptual_hash(image_file.read()) + + if image_binary_hash in request.seen_images: + return _build_duplicate_image_asset( + source_path=source_path, + cache_entry=request.seen_images[image_binary_hash], + timestamp=request.timestamp, + current_page_number=request.current_page_number, + ) + + relative_image_path = f"images/{request.image_name}{image_suffix}" + target_image_path = os.path.join( + request.image_dir, + f"{request.image_name}{image_suffix}", + ) + os.rename(source_path, target_image_path) + + image_index = f"image-{request.image_count}" + effective_summary = request.image_summary or request.last_context or None + image_summary_field = ( + f"{image_index}\n{effective_summary}" if effective_summary else image_index + ) + image_content = _build_image_content( + relative_image_path=relative_image_path, + summary=effective_summary, + ) + image_know_id = gen_str_codes(image_binary_hash) + row_values = _build_image_row_values( + content=image_content, + relative_path=relative_image_path, + summary=image_summary_field, + know_id=image_know_id, + timestamp=request.timestamp, + current_page_number=request.current_page_number, + ) + cache_entry = { + "relative_img_path": relative_image_path, + "img_content": image_content, + "img_summary_field": image_summary_field, + "temp_uid": image_know_id, + } + + deferred_task = None + if request.summary_image: + deferred_task = ImageDeferredSummaryTask( + row_index=request.row_index, + relative_path=relative_image_path, + image_dir=request.image_dir, + image_name=request.image_name, + image_suffix=image_suffix, + ) + + return MarkdownImageAsset( + content_item=image_content, + row_values=row_values, + cache_key=image_binary_hash, + cache_entry=cache_entry, + deferred_task=deferred_task, + should_advance_image_count=True, + ) + + +def build_markdown_image_name(*, image_count: int, last_context: str) -> str: + image_name_context = path_handle(last_context[:10], mode="clean_single") + return f"image-{str(image_count)}-{image_name_context}" + + +def resolve_workspace_image_path( + candidate_path: Path, workspace_path: Path, +) -> Path | None: + """Return the candidate only when it exists inside the current job workspace.""" + resolved_path = candidate_path.resolve(strict=False) + try: + resolved_path.relative_to(workspace_path) + except ValueError: + return None + return resolved_path if resolved_path.exists() else None + + +def resolve_markdown_image_source_path(output_dir: str, image_path: str) -> Path | None: + """Handle local absolute refs and container cwd-relative refs safely.""" + if not image_path: + return None + + workspace_path = Path(output_dir).resolve() + raw_path = Path(image_path).expanduser() + candidate_paths = ( + [raw_path] + if raw_path.is_absolute() + else [ + workspace_path / raw_path, + Path.cwd() / raw_path, + ] + ) + + for candidate_path in candidate_paths: + resolved_path = resolve_workspace_image_path(candidate_path, workspace_path) + if resolved_path is not None: + return resolved_path + + return None + + +def _build_duplicate_image_asset( + *, + source_path: Path, + cache_entry: dict[str, str], + timestamp: str, + current_page_number: int, +) -> MarkdownImageAsset: + row_values = _build_image_row_values( + content=cache_entry["img_content"], + relative_path=cache_entry["relative_img_path"], + summary=cache_entry["img_summary_field"], + know_id=cache_entry["temp_uid"], + timestamp=timestamp, + current_page_number=current_page_number, + ) + try: + source_path.unlink() + except OSError: + pass + logger.debug("Skipped duplicate image") + return MarkdownImageAsset( + content_item=cache_entry["img_content"], + row_values=row_values, + cache_key=None, + cache_entry=None, + deferred_task=None, + should_advance_image_count=False, + ) + + +def _build_image_content(*, relative_image_path: str, summary: str | None) -> str: + image_reference = build_chunk_ref(relative_image_path) + if summary: + return f"\n{summary}\n{image_reference}\n" + return f"\n{image_reference}\n" + + +def _build_image_row_values( + *, + content: str, + relative_path: str, + summary: str, + know_id: str, + timestamp: str, + current_page_number: int, +) -> ParserRowValues: + image_row = build_image_asset_row( + content=content, + relative_path=relative_path, + summary=summary, + know_id=know_id, + addtime=timestamp, + page_nums=str(current_page_number) if current_page_number > 0 else "", + ) + return cast(ParserRowValues, image_row.to_list()) + + +def _empty_asset(*, should_advance_image_count: bool) -> MarkdownImageAsset: + return MarkdownImageAsset( + content_item=None, + row_values=None, + cache_key=None, + cache_entry=None, + deferred_task=None, + should_advance_image_count=should_advance_image_count, + ) diff --git a/apps/worker/app/services/document_parser/markdown_parse_state.py b/apps/worker/app/services/document_parser/markdown_parse_state.py index eb8837f0f..51f97b954 100644 --- a/apps/worker/app/services/document_parser/markdown_parse_state.py +++ b/apps/worker/app/services/document_parser/markdown_parse_state.py @@ -8,6 +8,10 @@ import pandas as pd from app.services.document_parser.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.markdown_deferred_task import ( + MarkdownDeferredSummaryTask, + TextDeferredSummaryTask, +) from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder ParserRowValues = list[str | int] @@ -36,7 +40,7 @@ class MarkdownParseState: base_level: int | None = None path: str = "" path_counter: dict[str, int] = field(default_factory=dict) - deferred_llm_tasks: list[tuple[Any, ...]] = field(default_factory=list) + deferred_llm_tasks: list[MarkdownDeferredSummaryTask] = field(default_factory=list) seen_images: dict[str, dict[str, str]] = field(default_factory=dict) image_count: int = 1 table_count: int = 1 @@ -137,7 +141,7 @@ def append_plain_text(self, text: str) -> None: def append_row(self, row: ParserRowValues) -> None: self.rows.append(row) - def schedule_deferred_task(self, task: tuple[Any, ...]) -> None: + def schedule_deferred_task(self, task: MarkdownDeferredSummaryTask) -> None: self.deferred_llm_tasks.append(task) def collect_text_summary_tasks(self, summary_len: int) -> None: @@ -153,7 +157,9 @@ def collect_text_summary_tasks(self, summary_len: int) -> None: continue content = str(entry[0]) if len(content) > summary_len and not entry[4] and not entry[5]: - self.deferred_llm_tasks.append(("text", index, content)) + self.deferred_llm_tasks.append( + TextDeferredSummaryTask(row_index=index, content=content) + ) def to_dataframe(self) -> pd.DataFrame: rows_builder = ParsedRowsBuilder() diff --git a/apps/worker/app/services/document_parser/markdown_table_asset.py b/apps/worker/app/services/document_parser/markdown_table_asset.py new file mode 100644 index 000000000..a01e8ce97 --- /dev/null +++ b/apps/worker/app/services/document_parser/markdown_table_asset.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import cast + +from app.services.document_parser.html_parser import first_cols_rows_html +from app.services.document_parser.identifiers import gen_str_codes +from app.services.document_parser.inline_asset import build_table_asset_row +from app.services.document_parser.markdown_deferred_task import ( + MarkdownDeferredSummaryTask, + TableDeferredSummaryTask, +) +from app.services.document_parser.markdown_parse_state import ParserRowValues +from app.services.document_parser.table_text_parser import sanitize_table_name_from_header + +from shared.utils.chunk_refs import build_chunk_ref +from shared.utils.file_utils import path_handle + + +@dataclass(frozen=True) +class MarkdownTableAsset: + content_item: str + row_values: ParserRowValues + deferred_task: MarkdownDeferredSummaryTask | None + relative_path: str + + +@dataclass(frozen=True) +class MarkdownTableAssetRequest: + table_html: str + table_dir: str + table_count: int + timestamp: str + current_page_number: int + summary_table: bool + row_index: int + + +def build_markdown_table_asset( + request: MarkdownTableAssetRequest, +) -> MarkdownTableAsset: + first_row_text, _first_col_text = first_cols_rows_html(request.table_html) + table_index = f"table-{request.table_count}" + + raw_table_name = ( + sanitize_table_name_from_header(first_row_text) if first_row_text else "" + ) + table_name = _sanitize_table_file_stem( + f"table-{str(request.table_count)} {raw_table_name}" + ) + relative_table_path = f"tables/{table_name}.html" + table_ref = build_chunk_ref(relative_table_path) + table_content_item = f"\n{table_ref}\n" + table_path = os.path.join(request.table_dir, f"{table_name}.html") + _write_table_html(table_path=table_path, table_html=request.table_html) + + table_row = build_table_asset_row( + content=request.table_html, + relative_path=relative_table_path, + summary=table_index, + keywords="", + know_id=gen_str_codes((request.table_html + str(request.table_count))), + addtime=request.timestamp, + page_nums=str(request.current_page_number) + if request.current_page_number > 0 + else "", + ) + + deferred_task = None + if request.summary_table: + deferred_task = TableDeferredSummaryTask( + row_index=request.row_index, + table_html=request.table_html, + table_dir=request.table_dir, + table_name=table_name, + table_count=request.table_count - 1, + ) + + return MarkdownTableAsset( + content_item=table_content_item, + row_values=cast(ParserRowValues, table_row.to_list()), + deferred_task=deferred_task, + relative_path=relative_table_path, + ) + + +def _sanitize_table_file_stem(raw_name: str) -> str: + table_name = path_handle(raw_name, mode="clean_single") + if not isinstance(table_name, str) or not table_name: + raise ValueError(f"Failed to sanitize Markdown table name: {raw_name}") + return table_name + + +def _write_table_html(*, table_path: str, table_html: str) -> None: + table_html_with_border = table_html.replace("", "
").replace( + "
Path | None: - """Return the candidate only when it exists inside the current job workspace.""" - resolved_path = candidate_path.resolve(strict=False) - try: - resolved_path.relative_to(workspace_path) - except ValueError: - return None - return resolved_path if resolved_path.exists() else None - - -def resolve_markdown_image_source_path(output_dir: str, img_path: str) -> Path | None: - """Handle local absolute refs and container cwd-relative refs safely.""" - if not img_path: - return None - - workspace_path = Path(output_dir).resolve() - raw_path = Path(img_path).expanduser() - candidate_paths = ( - [raw_path] - if raw_path.is_absolute() - else [ - workspace_path / raw_path, - Path.cwd() / raw_path, - ] - ) - - for candidate_path in candidate_paths: - resolved_path = resolve_workspace_image_path(candidate_path, workspace_path) - if resolved_path is not None: - return resolved_path - - return None - - def find_surround_context(md_lines, lid): def is_skip(line): s = line.strip() @@ -365,112 +329,47 @@ def parse_md( else: # no path change, remain in the same hierarchy # a. handle lines containing images (LLM deferred to post-loop parallel batch) - img_name_context = path_handle(last_context[:10], mode="clean_single") - img_name = f"image-{str(parser_state.image_count)}-{img_name_context}" # Always skip inline LLM — vision calls are deferred to parallel batch imgs = detect_summary_img_md(line, last_context, output_dir, mode=False) - - for img_path, img_title, img_summary in imgs: - img_suffix = os.path.splitext(img_path)[-1] - update_img_path = os.path.join(img_dir, f"{img_name}{img_suffix}") - - # Check if source image file exists before renaming - source_path = resolve_markdown_image_source_path(output_dir, img_path) - if source_path is None or not source_path.exists(): - logger.warning(f"Image file not found, skipping rename: {img_path}") - parser_state.image_count += 1 - continue - - # Document-level dedup: perceptual hash for visual duplicates - with open(source_path, "rb") as f: - img_binary_hash = perceptual_hash(f.read()) - - if img_binary_hash in parser_state.seen_images: - cached = parser_state.seen_images[img_binary_hash] - parser_state.append_content_item(cached["img_content"]) - parser_state.append_row( - build_image_asset_row( - content=cached["img_content"], - relative_path=cached["relative_img_path"], - summary=cached["img_summary_field"], - know_id=cached["temp_uid"], - addtime=parser_state.timestamp, - page_nums=str(parser_state.current_page_number) - if parser_state.current_page_number > 0 - else "", - ).to_list() - ) - logger.debug( - f"Skipped duplicate image (hash={img_binary_hash[:12]}...)" + image_name = build_markdown_image_name( + image_count=parser_state.image_count, + last_context=last_context, + ) + + for img_path, _img_title, img_summary in imgs: + image_asset = build_markdown_image_asset( + MarkdownImageAssetRequest( + output_dir=output_dir, + image_dir=img_dir, + image_path=img_path, + image_name=image_name, + image_count=parser_state.image_count, + last_context=last_context, + image_summary=img_summary, + timestamp=parser_state.timestamp, + current_page_number=parser_state.current_page_number, + seen_images=parser_state.seen_images, + summary_image=bool(base_llm_paras["summary_image"]), + row_index=len(parser_state.rows), ) - # Remove unused source file since we reuse the cached image - try: - source_path.unlink() - except OSError: - pass - continue - - os.rename(source_path, update_img_path) - - # Image index (always present) - image_index = f"image-{parser_state.image_count}" - - # Fallback: LLM summary -> last_context -> None - effective_summary = img_summary or last_context or None - - # Deterministic know_id: use image binary hash - temp_uid = gen_str_codes(img_binary_hash) - relative_img_path = f"images/{img_name}{img_suffix}" - img_ref = build_chunk_ref(relative_img_path) - - # Build img_summary_field for df_list: image-n + optional summary - if effective_summary: - img_summary_field = f"{image_index}\n{effective_summary}" - else: - img_summary_field = image_index - - # Build image_ref for content: optional summary + image path ref - if effective_summary: - img_content = f"\n{effective_summary}\n{img_ref}\n" - else: - img_content = f"\n{img_ref}\n" - - parser_state.append_content_item(img_content) - - parser_state.append_row( - build_image_asset_row( - content=img_content, - relative_path=relative_img_path, - summary=img_summary_field, - know_id=temp_uid, - addtime=parser_state.timestamp, - page_nums=str(parser_state.current_page_number) - if parser_state.current_page_number > 0 - else "", - ).to_list() ) - - # Cache result for document-level dedup - parser_state.seen_images[img_binary_hash] = { - "relative_img_path": relative_img_path, - "img_content": img_content, - "img_summary_field": img_summary_field, - "temp_uid": temp_uid, - } - - if base_llm_paras["summary_image"]: - # Store img_dir, img_name, img_suffix for post-loop rename (mirrors table deferred task) - parser_state.schedule_deferred_task( - ( - "image", - len(parser_state.rows) - 1, - relative_img_path, - img_dir, - img_name, - img_suffix, - ) + if ( + image_asset.content_item is not None + and image_asset.row_values is not None + ): + parser_state.append_content_item(image_asset.content_item) + parser_state.append_row(image_asset.row_values) + if ( + image_asset.cache_key is not None + and image_asset.cache_entry is not None + ): + parser_state.seen_images[image_asset.cache_key] = ( + image_asset.cache_entry ) - parser_state.image_count += 1 + if image_asset.deferred_task is not None: + parser_state.schedule_deferred_task(image_asset.deferred_task) + if image_asset.should_advance_image_count: + parser_state.image_count += 1 # TODO for large and dense tables, such as "Epstein flight logs", # integrate tabula-py as an independent extraction path to solve VLM hallucinations and misplacement @@ -502,76 +401,21 @@ def parse_md( else: continue # Unknown form, skip - # Extract first row and first column for fallback file naming only - first_row_text, first_col_text = first_cols_rows_html(tb_str) - - # Table index (always present) - table_index = f"table-{parser_state.table_count}" - - # LLM title + keywords + summary deferred to post-loop parallel batch - llm_title = None - llm_summary = None - tb_keywords = "" - - # Build tb_summary for df_list: table-n + optional LLM summary - if llm_summary: - tb_summary = f"{table_index}\n{llm_summary}" - else: - tb_summary = table_index - - raw_tb_name = ( - sanitize_table_name_from_header(first_row_text) - if first_row_text - else "" - ) - # Use LLM title for filename when available, fallback to sanitized header - effective_name = llm_title if llm_title else raw_tb_name - tb_name = path_handle( - f"table-{str(parser_state.table_count)} {effective_name}", - mode="clean_single", - ) - temp_uid = gen_str_codes((tb_str + str(parser_state.table_count))) - - relative_tb_path = f"tables/{tb_name}.html" - tb_ref = build_chunk_ref(relative_tb_path) - - # Build table_ref for content: optional LLM summary + table path ref - if llm_summary: - parser_state.append_content_item(f"\n{llm_summary}\n{tb_ref}\n") - else: - parser_state.append_content_item(f"\n{tb_ref}\n") - tb_path = os.path.join(tb_dir, f"{tb_name}.html") - # Add border to HTML tables for consistent display - tb_str_with_border = tb_str.replace( - "
", "
" - ).replace("
0 - else "", - ).to_list() - ) - if base_llm_paras["summary_table"]: - parser_state.schedule_deferred_task( - ( - "table", - len(parser_state.rows) - 1, - tb_str, - tb_dir, - tb_name, - parser_state.table_count - 1, - ) + table_asset = build_markdown_table_asset( + MarkdownTableAssetRequest( + table_html=tb_str, + table_dir=tb_dir, + table_count=parser_state.table_count, + timestamp=parser_state.timestamp, + current_page_number=parser_state.current_page_number, + summary_table=bool(base_llm_paras["summary_table"]), + row_index=len(parser_state.rows), ) + ) + parser_state.append_content_item(table_asset.content_item) + parser_state.append_row(table_asset.row_values) + if table_asset.deferred_task is not None: + parser_state.schedule_deferred_task(table_asset.deferred_task) parser_state.table_lines = [] parser_state.table_count += 1 diff --git a/apps/worker/app/services/document_parser/table_frame_parser.py b/apps/worker/app/services/document_parser/table_frame_parser.py new file mode 100644 index 000000000..63f39feb8 --- /dev/null +++ b/apps/worker/app/services/document_parser/table_frame_parser.py @@ -0,0 +1,422 @@ +# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalOperand=false, reportOptionalSubscript=false, reportReturnType=false, reportOperatorIssue=false, reportIndexIssue=false, reportAssignmentType=false, reportGeneralTypeIssues=false +from __future__ import annotations + +import datetime +import os +import uuid +from collections import OrderedDict + +import numpy as np +import pandas as pd +from app.services.document_parser.dataframe_html_renderer import df2html +from app.services.document_parser.identifiers import gen_str_codes +from app.services.document_parser.path_helpers import flatten_dic2paths +from loguru import logger + +from shared.services.ai.prompt_service import build_prompt +from shared.services.ai.response_process_service import eval_response +from shared.utils.OpenAICompatibleClientSync import get_openai_client +from shared.utils.text_utils import remove_duplicates_orderkept + + +def parse_headers( + table_frame: pd.DataFrame, + paras: dict[str, object] | None = None, + header_window: int = 5, + smart_headers: bool = True, +) -> pd.DataFrame: + llm_parameters = paras or {"summary_table": False} + + def parse_headers_nonsmart(candidate_frame: pd.DataFrame) -> list[int]: + non_na_row = candidate_frame[candidate_frame.notna().any(axis=1)].head(1) + header_id = non_na_row.index[-1] if not non_na_row.empty else None + return list(range(header_id + 1)) + + if not pd.isna(table_frame.columns).all(): + table_frame.loc[-1] = table_frame.columns + table_frame.index = table_frame.index + 1 + table_frame = table_frame.sort_index() + table_frame.columns = [np.nan] * table_frame.shape[1] + + if llm_parameters["summary_table"] and smart_headers: + try: + table_sample = table_frame.head(header_window) + table_sample_html = df2html(table_sample) + prompt, _temperature, _top_p, _max_tokens = build_prompt( + task="detect-table-headers", + texts=table_sample_html, + query="", + paras=llm_parameters, + ) + + messages = [ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": prompt}, + ] + + context_task_id = gen_str_codes((str(uuid.uuid4()) + table_sample_html)) + + if os.getenv("LOCAL_DEBUG", "0") != "1": + from shared.services.redis.redis_sync_service import ( + SyncRedisServiceFactory, + ) + + redis_service = SyncRedisServiceFactory.get_service() + redis_service.set( + f"task:{context_task_id}:status", + "processing", + ttl=7200, + ) + + header_response = get_openai_client().chat_completion( + messages=messages, + timeout=60, + ) + parsed_response = eval_response(header_response) + if isinstance(parsed_response, dict): + answer = parsed_response.get("answer", []) + else: + answer = parsed_response if isinstance(parsed_response, list) else [] + + if not answer or len(answer) == 0: + logger.warning( + "AI returned empty list, cannot identify headers, falling back to traditional mode..." + ) + header_rows = parse_headers_nonsmart(table_frame) + else: + try: + header_id = answer[-1] + header_rows = list(range(header_id + 1)) + except Exception as exc: + logger.warning( + f"Failed to parse header row number: {exc}, falling back to traditional mode..." + ) + header_rows = parse_headers_nonsmart(table_frame) + + except Exception as exc: + logger.warning( + f"Smart header parsing failed: {exc}, falling back to traditional mode..." + ) + header_rows = parse_headers_nonsmart(table_frame) + else: + header_rows = parse_headers_nonsmart(table_frame) + + if len(header_rows) == 0 or (all(header is None for header in header_rows)): + logger.warning("No valid headers detected, fallback to using row 0 as header") + new_header = table_frame.iloc[0].ffill().bfill().tolist() + table_frame.columns = new_header + return table_frame.iloc[1:].reset_index(drop=True) + + if len(header_rows) > 1: + header_levels = [] + for header_index in range(0, len(header_rows)): + header_level = table_frame.iloc[header_index].ffill().bfill().tolist() + header_levels.append(header_level) + new_header = pd.MultiIndex.from_arrays(np.array(header_levels)) + else: + new_header = table_frame.iloc[header_rows[-1]].ffill().bfill().tolist() + + table_frame.columns = new_header + table_frame = table_frame.iloc[(header_rows[-1]) + 1 :] + return table_frame.reset_index(drop=True) + + +def parse_tb_keywords(table_frame: pd.DataFrame, kw_spit: str = ">>>") -> str: + def parse_single_level(columns: list[object], keywords: list[str]) -> list[str]: + column_texts = [str(column) for column in columns] + for column_text in column_texts: + if kw_spit in column_text: + keyword = column_text.split(">>>")[0] + else: + keyword = column_text + if keyword not in keywords: + keywords.append(column_text) + return list({keyword.strip() for keyword in keywords}) + + table_keywords: list[str] = [] + if isinstance(table_frame.columns, pd.MultiIndex): + multi_columns = table_frame.columns + columns_frame = pd.DataFrame( + multi_columns.tolist(), + columns=[f"level_{i}" for i in range(multi_columns.nlevels)], + ) + for level_index in range(multi_columns.nlevels): + level_keywords: list[str] = [] + level_keywords = parse_single_level( + columns_frame[f"level_{level_index}"].tolist(), + level_keywords, + ) + table_keywords.extend(level_keywords) + else: + table_keywords = parse_single_level(table_frame.columns.tolist(), table_keywords) + + table_keywords = remove_duplicates_orderkept(table_keywords) + table_keywords = [ + keyword + for keyword in table_keywords + if isinstance(keyword, str) + and keyword.strip() + and keyword.strip() != "nan" + and "Unnamed" not in keyword + ] + return ";".join(table_keywords) + + +def parse_tb_contents( + table_frame: pd.DataFrame, + parent_dic: dict[str, object] | None = None, + file_name: str = "", + sheet_name: str = "", + row_header_cols: int = 0, +) -> tuple[list[str], str]: + if parent_dic is None: + parent_dic = {} + + rendered_frame = table_frame.fillna("").infer_objects(copy=False) + table_html = df2html(rendered_frame, row_header_cols=row_header_cols) + + table_tree = tb_columns_to_tree(table_frame, parent_dic, file_name, sheet_name) + table_paths = flatten_dic2paths(table_tree) + return table_paths, table_html + + +def tb_columns_to_tree( + table_frame: pd.DataFrame, + parent_dic: dict[str, object], + file_name: str, + sheet_name: str, +) -> dict[str, object]: + if isinstance(table_frame.columns, pd.MultiIndex): + columns = pd.DataFrame(table_frame.columns.tolist()) + for level in range(columns.shape[1]): + columns[level] = process_duplicate_cols(columns[level]) + + new_columns = pd.MultiIndex.from_frame(columns) + tree_structure = multiindex_to_tree(new_columns) + else: + new_columns = process_duplicate_cols(table_frame.columns) + tree_structure = {column: {} for column in new_columns} + + table_frame.columns = new_columns + if (not file_name == "") and (not sheet_name == ""): + parent_dic[file_name][sheet_name] = tree_structure + elif not sheet_name == "": + parent_dic[sheet_name] = tree_structure + elif not file_name == "": + parent_dic[file_name] = tree_structure + else: + parent_dic = tree_structure + return parent_dic + + +def multiindex_to_tree(multiindex: pd.MultiIndex) -> dict[object, object]: + def tree() -> OrderedDict[object, object]: + return OrderedDict() + + root = tree() + for keys in multiindex: + current_level = root + for key in keys: + if key not in current_level: + current_level[key] = tree() + current_level = current_level[key] + + def convert_to_dict(value: object) -> object: + if isinstance(value, OrderedDict): + return {key: convert_to_dict(child) for key, child in value.items()} + return value + + return dict(convert_to_dict(root)) + + +def postprocess_tb(table_frame: pd.DataFrame, drop: bool = False) -> pd.DataFrame: + if drop: + was_range_index = isinstance(table_frame.index, pd.RangeIndex) + + source_row_columns = [ + column + for column in table_frame.columns + if (isinstance(column, tuple) and column[0] == "_src_row") + or column == "_src_row" + ] + if source_row_columns: + data_columns = [ + column for column in table_frame.columns if column not in source_row_columns + ] + mask = table_frame[data_columns].isna().all(axis=1) + table_frame = table_frame[~mask] + else: + table_frame = table_frame.dropna(how="all") + + if was_range_index: + table_frame = table_frame.reset_index(drop=True) + + cols_to_drop: list[int] = [] + for column_index, column in enumerate(table_frame.columns): + if table_frame.iloc[:, column_index].isna().all(): + has_meaningful_header = False + if isinstance(column, tuple): + for level in column: + if ( + level + and str(level).strip() + and str(level).strip() not in ["None", "nan", "NaN"] + ): + has_meaningful_header = True + break + elif ( + column + and str(column).strip() + and str(column).strip() not in ["None", "nan", "NaN"] + ): + has_meaningful_header = True + + if not has_meaningful_header: + cols_to_drop.append(column_index) + + if cols_to_drop: + cols_to_keep = [ + index + for index in range(len(table_frame.columns)) + if index not in cols_to_drop + ] + table_frame = table_frame.iloc[:, cols_to_keep] + + logger.debug(f"Dropped {len(cols_to_drop)} empty columns") + + if not isinstance(table_frame.index, pd.RangeIndex): + was_multiindex = isinstance(table_frame.columns, pd.MultiIndex) + column_level_count = table_frame.columns.nlevels if was_multiindex else 1 + existing_column_set = set(table_frame.columns) + + def make_padded(name: object) -> object: + if was_multiindex: + return (name,) + ("",) * (column_level_count - 1) + return name + + if isinstance(table_frame.index, pd.MultiIndex): + seen_counts: dict[object, int] = {} + deduped_names: list[object | None] = [] + for name in table_frame.index.names: + if name is None: + deduped_names.append(None) + continue + padded = make_padded(name) + if padded in existing_column_set or name in seen_counts: + deduped_names.append(None) + else: + deduped_names.append(name) + seen_counts[name] = seen_counts.get(name, 0) + 1 + table_frame.index.names = deduped_names + elif hasattr(table_frame.index, "name") and table_frame.index.name is not None: + padded = make_padded(table_frame.index.name) + if padded in existing_column_set: + table_frame.index.name = None + + table_frame = table_frame.reset_index() + + if was_multiindex: + new_columns = [] + for column in table_frame.columns: + if isinstance(column, str) and ( + column.startswith("level_") or column == "index" + ): + new_columns.append(tuple([""] * column_level_count)) + else: + new_columns.append(column) + table_frame.columns = pd.MultiIndex.from_tuples(new_columns) + else: + new_columns = [] + for column in table_frame.columns: + if isinstance(column, str) and ( + column.startswith("level_") or column == "index" + ): + new_columns.append("") + else: + new_columns.append(column) + table_frame.columns = new_columns + else: + table_frame.reset_index(drop=True, inplace=True) + + if isinstance(table_frame.columns, pd.MultiIndex): + new_levels = [] + for level_index in range(table_frame.columns.nlevels): + level_values = table_frame.columns.get_level_values(level_index) + cleaned = [ + str(value).replace("\n", "") if value is not None else "" + for value in level_values + ] + new_levels.append(cleaned) + table_frame.columns = pd.MultiIndex.from_arrays( + new_levels, + names=table_frame.columns.names, + ) + + new_levels = [] + for level_index in range(table_frame.columns.nlevels): + level_values = table_frame.columns.get_level_values(level_index) + cleaned = [np.nan if "Unnamed" in str(value) else value for value in level_values] + new_levels.append(cleaned) + table_frame.columns = pd.MultiIndex.from_arrays( + new_levels, + names=table_frame.columns.names, + ) + else: + table_frame.columns = [ + str(column).replace("\n", "") for column in table_frame.columns + ] + table_frame.columns = [ + np.nan if "Unnamed" in str(column) else column + for column in table_frame.columns + ] + + table_frame = table_frame.map( + lambda value: value.replace("\n", "") if isinstance(value, str) else value + ) + return process_datetime_cells(table_frame) + + +def process_datetime_cells(table_frame: pd.DataFrame) -> pd.DataFrame: + table_frame = table_frame.copy() + + def convert(value: object) -> object: + if isinstance(value, (pd.Timestamp, datetime.datetime)): + return value.strftime("%Y-%m-%d %H:%M:%S") + return value + + return table_frame.apply(lambda column: column.map(convert)) + + +def process_duplicate_cols(columns: object) -> list[object]: + column_counts: dict[object, int] = {} + new_columns: list[object] = [] + for column in columns: + if column in column_counts: + new_columns.append(f"{column}>>>{column_counts[column]}") + column_counts[column] += 1 + else: + new_columns.append(column) + column_counts[column] = 1 + return new_columns + + +def format_tb_scope(table_frame: pd.DataFrame, num: int) -> str: + if len(table_frame) > int(num * 3 + 1): + head_frame = table_frame.head(num) + tail_frame = table_frame.tail(num) + middle_frame = table_frame.iloc[num : len(table_frame) - num] + + if len(middle_frame) >= num: + mid_sample_frame = middle_frame.sample(n=num, random_state=42) + else: + mid_sample_frame = middle_frame + scope_frame = pd.concat( + objs=[head_frame, mid_sample_frame, tail_frame], + ignore_index=True, + ) + else: + scope_frame = table_frame + scope_frame = scope_frame.map( + lambda value: str(value).strip() if pd.notnull(value) else value + ) + return df2html(scope_frame) diff --git a/apps/worker/app/services/document_parser/table_parser.py b/apps/worker/app/services/document_parser/table_parser.py index 342200836..2c189c0c7 100755 --- a/apps/worker/app/services/document_parser/table_parser.py +++ b/apps/worker/app/services/document_parser/table_parser.py @@ -1,626 +1,24 @@ -# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalOperand=false, reportOptionalSubscript=false, reportReturnType=false -import datetime -import re -import threading -import uuid -from collections import OrderedDict +from __future__ import annotations -import numpy as np import pandas as pd -from app.services.document_parser.identifiers import gen_str_codes -from app.services.document_parser.path_helpers import flatten_dic2paths -from app.services.document_parser.dataframe_html_renderer import df2html -from bs4 import BeautifulSoup -from loguru import logger - -from shared.services.ai.prompt_service import build_prompt -from shared.services.ai.response_process_service import eval_response -from shared.utils.OpenAICompatibleClientSync import get_openai_client -from shared.utils.text_utils import remove_duplicates_orderkept - -# ── Table filename sanitizer ──────────────────────────── -# Max byte-safe filename length. Most filesystems cap at 255 bytes; we leave -# room for the "table-N " prefix (~10 chars) and ".html" suffix (5 chars). -_MAX_TABLE_NAME_CHARS = 80 - - -def sanitize_table_name_from_header(raw_header_text: str) -> str: - """Build a concise, filesystem-safe table name from raw first-row header text. - - Pipeline: - 1. Split by common delimiters (' | ', '_br_'/'__br_', '\\n') - 2. Strip whitespace, deduplicate (preserve order) - 3. Drop trivial single-character tokens (single CJK char, single digit, - single letter) — reuses ``_is_meaningful_token`` from shared text_utils - 4. Rejoin with spaces and cap at ``_MAX_TABLE_NAME_CHARS`` - - Args: - raw_header_text: The raw first-row text, often pipe-separated. - - Returns: - A cleaned string suitable for use in a filename (may be empty if all - fields were trivial). - """ - from shared.utils.text_utils import _is_meaningful_token - - if not raw_header_text: - return "" - - # 1. Split on common header delimiters - parts = re.split(r"\s*\|\s*|_+br_|\n", raw_header_text) - - # 2. Strip + deduplicate (order-preserved) - seen: set[str] = set() - unique: list[str] = [] - for p in parts: - p = p.strip() - if not p or p in seen: - continue - seen.add(p) - unique.append(p) - - # 3. Keep only meaningful fields (drop single-char noise) - meaningful = [f for f in unique if _is_meaningful_token(f)] - - # 4. Join and enforce length cap - result = " ".join(meaningful) - if len(result) > _MAX_TABLE_NAME_CHARS: - result = result[:_MAX_TABLE_NAME_CHARS].rstrip() - return result - - -g_tbl_lock = threading.Lock() - - -def identify_tables(line): - """Identify if a line contains a table. - - Note: For HTML tables, use merge_html_tables() from html_parser.py - to preprocess multi-line tables before calling this function. - """ - # HTML table: complete
...
in one line - html_tb_pattern = r".*?" - tables = re.findall(html_tb_pattern, line, re.DOTALL) - if bool(tables): - return True, "html", tables - - # MD table: lines starting and ending with | - if line.startswith("|") and line.endswith("|"): - return True, "md", [] - - return False, None, None - - -def df2md(tb_df: pd.DataFrame, *, index: bool = False, na_rep: str = "—") -> str: - """Convert DataFrame to Markdown table format with dynamic column widths. - - Note: Truncation should be done externally using truncate_text before calling this function. - - Args: - tb_df: Input DataFrame - index: Whether to include index column - na_rep: String to represent NA values - - Returns: - Markdown table string - """ - import unicodedata - - def get_display_width(text: str) -> int: - """eval width for both ASCII and Chinese""" - width = 0 - for char in text: - if unicodedata.east_asian_width(char) in ("F", "W"): - width += 2 - else: - width += 1 - return width - - def pad_to_width(text: str, target_width: int) -> str: - current_width = get_display_width(text) - padding = target_width - current_width - return text + " " * max(0, padding) - - df = tb_df.copy() - - # Handle index - if index: - df = df.reset_index() - - # Replace NA values - df = df.fillna(na_rep) - - # Convert all values to string - df = df.astype(str) - - # Calculate column widths based on actual display width (no truncation) - col_widths = {} - for col in df.columns: - header_width = get_display_width(str(col)) - max_content_width = max(df[col].apply(get_display_width)) if len(df) > 0 else 0 - col_widths[col] = max(header_width, max_content_width) - - # Build header row - header_cells = [pad_to_width(str(col), col_widths[col]) for col in df.columns] - header_line = "| " + " | ".join(header_cells) + " |" - - # Build separator row - separator_cells = ["-" * col_widths[col] for col in df.columns] - separator_line = "|-" + "-|-".join(separator_cells) + "-|" - - # Build data rows - data_lines = [] - for _, row in df.iterrows(): - cells = [pad_to_width(str(row[col]), col_widths[col]) for col in df.columns] - data_lines.append("| " + " | ".join(cells) + " |") - - # Combine all parts - lines = [header_line, separator_line] + data_lines - return "\n".join(lines) - - -def clean_html_tb(html: str) -> str: - soup = BeautifulSoup(html, "html.parser") - for row in soup.find_all("tr"): - seen = set() - unique_cells = [] - for cell in row.find_all("td", recursive=False): - content = cell.encode_contents() - if content not in seen: - seen.add(content) - unique_cells.append(cell) - row.clear() - for cell in unique_cells: - row.append(cell) - return soup.prettify() - - -def extract_tables_by_forms(tb_txt, form): - if form == "html": - return tb_txt - elif form == "md": - tb_df = pd.read_table( - pd.io.common.StringIO(tb_txt), sep="|", engine="python", on_bad_lines="skip" - ) - tb_df = tb_df.drop(columns=tb_df.columns[0]) # Drop extra leading column - tb_df = tb_df.drop(columns=tb_df.columns[-1]) # Drop extra trailing column - tb_df.columns = tb_df.columns.str.strip() # Clean up headers - # Filter out MD separator lines (e.g. "---", ":---:", "---:") - separator_pattern = r"^[\s\-:]+$" - tb_df = tb_df[ - ~tb_df.apply( - lambda row: row.astype(str).str.match(separator_pattern).all(), axis=1 - ) - ] - tb_strs = tb_df.to_html(index=False) - else: - tb_strs = None # UNDER DEVELOPMENT other forms of tables... - return tb_strs - - -def parse_headers(df_temp, paras=None, header_window=5, smart_headers=True): - def parse_headers_nonsmart(df_): - non_na_row = df_[df_.notna().any(axis=1)].head(1) - header_id = non_na_row.index[-1] if not non_na_row.empty else None - header_rows = list(range(header_id + 1)) - return header_rows - - if not pd.isna( - df_temp.columns - ).all(): # If columns are not all NaN, no need to add extra row - df_temp.loc[-1] = df_temp.columns - df_temp.index = df_temp.index + 1 - df_temp = df_temp.sort_index() - df_temp.columns = [np.nan] * df_temp.shape[1] - - if paras["summary_table"] and smart_headers: - try: - tb_small = df_temp.head(header_window) - tb_small_str = df2html(tb_small) - prompt, temperature, top_p, max_tokens = build_prompt( - task="detect-table-headers", texts=tb_small_str, query="", paras=paras - ) - - messages = [ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": prompt}, - ] - - ctx_task_id = gen_str_codes((str(uuid.uuid4()) + tb_small_str)) - - # Track task status via Redis (skip in LOCAL_DEBUG mode) - import os - - if os.getenv("LOCAL_DEBUG", "0") != "1": - from shared.services.redis.redis_sync_service import ( - SyncRedisServiceFactory, - ) - - redis_service = SyncRedisServiceFactory.get_service() - redis_service.set(f"task:{ctx_task_id}:status", "processing", ttl=7200) - - # Use unified AI service - header_res = get_openai_client().chat_completion( - messages=messages, timeout=60 - ) - header_res = eval_response(header_res) - # Extract answer field - if isinstance(header_res, dict): - answer = header_res.get("answer", []) - else: - answer = header_res if isinstance(header_res, list) else [] - - # Check if answer is empty list - if not answer or len(answer) == 0: - logger.warning( - "AI returned empty list, cannot identify headers, falling back to traditional mode..." - ) - header_rows = parse_headers_nonsmart(df_temp) - else: - try: - header_id = answer[-1] - header_rows = list(range(header_id + 1)) - except Exception as e: - logger.warning( - f"Failed to parse header row number: {e}, falling back to traditional mode..." - ) - header_rows = parse_headers_nonsmart(df_temp) - - except Exception as e: - logger.warning( - f"Smart header parsing failed: {e}, falling back to traditional mode..." - ) - header_rows = parse_headers_nonsmart(df_temp) - else: - header_rows = parse_headers_nonsmart(df_temp) - - # improve table structure based on header rows - if len(header_rows) == 0 or (all(h is None for h in header_rows)): - logger.warning("No valid headers detected, fallback to using row 0 as header") - new_header = df_temp.iloc[0].ffill().bfill().tolist() - df_temp.columns = new_header - df_temp = df_temp.iloc[1:].reset_index(drop=True) - return df_temp - elif len(header_rows) > 1: - head_lst = [] - for i in range(0, len(header_rows)): - temp_lst = df_temp.iloc[i].ffill().bfill().tolist() - head_lst.append(temp_lst) - new_header = pd.MultiIndex.from_arrays(np.array(head_lst)) - else: - new_header = df_temp.iloc[header_rows[-1]].ffill().bfill().tolist() - - df_temp.columns = new_header - df_temp = df_temp.iloc[(header_rows[-1]) + 1 :] - df_temp = df_temp.reset_index(drop=True) - return df_temp - - -def parse_tb_keywords( - tb_df, kw_spit=">>>" -): # Extract keywords from headers (can also add LLM extraction) - def parse_single_level_(cols, keywords): - cols = [str(c) for c in cols] - for col in cols: - if kw_spit in col: - tmp_kw = col.split(">>>")[0] - else: # May be first occurrence - tmp_kw = col - if tmp_kw not in keywords: - keywords.append(col) - keywords_a_level = list(set([k.strip() for k in keywords])) - return keywords_a_level - - tb_keywords = [] - if isinstance(tb_df.columns, pd.MultiIndex): - multi_cols = tb_df.columns - cols_df = pd.DataFrame( - multi_cols.tolist(), - columns=[f"level_{i}" for i in range(multi_cols.nlevels)], - ) - for i in range(multi_cols.nlevels): # Extract each level as list - level_kws = [] - level_kws = parse_single_level_(cols_df[f"level_{i}"].tolist(), level_kws) - tb_keywords.extend(level_kws) - else: - tb_keywords = parse_single_level_(tb_df.columns, tb_keywords) - - # Remove duplicates while preserving column order - tb_keywords = remove_duplicates_orderkept(tb_keywords) - tb_keywords = [ - k - for k in tb_keywords - if isinstance(k, str) - and k.strip() - and k.strip() != "nan" - and "Unnamed" not in k - ] - return ";".join(tb_keywords) - - -def parse_tb_contents( - df_temp, parent_dic=None, file_name="", sheet_name="", row_header_cols=0 -): - """Parse table contents and generate HTML. - - Args: - row_header_cols: Number of leftmost columns that are row headers (will be rendered as ) - """ - if parent_dic is None: - parent_dic = {} - - tb_res = df_temp.fillna("").infer_objects(copy=False) - tb_strs = df2html(tb_res, row_header_cols=row_header_cols) - - tb_tree = tb_columns_to_tree(df_temp, parent_dic, file_name, sheet_name) - tb_paths = flatten_dic2paths(tb_tree) - return tb_paths, tb_strs - - -def tb_columns_to_tree(df, parent_dic, file_name, sheet_name): - if isinstance(df.columns, pd.MultiIndex): - # Convert MultiIndex columns to a nested dictionary (tree-like structure) - columns = pd.DataFrame(df.columns.tolist()) - for level in range(columns.shape[1]): - columns[level] = process_duplicate_cols(columns[level]) - - new_columns = pd.MultiIndex.from_frame(columns) - tree_structure = multiindex_to_tree(new_columns) - else: - # If columns are not MultiIndex, convert them to a dictionary with empty dictionaries as values - new_columns = process_duplicate_cols(df.columns) - tree_structure = {col: {} for col in new_columns} - - df.columns = new_columns - if (not file_name == "") and (not sheet_name == ""): - parent_dic[file_name][sheet_name] = tree_structure - elif not sheet_name == "": - parent_dic[sheet_name] = tree_structure - elif not file_name == "": - parent_dic[file_name] = tree_structure - else: - parent_dic = tree_structure - return parent_dic - - -def multiindex_to_tree(multiindex): - """Convert a MultiIndex to a tree-like nested dictionary structure.""" - - def tree(): - return OrderedDict() - - root = tree() - for keys in multiindex: - current_level = root - for key in keys: - if key not in current_level: - current_level[key] = tree() - current_level = current_level[key] - - def convert_to_dict(d): - if isinstance(d, OrderedDict): - d = {k: convert_to_dict(v) for k, v in d.items()} - return d - - return convert_to_dict(root) - - -def postprocess_tb(df, drop=False): - if drop: - # Track if index was originally a simple RangeIndex (no semantic meaning) - # dropna(how='all') can turn RangeIndex into Int64Index by introducing gaps, - # which would incorrectly trigger the "preserve row index" logic below. - was_range_index = isinstance(df.index, pd.RangeIndex) - - # Drop rows where all data columns are empty (exclude _src_row from the check) - # _src_row is always non-null, so including it would prevent any row from being dropped. - src_row_cols = [ - c - for c in df.columns - if (isinstance(c, tuple) and c[0] == "_src_row") or c == "_src_row" - ] - if src_row_cols: - data_cols = [c for c in df.columns if c not in src_row_cols] - mask = df[data_cols].isna().all(axis=1) - df = df[~mask] - else: - df = df.dropna(how="all") - - # If index was originally RangeIndex, re-number it to avoid gaps - if was_range_index: - df = df.reset_index(drop=True) - - # Drop columns that are all empty AND have no meaningful header - # A column with a valid header should be preserved even if data is empty - cols_to_drop = [] - for col_idx, col in enumerate(df.columns): - # Check if all data values are NaN - # Use iloc to avoid ambiguity when MultiIndex has duplicate tuple keys - if df.iloc[:, col_idx].isna().all(): - # Check if the column header is meaningful - # For MultiIndex: check if any level has a non-empty meaningful value - # For simple index: check if the header is not None/empty - has_meaningful_header = False - if isinstance(col, tuple): - # MultiIndex column - check if any level has meaningful content - for level in col: - if ( - level - and str(level).strip() - and str(level).strip() not in ["None", "nan", "NaN"] - ): - has_meaningful_header = True - break - else: - # Simple column name - if ( - col - and str(col).strip() - and str(col).strip() not in ["None", "nan", "NaN"] - ): - has_meaningful_header = True - - # Only drop if header is not meaningful - if not has_meaningful_header: - cols_to_drop.append(col_idx) - - if cols_to_drop: - # Use positional indices to drop columns safely (avoids duplicate MultiIndex key issues) - cols_to_keep = [i for i in range(len(df.columns)) if i not in cols_to_drop] - df = df.iloc[:, cols_to_keep] - - logger.debug(f"Dropped {len(cols_to_drop)} empty columns") - - # Preserve meaningful row index (header columns) as regular columns - # Only drop=True if it's a simple RangeIndex (no semantic meaning) - if not isinstance(df.index, pd.RangeIndex): - # Remember if columns were MultiIndex before reset - was_multiindex = isinstance(df.columns, pd.MultiIndex) - n_levels = df.columns.nlevels if was_multiindex else 1 - - # Avoid name collision before reset_index(). - # Two collision sources: - # A) An index level name, when padded into a tuple by pandas, - # matches an existing column. - # B) Multiple index levels share the same name → pandas tries - # to insert duplicate columns (e.g. five levels all named - # one merged header repeated across five padded columns. - # Strategy: de-duplicate index.names so every level gets a unique - # column name during reset_index, then clean up afterwards. - existing_col_set = set(df.columns) - - def _make_padded(name): - """Simulate the column name pandas would create for this index level.""" - if was_multiindex: - return (name,) + ("",) * (n_levels - 1) - return name - - if isinstance(df.index, pd.MultiIndex): - seen_counts = {} # name → how many times seen so far - deduped = [] - for n in df.index.names: - if n is None: - deduped.append(None) - continue - padded = _make_padded(n) - # Collision with existing column OR with a previously-seen index name - if padded in existing_col_set or n in seen_counts: - deduped.append( - None - ) # let pandas auto-name it (level_0, level_1 …) - else: - deduped.append(n) - seen_counts[n] = seen_counts.get(n, 0) + 1 - df.index.names = deduped - elif hasattr(df.index, "name") and df.index.name is not None: - padded = _make_padded(df.index.name) - if padded in existing_col_set: - df.index.name = None - - df = df.reset_index() # Converts index to columns - - # Clean up auto-generated column names like 'index', 'level_0', 'level_1' - # For MultiIndex columns, we need to preserve the structure - if was_multiindex: - # Build new column tuples for the index columns - new_cols = [] - for col in df.columns: - if isinstance(col, str) and ( - col.startswith("level_") or col == "index" - ): - # Create a tuple with empty strings to match MultiIndex levels - new_cols.append(tuple([""] * n_levels)) - else: - new_cols.append(col) - df.columns = pd.MultiIndex.from_tuples(new_cols) - else: - # For simple columns - new_cols = [] - for col in df.columns: - if isinstance(col, str) and ( - col.startswith("level_") or col == "index" - ): - new_cols.append("") - else: - new_cols.append(col) - df.columns = new_cols - else: - df.reset_index(drop=True, inplace=True) - - # Clean column names - preserve MultiIndex structure if present - if isinstance(df.columns, pd.MultiIndex): - # For MultiIndex, clean each level's values while preserving structure - new_levels = [] - for level_idx in range(df.columns.nlevels): - level_vals = df.columns.get_level_values(level_idx) - cleaned = [ - str(v).replace("\n", "") if v is not None else "" for v in level_vals - ] - new_levels.append(cleaned) - df.columns = pd.MultiIndex.from_arrays(new_levels, names=df.columns.names) - # Also handle 'Unnamed' in MultiIndex - new_levels = [] - for level_idx in range(df.columns.nlevels): - level_vals = df.columns.get_level_values(level_idx) - cleaned = [np.nan if "Unnamed" in str(v) else v for v in level_vals] - new_levels.append(cleaned) - df.columns = pd.MultiIndex.from_arrays(new_levels, names=df.columns.names) - else: - df.columns = [ - str(col).replace("\n", "") for col in df.columns - ] # Replace '\n' in column headers - df.columns = [ - np.nan if "Unnamed" in str(col) else col for col in df.columns - ] # Replace Unnamed with nan - df = df.map( - lambda x: x.replace("\n", "") if isinstance(x, str) else x - ) # Replace '\n' in each cell - df = process_datetime_cells(df) - return df - - -def process_datetime_cells(df): - df = df.copy() - - def convert(x): - if isinstance(x, (pd.Timestamp, datetime.datetime)): - return x.strftime("%Y-%m-%d %H:%M:%S") - return x - - return df.apply(lambda col: col.map(convert)) - - -def process_duplicate_cols(columns): - col_count = {} - new_columns = [] - for col in columns: - if col in col_count: - new_columns.append(f"{col}>>>{col_count[col]}") - col_count[col] += 1 - else: - new_columns.append(col) - col_count[col] = 1 - return new_columns - - -def format_tb_scope(df, num): - if len(df) > int(num * 3 + 1): - # Get head and tail rows - head_df = df.head(num) - tail_df = df.tail(num) - # Middle portion excluding head and tail - middle_df = df.iloc[num : len(df) - num] - - if len(middle_df) >= num: - mid_sample_df = middle_df.sample(n=num, random_state=42) - else: # If middle has less than num rows, take all - mid_sample_df = middle_df - scope_df = pd.concat(objs=[head_df, mid_sample_df, tail_df], ignore_index=True) - else: - scope_df = df - scope_df = scope_df.applymap(lambda x: str(x).strip() if pd.notnull(x) else x) - scope_str = df2html(scope_df) - return scope_str +from app.services.document_parser.table_frame_parser import ( + format_tb_scope as format_tb_scope, + multiindex_to_tree as multiindex_to_tree, + parse_headers as parse_headers, + parse_tb_contents as parse_tb_contents, + parse_tb_keywords as parse_tb_keywords, + postprocess_tb as postprocess_tb, + process_datetime_cells as process_datetime_cells, + process_duplicate_cols as process_duplicate_cols, + tb_columns_to_tree as tb_columns_to_tree, +) +from app.services.document_parser.table_text_parser import ( + clean_html_tb as clean_html_tb, + df2md as df2md, + extract_tables_by_forms as extract_tables_by_forms, + identify_tables as identify_tables, + sanitize_table_name_from_header as sanitize_table_name_from_header, +) def parse_xlsx( @@ -628,7 +26,7 @@ def parse_xlsx( file_name: str, output_dir: str, baseurl: str, - base_llm_paras: dict | None = None, + base_llm_paras: dict[str, object] | None = None, window_h: int = 10, relative_root: str | None = None, use_precision_mode: bool = True, diff --git a/apps/worker/app/services/document_parser/table_text_parser.py b/apps/worker/app/services/document_parser/table_text_parser.py new file mode 100644 index 000000000..261bbe653 --- /dev/null +++ b/apps/worker/app/services/document_parser/table_text_parser.py @@ -0,0 +1,149 @@ +# pyright: reportArgumentType=false +from __future__ import annotations + +import io +import re +import unicodedata + +import pandas as pd +from bs4 import BeautifulSoup, Tag + +_MAX_TABLE_NAME_CHARS = 80 + + +def sanitize_table_name_from_header(raw_header_text: str) -> str: + """Build a concise, filesystem-safe table name from raw first-row header text.""" + from shared.utils.text_utils import _is_meaningful_token + + if not raw_header_text: + return "" + + parts = re.split(r"\s*\|\s*|_+br_|\n", raw_header_text) + + seen: set[str] = set() + unique: list[str] = [] + for part in parts: + part = part.strip() + if not part or part in seen: + continue + seen.add(part) + unique.append(part) + + meaningful = [field for field in unique if _is_meaningful_token(field)] + result = " ".join(meaningful) + if len(result) > _MAX_TABLE_NAME_CHARS: + result = result[:_MAX_TABLE_NAME_CHARS].rstrip() + return result + + +def identify_tables(line: str) -> tuple[bool, str | None, list[str] | None]: + """Identify whether one logical Markdown line contains a table.""" + html_table_pattern = r".*?" + tables = re.findall(html_table_pattern, line, re.DOTALL) + if bool(tables): + return True, "html", tables + + if line.startswith("|") and line.endswith("|"): + return True, "md", [] + + return False, None, None + + +def df2md(table_frame: pd.DataFrame, *, index: bool = False, na_rep: str = "—") -> str: + """Convert a DataFrame to a Markdown table while preserving display width.""" + + def get_display_width(text: str) -> int: + width = 0 + for character in text: + if unicodedata.east_asian_width(character) in ("F", "W"): + width += 2 + else: + width += 1 + return width + + def pad_to_width(text: str, target_width: int) -> str: + current_width = get_display_width(text) + padding = target_width - current_width + return text + " " * max(0, padding) + + table_frame = table_frame.copy() + + if index: + table_frame = table_frame.reset_index() + + table_frame = table_frame.fillna(na_rep).astype(str) + + column_widths: dict[object, int] = {} + for column in table_frame.columns: + header_width = get_display_width(str(column)) + max_content_width = ( + max(table_frame[column].apply(get_display_width)) + if len(table_frame) > 0 + else 0 + ) + column_widths[column] = max(header_width, max_content_width) + + header_cells = [ + pad_to_width(str(column), column_widths[column]) + for column in table_frame.columns + ] + header_line = "| " + " | ".join(header_cells) + " |" + + separator_cells = ["-" * column_widths[column] for column in table_frame.columns] + separator_line = "|-" + "-|-".join(separator_cells) + "-|" + + data_lines: list[str] = [] + for _, row in table_frame.iterrows(): + cells = [ + pad_to_width(str(row[column]), column_widths[column]) + for column in table_frame.columns + ] + data_lines.append("| " + " | ".join(cells) + " |") + + return "\n".join([header_line, separator_line, *data_lines]) + + +def clean_html_tb(html: str) -> str: + soup = BeautifulSoup(html, "html.parser") + for row in soup.find_all("tr"): + if not isinstance(row, Tag): + continue + seen: set[bytes] = set() + unique_cells: list[Tag] = [] + for cell in row.find_all("td", recursive=False): + if not isinstance(cell, Tag): + continue + content = cell.encode_contents() + if content not in seen: + seen.add(content) + unique_cells.append(cell) + row.clear() + for cell in unique_cells: + row.append(cell) + return str(soup.prettify()) + + +def extract_tables_by_forms(table_text: str, form: str) -> str | None: + if form == "html": + return table_text + + if form != "md": + return None + + table_frame = pd.read_table( + io.StringIO(table_text), + sep="|", + engine="python", + on_bad_lines="skip", + ) + table_frame = table_frame.iloc[:, 1:-1] + table_frame.columns = table_frame.columns.astype(str).str.strip() + + separator_pattern = r"^[\s\-:]+$" + table_frame = table_frame[ + ~table_frame.apply( + lambda row: row.astype(str).str.match(separator_pattern).all(), + axis=1, + ) + ] + return table_frame.to_html(index=False) diff --git a/apps/worker/app/services/document_parser/toc_hierarchy.py b/apps/worker/app/services/document_parser/toc_hierarchy.py index 693350a37..b88420e32 100644 --- a/apps/worker/app/services/document_parser/toc_hierarchy.py +++ b/apps/worker/app/services/document_parser/toc_hierarchy.py @@ -3,7 +3,7 @@ import pandas as pd from app.services.document_parser.layout_parser import hiearchy_llm from app.services.document_parser.stage_profiler import stage_timer -from app.services.document_parser.table_parser import df2md +from app.services.document_parser.table_text_parser import df2md from app.services.document_parser.text_helpers import normalize_md from loguru import logger from pandas import Index diff --git a/apps/worker/app/services/document_parser/toc_parser.py b/apps/worker/app/services/document_parser/toc_parser.py index 7a0229859..2efb1a4cd 100644 --- a/apps/worker/app/services/document_parser/toc_parser.py +++ b/apps/worker/app/services/document_parser/toc_parser.py @@ -15,7 +15,7 @@ from app.services.document_parser.toc_hierarchy import eval_toc_levels from app.services.document_parser.text_helpers import normalize_md, truncate_text_by_tokens from app.services.document_parser.stage_profiler import stage_timer -from app.services.document_parser.table_parser import df2md +from app.services.document_parser.table_text_parser import df2md from gevent.pool import Pool as GeventPool from loguru import logger diff --git a/apps/worker/tests/contract/test_document_parser_architecture_contract.py b/apps/worker/tests/contract/test_document_parser_architecture_contract.py index 0f3a377ee..f0d5d381a 100644 --- a/apps/worker/tests/contract/test_document_parser_architecture_contract.py +++ b/apps/worker/tests/contract/test_document_parser_architecture_contract.py @@ -393,6 +393,212 @@ def test_html_table_modules_separate_docx_and_dataframe_rendering( assert "3" in dataframe_html +def test_table_text_parser_owns_markdown_table_text_contract( + worker_contract_environment: None, +) -> None: + from app.services.document_parser.table_text_parser import ( + df2md, + extract_tables_by_forms, + identify_tables, + sanitize_table_name_from_header, + ) + + markdown_table = "\n".join( + [ + "| Product | Revenue |", + "| --- | ---: |", + "| Notebook | 42 |", + ] + ) + + is_table, table_form, _tables = identify_tables("| Product | Revenue |") + table_html = extract_tables_by_forms(markdown_table, form="md") + markdown_output = df2md(pd.DataFrame([{"City": "北京", "Value": 7}])) + + assert is_table is True + assert table_form == "md" + assert table_html is not None + assert "Product" in table_html + assert "Notebook" in table_html + assert "---" not in table_html + assert sanitize_table_name_from_header("A | Revenue | Revenue | 市场") == ( + "Revenue 市场" + ) + assert "| City | Value |" in markdown_output + assert "| 北京 | 7 |" in markdown_output + + +def test_table_frame_parser_owns_dataframe_table_contract( + worker_contract_environment: None, +) -> None: + from app.services.document_parser.table_frame_parser import ( + parse_tb_contents, + parse_tb_keywords, + postprocess_tb, + ) + + raw_frame = pd.DataFrame( + [["North\nAmerica", 42, None]], + columns=["Region\nName", "Revenue", None], + ) + + normalized_frame = postprocess_tb(raw_frame, drop=True) + paths, table_html = parse_tb_contents( + normalized_frame, + parent_dic={"budget.xlsx": {"Visible": {}}}, + file_name="budget.xlsx", + sheet_name="Visible", + ) + keywords = parse_tb_keywords(normalized_frame) + + assert normalized_frame.columns.tolist() == ["RegionName", "Revenue"] + assert "NorthAmerica" in table_html + assert "RegionName" in keywords + assert "Revenue" in keywords + assert "budget.xlsx/Visible/RegionName" in paths + assert "budget.xlsx/Visible/Revenue" in paths + + +def test_markdown_table_asset_module_owns_table_asset_contract( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.markdown_table_asset import ( + MarkdownTableAssetRequest, + build_markdown_table_asset, + ) + from app.services.document_parser.markdown_deferred_task import ( + TableDeferredSummaryTask, + ) + + table_dir = tmp_path / "tables" + table_dir.mkdir() + table_html = ( + "" + "
ProductRevenue
Notebook42
" + ) + + asset = build_markdown_table_asset( + MarkdownTableAssetRequest( + table_html=table_html, + table_dir=str(table_dir), + table_count=3, + timestamp="now", + current_page_number=9, + summary_table=True, + row_index=7, + ) + ) + + assert asset.content_item == f"\n[{asset.relative_path}]\n" + assert asset.row_values[1] == asset.relative_path + assert asset.row_values[2] == "table" + assert asset.row_values[5] == "table-3" + assert asset.row_values[10] == "9" + assert asset.deferred_task == TableDeferredSummaryTask( + row_index=7, + table_html=table_html, + table_dir=str(table_dir), + table_name=Path(asset.relative_path).stem, + table_count=2, + ) + assert "border='1'" in (tmp_path / asset.relative_path).read_text( + encoding="utf-8" + ) + + +def test_markdown_image_asset_module_owns_image_materialization_contract( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_parser.markdown_image_asset import ( + MarkdownImageAssetRequest, + build_markdown_image_name, + build_markdown_image_asset, + ) + from app.services.document_parser.markdown_deferred_task import ( + ImageDeferredSummaryTask, + ) + + image_dir = tmp_path / "images" + image_dir.mkdir() + source_image = tmp_path / "raw.png" + source_image.write_bytes(b"same pixels") + seen_images: dict[str, dict[str, str]] = {} + + image_asset = build_markdown_image_asset( + MarkdownImageAssetRequest( + output_dir=str(tmp_path), + image_dir=str(image_dir), + image_path=str(source_image), + image_name=build_markdown_image_name( + image_count=2, + last_context="Revenue Chart", + ), + image_count=2, + last_context="Revenue Chart", + image_summary="Sales by region", + timestamp="now", + current_page_number=8, + seen_images=seen_images, + summary_image=True, + row_index=4, + ) + ) + + assert image_asset.content_item is not None + assert image_asset.row_values is not None + assert image_asset.cache_key is not None + assert image_asset.cache_entry is not None + assert image_asset.should_advance_image_count is True + assert image_asset.row_values[1] == "images/image-2-Revenue Ch.png" + assert image_asset.row_values[2] == "image" + assert image_asset.row_values[5] == "image-2\nSales by region" + assert image_asset.row_values[10] == "8" + assert image_asset.deferred_task == ImageDeferredSummaryTask( + row_index=4, + relative_path="images/image-2-Revenue Ch.png", + image_dir=str(image_dir), + image_name="image-2-Revenue Ch", + image_suffix=".png", + ) + assert (tmp_path / "images" / "image-2-Revenue Ch.png").read_bytes() == ( + b"same pixels" + ) + assert not source_image.exists() + + seen_images[image_asset.cache_key] = image_asset.cache_entry + duplicate_source = tmp_path / "duplicate.png" + duplicate_source.write_bytes(b"same pixels") + + duplicate_asset = build_markdown_image_asset( + MarkdownImageAssetRequest( + output_dir=str(tmp_path), + image_dir=str(image_dir), + image_path=str(duplicate_source), + image_name=build_markdown_image_name( + image_count=3, + last_context="Other Chart", + ), + image_count=3, + last_context="Other Chart", + image_summary=None, + timestamp="now", + current_page_number=9, + seen_images=seen_images, + summary_image=True, + row_index=5, + ) + ) + + assert duplicate_asset.content_item == image_asset.content_item + assert duplicate_asset.row_values is not None + assert duplicate_asset.row_values[1] == "images/image-2-Revenue Ch.png" + assert duplicate_asset.deferred_task is None + assert duplicate_asset.should_advance_image_count is False + assert not duplicate_source.exists() + + def test_mineru_modules_separate_client_and_task_polling( worker_contract_environment: None, ) -> None: @@ -598,6 +804,11 @@ def test_markdown_deferred_summary_module_updates_rows_and_refs( MarkdownDeferredSummaryInput, apply_markdown_deferred_summaries, ) + from app.services.document_parser.markdown_deferred_task import ( + ImageDeferredSummaryTask, + TableDeferredSummaryTask, + TextDeferredSummaryTask, + ) image_dir = tmp_path / "images" table_dir = tmp_path / "tables" @@ -668,9 +879,21 @@ def fake_extract(text: str, **_kwargs: Any) -> tuple[str, str, str]: MarkdownDeferredSummaryInput( rows=rows, tasks=[ - ("image", 0, "images/image-3-old.png", str(image_dir), "image-3-old", ".png"), - ("table", 1, "
", str(table_dir), "table-0 old", 0), - ("text", 2, "long text"), + ImageDeferredSummaryTask( + row_index=0, + relative_path="images/image-3-old.png", + image_dir=str(image_dir), + image_name="image-3-old", + image_suffix=".png", + ), + TableDeferredSummaryTask( + row_index=1, + table_html="
", + table_dir=str(table_dir), + table_name="table-0 old", + table_count=0, + ), + TextDeferredSummaryTask(row_index=2, content="long text"), ], output_dir=str(tmp_path), ) From 8081faf51f7c3475bbce221766f3a7d3d71d3d64 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 17:34:10 +0000 Subject: [PATCH 39/40] refactor(api): deepen demo source materialization --- CONTEXT.md | 1 + .../api/app/services/demo_document_service.py | 308 +---------------- .../app/services/demo_source_materializer.py | 326 ++++++++++++++++++ .../contract/test_demo_documents_contract.py | 72 +++- 4 files changed, 406 insertions(+), 301 deletions(-) create mode 100644 apps/api/app/services/demo_source_materializer.py diff --git a/CONTEXT.md b/CONTEXT.md index f87d42697..44e050713 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -228,6 +228,7 @@ exceptions. - `app/api/v1/routes/demo.py` - `app/services/demo_document_service.py` +- `app/services/demo_source_materializer.py` ### Billing Workflow diff --git a/apps/api/app/services/demo_document_service.py b/apps/api/app/services/demo_document_service.py index d1bbd779c..99d735f75 100644 --- a/apps/api/app/services/demo_document_service.py +++ b/apps/api/app/services/demo_document_service.py @@ -1,46 +1,22 @@ -"""Demo Source materialization workflow.""" +"""Demo Source catalog facade.""" from __future__ import annotations -import shutil -import tempfile -from dataclasses import dataclass -from datetime import datetime, timezone -from hashlib import blake2b from pathlib import Path from typing import Any -from uuid import uuid4 -from sqlalchemy import func, select +from app.services.demo_source_catalog import DemoSourceCatalog +from app.services.demo_source_materializer import ( + DemoSourceMaterializer, + MaterializedDemoSource, +) from sqlalchemy.ext.asyncio import AsyncSession -from app.services.demo_source_catalog import DemoSourceCatalog, DemoSourceDefinition - -from shared.core.exceptions.domain_exceptions import ValidationException -from shared.models.database.demo_materialization import DemoMaterialization -from shared.models.database.document import Document -from shared.models.database.job import Job -from shared.models.database.job_result import JobResult -from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces from shared.services.retrieval.publication_service import RetrievalPublicationService -from shared.services.storage.result_storage import get_result_storage - - -@dataclass(frozen=True) -class MaterializedDemoSource: - """User-owned copy of one canonical demo source.""" - - demo_source_id: str - document_id: str - status: str - title: str - mime_type: str - size_bytes: int - chunk_count: int class DemoDocumentService: - """Serves canonical demo data and copies it into user namespaces.""" + """Serves canonical demo data and delegates user-copy side effects.""" def __init__( self, @@ -49,7 +25,10 @@ def __init__( publication_service: RetrievalPublicationService | None = None, ) -> None: self._catalog = catalog or DemoSourceCatalog() - self._publication_service = publication_service or RetrievalPublicationService() + self._materializer = DemoSourceMaterializer( + catalog=self._catalog, + publication_service=publication_service, + ) def get_catalog(self) -> dict[str, Any]: return self._catalog.get_catalog() @@ -100,270 +79,9 @@ async def materialize_sources( namespace: str, demo_source_ids: list[str], ) -> list[MaterializedDemoSource]: - selected_demo_source_ids = _deduplicate_source_ids(demo_source_ids) - if not selected_demo_source_ids: - raise ValidationException( - user_message="At least one demo source must be selected.", - violations=[ - { - "field": "demo_source_ids", - "description": "Select one or more demo source IDs.", - } - ], - ) - - selected_sources = [ - self._catalog.require_source(demo_source_id) - for demo_source_id in selected_demo_source_ids - ] - results: list[MaterializedDemoSource] = [] - for source in selected_sources: - result = await self._materialize_source( - db, - user_id=user_id, - namespace=namespace, - source=source, - ) - results.append(result) - - await db.commit() - await invalidate_retrieval_cache_namespaces( - user_id=user_id, - namespaces=[namespace], - ) - return results - - async def _materialize_source( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - source: DemoSourceDefinition, - ) -> MaterializedDemoSource: - await _lock_materialization_scope( + return await self._materializer.materialize_sources( db, user_id=user_id, namespace=namespace, - demo_source_id=source.demo_source_id, - ) - existing = await self._get_existing_materialization( - db, - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - ) - if existing is not None and await self._is_active_document( - db, - document_id=existing.document_id, - ): - return _materialized_source_payload( - source=source, - document_id=existing.document_id, - status="existing", - ) - - document_id = f"doc_{uuid4().hex[:12]}" - job_id = f"job_demo_{uuid4().hex[:12]}" - job_result_id = str(uuid4()) - timestamp = _utc_now() - result_bundle = _upload_demo_result_bundle( - job_id=job_id, - source_directory=self._catalog.source_directory(source), - ) - - db.add( - Job( - job_id=job_id, - user_id=user_id, - job_type="demo_materialization", - status="done", - source_type="demo", - webhook_enabled=False, - job_metadata={ - "document_id": document_id, - "namespace": namespace, - "source_type": "demo", - "source_file_name": source.title, - "demo_source_id": source.demo_source_id, - }, - version=0, - created_at=timestamp, - updated_at=timestamp, - credits_charged=0, - billing_status="skipped", - ) - ) - db.add( - JobResult( - id=job_result_id, - job_id=job_id, - delivery_mode="inline", - document_metadata={ - "source_file_name": source.title, - "demo_source_id": source.demo_source_id, - }, - inline_payload={"source": "canonical_demo"}, - result_s3_key=result_bundle["zip_key"], - result_size=result_bundle["zip_size"], - created_at=timestamp, - updated_at=timestamp, - ) - ) - await db.flush() - chunks = self._catalog.publication_chunks(source) - await db.run_sync( - lambda sync_db: self._publication_service.publish_document_state( - sync_db, - job_id=job_id, - job_result_id=job_result_id, - chunks=[dict(chunk) for chunk in chunks], - ) - ) - await db.run_sync( - lambda sync_db: self._publication_service.publish_document_graph( - sync_db, - job_id=job_id, - job_result_id=job_result_id, - ) - ) - await db.flush() - - if existing is None: - db.add( - DemoMaterialization( - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - document_id=document_id, - created_at=timestamp, - updated_at=timestamp, - ) - ) - else: - existing.document_id = document_id - existing.updated_at = timestamp - await db.flush() - return _materialized_source_payload( - source=source, - document_id=document_id, - status="created", + demo_source_ids=demo_source_ids, ) - - async def _get_existing_materialization( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - demo_source_id: str, - ) -> DemoMaterialization | None: - result = await db.execute( - select(DemoMaterialization) - .where(DemoMaterialization.user_id == user_id) - .where(DemoMaterialization.namespace == namespace) - .where(DemoMaterialization.demo_source_id == demo_source_id) - .with_for_update() - .limit(1) - ) - return result.scalar_one_or_none() - - async def _is_active_document( - self, - db: AsyncSession, - *, - document_id: str, - ) -> bool: - result = await db.execute( - select(Document.document_id) - .where(Document.document_id == document_id) - .where(Document.status == "active") - .limit(1) - ) - return result.scalar_one_or_none() is not None - - -def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: - selected: list[str] = [] - seen: set[str] = set() - for demo_source_id in demo_source_ids: - normalized = str(demo_source_id).strip() - if not normalized or normalized in seen: - continue - selected.append(normalized) - seen.add(normalized) - return selected - - -async def _lock_materialization_scope( - db: AsyncSession, - *, - user_id: str, - namespace: str, - demo_source_id: str, -) -> None: - lock_id = _materialization_lock_id( - user_id=user_id, - namespace=namespace, - demo_source_id=demo_source_id, - ) - await db.execute(select(func.pg_advisory_xact_lock(lock_id))) - - -def _materialization_lock_id( - *, - user_id: str, - namespace: str, - demo_source_id: str, -) -> int: - lock_key = f"{user_id}\0{namespace}\0{demo_source_id}" - digest = blake2b(lock_key.encode("utf-8"), digest_size=8).digest() - return int.from_bytes(digest, byteorder="big", signed=True) - - -def _materialized_source_payload( - *, - source: DemoSourceDefinition, - document_id: str, - status: str, -) -> MaterializedDemoSource: - return MaterializedDemoSource( - demo_source_id=source.demo_source_id, - document_id=document_id, - status=status, - title=source.title, - mime_type=source.mime_type, - size_bytes=source.size_bytes, - chunk_count=source.chunk_count, - ) - - -def _upload_demo_result_bundle( - *, - job_id: str, - source_directory: Path, -) -> dict[str, int | str]: - with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temp_directory: - zip_base_path = Path(temp_directory) / job_id - zip_file_path = Path( - shutil.make_archive( - str(zip_base_path), - "zip", - root_dir=source_directory, - ) - ) - zip_size = zip_file_path.stat().st_size - bundle = get_result_storage().upload( - job_id=job_id, - result_dir=str(source_directory), - zip_file_path=str(zip_file_path), - ) - - return { - "zip_key": bundle.zip_key, - "zip_size": zip_size, - } - - -def _utc_now() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) diff --git a/apps/api/app/services/demo_source_materializer.py b/apps/api/app/services/demo_source_materializer.py new file mode 100644 index 000000000..6b418ff3c --- /dev/null +++ b/apps/api/app/services/demo_source_materializer.py @@ -0,0 +1,326 @@ +"""Demo Source Materialization workflow.""" + +from __future__ import annotations + +import shutil +import tempfile +from dataclasses import dataclass +from datetime import datetime, timezone +from hashlib import blake2b +from pathlib import Path +from uuid import uuid4 + +from app.services.demo_source_catalog import DemoSourceCatalog, DemoSourceDefinition +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import ValidationException +from shared.models.database.demo_materialization import DemoMaterialization +from shared.models.database.document import Document +from shared.models.database.job import Job +from shared.models.database.job_result import JobResult +from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces +from shared.services.retrieval.publication_service import RetrievalPublicationService +from shared.services.storage.result_storage import get_result_storage + + +@dataclass(frozen=True) +class MaterializedDemoSource: + """User-owned copy of one canonical demo source.""" + + demo_source_id: str + document_id: str + status: str + title: str + mime_type: str + size_bytes: int + chunk_count: int + + +class DemoSourceMaterializer: + """Copies canonical Demo Sources into user-owned retrieval state.""" + + def __init__( + self, + *, + catalog: DemoSourceCatalog, + publication_service: RetrievalPublicationService | None = None, + ) -> None: + self._catalog = catalog + self._publication_service = publication_service or RetrievalPublicationService() + + async def materialize_sources( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + demo_source_ids: list[str], + ) -> list[MaterializedDemoSource]: + selected_demo_source_ids = _deduplicate_source_ids(demo_source_ids) + if not selected_demo_source_ids: + raise ValidationException( + user_message="At least one demo source must be selected.", + violations=[ + { + "field": "demo_source_ids", + "description": "Select one or more demo source IDs.", + } + ], + ) + + selected_sources = [ + self._catalog.require_source(demo_source_id) + for demo_source_id in selected_demo_source_ids + ] + results: list[MaterializedDemoSource] = [] + for source in selected_sources: + result = await self._materialize_source( + db, + user_id=user_id, + namespace=namespace, + source=source, + ) + results.append(result) + + await db.commit() + await invalidate_retrieval_cache_namespaces( + user_id=user_id, + namespaces=[namespace], + ) + return results + + async def _materialize_source( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + source: DemoSourceDefinition, + ) -> MaterializedDemoSource: + await _lock_materialization_scope( + db, + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + ) + existing = await self._get_existing_materialization( + db, + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + ) + if existing is not None and await self._is_active_document( + db, + document_id=existing.document_id, + ): + return _materialized_source_payload( + source=source, + document_id=existing.document_id, + status="existing", + ) + + document_id = f"doc_{uuid4().hex[:12]}" + job_id = f"job_demo_{uuid4().hex[:12]}" + job_result_id = str(uuid4()) + timestamp = _utc_now() + result_bundle = _upload_demo_result_bundle( + job_id=job_id, + source_directory=self._catalog.source_directory(source), + ) + + db.add( + Job( + job_id=job_id, + user_id=user_id, + job_type="demo_materialization", + status="done", + source_type="demo", + webhook_enabled=False, + job_metadata={ + "document_id": document_id, + "namespace": namespace, + "source_type": "demo", + "source_file_name": source.title, + "demo_source_id": source.demo_source_id, + }, + version=0, + created_at=timestamp, + updated_at=timestamp, + credits_charged=0, + billing_status="skipped", + ) + ) + db.add( + JobResult( + id=job_result_id, + job_id=job_id, + delivery_mode="inline", + document_metadata={ + "source_file_name": source.title, + "demo_source_id": source.demo_source_id, + }, + inline_payload={"source": "canonical_demo"}, + result_s3_key=result_bundle["zip_key"], + result_size=result_bundle["zip_size"], + created_at=timestamp, + updated_at=timestamp, + ) + ) + await db.flush() + chunks = self._catalog.publication_chunks(source) + await db.run_sync( + lambda sync_db: self._publication_service.publish_document_state( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + chunks=[dict(chunk) for chunk in chunks], + ) + ) + await db.run_sync( + lambda sync_db: self._publication_service.publish_document_graph( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + ) + ) + await db.flush() + + if existing is None: + db.add( + DemoMaterialization( + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + document_id=document_id, + created_at=timestamp, + updated_at=timestamp, + ) + ) + else: + existing.document_id = document_id + existing.updated_at = timestamp + await db.flush() + return _materialized_source_payload( + source=source, + document_id=document_id, + status="created", + ) + + async def _get_existing_materialization( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + demo_source_id: str, + ) -> DemoMaterialization | None: + result = await db.execute( + select(DemoMaterialization) + .where(DemoMaterialization.user_id == user_id) + .where(DemoMaterialization.namespace == namespace) + .where(DemoMaterialization.demo_source_id == demo_source_id) + .with_for_update() + .limit(1) + ) + return result.scalar_one_or_none() + + async def _is_active_document( + self, + db: AsyncSession, + *, + document_id: str, + ) -> bool: + result = await db.execute( + select(Document.document_id) + .where(Document.document_id == document_id) + .where(Document.status == "active") + .limit(1) + ) + return result.scalar_one_or_none() is not None + + +def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: + selected: list[str] = [] + seen: set[str] = set() + for demo_source_id in demo_source_ids: + normalized = str(demo_source_id).strip() + if not normalized or normalized in seen: + continue + selected.append(normalized) + seen.add(normalized) + return selected + + +async def _lock_materialization_scope( + db: AsyncSession, + *, + user_id: str, + namespace: str, + demo_source_id: str, +) -> None: + lock_id = _materialization_lock_id( + user_id=user_id, + namespace=namespace, + demo_source_id=demo_source_id, + ) + await db.execute(select(func.pg_advisory_xact_lock(lock_id))) + + +def _materialization_lock_id( + *, + user_id: str, + namespace: str, + demo_source_id: str, +) -> int: + lock_key = f"{user_id}\0{namespace}\0{demo_source_id}" + digest = blake2b(lock_key.encode("utf-8"), digest_size=8).digest() + return int.from_bytes(digest, byteorder="big", signed=True) + + +def _materialized_source_payload( + *, + source: DemoSourceDefinition, + document_id: str, + status: str, +) -> MaterializedDemoSource: + return MaterializedDemoSource( + demo_source_id=source.demo_source_id, + document_id=document_id, + status=status, + title=source.title, + mime_type=source.mime_type, + size_bytes=source.size_bytes, + chunk_count=source.chunk_count, + ) + + +def _upload_demo_result_bundle( + *, + job_id: str, + source_directory: Path, +) -> dict[str, int | str]: + with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temp_directory: + zip_base_path = Path(temp_directory) / job_id + zip_file_path = Path( + shutil.make_archive( + str(zip_base_path), + "zip", + root_dir=source_directory, + ) + ) + zip_size = zip_file_path.stat().st_size + bundle = get_result_storage().upload( + job_id=job_id, + result_dir=str(source_directory), + zip_file_path=str(zip_file_path), + ) + + return { + "zip_key": bundle.zip_key, + "zip_size": zip_size, + } + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc).replace(tzinfo=None) diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index a3b63b9ab..6975a378f 100644 --- a/apps/api/tests/contract/test_demo_documents_contract.py +++ b/apps/api/tests/contract/test_demo_documents_contract.py @@ -139,10 +139,10 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( monkeypatch.setenv("RETRIEVAL_AGENTIC_ENABLED", "false") async with developer_api_client_factory() as api_client: - import app.services.demo_document_service as demo_document_service + import app.services.demo_source_materializer as demo_source_materializer monkeypatch.setattr( - demo_document_service, + demo_source_materializer, "get_result_storage", lambda: fake_result_storage, ) @@ -271,6 +271,66 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( assert any(file_path.startswith("tables/") for file_path in uploaded_files) +@pytest.mark.asyncio +async def test_should_materialize_each_normalized_demo_source_once_per_request( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + fake_result_storage = FakeResultStorage() + + async with developer_api_client_factory() as api_client: + import app.services.demo_source_materializer as demo_source_materializer + + monkeypatch.setattr( + demo_source_materializer, + "get_result_storage", + lambda: fake_result_storage, + ) + response = await api_client.post( + "/api/v1/demo/materializations", + json={ + "namespace": "contract-demo-deduplicate", + "demo_source_ids": [ + DEMO_SOURCE_ID, + f" {DEMO_SOURCE_ID} ", + DEMO_SOURCE_ID, + "", + ], + }, + ) + + assert response.status_code == 200 + + body = cast(dict[str, Any], response.json()) + sources = cast(list[dict[str, Any]], body["sources"]) + assert [source["demo_source_id"] for source in sources] == [DEMO_SOURCE_ID] + + materialization_rows = await ContractDatabase.fetch_all( + """ + SELECT demo_source_id, document_id + FROM demo_materializations + WHERE user_id = 'local-dev-user' + AND namespace = 'contract-demo-deduplicate' + """, + ) + job_rows = await ContractDatabase.fetch_all( + """ + SELECT job_id + FROM jobs + WHERE user_id = 'local-dev-user' + AND job_metadata ->> 'namespace' = 'contract-demo-deduplicate' + AND job_metadata ->> 'demo_source_id' = :demo_source_id + """, + {"demo_source_id": DEMO_SOURCE_ID}, + ) + + assert len(materialization_rows) == 1 + assert len(job_rows) == 1 + assert len(fake_result_storage.raw_files_by_job_id) == 1 + + @pytest.mark.asyncio async def test_should_serialize_concurrent_first_demo_materialization( developer_api_client_factory: Callable[ @@ -281,10 +341,10 @@ async def test_should_serialize_concurrent_first_demo_materialization( fake_result_storage = FakeResultStorage() async with developer_api_client_factory() as api_client: - import app.services.demo_document_service as demo_document_service + import app.services.demo_source_materializer as demo_source_materializer monkeypatch.setattr( - demo_document_service, + demo_source_materializer, "get_result_storage", lambda: fake_result_storage, ) @@ -377,10 +437,10 @@ async def test_should_reject_mixed_demo_materialization_selection_before_upload( fake_result_storage = FakeResultStorage() async with developer_api_client_factory() as api_client: - import app.services.demo_document_service as demo_document_service + import app.services.demo_source_materializer as demo_source_materializer monkeypatch.setattr( - demo_document_service, + demo_source_materializer, "get_result_storage", lambda: fake_result_storage, ) From f278e1ef7948fd4bb3a0dd01934e448ebf6f6b83 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:51:03 +0000 Subject: [PATCH 40/40] refactor: remove dead code --- apps/api/app/core/response/ResponseCode.py | 31 ---- .../s3_events/signature_verification.py | 8 +- apps/api/scripts/add_credits.py | 1 - .../connect_builder/summary_builder.py | 10 - .../document_ingestion/processing_run.py | 9 +- .../services/document_parser/doc_parser.py | 2 - .../services/document_parser/layout_parser.py | 86 --------- .../document_parser/mineru_pdf_service.py | 7 + .../document_parser/pymupdf_subprocess.py | 4 +- .../shared-python/shared/core/async_utils.py | 32 ---- .../shared/core/celery_router.py | 10 +- packages/shared-python/shared/core/logging.py | 5 +- .../shared/core/state_machine/config.py | 92 ---------- .../storage/file_encryptor_service.py | 44 ----- .../shared/services/webhook/dispatcher.py | 1 - .../shared/utils/file_transfer.py | 171 ------------------ 16 files changed, 19 insertions(+), 494 deletions(-) delete mode 100644 packages/shared-python/shared/core/async_utils.py delete mode 100644 packages/shared-python/shared/core/state_machine/config.py delete mode 100755 packages/shared-python/shared/services/storage/file_encryptor_service.py delete mode 100644 packages/shared-python/shared/utils/file_transfer.py diff --git a/apps/api/app/core/response/ResponseCode.py b/apps/api/app/core/response/ResponseCode.py index 55bb13358..85ee75d9c 100644 --- a/apps/api/app/core/response/ResponseCode.py +++ b/apps/api/app/core/response/ResponseCode.py @@ -28,34 +28,3 @@ def get_all_as_dict(cls) -> Dict[int, str]: """Return all response codes and messages as a dictionary.""" return {member.code: member.msg for member in cls} - - -IS_TEST_MODE = False - - -if __name__ == "__main__": - if IS_TEST_MODE: - # 1. Access an enum member. - success_code = ResponseCode.SUCCESS - print(f"Member: {success_code}") - # Output: Member: ResponseCode.SUCCESS - - # 2. Access member attributes (code and msg). - print(f"Code: {success_code.code}, Message: {success_code.msg}") - # Output: Code: 200, Message: Operation succeeded - - fail_code = ResponseCode.FAIL - print(f"Code: {fail_code.code}, Message: {fail_code.msg}") - # Output: Code: 1, Message: Operation failed - - # 3. Iterate over all enum members. - print("\n--- All Response Codes ---") - for member in ResponseCode: - print(f"{member.name}: code={member.code}, msg='{member.msg}'") - - # 4. Call the class method to get a dictionary. - all_messages = ResponseCode.get_all_as_dict() - print("\n--- Dictionary Form ---") - import json - - print(json.dumps(all_messages, indent=2, ensure_ascii=False)) diff --git a/apps/api/app/services/s3_events/signature_verification.py b/apps/api/app/services/s3_events/signature_verification.py index 7665ef07d..86d918b23 100644 --- a/apps/api/app/services/s3_events/signature_verification.py +++ b/apps/api/app/services/s3_events/signature_verification.py @@ -5,11 +5,8 @@ def verify_sns_signature(request_body: bytes, signature: str, message: str) -> bool: - try: - return True - except Exception as exc: - logger.error(f"SNS signature verification failed: {exc}") - return False + del request_body, signature, message + return True def verify_minio_signature(auth_token: str, expected_token: str) -> bool: @@ -19,6 +16,7 @@ def verify_minio_signature(auth_token: str, expected_token: str) -> bool: def verify_oss_signature(request_body: bytes, headers: dict[str, str]) -> bool: + del request_body, headers try: from shared.core.config import settings diff --git a/apps/api/scripts/add_credits.py b/apps/api/scripts/add_credits.py index f468aa281..675f5ce26 100644 --- a/apps/api/scripts/add_credits.py +++ b/apps/api/scripts/add_credits.py @@ -33,7 +33,6 @@ def _bootstrap_python_path() -> None: from shared.models.database.tier_limit import TierLimit from shared.models.database.user import User from shared.models.database.user_balance import UserBalance -from shared.models.database.webhook import WebhookEvent # noqa: F401 from shared.services.billing import CreditsService from shared.utils.utc_now import utc_now_naive diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 6d1e507c8..7fcf8c6eb 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -28,16 +28,6 @@ NON_LLM_TOP_SUMMARY_MAX_SECTIONS = 20 NON_LLM_TOP_SUMMARY_MAX_DEPTH = 2 -_TREE_EXCLUDED_TITLES = {"root", "images", "tables"} -_TREE_TITLE_MAX_TOKENS_START = 20 -_TREE_TITLE_MAX_TOKENS_END = 5 -_TITLE_ENUM_PREFIXES = ( - "this section covers:", - "this section includes", - "this document covers:", - "this document includes", -) - # ─── LLM Interface ─────────────────────────────────────────────────────────── diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index 89dc4d765..789b9ee21 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -67,7 +67,7 @@ def execute(self, job_id: str, user_id: str | None) -> dict[str, object]: with RedisJobLock(job_context.redis_service, job_id): task_workspace_dir, input_dir, output_dir = _prepare_task_workspace(job_id) try: - return _run_parse_job( + result = _run_parse_job( job_id=job_id, job_context=job_context, lifecycle_service=lifecycle_service, @@ -78,12 +78,7 @@ def execute(self, job_id: str, user_id: str | None) -> dict[str, object]: finally: cleanup_task_workspace(task_workspace_dir) - raise WorkerHandlingException( - user_message="We could not complete document processing", - internal_message=( - f"Parse workflow exited without a result for job_id={job_id}" - ), - ) + return result def _prepare_task_workspace(job_id: str) -> tuple[str, str, str]: diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index 77238675d..d405fed56 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -440,8 +440,6 @@ def parse_docx( output_dir=None, filename="", file_url="", - start_text="", - end_text="", relative_root=None, ): doc_data = load_file_bytes(docx_path, file_url=file_url) diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py index 8d9e94ed5..b0490d576 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/layout_parser.py @@ -6,7 +6,6 @@ from app.services.document_parser.heading_candidates import ( filter_document_headings, filter_markdown_headings, - judge_by_conditions, postprocess_headings, ) from app.services.document_parser.heading_llm_executor import ( @@ -27,15 +26,6 @@ from app.services.document_parser.table_text_parser import df2md from gevent.pool import Pool as GeventPool -try: - from markitdown import MarkItDown -except ImportError: - # Fall back to a pass-through shim when markitdown is unavailable. - class MarkItDown: - def convert(self, content): - return content - - from loguru import logger from shared.core.config import settings @@ -93,31 +83,6 @@ def remove_isolated_nodes(tree): return remove_isolated_heading_nodes(tree) -# def if_no_pos_code(reason_str: str) -> bool: -# """ -# Check whether all pos_code values are zero. -# reason format: "POS [0, 0, ...] NEG [...]" -# """ -# if not reason_str or not isinstance(reason_str, str): -# return True - -# pos_match = re.search(r'POS\s*\[([^\]]*)\]', reason_str) -# if not pos_match: -# return True -# pos_content = pos_match.group(1) -# try: -# nums = [int(x.strip()) for x in pos_content.split(',') if x.strip()] -# return all(x == 0 for x in nums) -# except: -# return True - - - -def detect_outlines_md(line): - pos_code = judge_by_conditions(line) - any(x > 0 for x in pos_code) - - def format_toc_context_for_llm(toc_context) -> str: """Convert TOC hierarchy or structured payloads into compact LLM-friendly plain text.""" if not toc_context: @@ -711,54 +676,3 @@ def est_hierarchies_llm( output_dir=output_dir, csv_suffix=csv_suffix, ) - - -# def parse_outline_hier(markdown_text): -# lines = markdown_text.strip().splitlines() -# stack = [] -# root = [] -# for line in lines: -# line = line.replace('markdown', '') # handle possible unexpected outputs -# if not line.strip(): -# continue - -# stripped = line.lstrip() -# indent = len(line) - len(stripped) -# match = re.match(r"[-*+] (.+)", stripped) -# if not match: -# continue - -# title = match.group(1).strip() -# node = {"chapter": title, "children": [], 'serial': 1} -# level = indent // 2 # Two spaces per level, adjustable if needed. -# if level == 0: -# node['serial'] = len(root)+1 -# root.append(node) -# stack = [(level, node)] -# else: -# while stack and stack[-1][0] >= level: -# stack.pop() -# if stack: -# parent = stack[-1][1] -# node['serial'] = len(parent['children']) + 1 -# parent["children"].append(node) -# stack.append((level, node)) -# return root - - -# def outline_to_markdown(nodes, level=0, path=""): -# rows = [] -# def traverse(node_list, level, path_prefix): -# for node in node_list: -# split_char = settings.SPLIT_CHAR or "/" -# current_path = f"{path_prefix} {split_char} {node['chapter']}" if path_prefix else node['chapter'] -# rows.append({ -# "path": current_path, -# "title": node["chapter"], -# "thoughts": node.get("thoughts", "").strip(), -# "level": level -# }) -# if node.get("children"): -# traverse(node["children"], level + 1, current_path) -# traverse(nodes, level, path) -# return pd.DataFrame(rows) diff --git a/apps/worker/app/services/document_parser/mineru_pdf_service.py b/apps/worker/app/services/document_parser/mineru_pdf_service.py index 2025609ff..bd4acdbc4 100644 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ b/apps/worker/app/services/document_parser/mineru_pdf_service.py @@ -395,6 +395,8 @@ def parse_via_full( output_dir: str, s3_key: Optional[str] = None, ) -> None: + batch_id: str | None = None + token_id: str | None = None resolved_s3_key = resolve_mineru_source_s3_key( s3_key=s3_key, local_file_path=None if is_remote(pdf_url) else pdf_url, @@ -426,6 +428,11 @@ def parse_via_full( batch_id, upload_url, token_id = _request_upload_target(pdf_url, filename) _upload_file_to_mineru(pdf_url, filename, upload_url, token_id) + if batch_id is None or token_id is None: + raise MinerUServiceException( + internal_message="MinerU task setup completed without a batch id or token" + ) + poll_mineru_task( status_url=f"{settings.MINERU_URL}/extract-results/batch/{batch_id}", task_id=batch_id, diff --git a/apps/worker/app/services/document_parser/pymupdf_subprocess.py b/apps/worker/app/services/document_parser/pymupdf_subprocess.py index c98b4678c..88bec26af 100644 --- a/apps/worker/app/services/document_parser/pymupdf_subprocess.py +++ b/apps/worker/app/services/document_parser/pymupdf_subprocess.py @@ -24,6 +24,7 @@ from multiprocessing.process import BaseProcess from multiprocessing.queues import Queue as MultiprocessingQueue from threading import RLock +from typing import TYPE_CHECKING from app.core.runtime_limits import read_pymupdf_max_concurrent from loguru import logger @@ -33,7 +34,8 @@ TimeoutException, ) -from gevent.threadpool import ThreadPool as GeventThreadPool +if TYPE_CHECKING: + from gevent.threadpool import ThreadPool as GeventThreadPool # Default timeout for child processes (seconds) DEFAULT_TIMEOUT = 3000 diff --git a/packages/shared-python/shared/core/async_utils.py b/packages/shared-python/shared/core/async_utils.py deleted file mode 100644 index 1975d8b7a..000000000 --- a/packages/shared-python/shared/core/async_utils.py +++ /dev/null @@ -1,32 +0,0 @@ -import asyncio -from typing import Any, Coroutine, TypeVar - -T = TypeVar("T") - - -def run_async_task(coro: Coroutine[Any, Any, T]) -> T: - """ - Run an async task in a synchronous context, reusing the event loop if possible. - - This function attempts to get the current event loop. If it's closed or missing, - it creates a new one but DOES NOT close it after execution (unlike asyncio.run). - This allows long-lived async resources to persist across tasks. - - Args: - coro: The coroutine to run. - - Returns: - The return value of the coroutine. - """ - try: - loop = asyncio.get_event_loop() - if loop.is_closed(): - # Loop exists but closed - create new one - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - except RuntimeError: - # No loop in this thread - create new one - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - return loop.run_until_complete(coro) diff --git a/packages/shared-python/shared/core/celery_router.py b/packages/shared-python/shared/core/celery_router.py index fe2ca4502..7c8fd86fb 100644 --- a/packages/shared-python/shared/core/celery_router.py +++ b/packages/shared-python/shared/core/celery_router.py @@ -214,16 +214,8 @@ def get_queue_for_job(self, job_type: str, user_id: str) -> str: """ try: # TODO: temporary simplified path to avoid async work here. - priority_level = 1 # Default free-subscription level. - - # Choose the queue by task type and priority level. if job_type in ["kb_management", "kb_encoding"]: - if priority_level >= 9: - return "kb_high" - elif priority_level >= 5: - return "kb_medium" - else: - return "kb_low" + return "kb_low" elif job_type in ["ai_query", "user_auth", "urgent_document"]: return "ai_high_priority" elif job_type in ["document_processing"]: diff --git a/packages/shared-python/shared/core/logging.py b/packages/shared-python/shared/core/logging.py index c2633b806..bb246f18e 100644 --- a/packages/shared-python/shared/core/logging.py +++ b/packages/shared-python/shared/core/logging.py @@ -3,11 +3,12 @@ from contextlib import contextmanager from contextvars import ContextVar from enum import Enum -from typing import Any, Dict +from typing import TYPE_CHECKING, Any, Dict from loguru import logger -from logfire.types import ExceptionCallbackHelper +if TYPE_CHECKING: + from logfire.types import ExceptionCallbackHelper _log_context: ContextVar[Dict[str, Any]] = ContextVar("log_context", default={}) _DEFAULT_CONSOLE_FORMAT = ( diff --git a/packages/shared-python/shared/core/state_machine/config.py b/packages/shared-python/shared/core/state_machine/config.py deleted file mode 100644 index 2128df763..000000000 --- a/packages/shared-python/shared/core/state_machine/config.py +++ /dev/null @@ -1,92 +0,0 @@ -"""State-machine configuration.""" - -from dataclasses import dataclass, field -from typing import Dict - - -DEFAULT_STATE_TIMEOUTS: Dict[str, int] = { - "pending": 300, - "uploading": 600, - "processing": 1800, - "completed": 0, - "failed": 0, -} - - -@dataclass -class StateMachineConfig: - """State-machine settings.""" - - max_retries: int = 3 # Maximum retry count. - base_retry_delay: float = 0.1 # Base retry delay in seconds. - - # Timeout settings used with Redis Keyspace Notifications. - state_timeouts: Dict[str, int] = field( - default_factory=lambda: DEFAULT_STATE_TIMEOUTS.copy() - ) - - # Synchronization settings. - sync_batch_size: int = 100 # Batch size for sync work. - sync_interval: int = 300 # Sync interval in seconds. - - # Maintenance settings. - maintenance_interval: int = 3600 # Maintenance interval in seconds. - cleanup_interval: int = 1800 # Cleanup interval in seconds. - - # Redis Keyspace Notifications support. - enable_keyspace_notifications: bool = True # Enable Keyspace Notifications. - -# Default state-machine configuration. -DEFAULT_CONFIG = StateMachineConfig() - - -def get_state_machine_config() -> StateMachineConfig: - """Return the active state-machine configuration.""" - return DEFAULT_CONFIG - - -def update_state_machine_config( - *, - max_retries: int | None = None, - base_retry_delay: float | None = None, - state_timeouts: Dict[str, int] | None = None, - sync_batch_size: int | None = None, - sync_interval: int | None = None, - maintenance_interval: int | None = None, - cleanup_interval: int | None = None, - enable_keyspace_notifications: bool | None = None, -) -> StateMachineConfig: - """Update and return the active state-machine configuration.""" - global DEFAULT_CONFIG - current = DEFAULT_CONFIG - DEFAULT_CONFIG = StateMachineConfig( - max_retries=current.max_retries if max_retries is None else max_retries, - base_retry_delay=( - current.base_retry_delay - if base_retry_delay is None - else base_retry_delay - ), - state_timeouts=( - current.state_timeouts.copy() - if state_timeouts is None - else state_timeouts - ), - sync_batch_size=( - current.sync_batch_size if sync_batch_size is None else sync_batch_size - ), - sync_interval=current.sync_interval if sync_interval is None else sync_interval, - maintenance_interval=( - current.maintenance_interval - if maintenance_interval is None - else maintenance_interval - ), - cleanup_interval=( - current.cleanup_interval if cleanup_interval is None else cleanup_interval - ), - enable_keyspace_notifications=( - current.enable_keyspace_notifications - if enable_keyspace_notifications is None - else enable_keyspace_notifications - ), - ) - return DEFAULT_CONFIG diff --git a/packages/shared-python/shared/services/storage/file_encryptor_service.py b/packages/shared-python/shared/services/storage/file_encryptor_service.py deleted file mode 100755 index e542a705a..000000000 --- a/packages/shared-python/shared/services/storage/file_encryptor_service.py +++ /dev/null @@ -1,44 +0,0 @@ -import os -import pickle -from typing import Any - -from cryptography.fernet import Fernet - - -class FernetPickleEncryptor: - encrypt = False - - def __init__(self, key: bytes = b"nc1BPZSkNb7Oc82_Wo3QoZTmJCEnQtpKZ2n-Z5F4CwY="): - self.cipher = Fernet(key) - - def save_to_file(self, data: Any, file_path: str) -> None: - serialized_data = pickle.dumps(data) # Serialize the input payload. - encrypted_data = self.cipher.encrypt( - serialized_data - ) # Encrypt the serialized bytes. - with open(file_path, "wb") as f: - f.write(encrypted_data) - - def load_from_file(self, file_path: str) -> Any: - if not os.path.exists(file_path): - raise FileNotFoundError(f"File {file_path} does not exist.") - with open(file_path, "rb") as f: - encrypted_data = f.read() - decrypted_data = self.cipher.decrypt(encrypted_data) # Decrypt the file bytes. - loaded_data = pickle.loads(decrypted_data) # Deserialize the decrypted payload. - return loaded_data - - -encryptor = FernetPickleEncryptor() - -if __name__ == "__main__": - # 2. Encrypt the payload. - data = {"key": "value"} - # 3. Save the encrypted payload to a file. - file_path = "data.pkl" - encryptor.save_to_file(data, file_path) - print(f"Encrypted data saved to file: {file_path}") - - # 4. Load the encrypted payload back from disk. - decrypted_data = encryptor.load_from_file(file_path) - print("Decrypted data loaded from file:", decrypted_data) diff --git a/packages/shared-python/shared/services/webhook/dispatcher.py b/packages/shared-python/shared/services/webhook/dispatcher.py index bf6ada8e9..0f3c4b953 100644 --- a/packages/shared-python/shared/services/webhook/dispatcher.py +++ b/packages/shared-python/shared/services/webhook/dispatcher.py @@ -13,7 +13,6 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -# Use standard db context - run_async_task handles the loop reuse from shared.core.database import get_db_context from shared.core.exceptions.webhook_exceptions import WebhookDeliveryException from shared.models.database.webhook import WebhookEvent, WebhookEventStatus diff --git a/packages/shared-python/shared/utils/file_transfer.py b/packages/shared-python/shared/utils/file_transfer.py deleted file mode 100644 index d307e7995..000000000 --- a/packages/shared-python/shared/utils/file_transfer.py +++ /dev/null @@ -1,171 +0,0 @@ -""" -File Transfer Utilities - -Provides reliable file transfer operations for large files using temp files as buffers. -Uses httpx for proper total timeout enforcement. -""" - -import os -import tempfile -from typing import Dict, Optional -from urllib.parse import urlparse - -import httpx -from loguru import logger - -from shared.utils.http_clients import get_sync_client - - -class FileTransferError(Exception): - """Base exception for file transfer operations""" - - def __init__(self, message: str, status_code: Optional[int] = None): - super().__init__(message) - self.status_code = status_code - - -class DownloadError(FileTransferError): - """ - Download failed - typically a client error. - - The source file may be inaccessible, expired, or invalid. - Worker should raise a 4xx (client error) when catching this. - """ - - pass - - -class UploadError(FileTransferError): - """ - Upload failed - typically a server/service error. - - The target service (e.g., MinerU) may be unavailable or experiencing issues. - Worker should raise a 5xx (server error) when catching this. - """ - - pass - - -def stream_download_and_upload( - source_url: str, - target_url: str, - download_timeout: int = 300, - upload_timeout: int = 300, - chunk_size: int = 8192, - upload_method: str = "PUT", - upload_headers: Optional[Dict[str, str]] = None, - upload_retries: int = 3, -) -> httpx.Response: - """ - Download a file from source_url and upload to target_url using a temp file buffer. - - Uses httpx for proper total timeout enforcement. - Retries upload on failure since temp file is preserved on disk. - - Args: - source_url: URL to download the file from - target_url: URL to upload the file to - download_timeout: Total timeout for download in seconds - upload_timeout: Total timeout for upload in seconds - chunk_size: Chunk size for streaming download - upload_method: HTTP method for upload (PUT or POST) - upload_headers: Additional headers for upload request - upload_retries: Number of retry attempts for upload (default 3) - - Returns: - httpx.Response: The upload response - - Raises: - DownloadError: If download fails (source inaccessible) - UploadError: If upload fails after all retries - """ - # Create temp file manually for explicit cleanup control - tmp_fd, tmp_path = tempfile.mkstemp(suffix=".tmp") - source_host = urlparse(source_url).hostname or source_url[:60] - target_host = urlparse(target_url).hostname or target_url[:60] - - try: - # Phase 1: Download to temp file - logger.debug(f"Downloading from {source_url[:100]}...") - try: - client = get_sync_client() - with client.stream("GET", source_url, timeout=download_timeout) as response: - response.raise_for_status() - - downloaded_bytes = 0 - with os.fdopen(tmp_fd, "wb") as tmp_file: - for chunk in response.iter_bytes(chunk_size=chunk_size): - tmp_file.write(chunk) - downloaded_bytes += len(chunk) - - except httpx.TimeoutException as e: - raise DownloadError( - f"Download timed out: host={source_host}, timeout={download_timeout}s" - ) from e - except httpx.HTTPStatusError as e: - raise DownloadError( - f"Download failed: host={source_host}, status={e.response.status_code}" - ) from e - except httpx.RequestError as e: - raise DownloadError( - f"Download failed: host={source_host}, error={e}" - ) from e - - # Get file size - file_size = os.path.getsize(tmp_path) - logger.info(f"Downloaded {file_size} bytes to temp file") - - # Phase 2: Upload from temp file (with retries) - headers = upload_headers or {} - headers["Content-Length"] = str(file_size) - - last_error = None - for attempt in range(1, upload_retries + 1): - try: - logger.info( - f"Uploading {file_size} bytes (attempt {attempt}/{upload_retries}, timeout={upload_timeout}s)..." - ) - - # Stream directly from file without loading to memory - with open(tmp_path, "rb") as f: - client = get_sync_client() - if upload_method.upper() == "PUT": - upload_response = client.put( - target_url, - content=f, - headers=headers, - timeout=upload_timeout, - ) - else: - upload_response = client.post( - target_url, - content=f, - headers=headers, - timeout=upload_timeout, - ) - - logger.info(f"Upload completed: status={upload_response.status_code}") - return upload_response - - except (httpx.TimeoutException, httpx.RequestError) as e: - last_error = e - logger.warning( - f"Upload attempt {attempt} failed: host={target_host}, error={e}" - ) - if attempt < upload_retries: - logger.info("Retrying upload...") - continue - - # All retries exhausted - raise UploadError( - f"Upload failed: host={target_host}, attempts={upload_retries}, last_error={last_error}" - ) from last_error - - finally: - # Manual cleanup of temp file - if os.path.exists(tmp_path): - try: - os.remove(tmp_path) - logger.debug(f"Temp file cleaned up: {tmp_path}") - except OSError as e: - logger.warning(f"Failed to cleanup temp file {tmp_path}: {e}")